diff --git a/packages/cloud_firestore/cloud_firestore/android/src/main/java/io/flutter/plugins/firebase/firestore/utils/PipelineStageHandlers.java b/packages/cloud_firestore/cloud_firestore/android/src/main/java/io/flutter/plugins/firebase/firestore/utils/PipelineStageHandlers.java index 78222ecef7e9..3f8072cb19af 100644 --- a/packages/cloud_firestore/cloud_firestore/android/src/main/java/io/flutter/plugins/firebase/firestore/utils/PipelineStageHandlers.java +++ b/packages/cloud_firestore/cloud_firestore/android/src/main/java/io/flutter/plugins/firebase/firestore/utils/PipelineStageHandlers.java @@ -311,6 +311,8 @@ private Pipeline handleFindNearest( List vectorValue = (List) args.get("vector_value"); String distanceMeasureStr = (String) args.get("distance_measure"); Number limitObj = (Number) args.get("limit"); + String distanceField = + args.get("distance_field") instanceof String ? (String) args.get("distance_field") : null; if (distanceMeasureStr == null) { throw new IllegalArgumentException("'find_nearest' requires a 'distance_measure' argument"); @@ -327,13 +329,18 @@ private Pipeline handleFindNearest( } Field fieldExpr = Expression.field(vectorField); + boolean hasDistanceField = distanceField != null && !distanceField.isEmpty(); - if (limitObj != null) { + if (limitObj != null || hasDistanceField) { + FindNearestOptions options = new FindNearestOptions(); + if (limitObj != null) { + options = options.withLimit(limitObj.intValue()); + } + if (distanceField != null && !distanceField.isEmpty()) { + options = options.withDistanceField(distanceField); + } return pipeline.findNearest( - vectorField, - Expression.vector(vectorArray), - distanceMeasure, - new FindNearestOptions().withLimit(limitObj.intValue())); + vectorField, Expression.vector(vectorArray), distanceMeasure, options); } else { return pipeline.findNearest(fieldExpr, vectorArray, distanceMeasure); } diff --git a/packages/cloud_firestore/cloud_firestore/ios/cloud_firestore/Sources/cloud_firestore/FLTPipelineParser.m b/packages/cloud_firestore/cloud_firestore/ios/cloud_firestore/Sources/cloud_firestore/FLTPipelineParser.m index f416fbc2bd4d..e5ee9faef8de 100644 --- a/packages/cloud_firestore/cloud_firestore/ios/cloud_firestore/Sources/cloud_firestore/FLTPipelineParser.m +++ b/packages/cloud_firestore/cloud_firestore/ios/cloud_firestore/Sources/cloud_firestore/FLTPipelineParser.m @@ -1603,11 +1603,14 @@ + (__FIRStageBridge *)parseSearchStageWithArgs:(NSDictionary *)args return nil; } FIRVectorValue *queryVector = [[FIRVectorValue alloc] initWithArray:numbers]; + __FIRExprBridge *distanceFieldExpr = + (distanceField.length > 0) ? [[__FIRFieldBridge alloc] initWithName:distanceField] + : nil; stage = [[__FIRFindNearestStageBridge alloc] initWithField:embeddingField vectorValue:queryVector distanceMeasure:distanceMeasure limit:limit - distanceField:distanceField]; + distanceField:distanceFieldExpr]; } else { if (error) *error = parseError([NSString stringWithFormat:@"Unknown pipeline stage: %@", stageName]); diff --git a/packages/cloud_firestore/cloud_firestore/lib/src/pipeline.dart b/packages/cloud_firestore/cloud_firestore/lib/src/pipeline.dart index 55ca21d79f47..083689c0143c 100644 --- a/packages/cloud_firestore/cloud_firestore/lib/src/pipeline.dart +++ b/packages/cloud_firestore/cloud_firestore/lib/src/pipeline.dart @@ -406,6 +406,12 @@ class Pipeline { /// compare vectors (e.g. [DistanceMeasure.cosine]). Use [limit] to return /// only the first N documents. /// + /// [distanceField] is the Firestore field name added to each result. That + /// field holds the calculated distance from [vectorValue]. A dot in the + /// name separates path segments, as it does for any other Firestore field + /// name. Omit it to return the nearest documents without the numeric + /// distance. + /// /// Example: /// ```dart /// firestore.pipeline().collection('books').findNearest( @@ -413,6 +419,7 @@ class Pipeline { /// [0.1, 0.2, 0.3], /// DistanceMeasure.cosine, /// limit: 10, + /// distanceField: 'distance', /// ); /// ``` Pipeline findNearest( @@ -420,12 +427,14 @@ class Pipeline { List vectorValue, DistanceMeasure distanceMeasure, { int? limit, + String? distanceField, }) { final stage = _FindNearestStage( vectorField, vectorValue, distanceMeasure, limit: limit, + distanceField: distanceField, ); return Pipeline._( _firestore, diff --git a/packages/cloud_firestore/cloud_firestore/lib/src/pipeline_stage.dart b/packages/cloud_firestore/cloud_firestore/lib/src/pipeline_stage.dart index e6eaf8df20f5..01385ecfd2fb 100644 --- a/packages/cloud_firestore/cloud_firestore/lib/src/pipeline_stage.dart +++ b/packages/cloud_firestore/cloud_firestore/lib/src/pipeline_stage.dart @@ -179,12 +179,14 @@ final class _FindNearestStage extends PipelineStage { final List vectorValue; final DistanceMeasure distanceMeasure; final int? limit; + final String? distanceField; _FindNearestStage( this.vectorField, this.vectorValue, this.distanceMeasure, { this.limit, + this.distanceField, }); @override @@ -203,6 +205,10 @@ final class _FindNearestStage extends PipelineStage { if (limit != null) { map['args']['limit'] = limit; } + final outputField = distanceField; + if (outputField != null && outputField.isNotEmpty) { + map['args']['distance_field'] = outputField; + } return map; } } diff --git a/packages/cloud_firestore/cloud_firestore/pipeline_example/integration_test/pipeline/pipeline_find_nearest_e2e.dart b/packages/cloud_firestore/cloud_firestore/pipeline_example/integration_test/pipeline/pipeline_find_nearest_e2e.dart index 35d5b6e75062..4425b63e5cef 100644 --- a/packages/cloud_firestore/cloud_firestore/pipeline_example/integration_test/pipeline/pipeline_find_nearest_e2e.dart +++ b/packages/cloud_firestore/cloud_firestore/pipeline_example/integration_test/pipeline/pipeline_find_nearest_e2e.dart @@ -41,5 +41,27 @@ void runPipelineFindNearestTests() { } }, ); + + test('findNearest distanceField is present on each result', () async { + final snapshot = await firestore + .pipeline() + .collection('pipeline-e2e') + .where(Expression.field('test').equalValue('find-nearest')) + .findNearest( + Field('embedding'), + [0.1, 0.2, 0.3], + DistanceMeasure.cosine, + limit: 5, + distanceField: 'distance', + ) + .execute(); + expect(snapshot.result, isNotEmpty); + for (final result in snapshot.result) { + final data = result.data(); + expect(data, isNotNull); + expect(data!.containsKey('distance'), isTrue); + expect(data['distance'], isA()); + } + }); }); } diff --git a/packages/cloud_firestore/cloud_firestore/test/pipeline_stage_test.dart b/packages/cloud_firestore/cloud_firestore/test/pipeline_stage_test.dart index 41a6fba094db..14485f6e6097 100644 --- a/packages/cloud_firestore/cloud_firestore/test/pipeline_stage_test.dart +++ b/packages/cloud_firestore/cloud_firestore/test/pipeline_stage_test.dart @@ -396,6 +396,38 @@ void main() { final stage = pipeline.stages.last; expect(stage['args']['distance_measure'], 'dotProduct'); }); + + test('serializes findNearest with distanceField', () { + final pipeline = firestore.pipeline().collection('items').findNearest( + Field('embedding'), + [0.1, 0.2, 0.3], + DistanceMeasure.cosine, + distanceField: 'distance', + ); + final stage = pipeline.stages.last; + expect(stage['args']['distance_field'], 'distance'); + }); + + test('omits distance_field when distanceField is absent', () { + final pipeline = firestore.pipeline().collection('items').findNearest( + Field('embedding'), + [0.1, 0.2, 0.3], + DistanceMeasure.cosine, + ); + final stage = pipeline.stages.last; + expect(stage['args'].containsKey('distance_field'), isFalse); + }); + + test('omits distance_field when distanceField is empty', () { + final pipeline = firestore.pipeline().collection('items').findNearest( + Field('embedding'), + [0.1, 0.2, 0.3], + DistanceMeasure.cosine, + distanceField: '', + ); + final stage = pipeline.stages.last; + expect(stage['args'].containsKey('distance_field'), isFalse); + }); }); group('_SearchStage', () {