Fix another memory leak

felixwang9817 · felixwang9817 · commit a6232d964cf2 · 2022-07-13T17:45:03.000-07:00
Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/go/embedded/online_features.go b/go/embedded/online_features.go
@@ -241,14 +241,13 @@ func (s *OnlineFeatureService) GetOnlineFeatures(
 		s.tsColumnBuildersToRelease = append(s.tsColumnBuildersToRelease, tsColumnBuilder)
 		s.arraysToRelease = append(s.arraysToRelease, statusColumn)
 		s.arraysToRelease = append(s.arraysToRelease, tsColumn)
+		s.arraysToRelease = append(s.arraysToRelease, featureVector.Values)
 	}
 
 	result := array.NewRecord(arrow.NewSchema(outputFields, nil), outputColumns, int64(numRows))
 	s.resultsToRelease = append(s.resultsToRelease, result)
 
-	cdata.ExportArrowRecordBatch(result,
-		cdata.ArrayFromPtr(output.DataPtr),
-		cdata.SchemaFromPtr(output.SchemaPtr))
+	cdata.ExportArrowRecordBatch(result, cdata.ArrayFromPtr(output.DataPtr), cdata.SchemaFromPtr(output.SchemaPtr))
 
 	return nil
 }
diff --git a/go/internal/feast/onlineserving/serving.go b/go/internal/feast/onlineserving/serving.go
@@ -415,6 +415,8 @@ func KeepOnlyRequestedFeatures(
 	vectorsByName := make(map[string]*FeatureVector)
 	expectedVectors := make([]*FeatureVector, 0)
 
+	usedVectors := make(map[string]bool)
+
 	for _, vector := range vectors {
 		vectorsByName[vector.Name] = vector
 	}
@@ -438,6 +440,14 @@ func KeepOnlyRequestedFeatures(
 			return nil, fmt.Errorf("requested feature %s can't be retrieved", featureRef)
 		}
 		expectedVectors = append(expectedVectors, vectorsByName[qualifiedName])
+		usedVectors[qualifiedName] = true
+	}
+
+	// Free arrow arrays for vectors that were not used.
+	for _, vector := range vectors {
+		if _, ok := usedVectors[vector.Name]; !ok {
+			vector.Values.Release()
+		}
 	}
 
 	return expectedVectors, nil

Original file line number	Diff line number	Diff line change
`@@ -241,14 +241,13 @@ func (s *OnlineFeatureService) GetOnlineFeatures(`
`241`	`241`	`s.tsColumnBuildersToRelease = append(s.tsColumnBuildersToRelease, tsColumnBuilder)`
`242`	`242`	`s.arraysToRelease = append(s.arraysToRelease, statusColumn)`
`243`	`243`	`s.arraysToRelease = append(s.arraysToRelease, tsColumn)`
	`244`	`+ s.arraysToRelease = append(s.arraysToRelease, featureVector.Values)`
`244`	`245`	`}`
`245`	`246`
`246`	`247`	`result := array.NewRecord(arrow.NewSchema(outputFields, nil), outputColumns, int64(numRows))`
`247`	`248`	`s.resultsToRelease = append(s.resultsToRelease, result)`
`248`	`249`
`249`		`- cdata.ExportArrowRecordBatch(result,`
`250`		`- cdata.ArrayFromPtr(output.DataPtr),`
`251`		`- cdata.SchemaFromPtr(output.SchemaPtr))`
	`250`	`+ cdata.ExportArrowRecordBatch(result, cdata.ArrayFromPtr(output.DataPtr), cdata.SchemaFromPtr(output.SchemaPtr))`
`252`	`251`
`253`	`252`	`return nil`
`254`	`253`	`}`