Add oneof support for Protobuf via Confluent Schema Registry #16836

adamjshook · 2023-04-01T00:05:52Z

Description

This adds Kafka connector support for Protobuf oneOf data types by encoding them as a JSON object.

Additional context and related issues

Protobuf oneOf data types are a single field in the message that can hold one value of multiple fields (of varying types) per the schema definition. This adds support for these data types by encoding the entire oneOf field as a JSON object where each field name is a key in the JSON object where all values, except one, are null. The value that is set corresponds to the set value in the DynamicMessage that is processed in the message.

Release notes

( ) This is not user-visible or docs only and no release notes are required.
( ) Release notes are required, please propose a release note for me.
(X) Release notes are required, with the following suggested text:

# Kafka Connector
* Added support for `oneof` Protobuf types when using the Confluent table description provider

vlad-lyutenko · 2023-04-05T14:29:09Z

I am new for this topic, but I made review, general implementation looks good to me, I have only question about dependencies in pom. A little bit missed with them, why we comment some protobuf plugins?

adamjshook · 2023-04-05T15:06:58Z

I am new for this topic, but I made review, general implementation looks good to me

Thanks for the review, I'm hoping to clean things up today and open this up for review.

I have only question about dependencies in pom. A little bit missed with them, why we comment some protobuf plugins?

I originally added the protobuf plugins to generate a TestOneOf object for unit testing, however by the end of implementation I was able to do everything using a DynamicMessage. I commented out the plugins to make sure they weren't needed in any of my tests. The final PR will not use these plugins (I'll be deleting all the commented code).

hashhar

skimmed, mostly editorial comments

Impl looks good.

hashhar · 2023-04-05T15:17:05Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufColumnDecoder.java

@@ -69,43 +79,51 @@ public ProtobufColumnDecoder(DecoderColumnHandle columnHandle)
            this.columnType = columnHandle.getType();
            this.columnMapping = columnHandle.getMapping();
            this.columnName = columnHandle.getName();
+            this.oneOfs = columnHandle.getOneOfs()
+                    .map(list -> list.stream().sorted(String::compareTo).collect(toImmutableList()));


why do we need to sort?

I was sorting it based on this comment in JsonType. Maybe it isn't relevant if the built type is a Slice instead of a Block. (Back when I was using Block to build the JSON, it did require sorting as the comment suggested.)

this is very interesting, I had no idea the Trino JsonType required keys to be sorted. I now see in JDBC connectors we don't need to sort because it's done by JsonTypeUtils#jsonParse/JsonTypeUtils#toJsonValue for us.

Maybe we should encode this knowledge?

Since the JSON record we create will only have at most one field in it (since oneof can only have one set), sorting the column handles isn't necessary here.

