diff --git a/snippets/firestore/pipeline.py b/snippets/firestore/pipeline.py new file mode 100644 index 00000000..971d09fd --- /dev/null +++ b/snippets/firestore/pipeline.py @@ -0,0 +1,346 @@ +# 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. + +"""Pipeline snippets for Cloud Firestore.""" + +# pylint: disable=invalid-name + + +def load_test_data(db): + # [START pipeline_join_test_data] + # Load set of cities. + cities = db.collection("cities") + + cities.document("SF").set( + { + "name": "San Francisco", + "state": "CA", + "country": "USA", + } + ) + cities.document("LA").set( + { + "name": "Los Angeles", + "state": "CA", + "country": "USA", + } + ) + cities.document("DC").set( + { + "name": "Washington, D.C.", + "state": None, + "country": "USA", + } + ) + cities.document("TOK").set( + { + "name": "Tokyo", + "state": None, + "country": "Japan", + } + ) + + # Load restaurants in various cities. + sf_restaurants = db.collection("cities").document("SF").collection("restaurants") + la_restaurants = db.collection("cities").document("LA").collection("restaurants") + dc_restaurants = db.collection("cities").document("DC").collection("restaurants") + + rest1 = sf_restaurants.add( + { + "name": "Golden Gate Pizza", + "type": "pizza", + "owner_id": "Mario Rossi", + } + )[1] + rest2 = sf_restaurants.add( + { + "name": "Bay Area Burger", + "type": "burger", + "owner_id": "Sarah Jenkins", + } + )[1] + rest3 = sf_restaurants.add( + { + "name": "Sunset Taco", + "type": "mexican", + "owner_id": "Edward", + } + )[1] + + rest4 = la_restaurants.add( + { + "name": "Hollywood Sushi", + "type": "sushi", + "owner_id": "Ken Kenji", + } + )[1] + rest5 = la_restaurants.add( + { + "name": "Venice Pizza", + "type": "pizza", + "owner_id": "Luigi Romano", + } + )[1] + + rest6 = dc_restaurants.add( + { + "name": "Capitol Tacos", + "type": "mexican", + "owner_id": "Maria Garcia", + } + )[1] + rest7 = dc_restaurants.add( + { + "name": "Georgetown Coffee", + "type": "cafe", + "owner_id": "David Kim", + } + )[1] + + # Load collection of reviews. + reviews = db.collection("reviews") + + reviews.add({"restaurant": rest1, "rating": 5, "reviewer_id": "Alice"}) + reviews.add({"restaurant": rest1, "rating": 4, "reviewer_id": "Bob"}) + reviews.add({"restaurant": rest2, "rating": 4, "reviewer_id": "Charlie"}) + reviews.add({"restaurant": rest3, "rating": 5, "reviewer_id": "Diana"}) + reviews.add({"restaurant": rest3, "rating": 4, "reviewer_id": "Edward"}) + reviews.add({"restaurant": rest3, "rating": 4, "reviewer_id": "Fiona"}) + # rest4 has 0 reviews + reviews.add({"restaurant": rest5, "rating": 3, "reviewer_id": "George"}) + reviews.add({"restaurant": rest6, "rating": 5, "reviewer_id": "Hannah"}) + reviews.add({"restaurant": rest6, "rating": 4, "reviewer_id": "Ian"}) + reviews.add({"restaurant": rest7, "rating": 5, "reviewer_id": "Julia"}) + # [END pipeline_join_test_data] + + +def pipeline_join_lookup(db): + from google.cloud.firestore_v1.pipeline_expressions import Field, Variable + + # [START pipeline_join_lookup] + results = ( + db.pipeline() + .collection_group("reviews") + .define(Field.of("restaurant").as_("restaurant_name")) + .add_fields( + db.pipeline() + .collection_group("restaurants") + .where(Field.of("__name__").equal(Variable("restaurant_name"))) + .select("name", "type") + .to_scalar_expression() + .as_("restaurant") + ) + .execute() + ) + # [END pipeline_join_lookup] + return results + + +def pipeline_join_array(db): + from google.cloud.firestore_v1.pipeline_expressions import Field, Variable + + # [START pipeline_join_array] + results = ( + db.pipeline() + .collection_group("restaurants") + .where(Field.of("type").equal("pizza")) + .define(Field.of("__name__").as_("restaurant_name")) + .select( + Field.of("name"), + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("restaurant_name"))) + .select("rating", "reviewer_id") + .to_array_expression() + .as_("reviews"), + ) + .execute() + ) + # [END pipeline_join_array] + return results + + +def pipeline_join_aggregate(db): + from google.cloud.firestore_v1.pipeline_expressions import Field, Variable + + # [START pipeline_join_aggregate] + results = ( + db.pipeline() + .collection_group("restaurants") + .where(Field.of("type").equal("pizza")) + .define(Field.of("__name__").as_("restaurant_name")) + .select( + Field.of("name"), + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("restaurant_name"))) + .aggregate(Field.of("rating").average().as_("avg_rating")) + .to_scalar_expression() + .as_("avg_rating"), + ) + .execute() + ) + # [END pipeline_join_aggregate] + return results + + +def pipeline_join_limit(db): + from google.cloud.firestore_v1.pipeline_expressions import Field, Variable + + # [START pipeline_join_limit] + results = ( + db.pipeline() + .collection_group("restaurants") + .define(Field.of("__name__").as_("restaurant_name")) + .select( + Field.of("name"), + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("restaurant_name"))) + .sort(Field.of("rating").descending()) + .limit(2) + .select("rating", "reviewer_id") + .to_array_expression() + .as_("top_reviews"), + ) + .execute() + ) + # [END pipeline_join_limit] + return results + + +def pipeline_join_subcollection(db): + from google.cloud.firestore_v1.pipeline_source import PipelineSource + + # [START pipeline_join_subcollection] + results = ( + db.pipeline() + .collection("cities") + .add_fields( + PipelineSource.subcollection("restaurants") + .to_array_expression() + .length() + .as_("restaurant_count") + ) + .execute() + ) + # [END pipeline_join_subcollection] + return results + + +def pipeline_join_multi_field(db): + from google.cloud.firestore_v1.pipeline_expressions import Count, Field, Variable + + # [START pipeline_join_multi_field] + results = ( + db.pipeline() + .collection_group("restaurants") + .define( + Field.of("owner_id").as_("owner_id"), + Field.of("__name__").as_("__name__"), + ) + .where( + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("__name__"))) + .where(Field.of("reviewer_id").equal(Variable("owner_id"))) + .aggregate(Count().as_("c")) + .to_scalar_expression() + .greater_than(0) + ) + .execute() + ) + # [END pipeline_join_multi_field] + return results + + +def pipeline_join_anti(db): + from google.cloud.firestore_v1.pipeline_expressions import Count, Field, Variable + + # [START pipeline_join_anti] + results = ( + db.pipeline() + .collection_group("restaurants") + .define(Field.of("__name__").as_("restaurant_name")) + .where( + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("restaurant_name"))) + .aggregate(Count().as_("review_count")) + .to_scalar_expression() + .equal(0) + ) + .execute() + ) + # [END pipeline_join_anti] + return results + + +def pipeline_join_unnest(db): + from google.cloud.firestore_v1.pipeline_expressions import Field, Variable + + # [START pipeline_join_unnest] + results = ( + db.pipeline() + .collection_group("restaurants") + .where(Field.of("type").equal("pizza")) + .define(Field.of("__name__").as_("restaurant_name")) + .unnest( + db.pipeline() + .collection_group("reviews") + .where(Field.of("restaurant").equal(Variable("restaurant_name"))) + .select("rating", "reviewer_id") + .to_array_expression(), + alias="review", + ) + .execute() + ) + # [END pipeline_join_unnest] + return results + + +def pipeline_join_uncorrelated(db): + from google.cloud.firestore_v1.pipeline_expressions import Field + + # [START pipeline_join_uncorrelated] + results = ( + db.pipeline() + .collection("reviews") + # Average review rating is 4.3 + .where( + Field.of("rating").greater_than( + db.pipeline() + .collection("reviews") + .aggregate(Field.of("rating").average().as_("avg")) + .to_scalar_expression() + ) + ) + .select("rating", "reviewer_id") + .execute() + ) + # [END pipeline_join_uncorrelated] + return results + + +def pipeline_force_table_scan(db): + # [START pipeline_force_table_scan] + # Force Planner to only do a Full-Table Scan + results = ( + db.pipeline() + .collection_group("customers") + .limit(100) + .execute() + ) + # [END pipeline_force_table_scan] + return results