cashmand commented on code in PR #3221: URL: https://github.com/apache/parquet-java/pull/3221#discussion_r2100527909
########## parquet-avro/src/test/java/org/apache/parquet/avro/TestWriteVariant.java: ########## @@ -0,0 +1,517 @@ +/** + * 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.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).getValueRawBytes(); + } + + // Returns a value based on building with fixed metadata. + private static ByteBuffer variant(ByteBuffer metadata, Consumer<VariantBuilder> appendValue) { + VariantBuilder builder = new VariantBuilder(); + builder.setFixedMetadata(metadata); + appendValue.accept(builder); + return ByteBuffer.wrap(builder.valueWithoutMetadata()); + } + + 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 GroupType variantGroup = 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 MessageType readSchema = parquetSchema(variantGroup); + + private static final Schema VARIANT_SCHEMA = new AvroSchemaConverter().convert(variantGroup); + private static final Schema SCHEMA = new AvroSchemaConverter().convert(readSchema); + + 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()).getMetadataRawBytes(); + 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(); + }) + .getMetadataRawBytes(); + + 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.getMetadataRawBytes()); + vRecord.put(1, v.getValueRawBytes()); + 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(readSchema, readSchema); + + 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, readSchema); + + 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)); + if (!((GenericRecord) record.get(1)).get(1).equals(((GenericRecord) actual.get(1)).get(1))) { + assertTrue(false); + } + } + } + + @Test + public void testMixedShredding() throws IOException { + for (Variant v : VARIANTS) { + List<GenericRecord> expected = new ArrayList<>(); + for (int i = 0; i < VARIANTS.length; i++) { + expected.add(createRecord(i, VARIANTS[i])); + } + + MessageType writeSchema = shreddingSchema(v); + TestSchema testSchema = new TestSchema(writeSchema, readSchema); + + List<GenericRecord> actual = writeAndRead(testSchema, expected); + // TODO: CHECK RESULTS + } + } + + // Write schema contains the full shredding schema. Read schema should just be a value/metadata pair. + private static class TestSchema { + MessageType writeSchema; + MessageType readSchema; + + TestSchema(MessageType writeSchema, MessageType readSchema) { + this.writeSchema = writeSchema; + this.readSchema = readSchema; + } + } + + /** + * This is a custom Parquet writer builder that injects a specific Parquet schema and then uses + * the Avro object model. This ensures that the Parquet file's schema is exactly what was passed. + */ + private static class TestWriterBuilder extends ParquetWriter.Builder<GenericRecord, TestWriterBuilder> { + private MessageType schema = null; + + protected TestWriterBuilder(Path path) { + super(path); + } + + TestWriterBuilder withFileType(MessageType schema) { + this.schema = schema; + return self(); + } + + @Override + protected TestWriterBuilder self() { + return this; + } + + @Override + protected WriteSupport<GenericRecord> getWriteSupport(Configuration conf) { + return new AvroWriteSupport<>(schema, new AvroSchemaConverter().convert(schema), GenericData.get()); + } + } + + GenericRecord writeAndRead(TestSchema testSchema, GenericRecord record) throws IOException { + List<GenericRecord> result = writeAndRead(testSchema, Arrays.asList(record)); + assert (result.size() == 1); + return result.get(0); + } + + private List<GenericRecord> writeAndRead(TestSchema testSchema, List<GenericRecord> records) throws IOException { + File tmp = File.createTempFile(getClass().getSimpleName(), ".tmp"); + tmp.deleteOnExit(); + tmp.delete(); + Path path = new Path(tmp.getPath()); + + try (ParquetWriter<GenericRecord> writer = + new TestWriterBuilder(path).withFileType(testSchema.writeSchema).build()) { + for (GenericRecord record : records) { + writer.write(record); + } + } + + Configuration conf = new Configuration(); + // We need to set an explicit read schema because Avro wrote the shredding schema as the Avro + // schema in the write, and it will use that by default. If we write using a proper shredding + // writer, the Avro schema should just contain a <metadata, value> record, and we won't need this. Review Comment: Yes, deleted. -- 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]