hashhar · 2023-04-05T15:19:31Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufColumnDecoder.java

    {
        Object value = message;
        Optional<Descriptor> valueDescriptor = Optional.of(message.getDescriptorForType());
        for (String pathElement : Splitter.on('/').omitEmptyStrings().split(columnMapping)) {
            if (valueDescriptor.filter(descriptor -> descriptor.findFieldByName(pathElement) != null).isEmpty()) {
-                return null;
+                if (oneofColumns.isPresent()) {
+                    DynamicSliceOutput output = new DynamicSliceOutput(40);


can we come up with better estimates, e.g. based on field count * some avg length?

I made a change to calculate the size based on the string values of the field name, value, plus the extra characters. Looking at the data type test case, there wouldn't be any resizing.

hashhar · 2023-04-05T15:20:29Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/DecoderColumnHandle.java

@@ -33,4 +36,9 @@
    String getMapping();

    String getDataFormat();
+
+    default Optional<List<String>> getOneOfs()


nit: Let's consitently use oneof instead of oneOf since that's what the protobuf schema calls it.

hashhar · 2023-04-05T15:21:01Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufUtils.java

separate commit, just formatting changes

I'm just going to revert this; accidentally committed from when I was messing around in this class during implementation.

hashhar · 2023-04-05T15:21:43Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java

+        if (type instanceof JsonType) {
+            if (builder != null) {
+                type.writeObject(builder, value);
+                return null;
+            }
+            return (Block) value;
+        }


extract to a serializeJson method?

hashhar · 2023-04-05T15:22:14Z

lib/trino-record-decoder/src/test/java/io/trino/decoder/DecoderTestColumnHandle.java

@@ -70,6 +72,22 @@
     */
    private final boolean internal;

+    private final Optional<List<String>> oneOfs;
+
+    public DecoderTestColumnHandle(


all ctors should ideally go via the JsonCreator, less chance of mistakes

if this causes a large diff you can extract as the first commit with some reasons in commit message

hashhar · 2023-04-05T15:23:48Z

lib/trino-record-decoder/src/test/java/io/trino/decoder/DecoderTestColumnHandle.java

@@ -172,7 +199,8 @@ public boolean equals(Object obj)
                Objects.equals(this.formatHint, other.formatHint) &&
                Objects.equals(this.keyDecoder, other.keyDecoder) &&
                Objects.equals(this.hidden, other.hidden) &&
-                Objects.equals(this.internal, other.internal);
+                Objects.equals(this.internal, other.internal) &&
+                Objects.equals(this.oneOfs, other.oneOfs);


list equality depends on element order.

Does a protobuf schema consider differently ordered oneofs as different?
At-least within JSON the field order doesn't matter - so maybe this equality method should do the same?

Can't say the protobuf schema would treat different ordered oneof definitions as different objects. The order of the column names in this list matches the order of how the fields are defined in the schema (by name, not by tag number). Since order doesn't matter, and the final JSON object will have at most one field, how about we change this to a Set<String> instead of List? @hashhar

Set sounds good if we're only going to keep the single field.

hashhar · 2023-04-05T15:24:39Z

plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/KafkaColumnHandle.java

@@ -68,6 +71,22 @@
     */
    private final boolean internal;

+    private final Optional<List<String>> oneOfs;
+
+    // TODO Should we remove this constructor and update all the tests to use the one annotated with @JsonCreator


if it's a large diff, extract a commit

same applies to other ctor changes as well

adamjshook · 2023-04-05T16:37:56Z

.../io/trino/plugin/kafka/protobuf/TestKafkaProtobufWithSchemaRegistryMinimalFunctionality.java

+
+        assertThat(query(format("SELECT testOneOfColumn FROM %s", toDoubleQuoted(topic))))
+                .matches("""
+                        VALUES (JSON '{"stringColumn":"%s","integerColumn":null,"longColumn":null,"doubleColumn":null,"floatColumn":null,"booleanColumn":null}')


Do we want all possible columns in the oneof to be in the JSON string, leaving all other values as null? Or just a single column e.g. {"stringColumn":"value"}?

I think having a single field populated would be better because otherwise you can't differentiate that stringColumn=null because stringColumn was set to null or because some other field existed and stringColumn didn't exist.

Makes sense and I agree. I'll update to only write a field if it is present. Protobuf doesn't do normal null values, but I have added a test case for an empty message (which results in an empty JSON value {}).

hashhar

A test showcasing this doesn't work with the file-based table description supplier would be useful.

LGTM from me.

hashhar · 2023-04-11T14:38:58Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufColumnDecoder.java

 import static java.util.Objects.requireNonNull;

 public class ProtobufColumnDecoder
 {
+    private static final JsonFactory JSON_FACTORY = new JsonFactoryBuilder().disable(CANONICALIZE_FIELD_NAMES).build();
+    private static final Slice EMPTY_JSON = Slices.utf8Slice("{}");
+    private static final int EXTRA_JSON_CHARACTERS_LENGTH = 7;


nit: "{\"\":\"\"}".length (or a comment).

adamjshook · 2023-04-11T18:28:46Z

A test showcasing this doesn't work with the file-based table description supplier would be useful.

Added a test case asserting the column names are not found when using the file-based descriptor.

adamjshook · 2023-04-12T14:12:28Z

I've stumbled on a library com.google.protobuf:protobuf-java-util which has support for JSON-encoding protobuf oneof types. I'm going to flip this back to a draft because I think it may make the implementation a lot more simple.

adamjshook · 2023-04-13T15:31:52Z

Flipping this back to review for merge in the current state. We can follow up on any improvements later on.

hashhar · 2023-04-14T09:33:27Z

@mosabua / @jhlodin can you please review the last commit which includes documentation changes.

Praveen2112 · 2023-04-14T09:34:13Z

Flipping this back to review for merge in the current state. We can follow up on any improvements later on.
Can we have a GH issue to track the followup ?

Praveen2112

Overall looks good to me.

But I would like to see how it serializes/deserializes if one of oneOf field is of a message type - From this docs - https://protobuf.dev/programming-guides/proto3/#using-oneof it looks like map and array are not supported but row should be supported by oneOf

And if it is possible for us to determine the oneOf fields when reading the message, do we need to propagate these information when we create KafkaColumnHandle ?

Praveen2112 · 2023-04-14T09:57:10Z

...n/trino-kafka/src/main/java/io/trino/plugin/kafka/encoder/protobuf/ProtobufSchemaParser.java

+                        .collect(toImmutableList()));
+    }
+
+    private static List<FieldDescriptor> getFields(Descriptor descriptor)


nit: Can we return as KafkaTopicFieldDescription here so as to be unified with getOneofs.

Praveen2112 · 2023-04-14T10:04:26Z

plugin/trino-kafka/src/test/resources/protobuf/test_oneof.proto

+        uint64 longColumn = 3;
+        double doubleColumn = 4;
+        float floatColumn = 5;
+        bool booleanColumn = 6;


Can we also pass a message type and check on how DynamicMessage would be serialized.

Praveen2112 · 2023-04-14T10:10:43Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufColumnDecoder.java

@@ -144,4 +177,36 @@ private static Optional<Descriptor> getDescriptor(FieldDescriptor fieldDescripto
        }
        return Optional.empty();
    }
+
+    private static Object createOneofJson(DynamicMessage message, Set<String> oneofColumns)


Any specific reason on why we should pass the oneofColumns from the KafkaColumnHandle ? Why can't we use message.getDescriptorForType() get oneOf types from here and handle them ? In this case we don't have to change other decoder right ?

Yeah. You can collect oneof fields from FieldDescriptor:

.getContainingOneof().getName()

Great, thanks for the suggestion. I just couldn't find a way to get the names to where they needed to be, but this does the trick. This removes the need to include the oneof names in the column handles.

Praveen2112 · 2023-04-14T10:12:22Z

lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java

        throw new TrinoException(DECODER_CONVERSION_NOT_SUPPORTED, format("cannot decode object of '%s' as '%s' for column '%s'", value.getClass(), type, columnName));
    }

    @Nullable
-    private static Block serializeObject(BlockBuilder builder, Object value, Type type, String columnName)
+    private Block serializeObject(BlockBuilder builder, Object value, Type type, String columnName)


nit: These can be extracted to a new commit.

Praveen2112 · 2023-04-14T10:15:07Z

lib/trino-record-decoder/src/test/java/io/trino/decoder/protobuf/TestProtobufDecoder.java

+                "booleanColumn",
+                "numberColumn",
+                "timestampColumn",
+                "bytesColumn");


Can we add additional columnName with message type

Praveen2112 · 2023-04-14T10:15:32Z

lib/trino-record-decoder/src/test/java/io/trino/decoder/protobuf/TestProtobufDecoder.java

+                    .build();
+            assertOneof(message, testOneOfColumn, oneofColumnNames, Map.of("booleanColumn", booleanData));
+        }
+    }


for timestamp ?

mosabua

A couple of fixes for docs necessary. Good once the suggestions are applied.

mosabua · 2023-04-14T15:49:18Z