Uh oh!
There was an error while loading. Please reload this page.
- Notifications
You must be signed in to change notification settings - Fork 1.6k
PARQUET-1711: support recursive proto schemas by limiting recursion depth#995
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Uh oh!
There was an error while loading. Please reload this page.
Merged
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Jump to
Jump to file
Failed to load files.
Loading
Uh oh!
There was an error while loading. Please reload this page.
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
148 changes: 119 additions & 29 deletions
148 parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoSchemaConverter.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -18,11 +18,13 @@ | ||
| */ | ||
| package org.apache.parquet.proto; | ||
| import com.google.common.collect.ImmutableSetMultimap; | ||
| import com.google.protobuf.Descriptors; | ||
| import com.google.protobuf.Descriptors.FieldDescriptor; | ||
| import com.google.protobuf.Descriptors.FieldDescriptor.JavaType; | ||
| import com.google.protobuf.Message; | ||
| import com.twitter.elephantbird.util.Protobufs; | ||
| import org.apache.hadoop.conf.Configuration; | ||
| import org.apache.parquet.schema.LogicalTypeAnnotation; | ||
| import org.apache.parquet.schema.MessageType; | ||
| import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; | ||
| @@ -34,6 +36,7 @@ | ||
| import org.slf4j.LoggerFactory; | ||
| import java.util.List; | ||
| import javax.annotation.Nullable; | ||
| import static org.apache.parquet.schema.LogicalTypeAnnotation.enumType; | ||
| import static org.apache.parquet.schema.LogicalTypeAnnotation.listType; | ||
| @@ -47,49 +50,111 @@ | ||
| public class ProtoSchemaConverter { | ||
| private static final Logger LOG = LoggerFactory.getLogger(ProtoSchemaConverter.class); | ||
| public static final String PB_MAX_RECURSION = "parquet.proto.maxRecursion"; | ||
| private final boolean parquetSpecsCompliant; | ||
| // TODO: use proto custom options to override per field. | ||
| private final int maxRecursion; | ||
| /** | ||
| * Instantiate a schema converter to get the parquet schema corresponding to protobuf classes. | ||
| * Returns instances that are not parquetSpecsCompliant with a maxRecursion of 5. | ||
| */ | ||
| public ProtoSchemaConverter() { | ||
| this(false); | ||
| } | ||
| /** | ||
| * Instantiate a schema converter to get the parquet schema corresponding to protobuf classes. | ||
| * Returns instances limited to 5 levels of recursion depth. | ||
| * | ||
| * @param parquetSpecsCompliant If set to false, the parquet schema generated will be using the old | ||
| * schema style (prior to PARQUET-968) to provide backward-compatibility | ||
| * but which does not use LIST and MAP wrappers around collections as required | ||
| * by the parquet specifications. If set to true, specs compliant schemas are used. | ||
| */ | ||
| public ProtoSchemaConverter(boolean parquetSpecsCompliant) { | ||
| this(parquetSpecsCompliant, 5); | ||
| } | ||
| /** | ||
| * Instantiate a schema converter to get the parquet schema corresponding to protobuf classes. | ||
| * Returns instances that are not specs compliant and limited to 5 levels of recursion depth. | ||
| * | ||
| * @param config Hadoop configuration object to parrse parquetSpecsCompliant and maxRecursion settings. | ||
| */ | ||
| public ProtoSchemaConverter(Configuration config) { | ||
| this( | ||
| config.getBoolean(ProtoWriteSupport.PB_SPECS_COMPLIANT_WRITE, false), | ||
| config.getInt(PB_MAX_RECURSION, 5)); | ||
| } | ||
| /** | ||
| * Instantiate a schema converter to get the parquet schema corresponding to protobuf classes. | ||
| * | ||
| * @param parquetSpecsCompliant If set to false, the parquet schema generated will be using the old | ||
| * schema style (prior to PARQUET-968) to provide backward-compatibility | ||
| * but which does not use LIST and MAP wrappers around collections as required | ||
| * by the parquet specifications. If set to true, specs compliant schemas are used. | ||
| * @param maxRecursion The maximum recursion depth messages are allowed to go before terminating as | ||
| * bytes instead of their actual schema. | ||
| */ | ||
| public ProtoSchemaConverter(boolean parquetSpecsCompliant, int maxRecursion) { | ||
| this.parquetSpecsCompliant = parquetSpecsCompliant; | ||
| this.maxRecursion = maxRecursion; | ||
| } | ||
| /** | ||
| * Sets the maximum recursion depth for recursive schemas. | ||
| * | ||
| * @param config The hadoop configuration to be updated. | ||
| * @param maxRecursion The maximum recursion depth messages are allowed to go before terminating as | ||
| * bytes instead of their actual schema. | ||
| */ | ||
| public static void setMaxRecursion(Configuration config, int maxRecursion) { | ||
| config.setInt(PB_MAX_RECURSION, maxRecursion); | ||
| } | ||
| /** | ||
| * Converts a given protobuf message descriptor to a parquet schema. | ||
| * | ||
| * @param descriptor The protobuf message descriptor to convert. | ||
| * @return The parquet schema encoded as a MessageType. | ||
| */ | ||
| public MessageType convert(Descriptors.Descriptor descriptor) { | ||
| MessageType messageType = | ||
| convertFields(Types.buildMessage(), descriptor.getFields()) | ||
| // Remember classes seen with depths to avoid cycles. | ||
| int depth = 0; | ||
| ImmutableSetMultimap<String, Integer> seen = ImmutableSetMultimap.of(descriptor.getFullName(), depth); | ||
| LOG.trace("convert:\n{}", descriptor.toProto()); | ||
| MessageType messageType = convertFields(Types.buildMessage(), descriptor.getFields(), seen, depth) | ||
| .named(descriptor.getFullName()); | ||
| LOG.debug("Converter info:\n " + descriptor.toProto() + " was converted to \n" + messageType); | ||
| LOG.debug("Converter info:\n{}\n was converted to:\n{}", descriptor.toProto(), messageType); | ||
| return messageType; | ||
| } | ||
| /** | ||
| * Converts a given protobuf message class to a parquet schema. | ||
| * | ||
| * @param protobufClass The protobuf message class (e.g. MyMessage.class) to convert. | ||
| * @return The parquet schema encoded as a MessageType. | ||
| */ | ||
| public MessageType convert(Class<? extends Message> protobufClass) { | ||
| LOG.debug("Converting protocol buffer class \"" + protobufClass + "\" to parquet schema."); | ||
| LOG.debug("Converting protocol buffer class \"{}\" to parquet schema", protobufClass); | ||
| Descriptors.Descriptor descriptor = Protobufs.getMessageDescriptor(protobufClass); | ||
| return convert(descriptor); | ||
| } | ||
| /* Iterates over list of fields. **/ | ||
| private <T> GroupBuilder<T> convertFields(GroupBuilder<T> groupBuilder, List<FieldDescriptor> fieldDescriptors) { | ||
| private <T> GroupBuilder<T> convertFields(GroupBuilder<T> groupBuilder, List<FieldDescriptor> fieldDescriptors, ImmutableSetMultimap<String, Integer> seen, int depth) { | ||
| for (FieldDescriptor fieldDescriptor : fieldDescriptors) { | ||
| groupBuilder = | ||
| addField(fieldDescriptor, groupBuilder) | ||
| groupBuilder = addField(fieldDescriptor, groupBuilder, seen, depth) | ||
| .id(fieldDescriptor.getNumber()) | ||
| .named(fieldDescriptor.getName()); | ||
| } | ||
| return groupBuilder; | ||
| } | ||
| private Type.Repetition getRepetition(FieldDescriptor descriptor) { | ||
| private static Type.Repetition getRepetition(FieldDescriptor descriptor) { | ||
| if (descriptor.isRequired()) { | ||
| return Type.Repetition.REQUIRED; | ||
| } else if (descriptor.isRepeated()) { | ||
| @@ -99,9 +164,9 @@ private Type.Repetition getRepetition(FieldDescriptor descriptor) { | ||
| } | ||
| } | ||
| private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addField(FieldDescriptor descriptor, final GroupBuilder<T> builder) { | ||
| private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addField(FieldDescriptor descriptor, final GroupBuilder<T> builder, ImmutableSetMultimap<String, Integer> seen, int depth) { | ||
| if (descriptor.getJavaType() == JavaType.MESSAGE) { | ||
| return addMessageField(descriptor, builder); | ||
| return addMessageField(descriptor, builder, seen, depth); | ||
| } | ||
| ParquetType parquetType = getParquetType(descriptor); | ||
| @@ -113,7 +178,7 @@ private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addF | ||
| return builder.primitive(parquetType.primitiveType, getRepetition(descriptor)).as(parquetType.logicalTypeAnnotation); | ||
| } | ||
| private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addRepeatedPrimitive(PrimitiveTypeName primitiveType, | ||
| private static <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addRepeatedPrimitive(PrimitiveTypeName primitiveType, | ||
| LogicalTypeAnnotation logicalTypeAnnotation, | ||
| final GroupBuilder<T> builder) { | ||
| return builder | ||
| @@ -124,35 +189,61 @@ private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addR | ||
| .named("list"); | ||
| } | ||
| private <T> GroupBuilder<GroupBuilder<T>> addRepeatedMessage(FieldDescriptor descriptor, GroupBuilder<T> builder) { | ||
| GroupBuilder<GroupBuilder<GroupBuilder<GroupBuilder<T>>>> result = | ||
| builder | ||
| private <T> GroupBuilder<GroupBuilder<T>> addRepeatedMessage(FieldDescriptor descriptor, GroupBuilder<T> builder, ImmutableSetMultimap<String, Integer> seen, int depth) { | ||
| GroupBuilder<GroupBuilder<GroupBuilder<GroupBuilder<T>>>> result = builder | ||
| .group(Type.Repetition.OPTIONAL).as(listType()) | ||
| .group(Type.Repetition.REPEATED) | ||
| .group(Type.Repetition.OPTIONAL); | ||
| convertFields(result, descriptor.getMessageType().getFields()); | ||
| convertFields(result, descriptor.getMessageType().getFields(), seen, depth); | ||
| return result.named("element").named("list"); | ||
| } | ||
| private <T> GroupBuilder<GroupBuilder<T>> addMessageField(FieldDescriptor descriptor, final GroupBuilder<T> builder) { | ||
| private <T> Builder<? extends Builder<?, GroupBuilder<T>>, GroupBuilder<T>> addMessageField(FieldDescriptor descriptor, final GroupBuilder<T> builder, ImmutableSetMultimap<String, Integer> seen, int depth) { | ||
| // Prevent recursion by terminating with optional proto bytes. | ||
| depth += 1; | ||
| String typeName = getInnerTypeName(descriptor); | ||
| LOG.trace("addMessageField: {} type: {} depth: {}", descriptor.getFullName(), typeName, depth); | ||
| if (typeName != null) { | ||
| if (seen.get(typeName).size() > maxRecursion) { | ||
| return builder.primitive(BINARY, Type.Repetition.OPTIONAL).as((LogicalTypeAnnotation) null); | ||
| } | ||
| } | ||
| if (descriptor.isMapField() && parquetSpecsCompliant) { | ||
| // the old schema style did not include the MAP wrapper around map groups | ||
| return addMapField(descriptor, builder); | ||
| return addMapField(descriptor, builder, seen, depth); | ||
| } | ||
| seen = ImmutableSetMultimap.<String, Integer>builder().putAll(seen).put(typeName, depth).build(); | ||
jinyius marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| if (descriptor.isRepeated() && parquetSpecsCompliant) { | ||
| // the old schema style did not include the LIST wrapper around repeated messages | ||
| return addRepeatedMessage(descriptor, builder); | ||
| return addRepeatedMessage(descriptor, builder, seen, depth); | ||
| } | ||
| // Plain message | ||
| // Plain message. | ||
| GroupBuilder<GroupBuilder<T>> group = builder.group(getRepetition(descriptor)); | ||
| convertFields(group, descriptor.getMessageType().getFields()); | ||
| convertFields(group, descriptor.getMessageType().getFields(), seen, depth); | ||
| return group; | ||
| } | ||
| private <T> GroupBuilder<GroupBuilder<T>> addMapField(FieldDescriptor descriptor, final GroupBuilder<T> builder) { | ||
| @Nullable | ||
| private String getInnerTypeName(FieldDescriptor descriptor) { | ||
| if (descriptor.isMapField() && parquetSpecsCompliant) { | ||
| descriptor = descriptor.getMessageType().getFields().get(1); | ||
| } | ||
| if (descriptor.getJavaType() != JavaType.MESSAGE) { | ||
| LOG.trace("getInnerTypeName: {} => primitive", descriptor.getFullName()); | ||
| return null; | ||
| } | ||
| String name = descriptor.getMessageType().getFullName(); | ||
| LOG.trace("getInnerTypeName: {} => {}", descriptor.getFullName(), name); | ||
| return name; | ||
| } | ||
| private <T> GroupBuilder<GroupBuilder<T>> addMapField(FieldDescriptor descriptor, final GroupBuilder<T> builder, ImmutableSetMultimap<String, Integer> seen, int depth) { | ||
| List<FieldDescriptor> fields = descriptor.getMessageType().getFields(); | ||
| if (fields.size() != 2) { | ||
| throw new UnsupportedOperationException("Expected two fields for the map (key/value), but got: " + fields); | ||
| @@ -161,16 +252,16 @@ private <T> GroupBuilder<GroupBuilder<T>> addMapField(FieldDescriptor descriptor | ||
| ParquetType mapKeyParquetType = getParquetType(fields.get(0)); | ||
| GroupBuilder<GroupBuilder<GroupBuilder<T>>> group = builder | ||
| .group(Type.Repetition.OPTIONAL).as(mapType()) // only optional maps are allowed in Proto3 | ||
| .group(Type.Repetition.REPEATED) // key_value wrapper | ||
| .primitive(mapKeyParquetType.primitiveType, Type.Repetition.REQUIRED).as(mapKeyParquetType.logicalTypeAnnotation).named("key"); | ||
| .group(Type.Repetition.OPTIONAL).as(mapType()) // only optional maps are allowed in Proto3 | ||
| .group(Type.Repetition.REPEATED) // key_value wrapper | ||
| .primitive(mapKeyParquetType.primitiveType, Type.Repetition.REQUIRED).as(mapKeyParquetType.logicalTypeAnnotation).named("key"); | ||
| return addField(fields.get(1), group).named("value") | ||
| .named("key_value"); | ||
| return addField(fields.get(1), group, seen, depth) | ||
| .named("value") | ||
| .named("key_value"); | ||
| } | ||
| private ParquetType getParquetType(FieldDescriptor fieldDescriptor) { | ||
| private static ParquetType getParquetType(FieldDescriptor fieldDescriptor) { | ||
| JavaType javaType = fieldDescriptor.getJavaType(); | ||
| switch (javaType) { | ||
| case INT: return ParquetType.of(INT32); | ||
| @@ -203,5 +294,4 @@ public static ParquetType of(PrimitiveTypeName primitiveType) { | ||
| return of(primitiveType, null); | ||
| } | ||
| } | ||
| } | ||
Oops, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
would it make sense to consolidate seen and depth into a single data-structure that can be passed through and abstract some of the direct access to the multimap?
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
?
the
seenmap does encode the seen fields along with their depth as a single datastructure.depthbeing a separate arg is important b/c it's the current depth in the traversal, and is used to update the seen data structure.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
right, I was thinking of encapsulating this logic into its own class, so they can be recorded and updated together, to 1. Reduce additional parameters that have to be passed through.
2. Encapsulate the logic behind more mnemonic method names (e.g. AddRecursiveStep())
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
i'm not sure encapsulation helps with readability or protection in this case. they are really tracking different things, and should be understood by readers of the traversal code to know how each piece of state is used.