fix memory error in findnearst option

Author: cherylEnkidu

Firestore/Source/API/FIRPipelineBridge.mm

@@ -73,7 +73,9 @@
 using firebase::firestore::api::Union;
 using firebase::firestore::api::Unnest;
 using firebase::firestore::api::Where;
+using firebase::firestore::model::DeepClone;
 using firebase::firestore::model::FieldPath;
+using firebase::firestore::nanopb::MakeSharedMessage;
 using firebase::firestore::nanopb::SharedMessage;
 using firebase::firestore::util::ComparisonResult;
 using firebase::firestore::util::MakeCallback;
@@ -560,7 +562,7 @@ @implementation FIRFindNearestStageBridge {
   FIRVectorValue *_vectorValue;
   NSString *_distanceMeasure;
   NSNumber *_limit;
-  NSString *_Nullable _distanceField;
+  FIRExprBridge *_Nullable _distanceField;
   Boolean isUserDataRead;
   std::shared_ptr<FindNearestStage> cpp_find_nearest;
 }
@@ -569,7 +571,7 @@ - (id)initWithField:(FIRFieldBridge *)field
         vectorValue:(FIRVectorValue *)vectorValue
     distanceMeasure:(NSString *)distanceMeasure
               limit:(NSNumber *_Nullable)limit
