diff --git a/sink/src/main/java/com/datastax/oss/kafka/sink/KafkaStruct.java b/sink/src/main/java/com/datastax/oss/kafka/sink/KafkaStruct.java index 68d8d804..8e6d2453 100644 --- a/sink/src/main/java/com/datastax/oss/kafka/sink/KafkaStruct.java +++ b/sink/src/main/java/com/datastax/oss/kafka/sink/KafkaStruct.java @@ -17,6 +17,7 @@ import com.datastax.oss.common.sink.AbstractSchema; import com.datastax.oss.common.sink.AbstractStruct; +import org.apache.kafka.connect.data.Schema; import org.apache.kafka.connect.data.Struct; /** Wrapper for Kafka Structs. */ @@ -44,4 +45,8 @@ public Object get(String field) { public AbstractSchema schema() { return KafkaSchema.of(struct.schema()); } + + public Schema nativeSchema() { + return struct.schema(); + } } diff --git a/sink/src/main/java/com/datastax/oss/kafka/sink/codecs/StructToUDTCodec.java b/sink/src/main/java/com/datastax/oss/kafka/sink/codecs/StructToUDTCodec.java index 5672db42..e5955c9c 100644 --- a/sink/src/main/java/com/datastax/oss/kafka/sink/codecs/StructToUDTCodec.java +++ b/sink/src/main/java/com/datastax/oss/kafka/sink/codecs/StructToUDTCodec.java @@ -26,21 +26,38 @@ import com.datastax.oss.dsbulk.codecs.api.ConvertingCodec; import com.datastax.oss.dsbulk.codecs.api.ConvertingCodecFactory; import com.datastax.oss.kafka.sink.KafkaStruct; +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; import java.util.List; import java.util.Set; import java.util.stream.Collectors; +import org.apache.kafka.connect.data.Schema; import org.apache.kafka.connect.data.Struct; /** Codec to convert a Kafka {@link Struct} to a UDT. */ public class StructToUDTCodec extends ConvertingCodec { + private static final int MAX_PLANS_BY_SCHEMA = 1000; + private final ConvertingCodecFactory codecFactory; private final UserDefinedType definition; + private final int size; + private final List udtFieldNames; + private final List udtFieldTypes; + private final List udtFieldNamesInternal; + private final Cache plansBySchema = + Caffeine.newBuilder().maximumSize(MAX_PLANS_BY_SCHEMA).build(); StructToUDTCodec(ConvertingCodecFactory codecFactory, UserDefinedType cqlType) { super(codecFactory.getCodecRegistry().codecFor(cqlType), KafkaStruct.class); this.codecFactory = codecFactory; definition = cqlType; + udtFieldNames = definition.getFieldNames(); + udtFieldTypes = definition.getFieldTypes(); + size = udtFieldNames.size(); + assert (size == udtFieldTypes.size()); + udtFieldNamesInternal = + udtFieldNames.stream().map(CqlIdentifier::asInternal).collect(Collectors.toList()); } @Override @@ -49,8 +66,22 @@ public UdtValue externalToInternal(KafkaStruct external) { return null; } - int size = definition.getFieldNames().size(); - AbstractSchema schema = external.schema(); + StructToUdtPlan plan = + plansBySchema.get(external.nativeSchema(), schema -> createPlan(external.schema())); + UdtValue value = definition.newValue(); + for (FieldBinding binding : plan.bindings) { + Object o = binding.codec.externalToInternal(external.get(binding.fieldNameInternal)); + value = setValue(value, binding.fieldName, o, binding.codecInternalType); + } + return value; + } + + private static UdtValue setValue( + UdtValue value, CqlIdentifier fieldName, Object raw, GenericType targetType) { + return value.set(fieldName, raw, targetType); + } + + private StructToUdtPlan createPlan(AbstractSchema schema) { StructDataMetadata structMetadata = new StructDataMetadata(schema); Set structFieldNames = schema.fields().stream().map(AbstractField::name).collect(Collectors.toSet()); @@ -59,31 +90,52 @@ public UdtValue externalToInternal(KafkaStruct external) { String.format("Expecting %d fields, got %d", size, structFieldNames.size())); } - UdtValue value = definition.newValue(); - List fieldNames = definition.getFieldNames(); - List fieldTypes = definition.getFieldTypes(); - assert (fieldNames.size() == fieldTypes.size()); - + FieldBinding[] bindings = new FieldBinding[size]; for (int idx = 0; idx < size; idx++) { - CqlIdentifier udtFieldName = fieldNames.get(idx); - DataType udtFieldType = fieldTypes.get(idx); + CqlIdentifier udtFieldName = udtFieldNames.get(idx); + DataType udtFieldType = udtFieldTypes.get(idx); + String fieldNameInternal = udtFieldNamesInternal.get(idx); - if (!structFieldNames.contains(udtFieldName.asInternal())) { + if (!structFieldNames.contains(fieldNameInternal)) { throw new IllegalArgumentException( String.format( "Field %s in UDT %s not found in input struct", udtFieldName, definition.getName())); } + @SuppressWarnings("unchecked") GenericType fieldType = - (GenericType) - structMetadata.getFieldType(udtFieldName.asInternal(), udtFieldType); + (GenericType) structMetadata.getFieldType(fieldNameInternal, udtFieldType); ConvertingCodec fieldCodec = codecFactory.createConvertingCodec(udtFieldType, fieldType, false); - Object o = fieldCodec.externalToInternal(external.get(udtFieldName.asInternal())); - value = value.set(udtFieldName, o, fieldCodec.getInternalJavaType()); + bindings[idx] = new FieldBinding(udtFieldName, fieldNameInternal, fieldCodec); + } + return new StructToUdtPlan(bindings); + } + + private static final class StructToUdtPlan { + private final FieldBinding[] bindings; + + private StructToUdtPlan(FieldBinding[] bindings) { + this.bindings = bindings; + } + } + + private static final class FieldBinding { + private final CqlIdentifier fieldName; + private final String fieldNameInternal; + private final ConvertingCodec codec; + private final GenericType codecInternalType; + + private FieldBinding( + CqlIdentifier fieldName, String fieldNameInternal, ConvertingCodec codec) { + this.fieldName = fieldName; + this.fieldNameInternal = fieldNameInternal; + this.codec = codec; + @SuppressWarnings("unchecked") + GenericType internalType = (GenericType) codec.getInternalJavaType(); + this.codecInternalType = internalType; } - return value; } @Override