From 0a1484e220333134fb1e4d4f894ea119cc8ce825 Mon Sep 17 00:00:00 2001 From: sohailshk <122166523+sohailshk@users.noreply.github.com> Date: Thu, 12 Nov 2026 10:15:47 +0530 Subject: [PATCH] [Python] Do not modify input elements in the embedding output functions The output functions of MLTransform dict embeddings (_dict_output_fn), of the RAG embedding adapters (_add_embedding_fn) and of the Vertex AI multimodal dict adapter (_multimodal_dict_output_fn) assigned the embeddings onto the elements they were given. Those elements are shared with any other consumer of the same PCollection, which then saw the embedded values instead of the original ones. Copy the elements before assigning. --- CHANGES.md | 1 + .../apache_beam/ml/rag/embeddings/base.py | 6 +- .../ml/rag/embeddings/base_test.py | 10 +++ sdks/python/apache_beam/ml/transforms/base.py | 5 ++ .../apache_beam/ml/transforms/base_test.py | 10 +++ .../ml/transforms/embeddings/vertex_ai.py | 16 +++-- .../transforms/embeddings/vertex_ai_test.py | 65 +++++++++++++++++++ 7 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py diff --git a/CHANGES.md b/CHANGES.md index 61b1547e025f..fe28499a5588 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/python/apache_beam/ml/rag/embeddings/base.py b/sdks/python/apache_beam/ml/rag/embeddings/base.py index 1f85f1d89daf..9c779cf24e4c 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/base.py +++ b/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 diff --git a/sdks/python/apache_beam/ml/rag/embeddings/base_test.py b/sdks/python/apache_beam/ml/rag/embeddings/base_test.py index 61e211e415c3..e6d20dae0f2c 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/base_test.py +++ b/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): diff --git a/sdks/python/apache_beam/ml/transforms/base.py b/sdks/python/apache_beam/ml/transforms/base.py index f52e9df40566..76527fe89481 100644 --- a/sdks/python/apache_beam/ml/transforms/base.py +++ b/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): diff --git a/sdks/python/apache_beam/ml/transforms/base_test.py b/sdks/python/apache_beam/ml/transforms/base_test.py index c7b904c09444..a6b467bb6ef0 100644 --- a/sdks/python/apache_beam/ml/transforms/base_test.py +++ b/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'] diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py index f6c15eeac092..4af3645dd197 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py +++ 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 diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py new file mode 100644 index 000000000000..3a814b4db71e --- /dev/null +++ b/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()