[Python] Do not modify input elements in the embedding output functions

CHANGES.md

@@ -150,6 +150,7 @@
 * (Java/Python) The SQL `SqlTransformSchemaTransformProvider` configuration schema (`schematransform:org.apache.beam:sql_transform:v1`) now uses only portable types so it can be discovered from Python: `dialect` is a string, `tableproviders` is an array of strings, and the unused `parameters` field was removed. Invalid dialects and unknown table providers are rejected with a clear error ([#34613](https://github.com/apache/beam/issues/34613)).
 * (Python) An unknown pipeline option is no longer given the value of a later argument that does not directly follow it, which made `Pipeline.run()` and `get_all_options(retain_unknown_options=True)` fail with `expected one argument` for flags like `--unknown_flag --job_name=foo bar` ([#22434](https://github.com/apache/beam/issues/22434)).
 * (Python) Fixed `WriteToJson` silently rounding floats to 10 decimal places: it now defaults pandas' `double_precision` to 15, so values such as `0.1234567890123` and `1e-11` are written correctly. Floats that need 16-17 significant digits may still be rounded ([#41102](https://github.com/apache/beam/issues/41102)).
+* (Python) Fixed MLTransform embeddings created with `columns=`, the RAG embedding transforms and `VertexAIMultiModalEmbeddings` modifying their input elements, so that other transforms reading the same PCollection could see the embeddings ([#41233](https://github.com/apache/beam/issues/41233)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/python/apache_beam/ml/rag/embeddings/base.py

@@ -21,6 +21,7 @@
 are used by EmbeddingsManager to support various input types.
 """
 
+import dataclasses
 from collections.abc import Sequence
 
 from apache_beam.ml.rag.types import EmbeddableItem
@@ -63,6 +64,9 @@ def _add_embedding_fn(
     items: Sequence[EmbeddableItem],
     embeddings: Sequence[list[float]]) -> list[EmbeddableItem]:
   """Create Embeddings from items and embedding vectors."""
+  # Copy the items so the inputs, which other transforms may also consume, are
+  # not modified.
+  items = [dataclasses.replace(item) for item in items]
   for item, embedding in zip(items, embeddings):
     item.embedding = Embedding(dense_embedding=embedding)
-  return list(items)
+  return items

sdks/python/apache_beam/ml/rag/embeddings/base_test.py

@@ -89,6 +89,16 @@ def test_adapter_output_conversion(self):
     embeddings = adapter.output_fn(self.test_chunks, mock_embeddings)
     self.assertListEqual(embeddings, expected)
 
+  def test_adapter_output_does_not_modify_input(self):
+    adapter = create_text_adapter()
+
+    results = adapter.output_fn(self.test_chunks, [[0.1, 0.2], [0.3, 0.4]])
+
+    self.assertEqual(results[0].embedding.dense_embedding, [0.1, 0.2])
+    self.assertEqual(results[1].embedding.dense_embedding, [0.3, 0.4])
+    for chunk in self.test_chunks:
+      self.assertIsNone(chunk.embedding)
+
 
 class ImageEmbeddableItemTest(unittest.TestCase):
   def test_from_image_str(self):

sdks/python/apache_beam/ml/transforms/base.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import abc
+import copy
 import functools
 import logging
 import os
@@ -230,6 +231,10 @@ def _dict_output_fn(
 
   result = []
   for batch_idx, item in enumerate(batch):
+    if isinstance(item, dict):
+      # Copy so the input element, which other transforms may also consume,
+      # is not modified. copy.copy keeps the type of dict subclasses.
+      item = copy.copy(item)
     for col_idx, col in enumerate(columns):
       embedding_idx = batch_idx * len(columns) + col_idx
       if isinstance(item, dict):

sdks/python/apache_beam/ml/transforms/base_test.py

@@ -763,6 +763,16 @@ def test_dict_output_fn_normal(self):
     self.assertEqual(
         base._dict_output_fn(columns, input_list, embeddings), expected_output)
 
+  def test_dict_output_fn_does_not_modify_input(self):
+    input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]
+    columns = ['a']
+    embeddings = [1.1, 3.3]
+
+    output = base._dict_output_fn(columns, input_list, embeddings)
+
+    self.assertEqual(output, [{'a': 1.1, 'b': 2}, {'a': 3.3, 'b': 4}])
+    self.assertEqual(input_list, [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}])
+
   def test_dict_input_fn_on_list_inputs(self):
     input_list = [{'a': [1, 2, 10], 'b': 3}, {'a': [1], 'b': 5}]
     columns = ['a', 'b']

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py

@@ -19,6 +19,8 @@
 # Follow https://cloud.google.com/vertex-ai/docs/python-sdk/use-vertex-ai-python-sdk # pylint: disable=line-too-long
 # to install Vertex AI Python SDK.
 
+import copy
+import dataclasses
 import functools
 import logging
 from collections.abc import Callable
@@ -408,13 +410,19 @@ def _multimodal_dict_output_fn(
   results = []
   for batch_idx, item in enumerate(batch):
     mm_embedding = embeddings[batch_idx]
+    # Copy the dict and the nested objects so the input element, which other
+    # transforms may also consume, is not modified.
+    item = copy.copy(item)
     if image_column:
-      item[image_column].embedding = mm_embedding.image_embedding
+      item[image_column] = dataclasses.replace(
+          item[image_column], embedding=mm_embedding.image_embedding)
     if video_column:
-      item[video_column].embeddings = mm_embedding.video_embeddings
+      item[video_column] = dataclasses.replace(
+          item[video_column], embeddings=mm_embedding.video_embeddings)
     if text_column:
-      item[text_column].embedding = Embedding(
-          dense_embedding=mm_embedding.text_embedding)
+      item[text_column] = dataclasses.replace(
+          item[text_column],
+          embedding=Embedding(dense_embedding=mm_embedding.text_embedding))
     results.append(item)
   return results
 

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py

@@ -0,0 +1,65 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+
+import types
+import unittest
+
+import pytest
+
+from apache_beam.ml.rag.types import EmbeddableItem
+
+pytest.importorskip("vertexai", reason="Vertex AI dependencies not available")
+
+# pylint: disable=wrong-import-position
+from vertexai.vision_models import Image
+from vertexai.vision_models import Video
+from vertexai.vision_models import VideoSegmentConfig
+
+from apache_beam.ml.transforms.embeddings import vertex_ai
+
+
+class VertexAIMultiModalOutputFnTest(unittest.TestCase):
+  def test_output_fn_does_not_modify_input(self):
+    image = vertex_ai.VertexImage(image_content=Image(image_bytes=b'img'))
+    video = vertex_ai.VertexVideo(
+        video_content=Video(video_bytes=b'vid'), config=VideoSegmentConfig())
+    text = EmbeddableItem.from_text('hello')
+    item = {'image': image, 'video': video, 'text': text, 'id': 1}
+    video_embeddings = ['video-embedding']
+    response = types.SimpleNamespace(
+        image_embedding=[0.1, 0.2],
+        video_embeddings=video_embeddings,
+        text_embedding=[0.3, 0.4])
+
+    output = vertex_ai._multimodal_dict_output_fn(
+        'image', 'video', 'text', [item], [response])
+
+    self.assertEqual(len(output), 1)
+    self.assertEqual(output[0]['image'].embedding, [0.1, 0.2])
+    self.assertEqual(output[0]['video'].embeddings, video_embeddings)
+    self.assertEqual(output[0]['text'].embedding.dense_embedding, [0.3, 0.4])
+    self.assertEqual(output[0]['id'], 1)
+    # The input element and its nested objects are unchanged.
+    self.assertIs(item['image'], image)
+    self.assertIs(item['video'], video)
+    self.assertIs(item['text'], text)
+    self.assertIsNone(image.embedding)
+    self.assertIsNone(video.embeddings)
+    self.assertIsNone(text.embedding)
+
+
+if __name__ == '__main__':
+  unittest.main()