diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/YamlUtils.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/YamlUtils.java index 4382062b4619..e2b3f9a72aa2 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/YamlUtils.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/YamlUtils.java @@ -19,6 +19,8 @@ import static org.apache.beam.sdk.values.Row.toRow; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; import java.io.InputStream; import java.math.BigDecimal; import java.util.ArrayList; @@ -40,6 +42,8 @@ import org.yaml.snakeyaml.error.YAMLException; public class YamlUtils { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final Map> YAML_VALUE_PARSERS = ImmutableMap .) yamlValue, nestedSchema, convertNamesToCamelCase); } else if (fieldType.getTypeName() == Schema.TypeName.MAP) { return yamlValue; + } else if (fieldType.getTypeName() == Schema.TypeName.STRING) { + return toJsonString(field, yamlValue); } } @@ -198,6 +207,18 @@ public static String yamlStringFromMap(@Nullable Map map) { } } + private static String toJsonString(Field field, Object yamlValue) { + try { + return OBJECT_MAPPER.writeValueAsString(yamlValue); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException( + String.format( + "Failed to serialize YAML %s to JSON string for field '%s': %s", + yamlValue instanceof List ? "list" : "map", field.getName(), e.getMessage()), + e); + } + } + private static List findNonSerializableKeys(Map map) { List problematicKeys = new ArrayList<>(); Yaml yaml = new Yaml(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/YamlUtilsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/YamlUtilsTest.java index 05bc73370c7d..eefd880688a1 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/YamlUtilsTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/YamlUtilsTest.java @@ -288,4 +288,38 @@ public void testYamlStringFromMapWithValidMap() { org.junit.Assert.assertNotNull(yaml); org.junit.Assert.assertTrue(yaml.contains("string_key")); } + + @Test + public void testMapAndListToStringField() { + Schema schema = + Schema.builder().addStringField("map_field").addStringField("list_field").build(); + + String yaml = + "map_field:\n" + + " key: value\n" + + " nested:\n" + + " foo: 123\n" + + "list_field:\n" + + " - a\n" + + " - b\n"; + + Row row = YamlUtils.toBeamRow(yaml, schema); + assertEquals("{\"key\":\"value\",\"nested\":{\"foo\":123}}", row.getString("map_field")); + assertEquals("[\"a\",\"b\"]", row.getString("list_field")); + } + + @Test + public void testMapAndListToStringFieldWithCamelCase() { + Schema schema = Schema.builder().addStringField("mapField").addStringField("listField").build(); + + Map map = new java.util.HashMap<>(); + Map nestedMap = new java.util.HashMap<>(); + nestedMap.put("key", "value"); + map.put("map_field", nestedMap); + map.put("list_field", Arrays.asList("a", "b")); + + Row row = YamlUtils.toBeamRow(map, schema, true); + assertEquals("{\"key\":\"value\"}", row.getString("mapField")); + assertEquals("[\"a\",\"b\"]", row.getString("listField")); + } } diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java index 60cbe87d43d3..b89a8f6ca0a7 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java @@ -378,7 +378,24 @@ public void testBuildTransformWithManaged() { + "redistribute_num_keys: 10\n" + "allow_duplicates: false\n" + "offset_deduplication: true\n" - + "redistribute_by_record_key: true"); + + "redistribute_by_record_key: true", + "topic: topic_7\n" + + "bootstrap_servers: some bootstrap\n" + + "format: AVRO\n" + + "schema:\n" + + " type: record\n" + + " name: my_record\n" + + " fields:\n" + + " - name: bool\n" + + " type: boolean", + "topic: topic_8\n" + + "bootstrap_servers: some bootstrap\n" + + "format: JSON\n" + + "schema:\n" + + " type: \"object\"\n" + + " properties:\n" + + " value: { type: \"string\" }\n" + + "auto_offset_reset_config: earliest"); for (String config : configs) { // Kafka Read SchemaTransform gets built in ManagedSchemaTransformProvider's expand diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java index f76c26af9468..24f295910de7 100644 --- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java +++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaWriteSchemaTransformProviderTest.java @@ -273,15 +273,33 @@ public void testBuildTransformWithManaged() { "topic: topic_4\n" + "bootstrap_servers: some bootstrap\n" + "format: RAW\n" - + "with_gcp_adc: true"); + + "with_gcp_adc: true", + "topic: topic_5\n" + + "bootstrap_servers: some bootstrap\n" + + "format: AVRO\n" + + "schema:\n" + + " type: record\n" + + " name: my_record\n" + + " fields:\n" + + " - name: str\n" + + " type: string", + "topic: topic_6\n" + + "bootstrap_servers: some bootstrap\n" + + "format: JSON\n" + + "schema:\n" + + " type: \"object\"\n" + + " properties:\n" + + " value: { type: \"string\" }"); for (String config : configs) { + Schema inputSchema = + config.contains("format: JSON") + ? Schema.builder().addStringField("value").build() + : Schema.builder().addByteArrayField("bytes").build(); // Kafka Write SchemaTransform gets built in ManagedSchemaTransformProvider's expand Managed.write(Managed.KAFKA) .withConfig(YamlUtils.yamlStringToMap(config)) - .expand( - Pipeline.create() - .apply(Create.empty(Schema.builder().addByteArrayField("bytes").build()))); + .expand(Pipeline.create().apply(Create.empty(inputSchema))); } } diff --git a/sdks/python/apache_beam/transforms/external.py b/sdks/python/apache_beam/transforms/external.py index 33926281074d..956db6eaea15 100644 --- a/sdks/python/apache_beam/transforms/external.py +++ b/sdks/python/apache_beam/transforms/external.py @@ -22,6 +22,7 @@ import copy import functools import glob +import json import logging import re import subprocess @@ -257,6 +258,10 @@ def dict_to_row_recursive(field_type, py_value): key: dict_to_row_recursive(field_type.map_type.value_type, value) for key, value in py_value.items() } + elif (type_info == 'atomic_type' and + field_type.atomic_type == schema_pb2.STRING and + isinstance(py_value, (dict, list))): + return json.dumps(py_value) else: return py_value diff --git a/sdks/python/apache_beam/transforms/external_test.py b/sdks/python/apache_beam/transforms/external_test.py index 0af89367fcc2..05987e166819 100644 --- a/sdks/python/apache_beam/transforms/external_test.py +++ b/sdks/python/apache_beam/transforms/external_test.py @@ -45,6 +45,7 @@ from apache_beam.transforms import external from apache_beam.transforms.external import MANAGED_SCHEMA_TRANSFORM_IDENTIFIER from apache_beam.transforms.external import AnnotationBasedPayloadBuilder +from apache_beam.transforms.external import ExplicitSchemaTransformPayloadBuilder from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.transforms.external import JavaClassLookupPayloadBuilder from apache_beam.transforms.external import JavaExternalTransform @@ -524,6 +525,32 @@ def test_build_payload(self): self.assertEqual('bbb', schema_transform_config.object_field.str_sub_field) self.assertEqual(456, schema_transform_config.object_field.int_sub_field) + def test_explicit_payload_builder_with_dict_and_list_to_string_field(self): + schema = schema_pb2.Schema( + fields=[ + schema_pb2.Field( + name='str_field', + type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), + schema_pb2.Field( + name='list_field', + type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), + ]) + + payload_builder = ExplicitSchemaTransformPayloadBuilder( + identifier='dummy_id', + schema_proto=schema, + str_field={'foo': 'bar'}, + list_field=[1, 2, 'baz']) + payload = payload_builder.build() + + self.assertEqual('dummy_id', payload.identifier) + + coder = RowCoder(payload.configuration_schema) + config = coder.decode(payload.configuration_row) + + self.assertEqual('{"foo": "bar"}', config.str_field) + self.assertEqual('[1, 2, "baz"]', config.list_field) + class SchemaAwareExternalTransformTest(unittest.TestCase): class MockDiscoveryService: diff --git a/sdks/python/apache_beam/yaml/yaml_io_test.py b/sdks/python/apache_beam/yaml/yaml_io_test.py index 7276aa5f0728..dd1ce3ecaf85 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/sdks/python/apache_beam/yaml/yaml_io_test.py @@ -33,8 +33,10 @@ from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import schemas as schema_utils +from apache_beam.utils import subprocess_server from apache_beam.utils.timestamp import Timestamp from apache_beam.yaml.yaml_transform import YamlTransform +from apache_beam.yaml.yaml_transform import expand_pipeline try: import jsonschema @@ -1002,6 +1004,62 @@ def test_read_from_delta_cdc(self): self.assertEqual(transform._source, "delta_cdc") +@unittest.skipIf( + not os.path.exists( + subprocess_server.JavaJarServer.path_to_dev_beam_jar( + 'sdks:java:io:expansion-service:shadowJar')), + "Requires expansion service jars.") +class YamlKafkaTest(unittest.TestCase): + def test_read_from_kafka_json_schema_expansion(self): + # Regression test for https://github.com/apache/beam/issues/35186. + # Verifies that ReadFromKafka expands with a nested JSON schema map without + # UnsupportedOperationException. + p = beam.Pipeline( + options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) + expand_pipeline( + p, + ''' + pipeline: + type: chain + transforms: + - type: ReadFromKafka + name: ReadFromMyTopic + config: + format: JSON + schema: + type: "object" + properties: + value: { type: "string" } + topic: test + bootstrap_servers: kafka:9092 + auto_offset_reset_config: earliest + + - type: LogForTesting + ''') + + def test_read_from_kafka_avro_schema_expansion(self): + # Verifies that ReadFromKafka expands with a nested AVRO schema map without + # UnsupportedOperationException. + p = beam.Pipeline( + options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) + _ = p | YamlTransform( + ''' + type: ReadFromKafka + config: + topic: my-topic + bootstrap_servers: kafka:9092 + format: AVRO + schema: + type: record + name: my_record + fields: + - name: bool + type: boolean + ''') + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main()