diff --git a/sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java b/sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java index a6cebe1851de..6c31d5a02349 100644 --- a/sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java +++ b/sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java @@ -23,6 +23,7 @@ import java.io.Serializable; import java.util.HashMap; import java.util.Map; +import java.util.stream.Collectors; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.MapCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; @@ -318,14 +319,23 @@ protected Schema getRecordSchema() { SourceRecord sampledRecord = fn.getOneRecord(getConnectorConfiguration().getConfigurationMap()); fn.reset(); + Schema keySchema = + sampledRecord.keySchema() != null + ? KafkaConnectUtils.beamSchemaFromKafkaConnectSchema(sampledRecord.keySchema()) + : Schema.builder().build(); Schema valueSchema = KafkaConnectUtils.beamSchemaFromKafkaConnectSchema(sampledRecord.valueSchema()); return Schema.builder() .addFields(valueSchema.getFields()) - // TODO(https://github.com/apache/beam/issues/39557): - // Restore 'primaryKeyColumns' once Python can decode ARRAY - // schema options across the Java/Python cross-language boundary. + .setOptions( + Schema.Options.builder() + .setOption( + "primaryKeyColumns", + Schema.FieldType.array(Schema.FieldType.STRING), + keySchema.getFields().stream() + .map(Schema.Field::getName) + .collect(Collectors.toList()))) .build(); } diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py index 084ecc93581c..2fd3c22e1e58 100644 --- a/sdks/python/apache_beam/typehints/schemas.py +++ b/sdks/python/apache_beam/typehints/schemas.py @@ -486,27 +486,40 @@ def value_from_runner_api( self, type_proto: schema_pb2.FieldType, value_proto: schema_pb2.FieldValue): - if type_proto.WhichOneof("type_info") != "atomic_type": - # TODO: Allow other value types + type_info = type_proto.WhichOneof("type_info") + if type_info == "atomic_type": + return self.atomic_value_from_runner_api( + type_proto.atomic_type, value_proto.atomic_value) + elif type_info == "array_type": + element_type = type_proto.array_type.element_type + return [ + self.value_from_runner_api(element_type, element) + for element in value_proto.array_value.element + ] + else: raise ValueError( - "Encounterd option with unsupported type. Only " - f"atomic_type options are supported: {type_proto}") - - value = self.atomic_value_from_runner_api( - type_proto.atomic_type, value_proto.atomic_value) - return value + "Encountered option with unsupported type. Only atomic_type and " + f"array_type options are supported: {type_proto}") def value_to_runner_api(self, typing_proto: schema_pb2.FieldType, value): - if typing_proto.WhichOneof("type_info") != "atomic_type": - # TODO: Allow other value types + type_info = typing_proto.WhichOneof("type_info") + if type_info == "atomic_type": + return schema_pb2.FieldValue( + atomic_value=self.atomic_value_to_runner_api( + typing_proto.atomic_type, value)) + elif type_info == "array_type": + element_type = typing_proto.array_type.element_type + return schema_pb2.FieldValue( + array_value=schema_pb2.ArrayTypeValue( + element=[ + self.value_to_runner_api(element_type, element) + for element in value + ])) + else: raise ValueError( - "Only atomic_type option values are currently supported in Python. " - f"Got {value!r}, which maps to fieldtype {typing_proto!r}.") - - atomic_value = self.atomic_value_to_runner_api( - typing_proto.atomic_type, value) - value_proto = schema_pb2.FieldValue(atomic_value=atomic_value) - return value_proto + "Only atomic_type and array_type option values are currently " + f"supported in Python. Got {value!r}, which maps to fieldtype " + f"{typing_proto!r}.") def option_from_runner_api( self, option_proto: schema_pb2.Option) -> Tuple[str, Any]: @@ -524,7 +537,10 @@ def option_to_runner_api(self, option: Tuple[str, Any]) -> schema_pb2.Option: # Don't set type, value return schema_pb2.Option(name=name) - type_proto = self.typing_to_runner_api(type(value)) + from apache_beam.typehints import trivial_inference + + type_proto = self.typing_to_runner_api( + trivial_inference.instance_to_type(value)) value_proto = self.value_to_runner_api(type_proto, value) return schema_pb2.Option(name=name, type=type_proto, value=value_proto) diff --git a/sdks/python/apache_beam/typehints/schemas_test.py b/sdks/python/apache_beam/typehints/schemas_test.py index c2c21a7ce391..327fe7947ca5 100644 --- a/sdks/python/apache_beam/typehints/schemas_test.py +++ b/sdks/python/apache_beam/typehints/schemas_test.py @@ -257,6 +257,22 @@ def get_test_beam_fieldtype_protos(): value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( bytes=b'bytes!'))), + schema_pb2.Option( + name='a_string_array', + type=schema_pb2.FieldType( + array_type=schema_pb2.ArrayType( + element_type=schema_pb2.FieldType( + atomic_type=schema_pb2.STRING))), + value=schema_pb2.FieldValue( + array_value=schema_pb2.ArrayTypeValue( + element=[ + schema_pb2.FieldValue( + atomic_value=schema_pb2. + AtomicTypeValue(string='a')), + schema_pb2.FieldValue( + atomic_value=schema_pb2. + AtomicTypeValue(string='b')), + ]))), ]))), schema_pb2.FieldType( row_type=schema_pb2.RowType(