Skip to content

Commit 835daed

Browse files
cashmandcashmand
committed
More code review feedback
1 parent 22f90f4 commit 835daed

3 files changed

Lines changed: 83 additions & 52 deletions

File tree

‎parquet-avro/src/main/java/org/apache/parquet/avro/AvroWriteSupport.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ private void writeVariantFields(GroupType schema, Schema avroSchema, Object reco
208208
break;
209209
}
210210
Type fieldType = fields.get(index);
211-
if (fieldType.getName() == "value") {
211+
if (fieldType.getName().equals("value")) {
212212
Object valueObj = model.getField(record, avroField.name(), index);
213213
if (valueObj instanceof byte[]) {
214214
value = ByteBuffer.wrap((byte[]) valueObj);
@@ -218,7 +218,7 @@ private void writeVariantFields(GroupType schema, Schema avroSchema, Object reco
218218
}
219219
value = (ByteBuffer) valueObj;
220220
}
221-
} else if (fieldType.getName() == "metadata") {
221+
} else if (fieldType.getName().equals("metadata")) {
222222
Object metadataObj = model.getField(record, avroField.name(), index);
223223
if (!(metadataObj instanceof ByteBuffer)) {
224224
throw new RuntimeException("Variant metadata must be a ByteBuffer: " + schema.getName());

‎parquet-avro/src/test/java/org/apache/parquet/avro/TestWriteVariant.java‎

Lines changed: 43 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,16 @@
1818
*/
1919
package org.apache.parquet.avro;
2020

21-
import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.*;
22-
import static org.junit.Assert.*;
21+
import static org.junit.Assert.assertEquals;
2322

2423
import java.io.File;
2524
import java.io.IOException;
2625
import java.math.BigDecimal;
2726
import java.nio.ByteBuffer;
28-
import java.util.*;
27+
import java.util.Arrays;
28+
import java.util.ArrayList;
29+
import java.util.List;
30+
import java.util.UUID;
2931
import java.util.function.Consumer;
3032
import org.apache.avro.Schema;
3133
import org.apache.avro.generic.GenericData;
@@ -35,7 +37,11 @@
3537
import org.apache.parquet.DirectWriterTest;
3638
import org.apache.parquet.hadoop.ParquetWriter;
3739
import org.apache.parquet.hadoop.api.WriteSupport;
38-
import org.apache.parquet.schema.*;
40+
import org.apache.parquet.schema.GroupType;
41+
import org.apache.parquet.schema.LogicalTypeAnnotation;
42+
import org.apache.parquet.schema.MessageType;
43+
import org.apache.parquet.schema.Type;
44+
import org.apache.parquet.schema.Types;
3945
import org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit;
4046
import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName;
4147
import org.apache.parquet.variant.ImmutableMetadata;
@@ -87,7 +93,7 @@ private static ByteBuffer variant(String s) {
8793

8894
private static MessageType parquetSchema(GroupType variantGroup) {
8995
return Types.buildMessage()
90-
.required(INT32)
96+
.required(PrimitiveTypeName.INT32)
9197
.named("id")
9298
.addField(variantGroup)
9399
.named("table");
@@ -359,18 +365,18 @@ public void testShreddedValues() throws IOException {
359365

360366
@Test
361367
public void testMixedShredding() throws IOException {
362-
for (Variant valueForSchema : VARIANTS) {
363-
List<GenericRecord> expected = new ArrayList<>();
364-
for (int i = 0; i < VARIANTS.length; i++) {
365-
expected.add(createRecord(i, VARIANTS[i]));
366-
}
368+
List<GenericRecord> expected = new ArrayList<>();
369+
for (int i = 0; i < VARIANTS.length; i++) {
370+
expected.add(createRecord(i, VARIANTS[i]));
371+
}
367372

373+
for (Variant valueForSchema : VARIANTS) {
368374
MessageType writeSchema = shreddingSchema(valueForSchema);
369375
TestSchema testSchema = new TestSchema(writeSchema, READ_SCHEMA);
370376

371377
List<GenericRecord> actual = writeAndRead(testSchema, expected);
372-
assertEquals(actual.size(), VARIANTS.length);
373-
for (int i = 0; i < VARIANTS.length; i++) {
378+
assertEquals(actual.size(), expected.size());
379+
for (int i = 0; i < expected.size(); i++) {
374380
Variant actualV =
375381
new Variant((ByteBuffer) ((GenericRecord) actual.get(i).get(1)).get(1), (ByteBuffer)
376382
((GenericRecord) actual.get(i).get(1)).get(0));
@@ -419,7 +425,7 @@ protected WriteSupport<GenericRecord> getWriteSupport(Configuration conf) {
419425

420426
GenericRecord writeAndRead(TestSchema testSchema, GenericRecord record) throws IOException {
421427
List<GenericRecord> result = writeAndRead(testSchema, Arrays.asList(record));
422-
assert (result.size() == 1);
428+
assertEquals(result.size(), 1);
423429
return result.get(0);
424430
}
425431

@@ -456,9 +462,9 @@ private static MessageType shreddingSchema(Variant v) {
456462
Type shreddedType = shreddedType(v);
457463
Types.GroupBuilder<GroupType> partialType = Types.buildGroup(Type.Repetition.OPTIONAL)
458464
.as(LogicalTypeAnnotation.variantType((byte) 1))
459-
.required(BINARY)
465+
.required(PrimitiveTypeName.BINARY)
460466
.named("metadata")
461-
.optional(BINARY)
467+
.optional(PrimitiveTypeName.BINARY)
462468
.named("value");
463469
Type variantType;
464470
if (shreddedType == null) {
@@ -467,7 +473,7 @@ private static MessageType shreddingSchema(Variant v) {
467473
variantType = partialType.addField(shreddedType).named("var");
468474
}
469475
return Types.buildMessage()
470-
.required(INT32)
476+
.required(PrimitiveTypeName.INT32)
471477
.named("id")
472478
.addField(variantType)
473479
.named("table");
@@ -477,12 +483,12 @@ private static GroupType shreddedGroup(Variant v, String name) {
477483
Type shreddedType = shreddedType(v);
478484
if (shreddedType == null) {
479485
return Types.buildGroup(Type.Repetition.OPTIONAL)
480-
.optional(BINARY)
486+
.optional(PrimitiveTypeName.BINARY)
481487
.named("value")
482488
.named(name);
483489
} else {
484490
return Types.buildGroup(Type.Repetition.OPTIONAL)
485-
.optional(BINARY)
491+
.optional(PrimitiveTypeName.BINARY)
486492
.named("value")
487493
.addField(shreddedType)
488494
.named(name);
@@ -497,67 +503,67 @@ private static Type shreddedType(Variant v) {
497503
case NULL:
498504
return null;
499505
case BOOLEAN:
500-
return Types.optional(BOOLEAN).named("typed_value");
506+
return Types.optional(PrimitiveTypeName.BOOLEAN).named("typed_value");
501507
case BYTE:
502-
return Types.optional(INT32)
508+
return Types.optional(PrimitiveTypeName.INT32)
503509
.as(LogicalTypeAnnotation.intType(8))
504510
.named("typed_value");
505511
case SHORT:
506-
return Types.optional(INT32)
512+
return Types.optional(PrimitiveTypeName.INT32)
507513
.as(LogicalTypeAnnotation.intType(16))
508514
.named("typed_value");
509515
case INT:
510-
return Types.optional(INT32).named("typed_value");
516+
return Types.optional(PrimitiveTypeName.INT32).named("typed_value");
511517
case LONG:
512-
return Types.optional(INT64).named("typed_value");
518+
return Types.optional(PrimitiveTypeName.INT64).named("typed_value");
513519
case FLOAT:
514-
return Types.optional(FLOAT).named("typed_value");
520+
return Types.optional(PrimitiveTypeName.FLOAT).named("typed_value");
515521
case DOUBLE:
516-
return Types.optional(DOUBLE).named("typed_value");
522+
return Types.optional(PrimitiveTypeName.DOUBLE).named("typed_value");
517523
case DECIMAL4:
518-
return Types.optional(INT32)
524+
return Types.optional(PrimitiveTypeName.INT32)
519525
.as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 9))
520526
.named("typed_value");
521527
case DECIMAL8:
522-
return Types.optional(INT64)
528+
return Types.optional(PrimitiveTypeName.INT64)
523529
.as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 18))
524530
.named("typed_value");
525531
case DECIMAL16:
526-
return Types.optional(BINARY)
532+
return Types.optional(PrimitiveTypeName.BINARY)
527533
.as(LogicalTypeAnnotation.decimalType(v.getDecimal().scale(), 38))
528534
.named("typed_value");
529535
case DATE:
530-
return Types.optional(INT32)
536+
return Types.optional(PrimitiveTypeName.INT32)
531537
.as(LogicalTypeAnnotation.dateType())
532538
.named("typed_value");
533539
case TIMESTAMP_TZ:
534-
return Types.optional(INT64)
540+
return Types.optional(PrimitiveTypeName.INT64)
535541
.as(LogicalTypeAnnotation.timestampType(true, TimeUnit.MICROS))
536542
.named("typed_value");
537543
case TIMESTAMP_NTZ:
538-
return Types.optional(INT64)
544+
return Types.optional(PrimitiveTypeName.INT64)
539545
.as(LogicalTypeAnnotation.timestampType(false, TimeUnit.MICROS))
540546
.named("typed_value");
541547
case BINARY:
542-
return Types.optional(BINARY).named("typed_value");
548+
return Types.optional(PrimitiveTypeName.BINARY).named("typed_value");
543549
case STRING:
544-
return Types.optional(BINARY)
550+
return Types.optional(PrimitiveTypeName.BINARY)
545551
.as(LogicalTypeAnnotation.stringType())
546552
.named("typed_value");
547553
case TIME:
548-
return Types.optional(INT64)
554+
return Types.optional(PrimitiveTypeName.INT64)
549555
.as(LogicalTypeAnnotation.timeType(false, TimeUnit.MICROS))
550556
.named("typed_value");
551557
case TIMESTAMP_NANOS_TZ:
552-
return Types.optional(INT64)
558+
return Types.optional(PrimitiveTypeName.INT64)
553559
.as(LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS))
554560
.named("typed_value");
555561
case TIMESTAMP_NANOS_NTZ:
556-
return Types.optional(INT64)
562+
return Types.optional(PrimitiveTypeName.INT64)
557563
.as(LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS))
558564
.named("typed_value");
559565
case UUID:
560-
return Types.optional(FIXED_LEN_BYTE_ARRAY)
566+
return Types.optional(PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY)
561567
.as(LogicalTypeAnnotation.uuidType())
562568
.named("typed_value");
563569
case OBJECT:

‎parquet-variant/src/main/java/org/apache/parquet/variant/VariantValueWriter.java‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.parquet.variant;
2020

2121
import java.nio.ByteBuffer;
22+
import org.apache.parquet.Preconditions;
2223
import org.apache.parquet.io.api.Binary;
2324
import org.apache.parquet.io.api.RecordConsumer;
2425
import org.apache.parquet.schema.GroupType;
@@ -36,6 +37,9 @@ public class VariantValueWriter {
3637
private ImmutableMetadata metadata = null;
3738
private RecordConsumer recordConsumer;
3839

40+
private static final String LIST_REPEATED_NAME = "list";
41+
private static final String LIST_ELEMENT_NAME = "element";
42+
3943
VariantValueWriter(RecordConsumer recordConsumer, ByteBuffer metadata) {
4044
this.recordConsumer = recordConsumer;
4145
this.metadataBuffer = metadata;
@@ -118,22 +122,24 @@ private boolean isTypeCompatible(Variant.Type variantType, Type typedValueField)
118122
case BYTE:
119123
return primitiveTypeName == PrimitiveType.PrimitiveTypeName.INT32
120124
&& logicalType instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation
125+
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).isSigned()
121126
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).getBitWidth() == 8;
122127
case SHORT:
123128
return primitiveTypeName == PrimitiveType.PrimitiveTypeName.INT32
124129
&& logicalType instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation
130+
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).isSigned()
125131
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).getBitWidth() == 16;
126132
case INT:
127133
return primitiveTypeName == PrimitiveType.PrimitiveTypeName.INT32
128134
&& (logicalType == null
129135
|| (logicalType instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation
130-
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType)
131-
.getBitWidth()
132-
== 32));
136+
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).isSigned()
137+
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).getBitWidth() == 32));
133138
case LONG:
134139
return primitiveTypeName == PrimitiveType.PrimitiveTypeName.INT64
135140
&& (logicalType == null
136-
|| logicalType instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation);
141+
|| (logicalType instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation
142+
&& ((LogicalTypeAnnotation.IntLogicalTypeAnnotation) logicalType).isSigned()));
137143
case FLOAT:
138144
return primitiveTypeName == PrimitiveType.PrimitiveTypeName.FLOAT;
139145
case DOUBLE:
@@ -250,31 +256,50 @@ private void writeScalarValue(RecordConsumer recordConsumer, Variant variant, Pr
250256
}
251257

252258
private void writeArrayValue(RecordConsumer recordConsumer, Variant variant, GroupType arrayType) {
253-
if (variant.getType() != Variant.Type.ARRAY) {
254-
throw new IllegalArgumentException("Expected array type but got: " + variant.getType());
259+
Preconditions.checkArgument(variant.getType() == Variant.Type.ARRAY,
260+
"Cannot write variant type " + variant.getType() + " as array");
261+
262+
// Validate that it's a 3-level array.
263+
if (arrayType.getFieldCount() != 1 ||
264+
arrayType.getRepetition() == Type.Repetition.REPEATED ||
265+
arrayType.getType(0).isPrimitive() ||
266+
!arrayType.getFieldName(0).equals(LIST_REPEATED_NAME)) {
267+
throw new IllegalArgumentException("Variant list must be a three-level list structure: " + arrayType);
255268
}
256269

257270
// Get the element type from the array schema
258-
GroupType listType = arrayType.getType(0).asGroupType();
259-
Type elementType = listType.getType(0);
271+
GroupType repeatedType = arrayType.getType(0).asGroupType();
272+
273+
if (repeatedType.getFieldCount() != 1 ||
274+
repeatedType.getRepetition() != Type.Repetition.REPEATED ||
275+
repeatedType.getType(0).isPrimitive() ||
276+
!repeatedType.getFieldName(0).equals(LIST_ELEMENT_NAME)) {
277+
throw new IllegalArgumentException("Variant list must be a three-level list structure: " + arrayType);
278+
}
279+
280+
GroupType elementType = repeatedType.getType(0).asGroupType();
260281

282+
// List field, annotated as LIST
261283
recordConsumer.startGroup();
262284
int numElements = variant.numArrayElements();
285+
// Can only call startField if there is at least one element.
263286
if (numElements > 0) {
264-
recordConsumer.startField(arrayType.getFieldName(0), 0);
287+
recordConsumer.startField(LIST_REPEATED_NAME, 0);
265288
// Write each array element
266289
for (int i = 0; i < numElements; i++) {
290+
// Repeated group.
267291
recordConsumer.startGroup();
268-
recordConsumer.startField("element", 0);
292+
recordConsumer.startField(LIST_ELEMENT_NAME, 0);
269293

294+
// Element group. Can never be null for shredded Variant.
270295
recordConsumer.startGroup();
271-
write(elementType.asGroupType(), variant.getElementAtIndex(i));
296+
write(elementType, variant.getElementAtIndex(i));
272297
recordConsumer.endGroup();
273298

274-
recordConsumer.endField("element", 0);
299+
recordConsumer.endField(LIST_ELEMENT_NAME, 0);
275300
recordConsumer.endGroup();
276301
}
277-
recordConsumer.endField(arrayType.getFieldName(0), 0);
302+
recordConsumer.endField(LIST_REPEATED_NAME, 0);
278303
}
279304
recordConsumer.endGroup();
280305
}

0 commit comments

Comments
 (0)