-      distanceField:(NSString *_Nullable)distanceField {
+      distanceField:(FIRExprBridge *_Nullable)distanceField {
   self = [super init];
   if (self) {
     _field = field;
@@ -584,21 +586,16 @@ - (id)initWithField:(FIRFieldBridge *)field
 
 - (std::shared_ptr<api::Stage>)cppStageWithReader:(FSTUserDataReader *)reader {
   if (!isUserDataRead) {
-    std::unordered_map<std::string,
-                       nanopb::SharedMessage<firebase::firestore::google_firestore_v1_Value>>
-        optional_value;
+    std::unordered_map<std::string, firebase::firestore::google_firestore_v1_Value> optional_value;
     if (_limit) {
-      optional_value.emplace(
-          std::make_pair(std::string("limit"),
-                         nanopb::SharedMessage<firebase::firestore::google_firestore_v1_Value>(
-                             [reader parsedQueryValue:_limit])));
+      optional_value.emplace(std::make_pair(
+          std::string("limit"), *DeepClone(*[reader parsedQueryValue:_limit]).release()));
     }
 
     if (_distanceField) {
+      std::shared_ptr<Expr> cpp_distance_field = [_distanceField cppExprWithReader:reader];
       optional_value.emplace(
-          std::make_pair(std::string("distance_field"),
-                         nanopb::SharedMessage<firebase::firestore::google_firestore_v1_Value>(
-                             [reader parsedQueryValue:_distanceField])));
+          std::make_pair(std::string("distance_field"), cpp_distance_field->to_proto()));
     }
 
     FindNearestStage::DistanceMeasure::Measure measure_enum;

Firestore/Source/Public/FirebaseFirestore/FIRPipelineBridge.h

@@ -160,7 +160,7 @@ NS_SWIFT_NAME(FindNearestStageBridge)
         vectorValue:(FIRVectorValue *)vectorValue
     distanceMeasure:(NSString *)distanceMeasure
               limit:(NSNumber *_Nullable)limit
-      distanceField:(NSString *_Nullable)distanceField;
+      distanceField:(FIRExprBridge *_Nullable)distanceField;
 @end
 
 NS_SWIFT_SENDABLE

Firestore/Swift/Source/SwiftAPI/Stages.swift

@@ -262,7 +262,7 @@ class FindNearest: Stage {
       vectorValue: VectorValue(vectorValue),
       distanceMeasure: distanceMeasure.kind.rawValue,
       limit: limit as NSNumber?,
-      distanceField: distanceField
+      distanceField: distanceField.map { Field($0).toBridge() } ?? nil
     )
   }
 }

Firestore/Swift/Tests/Integration/PipelineTests.swift

@@ -1508,57 +1508,57 @@ class PipelineIntegrationTests: FSTIntegrationTestCase {
 
     TestHelper.compare(pipelineSnapshot: snapshot, expected: expectedResults, enforceOrder: false)
   }
-  
-  func testFindNearest() async throws {
-      let collRef = collectionRef(withDocuments: bookDocs)
-      let db = collRef.firestore
-
-      let measures: [DistanceMeasure] = [.euclidean, .dotProduct, .cosine]
-      let expectedResults: [[String: Sendable]] = [
-        ["title": "The Hitchhiker's Guide to the Galaxy"],
-        ["title": "One Hundred Years of Solitude"],
-        ["title": "The Handmaid's Tale"],
-      ]
-
-      for measure in measures {
-        let pipeline = db.pipeline()
-          .collection(collRef.path)
-          .findNearest(
-            field: Field("embedding"),
-            vectorValue: [10, 1, 3, 1, 2, 1, 1, 1, 1, 1],
-            distanceMeasure: measure, limit: 3
-          )
-          .select("title")
-        let snapshot = try await pipeline.execute()
-        TestHelper.compare(pipelineSnapshot: snapshot, expected: expectedResults, enforceOrder: true)
-      }
-    }
 
-    func testFindNearestWithDistance() async throws {
-      let collRef = collectionRef(withDocuments: bookDocs)
-      let db = collRef.firestore
+  func testFindNearest() async throws {
+    let collRef = collectionRef(withDocuments: bookDocs)
+    let db = collRef.firestore
 
-      let expectedResults: [[String: Sendable]] = [
-        [
-          "title": "The Hitchhiker's Guide to the Galaxy",
-          "computedDistance": 1.0,
-        ],
-        [
-          "title": "One Hundred Years of Solitude",
-          "computedDistance": 12.041594578792296,
-        ],
-      ]
+    let measures: [DistanceMeasure] = [.euclidean, .dotProduct, .cosine]
+    let expectedResults: [[String: Sendable]] = [
+      ["title": "The Hitchhiker's Guide to the Galaxy"],
+      ["title": "One Hundred Years of Solitude"],
+      ["title": "The Handmaid's Tale"],
+    ]
 
+    for measure in measures {
       let pipeline = db.pipeline()
         .collection(collRef.path)
         .findNearest(
           field: Field("embedding"),
-          vectorValue: [10, 1, 2, 1, 1, 1, 1, 1, 1, 1],
-          distanceMeasure: .euclidean, limit: 2,
-          distanceField: "computedDistance"
+          vectorValue: [10, 1, 3, 1, 2, 1, 1, 1, 1, 1],
+          distanceMeasure: measure, limit: 3
         )
-        .select("title", "computedDistance")
+        .select("title")
       let snapshot = try await pipeline.execute()
-      TestHelper.compare(pipelineSnapshot: snapshot, expected: expectedResults, enforceOrder: false)
+      TestHelper.compare(pipelineSnapshot: snapshot, expected: expectedResults, enforceOrder: true)
     }
+  }
+
+  func testFindNearestWithDistance() async throws {
+    let collRef = collectionRef(withDocuments: bookDocs)
+    let db = collRef.firestore
+
+    let expectedResults: [[String: Sendable]] = [
+      [
+        "title": "The Hitchhiker's Guide to the Galaxy",
+        "computedDistance": 1.0,
+      ],
+      [
+        "title": "One Hundred Years of Solitude",
+        "computedDistance": 12.041594578792296,
+      ],
+    ]
+
+    let pipeline = db.pipeline()
+      .collection(collRef.path)
+      .findNearest(
+        field: Field("embedding"),
+        vectorValue: [10, 1, 2, 1, 1, 1, 1, 1, 1, 1],
+        distanceMeasure: .euclidean, limit: 2,
+        distanceField: "computedDistance"
+      )
+      .select("title", "computedDistance")
+    let snapshot = try await pipeline.execute()
+    TestHelper.compare(pipelineSnapshot: snapshot, expected: expectedResults, enforceOrder: false)
+  }
 }

Firestore/core/src/api/stages.cc

@@ -198,11 +198,9 @@ google_firestore_v1_Pipeline_Stage FindNearestStage::to_proto() const {
 
   nanopb::SetRepeatedField(
       &result.options, &result.options_count, options_,
-      [](const std::pair<std::string,
-                         nanopb::SharedMessage<google_firestore_v1_Value>>&
-             entry) {
+      [](const std::pair<std::string, google_firestore_v1_Value>& entry) {
         return _google_firestore_v1_Pipeline_Stage_OptionsEntry{
-            nanopb::MakeBytesArray(entry.first), *DeepClone(*entry.second).release()};
+            nanopb::MakeBytesArray(entry.first), entry.second};
       });
 
   return result;

Firestore/core/src/api/stages.h

@@ -150,13 +150,11 @@ class FindNearestStage : public Stage {
       std::shared_ptr<Expr> property,
       nanopb::SharedMessage<google_firestore_v1_Value> vector,
       DistanceMeasure distance_measure,
-      std::unordered_map<std::string,
-                         nanopb::SharedMessage<google_firestore_v1_Value>>
-          options)
+      std::unordered_map<std::string, google_firestore_v1_Value> options)
       : property_(std::move(property)),
         vector_(std::move(vector)),
         distance_measure_(distance_measure),
-        options_(options) {
+        options_(std::move(options)) {
   }
 
   ~FindNearestStage() override = default;
@@ -167,9 +165,7 @@ class FindNearestStage : public Stage {
   std::shared_ptr<Expr> property_;
   nanopb::SharedMessage<google_firestore_v1_Value> vector_;
   DistanceMeasure distance_measure_;
-  std::unordered_map<std::string,
-                     nanopb::SharedMessage<google_firestore_v1_Value>>
-      options_;
+  std::unordered_map<std::string, google_firestore_v1_Value> options_;
 };
 
 class LimitStage : public Stage {