[Docs] Beam website doc for vertex ai enrichment handler (apache#30692)

* add vetex ai page * update beam side docs for vertex ai enrichment * add links to release doc * update links, correct table formatting * Update website/www/site/content/en/documentation/transforms/python/elementwise/enrichment-bigtable.md Co-authored-by: Rebecca Szper <[email protected]> * Update website/www/site/content/en/documentation/transforms/python/elementwise/enrichment-vertexai.md Co-authored-by: Rebecca Szper <[email protected]> * Apply suggestions from code review Co-authored-by: Rebecca Szper <[email protected]> --------- Co-authored-by: Rebecca Szper <[email protected]>
acrites · Mar 26, 2024 · 4af19ff · 4af19ff
1 parent 488d2a1
commit 4af19ff
Show file tree

Hide file tree

Showing 7 changed files with 275 additions and 28 deletions.
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py
@@ -49,3 +49,70 @@ def enrichment_with_bigtable():
         | "Enrich W/ BigTable" >> Enrichment(bigtable_handler)
         | "Print" >> beam.Map(print))
   # [END enrichment_with_bigtable]
+
+
+def enrichment_with_vertex_ai():
+  # [START enrichment_with_vertex_ai]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store \
+    import VertexAIFeatureStoreEnrichmentHandler
+
+  project_id = 'apache-beam-testing'
+  location = 'us-central1'
+  api_endpoint = f"{location}-aiplatform.googleapis.com"
+  data = [
+      beam.Row(user_id='2963', product_id=14235, sale_price=15.0),
+      beam.Row(user_id='21422', product_id=11203, sale_price=12.0),
+      beam.Row(user_id='20592', product_id=8579, sale_price=9.0),
+  ]
+
+  vertex_ai_handler = VertexAIFeatureStoreEnrichmentHandler(
+      project=project_id,
+      location=location,
+      api_endpoint=api_endpoint,
+      feature_store_name="vertexai_enrichment_example",
+      feature_view_name="users",
+      row_key="user_id",
+  )
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_vertex_ai]
+
+
+def enrichment_with_vertex_ai_legacy():
+  # [START enrichment_with_vertex_ai_legacy]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store \
+    import VertexAIFeatureStoreLegacyEnrichmentHandler
+
+  project_id = 'apache-beam-testing'
+  location = 'us-central1'
+  api_endpoint = f"{location}-aiplatform.googleapis.com"
+  data = [
+      beam.Row(entity_id="movie_01", title='The Shawshank Redemption'),
+      beam.Row(entity_id="movie_02", title="The Shining"),
+      beam.Row(entity_id="movie_04", title='The Dark Knight'),
+  ]
+
+  vertex_ai_handler = VertexAIFeatureStoreLegacyEnrichmentHandler(
+      project=project_id,
+      location=location,
+      api_endpoint=api_endpoint,
+      entity_type_id='movies',
+      feature_store_id="movie_prediction_unique",
+      feature_ids=["title", "genres"],
+      row_key="entity_id",
+  )
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_vertex_ai_legacy]
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
@@ -25,7 +25,9 @@
 
 # pylint: disable=unused-import
 try:
-  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_bigtable
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_bigtable, \
+  enrichment_with_vertex_ai_legacy
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_vertex_ai
   from apache_beam.io.requestresponse import RequestResponseIO
 except ImportError:
   raise unittest.SkipTest('RequestResponseIO dependencies are not installed')
@@ -40,6 +42,24 @@ def validate_enrichment_with_bigtable():
   return expected
 
 
