rdblue commented on code in PR #3221: URL: https://github.com/apache/parquet-java/pull/3221#discussion_r2130687824
########## 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 -> { + VariantArrayBuilder ab = b.startArray(); + VariantObjectBuilder ob = ab.startObject(); + ob.appendKey("a"); + ob.appendNull(); + ob.appendKey("d"); + ob.appendString("iceberg"); + ab.endObject(); + ab.appendInt(123); + VariantObjectBuilder ob2 = ab.startObject(); + ob2.appendKey("c"); + ob2.appendString("hello"); + ob2.appendKey("d"); + ob2.appendDate(12345); + ab.endObject(); + b.endArray(); + }); + + // Change one field name and one type in the first element to change the schema. + SIMILAR_ARRAY = variant(TEST_METADATA, b -> { + VariantArrayBuilder ab = b.startArray(); + VariantObjectBuilder ob = ab.startObject(); + ob.appendKey("c"); + ob.appendString("iceberg"); + ob.appendKey("a"); + ob.appendString("parquet"); + ab.endObject(); + ab.appendInt(123); + VariantObjectBuilder ob2 = ab.startObject(); + ob2.appendKey("c"); + ob2.appendString("hello"); + ob2.appendKey("d"); + ob2.appendDate(12345); + ab.endObject(); + b.endArray(); + }); + + EMPTY_OBJECT = variant(TEST_METADATA, b -> { + b.startObject(); + b.endObject(); + }); + + VARIANTS = new Variant[] { + fullVariant(b -> b.appendNull()), + fullVariant(b -> b.appendBoolean(true)), + fullVariant(b -> b.appendBoolean(false)), + fullVariant(b -> b.appendByte((byte) 34)), + fullVariant(b -> b.appendByte((byte) -34)), + fullVariant(b -> b.appendShort((byte) 1234)), + fullVariant(b -> b.appendShort((byte) -1234)), + fullVariant(b -> b.appendInt(12345)), + fullVariant(b -> b.appendInt(-12345)), + fullVariant(b -> b.appendLong(9876543210L)), + fullVariant(b -> b.appendLong(-9876543210L)), + fullVariant(b -> b.appendFloat(10.11F)), + fullVariant(b -> b.appendFloat(-10.11F)), + fullVariant(b -> b.appendDouble(14.3D)), + fullVariant(b -> b.appendDouble(-14.3D)), + new Variant(EMPTY_OBJECT, EMPTY_METADATA), + new Variant(TEST_OBJECT, TEST_METADATA), + new Variant(SIMILAR_OBJECT, TEST_METADATA), + new Variant(TEST_ARRAY, TEST_METADATA), + new Variant(SIMILAR_ARRAY, TEST_METADATA), + fullVariant(b -> b.appendDate(12345)), + fullVariant(b -> b.appendDate(-12345)), + fullVariant(b -> b.appendTimestampTz(1234567890L)), + fullVariant(b -> b.appendTimestampTz(-1234567890L)), + fullVariant(b -> b.appendTimestampNtz(1234567890L)), + fullVariant(b -> b.appendTimestampNtz(-1234567890L)), + fullVariant(b -> b.appendDecimal(new BigDecimal("123456.789"))), // decimal4 + fullVariant(b -> b.appendDecimal(new BigDecimal("-123456.789"))), // decimal4 + fullVariant(b -> b.appendDecimal(new BigDecimal("123456789.987654321"))), // decimal8 + fullVariant(b -> b.appendDecimal(new BigDecimal("-123456789.987654321"))), // decimal8 + fullVariant(b -> b.appendDecimal(new BigDecimal("9876543210.123456789"))), // decimal16 + fullVariant(b -> b.appendDecimal(new BigDecimal("-9876543210.123456789"))), // decimal16 + fullVariant(b -> b.appendBinary(ByteBuffer.wrap(new byte[] {0x0a, 0x0b, 0x0c, 0x0d}))), + fullVariant(b -> b.appendString("iceberg")), + fullVariant(b -> b.appendTime(1234567890)), + fullVariant(b -> b.appendTimestampNanosTz(1234567890L)), + fullVariant(b -> b.appendTimestampNanosTz(-1234567890L)), + fullVariant(b -> b.appendTimestampNanosNtz(1234567890L)), + fullVariant(b -> b.appendTimestampNanosNtz(-1234567890L)), + fullVariant(b -> b.appendUUID(UUID.fromString("f24f9b64-81fa-49d1-b74e-8c09a6e31c56"))) + }; + } + + /** + * Create a record containing a Variant value using the standard schema. + * @return + */ + GenericRecord createRecord(int i, Variant v) { + GenericRecord vRecord = new GenericData.Record(VARIANT_SCHEMA); + vRecord.put(0, v.getMetadataBuffer()); + vRecord.put(1, v.getValueBuffer()); + GenericRecord record = new GenericData.Record(SCHEMA); + record.put(0, i); + record.put(1, vRecord); + return record; + } + + // Tests in this file are based on Iceberg's TestVariantWriters suite. + @Test + public void testUnshreddedValues() throws IOException { + for (Variant v : VARIANTS) { + GenericRecord record = createRecord(1, v); + TestSchema testSchema = new TestSchema(READ_SCHEMA, READ_SCHEMA); + + GenericRecord actual = writeAndRead(testSchema, record); + + assertEquals(record.get(0), actual.get(0)); + assertEquals(((GenericRecord) record.get(1)).get(0), ((GenericRecord) actual.get(1)).get(0)); + assertEquals(((GenericRecord) record.get(1)).get(1), ((GenericRecord) actual.get(1)).get(1)); + } + } + + @Test + public void testShreddedValues() throws IOException { + for (Variant v : VARIANTS) { + GenericRecord record = createRecord(1, v); + MessageType writeSchema = shreddingSchema(v); + TestSchema testSchema = new TestSchema(writeSchema, READ_SCHEMA); + + GenericRecord actual = writeAndRead(testSchema, record); + assertEquals(record.get(0), actual.get(0)); + Variant actualV = new Variant((ByteBuffer) ((GenericRecord) actual.get(1)).get(1), (ByteBuffer) + ((GenericRecord) actual.get(1)).get(0)); + AvroTestUtil.assertEquivalent(v, actualV); + } + } + + @Test + public void testMixedShredding() throws IOException { + for (Variant valueForSchema : VARIANTS) { + List<GenericRecord> expected = new ArrayList<>(); Review Comment: This can be moved outside the `valueForSchema` loop because it doesn't change. -- 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]
