Add subquery support in pipeline (#15963)

Co-authored-by: Nick Cooke <36927374+ncooke3@users.noreply.github.com>
Co-authored-by: Yvonne Yu <yyyu@google.com>
Co-authored-by: Jeff Scudder <jscudder@google.com>

Author: 4 people

Firestore/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [feature] Added support for Firestore Pipeline Subqueries and scope bridging with `Subcollection`, `define()`, `toArrayExpression()`, `toScalarExpression()`, `Variable`, and `CurrentDocument` APIs.
+
 - [feature] Added support for Pipeline expressions `nor` and `switchOn`. (#15943)
 
 # 12.11.0

Firestore/Example/Firestore.xcodeproj/project.pbxproj

@@ -52,6 +52,10 @@ NS_ASSUME_NONNULL_BEGIN
 - (id)initWithCppStage:(std::shared_ptr<const firebase::firestore::api::CollectionSource>)stage;
 @end
 
+@interface FIRSubcollectionSourceStageBridge (Internal)
+- (id)initWithCppStage:(std::shared_ptr<const firebase::firestore::api::SubcollectionSource>)stage;
+@end
+
 @interface FIRDatabaseSourceStageBridge (Internal)
 - (id)initWithCppStage:(std::shared_ptr<const firebase::firestore::api::DatabaseSource>)stage;
 @end

Firestore/Source/API/FIRPipelineBridge+Internal.h

@@ -64,6 +64,7 @@
 using firebase::firestore::api::CollectionSource;
 using firebase::firestore::api::Constant;
 using firebase::firestore::api::DatabaseSource;
+using firebase::firestore::api::DefineStage;
 using firebase::firestore::api::DistinctStage;
 using firebase::firestore::api::DocumentChange;
 using firebase::firestore::api::DocumentReference;
@@ -88,8 +89,10 @@
 using firebase::firestore::api::SelectStage;
 using firebase::firestore::api::SnapshotMetadata;
 using firebase::firestore::api::SortStage;
+using firebase::firestore::api::SubcollectionSource;
 using firebase::firestore::api::Union;
 using firebase::firestore::api::Unnest;
+using firebase::firestore::api::Variable;
 using firebase::firestore::api::Where;
 using firebase::firestore::core::EventListener;
 using firebase::firestore::core::ViewSnapshot;
@@ -148,6 +151,30 @@ - (NSString *)field_name {
 
 @end
 
+@implementation FIRVariableBridge {
+  std::shared_ptr<Variable> cpp_variable;
+  NSString *_name;
+  Boolean isUserDataRead;
+}
+
+- (id)initWithName:(NSString *)name {
+  self = [super init];
+  if (self) {
+    _name = name;
+    isUserDataRead = NO;
+  }
+  return self;
+}
+
+- (std::shared_ptr<api::Expr>)cppExprWithReader:(FSTUserDataReader *)reader {
+  if (!isUserDataRead) {
+    cpp_variable = std::make_shared<Variable>(MakeString(_name));
+  }
+  isUserDataRead = YES;
+  return cpp_variable;
+}
+@end
+
 @implementation FIRConstantBridge {
   std::shared_ptr<Constant> cpp_constant;
   id _input;
@@ -302,6 +329,27 @@ - (NSString *)name {
 }
 @end
 
+@implementation FIRSubcollectionSourceStageBridge {
+  std::shared_ptr<SubcollectionSource> cpp_subcollection_source;
+}
+
+- (id)initWithPath:(NSString *)path {
+  self = [super init];
+  if (self) {
+    cpp_subcollection_source = std::make_shared<SubcollectionSource>(MakeString(path));
+  }
+  return self;
+}
+
+- (std::shared_ptr<api::Stage>)cppStageWithReader:(FSTUserDataReader *)reader {
+  return cpp_subcollection_source;
+}
+
+- (NSString *)name {
+  return @"subcollection";
+}
+@end
+
 @implementation FIRDatabaseSourceStageBridge {
   std::shared_ptr<DatabaseSource> cpp_database_source;
 }
@@ -617,6 +665,38 @@ - (NSString *)name {
 }
 @end
 
+@implementation FIRDefineStageBridge {
+  NSDictionary<NSString *, FIRExprBridge *> *_variables;
+  Boolean isUserDataRead;
+  std::shared_ptr<DefineStage> cpp_define;
+}
+
+- (id)initWithVariables:(NSDictionary<NSString *, FIRExprBridge *> *)variables {
+  self = [super init];
+  if (self) {
+    _variables = variables;
+    isUserDataRead = NO;
+  }
+  return self;
+}
+
+- (std::shared_ptr<api::Stage>)cppStageWithReader:(FSTUserDataReader *)reader {
+  if (!isUserDataRead) {
+    std::unordered_map<std::string, std::shared_ptr<Expr>> cpp_fields;
+    for (NSString *key in _variables) {
+      cpp_fields[MakeString(key)] = [_variables[key] cppExprWithReader:reader];
+    }
+    cpp_define = std::make_shared<DefineStage>(std::move(cpp_fields));
+  }
+  isUserDataRead = YES;
+  return cpp_define;
+}
+
+- (NSString *)name {
+  return @"let";
+}
+@end
+
 @implementation FIRDistinctStageBridge {
   NSDictionary<NSString *, FIRExprBridge *> *_groups;
   Boolean isUserDataRead;
@@ -1205,7 +1285,30 @@ - (id)initWithCppChange:(api::PipelineResultChange)change db:(std::shared_ptr<ap
 
   return self;
 }
+@end
 
+@implementation FIRPipelineExprBridge {
+  NSArray<FIRStageBridge *> *_stages;
+  Boolean isUserDataRead;
+  std::vector<std::shared_ptr<api::Stage>> cpp_stages;
+}
+- (id)initWithStages:(NSArray<FIRStageBridge *> *)stages {
+  self = [super init];
+  if (self) {
+    _stages = stages;
+    isUserDataRead = NO;
+  }
+  return self;
+}
+- (std::shared_ptr<api::Expr>)cppExprWithReader:(FSTUserDataReader *)reader {
+  if (!isUserDataRead) {
+    for (FIRStageBridge *stage in _stages) {
+      cpp_stages.push_back([stage cppStageWithReader:reader]);
+    }
+    isUserDataRead = YES;
+  }
+  return std::make_shared<api::PipelineExpr>(cpp_stages);
+}
 @end
 
 @implementation FIRPipelineBridge {
@@ -1240,7 +1343,7 @@ - (void)executeWithCompletion:(void (^)(__FIRPipelineSnapshotBridge *_Nullable r
   if (!isUserDataRead) {
     std::vector<std::shared_ptr<firebase::firestore::api::Stage>> cpp_stages;
     for (FIRStageBridge *stage in _stages) {
-      cpp_stages.push_back([stage cppStageWithReader:firestore.dataReader]);
+      cpp_stages.push_back([stage cppStageWithReader:reader]);
     }
     cpp_pipeline = std::make_shared<Pipeline>(cpp_stages, firestore.wrapped);
   }

Firestore/Source/API/FIRPipelineBridge.mm

@@ -42,6 +42,12 @@ NS_SWIFT_NAME(FieldBridge)
 - (NSString *)field_name;
 @end
 
+NS_SWIFT_SENDABLE
+NS_SWIFT_NAME(VariableBridge)
+@interface FIRVariableBridge : FIRExprBridge
+- (id)initWithName:(NSString *)name;
+@end
+
 NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(ConstantBridge)
 @interface FIRConstantBridge : FIRExprBridge
@@ -79,6 +85,13 @@ NS_SWIFT_NAME(CollectionSourceStageBridge)
 - (id)initWithRef:(FIRCollectionReference *)ref firestore:(FIRFirestore *)db;
 @end
 
+NS_SWIFT_SENDABLE
+NS_SWIFT_NAME(SubcollectionSourceStageBridge)
+@interface FIRSubcollectionSourceStageBridge : FIRStageBridge
+
+- (id)initWithPath:(NSString *)path;
+@end
+
 NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(DatabaseSourceStageBridge)
 @interface FIRDatabaseSourceStageBridge : FIRStageBridge
@@ -139,6 +152,12 @@ NS_SWIFT_NAME(SelectStageBridge)
 - (id)initWithSelections:(NSDictionary<NSString *, FIRExprBridge *> *)selections;
 @end
 
+NS_SWIFT_SENDABLE
+NS_SWIFT_NAME(DefineStageBridge)
+@interface FIRDefineStageBridge : FIRStageBridge
+- (id)initWithVariables:(NSDictionary<NSString *, FIRExprBridge *> *)variables;
+@end
+
 NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(DistinctStageBridge)
 @interface FIRDistinctStageBridge : FIRStageBridge
@@ -266,6 +285,12 @@ NS_SWIFT_NAME(PipelineBridge)
 + (NSArray<FIRStageBridge *> *)createStageBridgesFromQuery:(FIRQuery *)query;
 @end
 
+NS_SWIFT_SENDABLE
+NS_SWIFT_NAME(PipelineExprBridge)
+@interface FIRPipelineExprBridge : FIRExprBridge
+- (id)initWithStages:(NSArray<FIRStageBridge *> *)stages;
+@end
+
 NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(__RealtimePipelineSnapshotBridge)
 @interface __FIRRealtimePipelineSnapshotBridge : NSObject

Firestore/Source/Public/FirebaseFirestore/FIRPipelineBridge.h

@@ -11,13 +11,13 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 extension Expression {
-  /// Returns the internal error message. It is overridden in specific expression implementations
-  /// (like FunctionExpression) but defaults to `nil` for others.
+  /// Returns the internal error message. It is dynamically dispatched
+  /// to specific expression implementations (like FunctionExpression), and returns `nil` for
+  /// others.
   /// This design is to support pipeline conversion to expression.
   var errorMessage: String? {
-    return nil
+    return Helper.errorMessage(for: self)
   }
 
   func toBridge() -> ExprBridge {
@@ -1190,4 +1190,14 @@ public extension Expression {
   func type() -> FunctionExpression {
     return FunctionExpression(functionName: "type", args: [self])
   }
+
+  /// Creates an expression that accesses a field on this expression using a string key.
+  func getField(_ key: String) -> FunctionExpression {
+    return FunctionExpression(functionName: "get_field", args: [self, Constant(key)])
+  }
+
+  /// Creates an expression that accesses a field on this expression using a dynamic key expression.
+  func getField(_ expression: Expression) -> FunctionExpression {
+    return FunctionExpression(functionName: "get_field", args: [self, expression])
+  }
 }

Firestore/Swift/Source/ExpressionImplementation.swift

@@ -26,6 +26,17 @@ enum Helper {
     }
   }
 
+  static func errorMessage(for expr: Expression) -> String? {
+    if let funcExpr = expr as? FunctionExpression {
+      return funcExpr.errorMessage
+    } else if let pipelineExpr = expr as? PipelineExpression {
+      return pipelineExpr.errorMessage
+    } else if let boolFuncExpr = expr as? BooleanFunctionExpression {
+      return boolFuncExpr.errorMessage
+    }
+    return nil
+  }
+
   static func sendableToExpr(_ value: Sendable?) -> Expression {
     guard let value else {
       return Constant.nil

Firestore/Swift/Source/Helper/PipelineHelper.swift

@@ -55,6 +55,16 @@ class CollectionSource: Stage {
   }
 }
 
+@available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
+class SubcollectionStage: Stage {
+  let name: String = "subcollection"
+  let bridge: StageBridge
+
+  init(path: String) {
+    bridge = SubcollectionSourceStageBridge(path: path)
+  }
+}
+
 @available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
 class CollectionGroupSource: Stage {
   let name: String = "collection_group"
@@ -192,6 +202,25 @@ class RemoveFieldsStage: Stage {
   }
 }
 
+@available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
+class Define: Stage {
+  let name: String = "let"
+  let bridge: StageBridge
+  let errorMessage: String?
+
+  init(variables: [Selectable]) {
+    let (exprMap, error) = Helper.selectablesToMap(selectables: variables)
+    if let error = error {
+      errorMessage = error.localizedDescription
+      bridge = DefineStageBridge(variables: [:])
+    } else {
+      errorMessage = nil
+      let bridgeVariables = exprMap.mapValues { $0.toBridge() }
+      bridge = DefineStageBridge(variables: bridgeVariables)
+    }
+  }
+}
+
 @available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
 class Select: Stage {
   let name: String = "select"
@@ -355,7 +384,7 @@ class Union: Stage {
 
   init(other: Pipeline) {
     self.other = other
-    bridge = UnionStageBridge(other: other.bridge)
+    bridge = UnionStageBridge(other: other.pipelineBridge)
     errorMessage = other.errorMessage
   }
 }

Firestore/Swift/Source/Stages.swift

@@ -0,0 +1,31 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// An expression that represents the current document being processed.
+///
+/// Example:
+/// ```swift
+/// // Define the current document as a variable "doc"
+/// firestore.pipeline().collection("books")
+///     .define([CurrentDocument().as("doc")])
+///     // Access a field from the defined document variable
+///     .select([Variable("doc").getField("title")])
+/// ```
+public struct CurrentDocument: Expression, BridgeWrapper {
+  let bridge: ExprBridge
+
+  public init() {
+    bridge = FunctionExprBridge(name: "current_document", args: [])
+  }
+}

Firestore/Swift/Source/SwiftAPI/Pipeline/Expressions/CurrentDocument.swift

@@ -1360,6 +1360,18 @@ public protocol Expression: Sendable {
 
   // MARK: Map Operations
 
+  /// Creates an expression that accesses a field on this expression using a string key.
+  ///
+  /// - Parameter key: The string key to access.
+  /// - Returns: A new `FunctionExpression` representing the field access.
+  func getField(_ key: String) -> FunctionExpression
+
+  /// Creates an expression that accesses a field on this expression using a dynamic key expression.
+  ///
+  /// - Parameter expression: The expression evaluating to a key to access.
+  /// - Returns: A new `FunctionExpression` representing the field access.
+  func getField(_ expression: Expression) -> FunctionExpression
+
   /// Accesses a value from a map (object) field using the provided literal string key.
   /// Assumes `self` evaluates to a Map.
   ///