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()