rdblue commented on code in PR #3221: URL: https://github.com/apache/parquet-java/pull/3221#discussion_r2130673927
########## parquet-avro/src/test/java/org/apache/parquet/avro/TestWriteVariant.java: ########## @@ -0,0 +1,518 @@ +/** + * 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 + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * 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.avro; + +import static org.apache.parquet.avro.AvroTestUtil.*; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.*; +import static org.junit.Assert.*; + +import java.io.File; +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.*; +import java.util.function.Consumer; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.DirectWriterTest; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.api.WriteSupport; +import org.apache.parquet.schema.*; +import org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.variant.ImmutableMetadata; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantArrayBuilder; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantObjectBuilder; +import org.junit.Test; + +public class TestWriteVariant extends DirectWriterTest { + + // Construct a variant, and return the value binary, dropping metadata. + private static Variant fullVariant(Consumer<VariantBuilder> appendValue) { + VariantBuilder builder = new VariantBuilder(); + appendValue.accept(builder); + return builder.build(); + } + + // Return only the byte[], which is usually all we want. + private static ByteBuffer variant(Consumer<VariantBuilder> appendValue) { + return fullVariant(appendValue).getValueBuffer(); + } + + // Returns a value based on building with fixed metadata. + private static ByteBuffer variant(ByteBuffer metadata, Consumer<VariantBuilder> appendValue) { + VariantBuilder builder = new VariantBuilder(new ImmutableMetadata(metadata)); + appendValue.accept(builder); + return builder.build().getValueBuffer(); + } + + private static ByteBuffer variant(int val) { + return variant(b -> b.appendInt(val)); + } + + private static ByteBuffer variant(long val) { + return variant(b -> b.appendLong(val)); + } + + private static ByteBuffer variant(String s) { + return variant(b -> b.appendString(s)); + } + + private static final GroupType UNSHREDDED_GROUP = Types.buildGroup(Type.Repetition.REQUIRED) + .as(LogicalTypeAnnotation.variantType((byte) 1)) + .required(PrimitiveTypeName.BINARY) + .named("metadata") + .required(PrimitiveTypeName.BINARY) + .named("value") + .named("var"); + + private static MessageType parquetSchema(GroupType variantGroup) { + return Types.buildMessage() + .required(INT32) + .named("id") + .addField(variantGroup) + .named("table"); + } + + private static final MessageType READ_SCHEMA = parquetSchema(UNSHREDDED_GROUP); + + private static final Schema VARIANT_SCHEMA = new AvroSchemaConverter().convert(UNSHREDDED_GROUP); + private static final Schema SCHEMA = new AvroSchemaConverter().convert(READ_SCHEMA); + + private ByteBuffer TEST_METADATA; + private ByteBuffer TEST_OBJECT; + private ByteBuffer SIMILAR_OBJECT; + private ByteBuffer TEST_ARRAY; + private ByteBuffer SIMILAR_ARRAY; + private ByteBuffer EMPTY_OBJECT; + private ByteBuffer EMPTY_METADATA = fullVariant(b -> b.appendNull()).getMetadataBuffer(); + private Variant[] VARIANTS; + + public TestWriteVariant() throws Exception { + TEST_METADATA = fullVariant(b -> { + VariantObjectBuilder ob = b.startObject(); + ob.appendKey("a"); + ob.appendNull(); + ob.appendKey("b"); + ob.appendNull(); + ob.appendKey("c"); + ob.appendNull(); + ob.appendKey("d"); + ob.appendNull(); + ob.appendKey("e"); + ob.appendNull(); + b.endObject(); + }) + .getMetadataBuffer(); + + TEST_OBJECT = variant(TEST_METADATA, b -> { + VariantObjectBuilder ob = b.startObject(); + ob.appendKey("a"); + ob.appendNull(); + ob.appendKey("d"); + ob.appendString("iceberg"); + b.endObject(); + }); + + SIMILAR_OBJECT = variant(TEST_METADATA, b -> { + VariantObjectBuilder ob = b.startObject(); + ob.appendKey("a"); + ob.appendInt(123456789); + ob.appendKey("c"); + ob.appendString("string"); + b.endObject(); + }); + + // The first array element defines the schema. + TEST_ARRAY = variant(TEST_METADATA, b -> { Review Comment: Looks like this deviates from the Iceberg array tests, which use an array of strings, an array of strings and an int, an array of arrays of strings, and then some nested objects. Why does this skip the simpler array cases? -- 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]