+def validate_enrichment_with_vertex_ai():
+  expected = '''[START enrichment_with_vertex_ai]
+Row(user_id='2963', product_id=14235, sale_price=15.0, age=29.0, gender='1', state='97', country='2')
+Row(user_id='21422', product_id=11203, sale_price=12.0, age=36.0, state='184', gender='1', country='5')
+Row(user_id='20592', product_id=8579, sale_price=9.0, age=30.0, state='86', gender='1', country='4')
+  [END enrichment_with_vertex_ai]'''.splitlines()[1:-1]
+  return expected
+
+
+def validate_enrichment_with_vertex_ai_legacy():
+  expected = '''[START enrichment_with_vertex_ai_legacy]
+Row(entity_id='movie_01', title='The Shawshank Redemption', genres='Drama')
+Row(entity_id='movie_02', title='The Shining', genres='Horror')
+Row(entity_id='movie_04', title='The Dark Knight', genres='Action')
+  [END enrichment_with_vertex_ai_legacy]'''.splitlines()[1:-1]
+  return expected
+
+
 @mock.patch('sys.stdout', new_callable=StringIO)
 class EnrichmentTest(unittest.TestCase):
   def test_enrichment_with_bigtable(self, mock_stdout):
@@ -48,6 +68,21 @@ def test_enrichment_with_bigtable(self, mock_stdout):
     expected = validate_enrichment_with_bigtable()
     self.assertEqual(output, expected)
 
