rdblue commented on code in PR #3221: URL: https://github.com/apache/parquet-java/pull/3221#discussion_r2155717992
########## parquet-variant/src/main/java/org/apache/parquet/variant/VariantValueWriter.java: ########## @@ -0,0 +1,372 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.parquet.variant; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import org.apache.parquet.Preconditions; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.io.api.RecordConsumer; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; + +/** + * Class to write Variant values to a shredded schema. + */ +public class VariantValueWriter { + private ByteBuffer metadataBuffer; + // We defer initializing the ImmutableMetata until it's needed. It has some construction cost, and if all + // object fields are shredded into typed_value, it will never be used. + private ImmutableMetadata metadata = null; + private RecordConsumer recordConsumer; + + private static final String LIST_REPEATED_NAME = "list"; + private static final String LIST_ELEMENT_NAME = "element"; + + VariantValueWriter(RecordConsumer recordConsumer, ByteBuffer metadata) { + this.recordConsumer = recordConsumer; + this.metadataBuffer = metadata; + } + + Metadata getMetadata() { + if (metadata == null) { + metadata = new ImmutableMetadata(metadataBuffer); + } + return metadata; + } + + /** + * Write a Variant value to a shredded schema. + */ + public static void write(RecordConsumer recordConsumer, GroupType schema, Variant value) { + recordConsumer.startGroup(); + int metadataIndex = schema.getFieldIndex("metadata"); + recordConsumer.startField("metadata", metadataIndex); + recordConsumer.addBinary(Binary.fromConstantByteBuffer(value.getMetadataBuffer())); + recordConsumer.endField("metadata", metadataIndex); + VariantValueWriter writer = new VariantValueWriter(recordConsumer, value.getMetadataBuffer()); + writer.write(schema, value); + recordConsumer.endGroup(); + } + + /** + * Write a Variant value to a shredded schema. The caller is responsible for calling startGroup() + * and endGroup(), and writing metadata if this is the top level of the Variant group. + */ + void write(GroupType schema, Variant value) { + Type typedValueField = null; + if (schema.containsField("typed_value")) { + typedValueField = schema.getType("typed_value"); + } + + Variant.Type variantType = value.getType(); + + // Handle typed_value if present + if (isTypeCompatible(variantType, typedValueField, value)) { + int typedValueIdx = schema.getFieldIndex("typed_value"); + recordConsumer.startField("typed_value", typedValueIdx); + ByteBuffer residual = null; + if (typedValueField.isPrimitive()) { + writeScalarValue(recordConsumer, value, typedValueField.asPrimitiveType()); + } else if (typedValueField.getLogicalTypeAnnotation() + instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation) { + writeArrayValue(recordConsumer, value, typedValueField.asGroupType()); + } else { + residual = writeObjectValue(recordConsumer, value, typedValueField.asGroupType()); + } + recordConsumer.endField("typed_value", typedValueIdx); + + if (residual != null) { + int valueIdx = schema.getFieldIndex("value"); + recordConsumer.startField("value", valueIdx); + recordConsumer.addBinary(Binary.fromConstantByteBuffer(residual)); + recordConsumer.endField("value", valueIdx); + } + } else { + int valueIdx = schema.getFieldIndex("value"); + recordConsumer.startField("value", valueIdx); + recordConsumer.addBinary(Binary.fromReusedByteBuffer(value.getValueBuffer())); + recordConsumer.endField("value", valueIdx); + } + } + + // Return true if the logical type is a decimal with the same scale as the provided value, with enough + // precision to hold the value. The provided value must be a decimal. + private boolean compatibleDecimalType(Variant value, LogicalTypeAnnotation logicalType) { Review Comment: This can be static, right? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
