release arrow context before any return

ExpediaGroup · Oct 9, 2023 · bb36c81 · bb36c81
1 parent 9a9577b
commit bb36c81
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/go/internal/feast/transformation/transformation.go b/go/internal/feast/transformation/transformation.go
@@ -56,13 +56,15 @@ func AugmentResponseWithOnDemandTransforms(
 		for name, values := range requestData {
 			requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
 			if err != nil {
+				ReleaseArrowContext(requestContextArrow)
 				return nil, err
 			}
 		}
 
 		for name, values := range entityRows {
 			requestContextArrow[name], err = types.ProtoValuesToArrowArray(values.Val, arrowMemory, numRows)
 			if err != nil {
+				ReleaseArrowContext(requestContextArrow)
 				return nil, err
 			}
 		}
@@ -80,19 +82,24 @@ func AugmentResponseWithOnDemandTransforms(
 			fullFeatureNames,
 		)
 		if err != nil {
+			ReleaseArrowContext(requestContextArrow)
 			return nil, err
 		}
 		result = append(result, onDemandFeatures...)
 
-		// Release memory used by requestContextArrow
-		for _, arrowArray := range requestContextArrow {
-			arrowArray.Release()
-		}
+		ReleaseArrowContext(requestContextArrow)
 	}
 
 	return result, nil
 }
 
+func ReleaseArrowContext(requestContextArrow map[string]arrow.Array) {
+	// Release memory used by requestContextArrow
+	for _, arrowArray := range requestContextArrow {
+		arrowArray.Release()
+	}
+}
+
 func CallTransformations(
 	featureView *model.OnDemandFeatureView,
 	retrievedFeatures map[string]arrow.Array,