+  def test_enrichment_with_vertex_ai(self, mock_stdout):
+    enrichment_with_vertex_ai()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_vertex_ai()
+
+    for i in range(len(expected)):
+      self.assertEqual(set(output[i].split(',')), set(expected[i].split(',')))
+
+  def test_enrichment_with_vertex_ai_legacy(self, mock_stdout):
+    enrichment_with_vertex_ai_legacy()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_vertex_ai_legacy()
+    self.maxDiff = None
+    self.assertEqual(output, expected)
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/...e/content/en/documentation/transforms/python/elementwise/enrichment-bigtable.md b/...e/content/en/documentation/transforms/python/elementwise/enrichment-bigtable.md
@@ -0,0 +1,62 @@
+---
+title: "Enrichment with Bigtable"
+---
+<!--
+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.
+-->
+
+# Use Bigtable to enrich data
+
+{{< localstorage language language-py >}}
+
+<table>
+  <tr>
+    <td>
+      <a>
+      {{< button-pydoc path="apache_beam.transforms.enrichment_handlers.bigtable" class="BigTableEnrichmentHandler" >}}
+      </a>
+   </td>
+  </tr>
+</table>
+
+In Apache Beam 2.54.0 and later versions, the enrichment transform includes a built-in enrichment handler for [Bigtable](https://cloud.google.com/bigtable/docs/overview).
+The following example demonstrates how to create a pipeline that use the enrichment transform with the [`BigTableEnrichmentHandler`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#apache_beam.transforms.enrichment_handlers.bigtable.BigTableEnrichmentHandler) handler.
+
+The data stored in the Bigtable cluster uses the following format:
+
+{{< table >}}
+|  Row key    |  product:product_id  |  product:product_name  |  product:product_stock  |
+|:-----------:|:--------------------:|:----------------------:|:-----------------------:|
+|     1       |          1           |        pixel 5         |            2            |
+|     2       |          2           |        pixel 6         |            4            |
+|     3       |          3           |        pixel 7         |           20            |
+|     4       |          4           |        pixel 8         |           10            |
+{{< /table >}}
+
+
+{{< highlight language="py" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py" enrichment_with_bigtable >}}
+{{</ highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py" enrichment_with_bigtable >}}
+{{< /highlight >}}
+
+## Related transforms
+
+Not applicable.
+
+{{< button-pydoc path="apache_beam.transforms.enrichment_handlers.bigtable" class="BigTableEnrichmentHandler" >}}
diff --git a/...e/content/en/documentation/transforms/python/elementwise/enrichment-vertexai.md b/...e/content/en/documentation/transforms/python/elementwise/enrichment-vertexai.md
@@ -0,0 +1,89 @@
+---
+title: "Enrichment with Vertex AI Feature Store"
+---
+<!--
+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.
+-->
+
+# Enrichment with Google Cloud Vertex AI Feature Store
+
+{{< localstorage language language-py >}}
+
+<table>
+  <tr>
+    <td>
+      <a>
+      {{< button-pydoc path="apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store" class="VertexAIFeatureStoreEnrichmentHandler" >}}
+      </a>
+   </td>
+  </tr>
+</table>
+
+
+In Apache Beam 2.55.0 and later versions, the enrichment transform includes a built-in enrichment handler for [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore).
+The following example demonstrates how to create a pipeline that use the enrichment transform with the [`VertexAIFeatureStoreEnrichmentHandler`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.html#apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreEnrichmentHandler) handler and the [`VertexAIFeatureStoreLegacyEnrichmentHandler`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.html#apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreLegacyEnrichmentHandler) handler.
+
+## Example 1: Enrichment with Vertex AI Feature Store
+
+The precomputed feature values stored in Vertex AI Feature Store uses the following format:
+
+{{< table >}}
+| user_id  | age  | gender  | state | country |
+|:--------:|:----:|:-------:|:-----:|:-------:|
+|  21422   |  12  |    0    |   0   |    0    |
+|   2963   |  12  |    1    |   1   |    1    |
+|  20592   |  12  |    1    |   2   |    2    |
+|  76538   |  12  |    1    |   3   |    0    |
+{{< /table >}}
+
+
+{{< highlight language="py" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py" enrichment_with_vertex_ai >}}
+{{</ highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py" enrichment_with_vertex_ai >}}
+{{< /highlight >}}
+
+## Example 2: Enrichment with Vertex AI Feature Store (legacy)
+
+The precomputed feature values stored in Vertex AI Feature Store (Legacy) use the following format:
+
+{{< table >}}
+| entity_id | title                    | genres  |
+|:----------|:-------------------------|:--------|
+| movie_01  | The Shawshank Redemption | Drama   |
+| movie_02  | The Shining              | Horror  |
+| movie_04  | The Dark Knight          | Action  |
+{{< /table >}}
+
+{{< highlight language="py" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py" enrichment_with_vertex_ai_legacy >}}
+{{</ highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py" enrichment_with_vertex_ai_legacy >}}
+{{< /highlight >}}
+
+
+## Related transforms
+
+Not applicable.
+
+{{< button-pydoc path="apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store" class="VertexAIFeatureStoreEnrichmentHandler" >}}
diff --git a/...e/www/site/content/en/documentation/transforms/python/elementwise/enrichment.md b/...e/www/site/content/en/documentation/transforms/python/elementwise/enrichment.md
@@ -23,7 +23,7 @@ limitations under the License.
   <tr>
     <td>
       <a>
-      {{< button-pydoc path="apache_beam.transforms" class="Enrichment" >}}
+      {{< button-pydoc path="apache_beam.transforms.enrichment" class="Enrichment" >}}
       </a>
    </td>
   </tr>
@@ -32,35 +32,22 @@ limitations under the License.
 
 The enrichment transform lets you dynamically enrich data in a pipeline by doing a key-value lookup to a remote service. The transform uses [`RequestResponeIO`](https://beam.apache.org/releases/pydoc/current/apache_beam.io.requestresponseio.html#apache_beam.io.requestresponseio.RequestResponseIO) internally. This feature uses client-side throttling to ensure that the remote service isn't overloaded with requests. If service-side errors occur, like `TooManyRequests` and `Timeout` exceptions, it retries the requests by using exponential backoff.
 
-In Apache Beam 2.54.0 and later versions, the transform includes a built-in enrichment handler for [Bigtable](https://cloud.google.com/bigtable/docs/overview).
+This transform is available in Apache Beam 2.54.0 and later versions.
 
-## Use Bigtable to enrich data
+## Examples
 
-The following example demonstrates how to create a pipeline that use the enrichment transform with [`BigTableEnrichmentHandler`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#apache_beam.transforms.enrichment_handlers.bigtable.BigTableEnrichmentHandler).
+The following examples demonstrate how to create a pipeline that use the enrichment transform to enrich data from external services.
 
-The data stored in the Bigtable cluster uses the following format:
-
-|  Row key  |  product:product_id  |  product:product_name  |  product:product_stock  |
-|:---------:|:--------------------:|:----------------------:|:-----------------------:|
-|     1     |          1           |        pixel 5         |            2            |
-|     2     |          2           |        pixel 6         |            4            |
-|     3     |          3           |        pixel 7         |           20            |
-|     4     |          4           |        pixel 8         |           10            |
-
-
-{{< highlight language="py" >}}
-{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py" enrichment_with_bigtable >}}
-{{</ highlight >}}
-
-{{< paragraph class="notebook-skip" >}}
-Output:
-{{< /paragraph >}}
-{{< highlight class="notebook-skip" >}}
-{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py" enrichment_with_bigtable >}}
-{{< /highlight >}}
+{{< table >}}
+| Service                            | Example                                                                                                                                                                      |
+|:-----------------------------------|:-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Cloud Bigtable                     | [Enrichment with Bigtable](/documentation/transforms/python/elementwise/enrichment-bigtable/#example)                                                                        |
+| Vertex AI Feature Store            | [Enrichment with Vertex AI Feature Store](/documentation/transforms/python/elementwise/enrichment-vertexai/#example-1-enrichment-with-vertex-ai-feature-store)               |
+| Vertex AI Feature Store (Legacy)   | [Enrichment with Legacy Vertex AI Feature Store](/documentation/transforms/python/elementwise/enrichment-vertexai/#example-2-enrichment-with-vertex-ai-feature-store-legacy) |
+{{< /table >}}
 
 ## Related transforms
 
 Not applicable.
 
-{{< button-pydoc path="apache_beam.transforms" class="Enrichment" >}}
+{{< button-pydoc path="apache_beam.transforms.enrichment" class="Enrichment" >}}
diff --git a/...www/site/content/en/documentation/transforms/python/elementwise/runinference.md b/...www/site/content/en/documentation/transforms/python/elementwise/runinference.md
@@ -23,7 +23,7 @@ limitations under the License.
   <tr>
     <td>
       <a>
-      {{< button-pydoc path="apache_beam.ml.inference" class="RunInference" >}}
+      {{< button-pydoc path="apache_beam.ml.inference.base" class="RunInference" >}}
       </a>
    </td>
   </tr>

diff --git a/website/www/site/layouts/partials/section-menu/en/documentation.html b/website/www/site/layouts/partials/section-menu/en/documentation.html
@@ -289,7 +289,14 @@
           <span class="section-nav-list-title">Element-wise</span>
 
           <ul class="section-nav-list">
-            <li><a href="/documentation/transforms/python/elementwise/enrichment/">Enrichment</a></li>
+            <li class="section-nav-item--collapsible">
+              <span class="section-nav-list-title">Enrichment</span>
+              <ul class="section-nav-list">
+                <li><a href="/documentation/transforms/python/elementwise/enrichment/">Overview</a></li>
+                <li><a href="/documentation/transforms/python/elementwise/enrichment-bigtable/">Bigtable example</a></li>
+                <li><a href="/documentation/transforms/python/elementwise/enrichment-vertexai/">Vertex AI Feature Store examples</a></li>
+              </ul>
+              </li>
             <li><a href="/documentation/transforms/python/elementwise/filter/">Filter</a></li>
             <li><a href="/documentation/transforms/python/elementwise/flatmap/">FlatMap</a></li>
             <li><a href="/documentation/transforms/python/elementwise/keys/">Keys</a></li>
-Original file line number
+Diff line change
@@ Expand Up / @@ -23,7 +23,7 @@ limitations under the License. @@
       <tr>
         <td>
           <a>
-          {{< button-pydoc path="apache_beam.ml.inference" class="RunInference" >}}
+          {{< button-pydoc path="apache_beam.ml.inference.base" class="RunInference" >}}
           </a>
        </td>
       </tr>
@@ Expand Down @@