Skip to content

Commit 1d3b6b9

Browse files
claudevdmclaudevdm
authored
Created using Colab (#35499)
1 parent da6df98 commit 1d3b6b9

1 file changed

Lines changed: 108 additions & 79 deletions

File tree

‎examples/notebooks/beam-ml/alloydb_product_catalog_embeddings.ipynb‎

Lines changed: 108 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,32 @@
103103
"id": "z2eAyRECIP3z"
104104
},
105105
"source": [
106+
"# Connecting Apache Beam to AlloyDB\n",
107+
"\n",
108+
"Beam uses the [AlloyDB Java Connector](https://github.com/GoogleCloudPlatform/alloydb-java-connector) to establish a connection to your database.\n",
109+
"\n",
110+
"AlloyDB Java Connector defaults to connecting via private IP, but can be overridden via the [AlloyDBLanguageConnectorConfig](https://beam.apache.org/releases/pydoc/current/apache_beam.ml.rag.ingestion.alloydb.html#apache_beam.ml.rag.ingestion.alloydb.AlloyDBLanguageConnectorConfig) to connect via public IP or PSC\n",
111+
"\n",
112+
"## Running on Dataflow\n",
113+
"To connect from Dataflow workers to an AlloyDB instance using private IP:\n",
114+
"- [Set up AlloyDB with Private Services Access](https://cloud.google.com/alloydb/docs/about-private-services-access)\n",
115+
"- [Run on Dataflow](#scrollTo=Quick_Start_Run_on_Dataflow) on the same VPC network as the AlloyDB instance\n",
116+
"\n",
117+
"Alternatively, dataflow can also connect to an AlloyDB instance that is set up with [Private Service Connect](https://cloud.google.com/alloydb/docs/about-private-service-connect)\n",
118+
"\n",
119+
"## Running on DirectRunner (for Development/Testing)\n",
120+
"For faster iteration during development, it's often convenient to use the DirectRunner. The connection method depends on your setup:\n",
121+
"\n",
122+
"- For simplicity and quick setup (with Public IP): If you are developing on a machine (like a local laptop or a VM) that is not on the same VPC network as your AlloyDB instance, the easiest way to connect is to enable public IP on a sandbox AlloyDB instance.\n",
123+
"- For a more secure setup (with Private IP): The recommended approach is to use a development VM that is on the same subnetwork as your AlloyDB cluster. In this scenario, you can connect using the instance's private IP, and a public IP is not needed.\n",
124+
"\n",
106125
"# Setup and Prerequisites\n",
107126
"\n",
108127
"This example requires:\n",
109-
"1. An AlloyDB instance with pgvector extension and PUBLIC IP enabled\n",
110-
"2. Apache Beam 2.64.0 or later"
128+
"1. An AlloyDB instance with pgvector extension\n",
129+
"2. An AlloyDB instance PUBLIC IP enabled if running on DirectRunner (Local Execution) or PRIVATE IP if running on Dataflow\n",
130+
"3. Apache Beam 2.66.0 or later\n",
131+
"\n"
111132
]
112133
},
113134
{
@@ -130,11 +151,22 @@
130151
"outputs": [],
131152
"source": [
132153
"# Apache Beam with GCP support\n",
133-
"!pip install apache_beam[gcp]>=v2.64.0 --quiet\n",
154+
"!pip install apache_beam[gcp]>=2.66.0\n",
134155
"# Huggingface sentence-transformers for embedding models\n",
135156
"!pip install sentence-transformers --quiet"
136157
]
137158
},
159+
{
160+
"cell_type": "code",
161+
"source": [
162+
"!pip show apache-beam"
163+
],
164+
"metadata": {
165+
"id": "2FlMPmA0IUuv"
166+
},
167+
"execution_count": null,
168+
"outputs": []
169+
},
138170
{
139171
"cell_type": "markdown",
140172
"metadata": {
@@ -165,9 +197,10 @@
165197
"\n",
166198
"To connect to AlloyDB, you'll need:\n",
167199
"1. GCP project ID where the AlloyDB instance is located\n",
168-
"2. The AlloyDB instance URI\n",
169-
"3. Database credentials\n",
170-
"4. The pgvector extension enabled in your database\n",
200+
"2. The AlloyDB instance URI. This is the fully qualified name of the AlloyDB instance found in the google cloud console under AlloyDB Cluster > Connectivity > Connection URI.\n",
201+
"3. Database name. This is the name of the postgres database within your AlloyDB instance. The default database name is postgres.\n",
202+
"4. Database credentials\n",
203+
"5. The pgvector extension enabled in your database\n",
171204
"\n",
172205
"Replace these placeholder values with your actual AlloyDB connection details:"
173206
]
@@ -182,7 +215,7 @@
182215
"source": [
183216
"PROJECT_ID = \"\" # @param {type:'string'}\n",
184217
"\n",
185-
"INSTANCE_URI = \"\" # @param {type:'string'}\n",
218+
"INSTANCE_URI = \"projects/<project>/locations/<region>/clusters/<cluster_name>/instances/<instance_name>\" # @param {type:'string'}\n",
186219
"\n",
187220
"DB_NAME = \"postgres\" # @param {type:'string'}\n",
188221
"\n",
@@ -459,7 +492,6 @@
459492
"cell_type": "code",
460493
"execution_count": null,
461494
"metadata": {
462-
"cellView": "form",
463495
"id": "7_J__S8JOwJ_"
464496
},
465497
"outputs": [],
@@ -534,10 +566,10 @@
534566
"## Importing Pipeline Components\n",
535567
"\n",
536568
"We import the following for configuring our embedding ingestion pipeline:\n",
537-
"- `Chunk`, the structured input for generating and ingesting embeddings\n",
538-
"- `AlloyDBConnectionConfig` for configuring database connection information\n",
539-
"- `AlloyDBVectorWriterConfig` for configuring write behavior like schema mapping and conflict resolution\n",
540-
"- `AlloyDBLanguageConnectorConfig` to connect using the [AlloyDB language connector](https://cloud.google.com/alloydb/docs/connect-language-connectors)"
569+
"- `apache_beam.ml.rag.types.Chunk`, the structured input for generating and ingesting embeddings\n",
570+
"- `apache_beam.ml.rag.ingestion.alloydb.AlloyDBVectorWriterConfig` for configuring write behavior like schema mapping and conflict resolution\n",
571+
"- `apache_beam.ml.rag.ingestion.alloydb.AlloyDBLanguageConnectorConfig` to connect using the [AlloyDB language connector](https://cloud.google.com/alloydb/docs/connect-language-connectors)\n",
572+
"- `apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform` to perform the write step using AlloyDB configs"
541573
]
542574
},
543575
{
@@ -548,20 +580,23 @@
548580
},
549581
"outputs": [],
550582
"source": [
551-
"# Embedding-specific imports\n",
552-
"from apache_beam.ml.rag.ingestion.alloydb import (\n",
553-
" AlloyDBVectorWriterConfig,\n",
554-
" AlloyDBConnectionConfig,\n",
555-
" AlloyDBLanguageConnectorConfig\n",
556-
")\n",
583+
"# AlloyDB imports\n",
584+
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBLanguageConnectorConfig\n",
585+
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig\n",
586+
"\n",
587+
"\n",
557588
"from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
558589
"from apache_beam.ml.rag.types import Chunk, Content\n",
559590
"from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
560591
"\n",
561592
"# Apache Beam core\n",
562593
"import apache_beam as beam\n",
563594
"from apache_beam.options.pipeline_options import PipelineOptions\n",
564-
"from apache_beam.ml.transforms.base import MLTransform"
595+
"from apache_beam.ml.transforms.base import MLTransform\n",
596+
"\n",
597+
"# JDBC and Postgres utilities\n",
598+
"from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig\n",
599+
"from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder, ConflictResolution"
565600
]
566601
},
567602
{
@@ -773,17 +808,15 @@
773808
"outputs": [],
774809
"source": [
775810
"# Configure the language connector so we can connect securly\n",
776-
"language_connector_config = AlloyDBLanguageConnectorConfig(\n",
777-
" database_name=DB_NAME, instance_name=INSTANCE_URI, ip_type=\"PUBLIC\"\n",
778-
")\n",
779-
"# Configure the AlloyDBConnectionConfig with language connector\n",
780-
"connection_config = AlloyDBConnectionConfig.with_language_connector(\n",
781-
" connector_options=language_connector_config,\n",
811+
"connector_config = AlloyDBLanguageConnectorConfig(\n",
782812
" username=DB_USER,\n",
783-
" password=DB_PASSWORD\n",
813+
" password=DB_PASSWORD,\n",
814+
" database_name=DB_NAME,\n",
815+
" instance_name=INSTANCE_URI,\n",
816+
" ip_type=\"PUBLIC\"\n",
784817
")\n",
785818
"alloydb_writer_config = AlloyDBVectorWriterConfig(\n",
786-
" connection_config=connection_config,\n",
819+
" connection_config=connector_config,\n",
787820
" table_name=table_name\n",
788821
")"
789822
]
@@ -941,7 +974,7 @@
941974
"from apache_beam.ml.transforms.base import MLTransform\n",
942975
"from apache_beam.ml.rag.types import Chunk, Content\n",
943976
"from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
944-
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig, AlloyDBConnectionConfig, AlloyDBLanguageConnectorConfig\n",
977+
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig, AlloyDBLanguageConnectorConfig\n",
945978
"from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
946979
"from apache_beam.options.pipeline_options import SetupOptions\n",
947980
"\n",
@@ -1020,14 +1053,13 @@
10201053
" )\n",
10211054
" | 'Write to AlloyDB' >> VectorDatabaseWriteTransform(\n",
10221055
" AlloyDBVectorWriterConfig(\n",
1023-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
1024-
" AlloyDBLanguageConnectorConfig(\n",
1025-
" database_name=known_args.alloydb_database, instance_name=known_args.instance_uri\n",
1026-
" ),\n",
1056+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
10271057
" username=known_args.alloydb_username,\n",
1028-
" password=known_args.alloydb_password\n",
1029-
" ),\n",
1030-
" table_name=known_args.alloydb_table\n",
1058+
" password=known_args.alloydb_password,\n",
1059+
" database_name=known_args.alloydb_database,\n",
1060+
" instance_name=known_args.instance_uri\n",
1061+
" ),\n",
1062+
" table_name=known_args.alloydb_table\n",
10311063
" )\n",
10321064
" )\n",
10331065
" )\n",
@@ -1109,7 +1141,7 @@
11091141
"outputs": [],
11101142
"source": [
11111143
"import os\n",
1112-
"BUCKET_NAME = '' # @param {type:'string'}\n",
1144+
"BUCKET_NAME = 'gs://' # @param {type:'string'}\n",
11131145
"REGION = 'us-central1' # @param {type:'string'}\n",
11141146
"os.environ['BUCKET_NAME'] = BUCKET_NAME\n",
11151147
"os.environ['REGION'] = REGION\n",
@@ -1286,6 +1318,8 @@
12861318
"\n",
12871319
"ColumnSpec specifies how to map data to a database column. For example:\n",
12881320
"```python\n",
1321+
"from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder\n",
1322+
"\n",
12891323
"ColumnSpec(\n",
12901324
" column_name=\"price\", # Database column\n",
12911325
" python_type=float, # Python Type for the value\n",
@@ -1392,8 +1426,8 @@
13921426
},
13931427
"outputs": [],
13941428
"source": [
1395-
"from apache_beam.ml.rag.ingestion.alloydb import ColumnSpec\n",
1396-
"from apache_beam.ml.rag.ingestion.alloydb import ColumnSpecsBuilder\n",
1429+
"from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder\n",
1430+
"from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec\n",
13971431
"from datetime import datetime\n",
13981432
"\n",
13991433
"column_specs = (\n",
@@ -1493,14 +1527,12 @@
14931527
" .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
14941528
" | 'Write to AlloyDB' >> VectorDatabaseWriteTransform(\n",
14951529
" AlloyDBVectorWriterConfig(\n",
1496-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
1497-
" connector_options=AlloyDBLanguageConnectorConfig(\n",
1498-
" database_name=DB_NAME,\n",
1499-
" instance_name=INSTANCE_URI,\n",
1500-
" ip_type=\"PUBLIC\"\n",
1501-
" ),\n",
1530+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
15021531
" username=DB_USER,\n",
1503-
" password=DB_PASSWORD\n",
1532+
" password=DB_PASSWORD,\n",
1533+
" database_name=DB_NAME,\n",
1534+
" instance_name=INSTANCE_URI,\n",
1535+
" ip_type=\"PUBLIC\"\n",
15041536
" ),\n",
15051537
" table_name=table_name,\n",
15061538
" column_specs=column_specs\n",
@@ -1644,9 +1676,9 @@
16441676
"source": [
16451677
"from apache_beam.ml.rag.ingestion.alloydb import (\n",
16461678
" AlloyDBVectorWriterConfig,\n",
1647-
" AlloyDBConnectionConfig,\n",
1648-
" ConflictResolution\n",
1679+
" AlloyDBLanguageConnectorConfig,\n",
16491680
")\n",
1681+
"from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution\n",
16501682
"\n",
16511683
"# Define how to handle conflicts - update all fields when ID matches\n",
16521684
"conflict_resolution = ConflictResolution(\n",
@@ -1657,14 +1689,12 @@
16571689
"\n",
16581690
"# Create writer config with conflict resolution\n",
16591691
"alloydb_writer_config = AlloyDBVectorWriterConfig(\n",
1660-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
1661-
" connector_options=AlloyDBLanguageConnectorConfig(\n",
1662-
" database_name=DB_NAME,\n",
1663-
" instance_name=INSTANCE_URI,\n",
1664-
" ip_type=\"PUBLIC\"\n",
1665-
" ),\n",
1692+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
16661693
" username=DB_USER,\n",
1667-
" password=DB_PASSWORD\n",
1694+
" password=DB_PASSWORD,\n",
1695+
" database_name=DB_NAME,\n",
1696+
" instance_name=INSTANCE_URI,\n",
1697+
" ip_type=\"PUBLIC\"\n",
16681698
" ),\n",
16691699
" table_name=table_name,\n",
16701700
" conflict_resolution=conflict_resolution,\n",
@@ -2000,14 +2030,12 @@
20002030
"\n",
20012031
"# Configure database writer\n",
20022032
"alloydb_writer_config = AlloyDBVectorWriterConfig(\n",
2003-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
2004-
" connector_options=AlloyDBLanguageConnectorConfig(\n",
2005-
" database_name=DB_NAME,\n",
2006-
" instance_name=INSTANCE_URI,\n",
2007-
" ip_type=\"PUBLIC\"\n",
2008-
" ),\n",
2033+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
20092034
" username=DB_USER,\n",
2010-
" password=DB_PASSWORD\n",
2035+
" password=DB_PASSWORD,\n",
2036+
" database_name=DB_NAME,\n",
2037+
" instance_name=INSTANCE_URI,\n",
2038+
" ip_type=\"PUBLIC\"\n",
20112039
" ),\n",
20122040
" table_name=table_name,\n",
20132041
" column_specs=(\n",
@@ -2039,6 +2067,8 @@
20392067
" table_name=table_name,\n",
20402068
" driver_class_name=\"org.postgresql.Driver\",\n",
20412069
" jdbc_url=AlloyDBLanguageConnectorConfig(\n",
2070+
" username=DB_USER,\n",
2071+
" password=DB_PASSWORD,\n",
20422072
" database_name=DB_NAME,\n",
20432073
" instance_name=INSTANCE_URI,\n",
20442074
" ip_type=\"PUBLIC\"\n",
@@ -2253,14 +2283,12 @@
22532283
" )\n",
22542284
" | 'Write to AlloyDB' >> VectorDatabaseWriteTransform(\n",
22552285
" AlloyDBVectorWriterConfig(\n",
2256-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
2257-
" connector_options=AlloyDBLanguageConnectorConfig(\n",
2258-
" database_name=DB_NAME,\n",
2259-
" instance_name=INSTANCE_URI,\n",
2260-
" ip_type=\"PUBLIC\"\n",
2261-
" ),\n",
2286+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
22622287
" username=DB_USER,\n",
2263-
" password=DB_PASSWORD\n",
2288+
" password=DB_PASSWORD,\n",
2289+
" database_name=DB_NAME,\n",
2290+
" instance_name=INSTANCE_URI,\n",
2291+
" ip_type=\"PUBLIC\"\n",
22642292
" ),\n",
22652293
" table_name=table_name\n",
22662294
" )\n",
@@ -2443,7 +2471,7 @@
24432471
"\n",
24442472
"\n",
24452473
"\n",
2446-
"BUCKET_NAME = '' # @param {type:'string'}\n",
2474+
"BUCKET_NAME = 'gs://' # @param {type:'string'}\n",
24472475
"dataflow_gcs_location = \"gs://%s/dataflow\" % BUCKET_NAME\n",
24482476
"\n",
24492477
"# The Dataflow staging location. This location is used to stage the Dataflow pipeline and the SDK binary.\n",
@@ -2517,7 +2545,11 @@
25172545
"from apache_beam.ml.transforms.base import MLTransform\n",
25182546
"from apache_beam.ml.rag.types import Chunk, Content\n",
25192547
"from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
2520-
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig, AlloyDBConnectionConfig, ConflictResolution\n",
2548+
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig\n",
2549+
"from apache_beam.ml.rag.ingestion.alloydb import AlloyDBLanguageConnectorConfig\n",
2550+
"\n",
2551+
"from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution\n",
2552+
"\n",
25212553
"from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
25222554
"from apache_beam.transforms.window import FixedWindows\n",
25232555
"\n",
@@ -2545,13 +2577,11 @@
25452577
" .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
25462578
" | \"Write to AlloyDB\" >> VectorDatabaseWriteTransform(\n",
25472579
" AlloyDBVectorWriterConfig(\n",
2548-
" connection_config=AlloyDBConnectionConfig.with_language_connector(\n",
2549-
" connector_options=AlloyDBLanguageConnectorConfig(\n",
2550-
" database_name=DB_NAME,\n",
2551-
" instance_name=INSTANCE_URI\n",
2552-
" ),\n",
2580+
" connection_config=AlloyDBLanguageConnectorConfig(\n",
25532581
" username=DB_USER,\n",
2554-
" password=DB_PASSWORD\n",
2582+
" password=DB_PASSWORD,\n",
2583+
" database_name=DB_NAME,\n",
2584+
" instance_name=INSTANCE_URI\n",
25552585
" ),\n",
25562586
" table_name=table_name,\n",
25572587
" conflict_resolution=ConflictResolution(\n",
@@ -2755,8 +2785,7 @@
27552785
"collapsed_sections": [
27562786
"mcZATJbaOec0"
27572787
],
2758-
"provenance": [],
2759-
"toc_visible": true
2788+
"provenance": []
27602789
},
27612790
"kernelspec": {
27622791
"display_name": "Python 3",
@@ -2768,4 +2797,4 @@
27682797
},
27692798
"nbformat": 4,
27702799
"nbformat_minor": 0
2771-
}
2800+
}

0 commit comments

Comments
 (0)