Skip to content

Commit 22f90f4

Browse files
cashmandcashmand
committed
Code review feedback
1 parent 783acdf commit 22f90f4

5 files changed

Lines changed: 121 additions & 36 deletions

File tree

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

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -213,15 +213,17 @@ private void writeVariantFields(GroupType schema, Schema avroSchema, Object reco
213213
if (valueObj instanceof byte[]) {
214214
value = ByteBuffer.wrap((byte[]) valueObj);
215215
} else {
216+
if (!(valueObj instanceof ByteBuffer)) {
217+
throw new RuntimeException("Variant value must be a ByteBuffer: " + schema.getName());
218+
}
216219
value = (ByteBuffer) valueObj;
217220
}
218221
} else if (fieldType.getName() == "metadata") {
219222
Object metadataObj = model.getField(record, avroField.name(), index);
220-
if (metadataObj instanceof byte[]) {
221-
metadata = ByteBuffer.wrap((byte[]) metadataObj);
222-
} else {
223-
metadata = (ByteBuffer) metadataObj;
223+
if (!(metadataObj instanceof ByteBuffer)) {
224+
throw new RuntimeException("Variant metadata must be a ByteBuffer: " + schema.getName());
224225
}
226+
metadata = (ByteBuffer) metadataObj;
225227
} else {
226228
binarySchema = false;
227229
break;
@@ -231,10 +233,7 @@ private void writeVariantFields(GroupType schema, Schema avroSchema, Object reco
231233
if (binarySchema) {
232234
VariantValueWriter.write(recordConsumer, schema, new Variant(value, metadata));
233235
} else {
234-
// If the schema was something other than value and metaadata, treat the value as a non-variant record.
235-
recordConsumer.startGroup();
236-
writeRecordFields(schema, avroSchema, record);
237-
recordConsumer.endGroup();
236+
throw new RuntimeException("Invalid Avro schema for Variant logical type: " + schema.getName());
238237
}
239238
}
240239

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ public static Configuration conf(String name, boolean value) {
133133

134134
/**
135135
* Assert that to Variant values are logically equivalent.
136-
* E.g. object fields may be ordered differently in the binary.
136+
* E.g. values in an object may be ordered differently in the binary.
137137
*/
138138
static void assertEquivalent(Variant expected, Variant actual) {
139139
Assert.assertEquals(expected.getType(), actual.getType());

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

Lines changed: 91 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
*/
1919
package org.apache.parquet.avro;
2020

21-
import static org.apache.parquet.avro.AvroTestUtil.*;
2221
import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.*;
2322
import static org.junit.Assert.*;
2423

@@ -48,7 +47,6 @@
4847

4948
public class TestWriteVariant extends DirectWriterTest {
5049

51-
// Construct a variant, and return the value binary, dropping metadata.
5250
private static Variant fullVariant(Consumer<VariantBuilder> appendValue) {
5351
VariantBuilder builder = new VariantBuilder();
5452
appendValue.accept(builder);
@@ -103,8 +101,14 @@ private static MessageType parquetSchema(GroupType variantGroup) {
103101
private ByteBuffer TEST_METADATA;
104102
private ByteBuffer TEST_OBJECT;
105103
private ByteBuffer SIMILAR_OBJECT;
106-
private ByteBuffer TEST_ARRAY;
107-
private ByteBuffer SIMILAR_ARRAY;
104+
private ByteBuffer EMPTY_ARRAY;
105+
private ByteBuffer STRING_ARRAY;
106+
private ByteBuffer MIXED_ARRAY;
107+
private ByteBuffer NESTED_ARRAY;
108+
private ByteBuffer MIXED_NESTED_ARRAY;
109+
private ByteBuffer OBJECT_IN_ARRAY;
110+
private ByteBuffer MIXED_OBJECT_IN_ARRAY;
111+
private ByteBuffer SIMILAR_OBJECT_IN_ARRAY;
108112
private ByteBuffer EMPTY_OBJECT;
109113
private ByteBuffer EMPTY_METADATA = fullVariant(b -> b.appendNull()).getMetadataBuffer();
110114
private Variant[] VARIANTS;
@@ -144,8 +148,79 @@ public TestWriteVariant() throws Exception {
144148
b.endObject();
145149
});
146150

151+
EMPTY_ARRAY = variant(b -> {
152+
b.startArray();
153+
b.endArray();
154+
});
155+
156+
STRING_ARRAY = variant(b -> {
157+
VariantArrayBuilder ab = b.startArray();
158+
ab.appendString("parquet");
159+
b.endArray();
160+
});
161+
162+
MIXED_ARRAY = variant(b -> {
163+
VariantArrayBuilder ab = b.startArray();
164+
ab.appendString("parquet");
165+
ab.appendString("string");
166+
ab.appendInt(34);
167+
b.endArray();
168+
});
169+
170+
NESTED_ARRAY = variant(b -> {
171+
VariantArrayBuilder ab = b.startArray();
172+
VariantArrayBuilder inner1 = ab.startArray();
173+
inner1.appendString("parquet");
174+
inner1.appendString("string");
175+
ab.endArray();
176+
VariantArrayBuilder inner2 = ab.startArray();
177+
inner2.appendString("parquet");
178+
inner2.appendString("string");
179+
ab.endArray();
180+
b.endArray();
181+
});
182+
183+
MIXED_NESTED_ARRAY = variant(b -> {
184+
VariantArrayBuilder ab = b.startArray();
185+
VariantArrayBuilder inner1 = ab.startArray();
186+
inner1.appendString("parquet");
187+
inner1.appendString("string");
188+
inner1.appendInt(34);
189+
ab.endArray();
190+
VariantArrayBuilder inner2 = ab.startArray();
191+
inner2.appendInt(34);
192+
inner2.appendNull();
193+
ab.endArray();
194+
ab.startArray();
195+
ab.endArray();
196+
VariantArrayBuilder inner4 = ab.startArray();
197+
inner4.appendString("parquet");
198+
inner4.appendString("string");
199+
inner4.appendInt(34);
200+
ab.endArray();
201+
b.endArray();
202+
});
203+
147204
// The first array element defines the schema.
148-
TEST_ARRAY = variant(TEST_METADATA, b -> {
205+
OBJECT_IN_ARRAY = variant(TEST_METADATA, b -> {
206+
VariantArrayBuilder ab = b.startArray();
207+
VariantObjectBuilder ob = ab.startObject();
208+
ob.appendKey("a");
209+
ob.appendNull();
210+
ob.appendKey("d");
211+
ob.appendString("iceberg");
212+
ab.endObject();
213+
ab.appendInt(123);
214+
VariantObjectBuilder ob2 = ab.startObject();
215+
ob2.appendKey("c");
216+
ob2.appendString("hello");
217+
ob2.appendKey("d");
218+
ob2.appendDate(12345);
219+
ab.endObject();
220+
b.endArray();
221+
});
222+
223+
MIXED_OBJECT_IN_ARRAY = variant(TEST_METADATA, b -> {
149224
VariantArrayBuilder ab = b.startArray();
150225
VariantObjectBuilder ob = ab.startObject();
151226
ob.appendKey("a");
@@ -160,11 +235,13 @@ public TestWriteVariant() throws Exception {
160235
ob2.appendKey("d");
161236
ob2.appendDate(12345);
162237
ab.endObject();
238+
ab.appendString("parquet");
239+
ab.appendInt(34);
163240
b.endArray();
164241
});
165242

166243
// Change one field name and one type in the first element to change the schema.
167-
SIMILAR_ARRAY = variant(TEST_METADATA, b -> {
244+
SIMILAR_OBJECT_IN_ARRAY = variant(TEST_METADATA, b -> {
168245
VariantArrayBuilder ab = b.startArray();
169246
VariantObjectBuilder ob = ab.startObject();
170247
ob.appendKey("c");
@@ -206,8 +283,14 @@ public TestWriteVariant() throws Exception {
206283
new Variant(EMPTY_OBJECT, EMPTY_METADATA),
207284
new Variant(TEST_OBJECT, TEST_METADATA),
208285
new Variant(SIMILAR_OBJECT, TEST_METADATA),
209-
new Variant(TEST_ARRAY, TEST_METADATA),
210-
new Variant(SIMILAR_ARRAY, TEST_METADATA),
286+
new Variant(EMPTY_ARRAY, EMPTY_METADATA),
287+
new Variant(STRING_ARRAY, EMPTY_METADATA),
288+
new Variant(MIXED_ARRAY, EMPTY_METADATA),
289+
new Variant(NESTED_ARRAY, EMPTY_METADATA),
290+
new Variant(MIXED_NESTED_ARRAY, EMPTY_METADATA),
291+
new Variant(OBJECT_IN_ARRAY, TEST_METADATA),
292+
new Variant(MIXED_OBJECT_IN_ARRAY, TEST_METADATA),
293+
new Variant(SIMILAR_OBJECT_IN_ARRAY, TEST_METADATA),
211294
fullVariant(b -> b.appendDate(12345)),
212295
fullVariant(b -> b.appendDate(-12345)),
213296
fullVariant(b -> b.appendTimestampTz(1234567890L)),
@@ -233,7 +316,6 @@ public TestWriteVariant() throws Exception {
233316

234317
/**
235318
* Create a record containing a Variant value using the standard schema.
236-
* @return
237319
*/
238320
GenericRecord createRecord(int i, Variant v) {
239321
GenericRecord vRecord = new GenericData.Record(VARIANT_SCHEMA);

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -850,18 +850,17 @@ static HashMap<String, Integer> getMetadataMap(ByteBuffer metadata) {
850850
}
851851

852852
/**
853-
* Computes the actual size (in bytes) of the Variant value at `value[pos...]`
854-
* @param value The Variant value
855-
* @return The size (in bytes) of the Variant value
853+
* Computes the actual size (in bytes) of the Variant value.
854+
* @param value The Variant value binary
855+
* @return The size (in bytes) of the Variant value, including the header byte
856856
*/
857857
public static int valueSize(ByteBuffer value) {
858858
int pos = value.position();
859-
checkIndex(pos, value.limit());
860859
int basicType = value.get(pos) & BASIC_TYPE_MASK;
861-
int typeInfo = (value.get(pos) >> BASIC_TYPE_BITS) & PRIMITIVE_TYPE_MASK;
862860
switch (basicType) {
863861
case SHORT_STR:
864-
return 1 + typeInfo;
862+
int stringSize = (value.get(pos) >> BASIC_TYPE_BITS) & PRIMITIVE_TYPE_MASK;
863+
return 1 + stringSize;
865864
case OBJECT: {
866865
VariantUtil.ObjectInfo info = VariantUtil.getObjectInfo(slice(value, pos));
867866
return info.dataStartOffset
@@ -878,7 +877,8 @@ public static int valueSize(ByteBuffer value) {
878877
pos + info.offsetStartOffset + info.numElements * info.offsetSize,
879878
info.offsetSize);
880879
}
881-
default:
880+
default: {
881+
int typeInfo = (value.get(pos) >> BASIC_TYPE_BITS) & PRIMITIVE_TYPE_MASK;
882882
switch (typeInfo) {
883883
case NULL:
884884
case TRUE:
@@ -915,6 +915,7 @@ public static int valueSize(ByteBuffer value) {
915915
throw new UnsupportedOperationException(
916916
String.format("Unknown type in Variant. primitive type: %d", typeInfo));
917917
}
918+
}
918919
}
919920
}
920921
}

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

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -259,20 +259,23 @@ private void writeArrayValue(RecordConsumer recordConsumer, Variant variant, Gro
259259
Type elementType = listType.getType(0);
260260

261261
recordConsumer.startGroup();
262-
recordConsumer.startField(arrayType.getFieldName(0), 0);
263-
// Write each array element
264-
for (int i = 0; i < variant.numArrayElements(); i++) {
265-
recordConsumer.startGroup();
266-
recordConsumer.startField("element", 0);
262+
int numElements = variant.numArrayElements();
263+
if (numElements > 0) {
264+
recordConsumer.startField(arrayType.getFieldName(0), 0);
265+
// Write each array element
266+
for (int i = 0; i < numElements; i++) {
267+
recordConsumer.startGroup();
268+
recordConsumer.startField("element", 0);
267269

268-
recordConsumer.startGroup();
269-
write(elementType.asGroupType(), variant.getElementAtIndex(i));
270-
recordConsumer.endGroup();
270+
recordConsumer.startGroup();
271+
write(elementType.asGroupType(), variant.getElementAtIndex(i));
272+
recordConsumer.endGroup();
271273

272-
recordConsumer.endField("element", 0);
273-
recordConsumer.endGroup();
274+
recordConsumer.endField("element", 0);
275+
recordConsumer.endGroup();
276+
}
277+
recordConsumer.endField(arrayType.getFieldName(0), 0);
274278
}
275-
recordConsumer.endField(arrayType.getFieldName(0), 0);
276279
recordConsumer.endGroup();
277280
}
278281

0 commit comments

Comments
 (0)