chore: Release column vectors and their arrays correctly (#19496)

Signed-off-by: Christian Haudum <[email protected]>

Author: chaudum

pkg/engine/internal/executor/expressions.go

@@ -259,6 +259,7 @@ var _ ColumnVector = (*Array)(nil)
 
 // ToArray implements ColumnVector.
 func (a *Array) ToArray() arrow.Array {
+	a.array.Retain()
 	return a.array
 }
 

pkg/engine/internal/executor/expressions_test.go

@@ -216,6 +216,7 @@ func TestEvaluateBinaryExpression(t *testing.T) {
 func collectBooleanColumnVector(vec ColumnVector) []bool {
 	res := make([]bool, 0, vec.Len())
 	arr := vec.ToArray().(*array.Boolean)
+	defer arr.Release()
 	for i := range int(vec.Len()) {
 		res = append(res, arr.Value(i))
 	}
@@ -326,8 +327,10 @@ null,null,null`
 		colVec, err := e.eval(colExpr, record)
 		require.NoError(t, err)
 		require.IsType(t, &CoalesceVector{}, colVec)
+		defer colVec.Release()
 
 		arr := colVec.ToArray()
+		defer arr.Release()
 		require.IsType(t, &array.String{}, arr)
 		stringArr := arr.(*array.String)
 

pkg/engine/internal/executor/filter.go

@@ -24,20 +24,21 @@ func NewFilterPipeline(filter *physical.Filter, input Pipeline, evaluator expres
 		cols := make([]*array.Boolean, 0, len(filter.Predicates))
 
 		for i, pred := range filter.Predicates {
-			res, err := evaluator.eval(pred, batch)
+			vec, err := evaluator.eval(pred, batch)
 			if err != nil {
 				return nil, err
 			}
-			data := res.ToArray()
+			defer vec.Release()
 
+			arr := vec.ToArray()
 			// boolean filters are only used for filtering; they're not returned
 			// and must be released
-			defer data.Release()
+			defer arr.Release()
 
-			if data.DataType().ID() != arrow.BOOL {
-				return nil, fmt.Errorf("predicate %d returned non-boolean type %s", i, data.DataType())
+			if arr.DataType().ID() != arrow.BOOL {
+				return nil, fmt.Errorf("predicate %d returned non-boolean type %s", i, arr.DataType())
 			}
-			casted := data.(*array.Boolean)
+			casted := arr.(*array.Boolean)
 			cols = append(cols, casted)
 		}
 

pkg/engine/internal/executor/functions.go

@@ -161,6 +161,7 @@ type arrayType[T comparable] interface {
 	IsNull(int) bool
 	Value(int) T
 	Len() int
+	Release()
 }
 
 // genericFunction is a struct that implements the [BinaryFunction] interface methods
@@ -179,11 +180,13 @@ func (f *genericFunction[E, T]) Evaluate(lhs ColumnVector, rhs ColumnVector) (Co
 	if !ok {
 		return nil, arrow.ErrType
 	}
+	defer lhsArr.Release()
 
 	rhsArr, ok := rhs.ToArray().(E)
 	if !ok {
 		return nil, arrow.ErrType
 	}
+	defer rhsArr.Release()
 
 	mem := memory.NewGoAllocator()
 	builder := array.NewBooleanBuilder(mem)

pkg/engine/internal/executor/functions_test.go

@@ -120,6 +120,8 @@ func createFloat64Array(mem memory.Allocator, values []float64, nulls []bool) *A
 // Helper function to extract boolean values from result
 func extractBoolValues(result ColumnVector) []bool {
 	arr := result.ToArray().(*array.Boolean)
+	defer arr.Release()
+
 	values := make([]bool, arr.Len())
 	for i := 0; i < arr.Len(); i++ {
 		if arr.IsNull(i) {

pkg/engine/internal/executor/project.go

@@ -40,6 +40,7 @@ func NewProjectPipeline(input Pipeline, columns []physical.ColumnExpression, eva
 			if err != nil {
 				return nil, err
 			}
+			defer vec.Release()
 			ident := semconv.NewIdentifier(columnNames[i], vec.ColumnType(), vec.Type())
 			fields = append(fields, semconv.FieldFromIdent(ident, true))
 			arr := vec.ToArray()

pkg/engine/internal/executor/range_aggregation.go

@@ -163,6 +163,7 @@ func (r *rangeAggregationPipeline) read(ctx context.Context) (arrow.Record, erro
 				if err != nil {
 					return nil, err
 				}
+				defer vec.Release()
 
 				if vec.Type() != types.Loki.String {
 					return nil, fmt.Errorf("unsupported datatype for partitioning %s", vec.Type())
@@ -179,6 +180,7 @@ func (r *rangeAggregationPipeline) read(ctx context.Context) (arrow.Record, erro
 			if err != nil {
 				return nil, err
 			}
+			defer tsVec.Release()
 			tsCol := tsVec.ToArray().(*array.Timestamp)
 			defer tsCol.Release()
 

pkg/engine/internal/executor/sortmerge.go

@@ -154,11 +154,14 @@ loop:
 		if err != nil {
 			return nil, err
 		}
+		defer col.Release()
+
 		tsCol, ok := col.ToArray().(*array.Timestamp)
 		if !ok {
 			return nil, errors.New("column is not a timestamp column")
 		}
 		ts := tsCol.Value(int(p.offsets[i]))
+		tsCol.Release()
 
 		// Populate slices for sorting
 		inputIndexes = append(inputIndexes, i)
@@ -199,11 +202,13 @@ loop:
 	if err != nil {
 		return nil, err
 	}
+	defer col.Release()
 	// We assume the column is a Uint64 array
 	tsCol, ok := col.ToArray().(*array.Timestamp)
 	if !ok {
 		return nil, errors.New("column is not a timestamp column")
 	}
+	defer tsCol.Release()
 
 	// Calculate start/end of the sub-slice of the record
 	start := p.offsets[j]

pkg/engine/internal/executor/sortmerge_test.go

@@ -80,7 +80,9 @@ func TestSortMerge(t *testing.T) {
 
 			tsCol, err := c.evaluator.eval(merge.Column, batch)
 			require.NoError(t, err)
+			defer tsCol.Release()
 			arr := tsCol.ToArray().(*array.Timestamp)
+			defer arr.Release()
 
 			timestamps = append(timestamps, arr.Values()...)
 			batches++
@@ -126,8 +128,10 @@ func TestSortMerge(t *testing.T) {
 			}
 
 			tsCol, err := c.evaluator.eval(merge.Column, batch)
+			defer tsCol.Release()
 			require.NoError(t, err)
 			arr := tsCol.ToArray().(*array.Timestamp)
+			defer arr.Release()
 
 			timestamps = append(timestamps, arr.Values()...)
 			batches++

pkg/engine/internal/executor/vector_aggregate.go

@@ -102,6 +102,7 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 			if err != nil {
 				return nil, err
 			}
+			defer tsVec.Release()
 			tsCol := tsVec.ToArray().(*array.Timestamp)
 			defer tsCol.Release()
 
@@ -110,6 +111,7 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 			if err != nil {
 				return nil, err
 			}
+			defer valueVec.Release()
 			valueArr := valueVec.ToArray().(*array.Float64)
 			defer valueArr.Release()
 
@@ -121,6 +123,7 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 				if err != nil {
 					return nil, err
 				}
+				defer vec.Release()
 
 				if vec.Type() != types.Loki.String {
 					return nil, fmt.Errorf("unsupported datatype for grouping %s", vec.Type())