diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/List3Levels.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/List3Levels.java index 297b063320..f1286d5a6f 100644 --- a/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/List3Levels.java +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/List3Levels.java @@ -42,7 +42,8 @@ class List3Levels { * @param list the Parquet List */ public List3Levels(GroupType list) { - if (list.getLogicalTypeAnnotation() != LogicalTypeAnnotation.listType() + LogicalTypeAnnotation logicalType = list.getLogicalTypeAnnotation(); + if (!(logicalType instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation) || list.getFields().size() != 1) { throw new IllegalArgumentException("invalid list type: " + list); } diff --git a/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/SchemaConverter.java b/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/SchemaConverter.java index e779b6b8c2..561001b850 100644 --- a/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/SchemaConverter.java +++ b/parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/SchemaConverter.java @@ -486,6 +486,21 @@ public Optional visit( return of(new ListTypeMapping(arrowField, list3Levels, child)); } + @Override + public Optional visit( + LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation vectorLogicalType) { + List3Levels list3Levels = new List3Levels(type); + Type element = list3Levels.getElement(); + TypeMapping child = fromParquet(element, null, element.getRepetition()); + ArrowType.FixedSizeList fixedSizeList = + new ArrowType.FixedSizeList(vectorLogicalType.getNumElements()); + Field arrowField = new Field( + name, + new FieldType(type.isRepetition(OPTIONAL), fixedSizeList, null), + Collections.singletonList(child.getArrowField())); + return of(new ListTypeMapping(arrowField, list3Levels, child)); + } + @Override public Optional visit( LogicalTypeAnnotation.MapLogicalTypeAnnotation mapLogicalType) { diff --git a/parquet-arrow/src/test/java/org/apache/parquet/arrow/schema/TestSchemaConverter.java b/parquet-arrow/src/test/java/org/apache/parquet/arrow/schema/TestSchemaConverter.java index 4d8a21a8d0..f697b53a29 100644 --- a/parquet-arrow/src/test/java/org/apache/parquet/arrow/schema/TestSchemaConverter.java +++ b/parquet-arrow/src/test/java/org/apache/parquet/arrow/schema/TestSchemaConverter.java @@ -55,13 +55,34 @@ import org.apache.parquet.schema.GroupType; import org.apache.parquet.schema.LogicalTypeAnnotation; import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; import org.apache.parquet.schema.Types; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; /** * @see SchemaConverter */ public class TestSchemaConverter { + @ParameterizedTest + @ValueSource(strings = {"required", "optional"}) + public void testVectorAsFixedSizeList(String repetition) { + String children = "repeated group list { required float element; }"; + MessageType vector = MessageTypeParser.parseMessageType( + "message m { " + repetition + " group v (FINITE_VECTOR(5)) { " + children + " } }"); + SchemaMapping mapping = converter.fromParquet(vector); + Field field = mapping.getArrowSchema().findField("v"); + assertThat(field.getType()).isEqualTo(new ArrowType.FixedSizeList(5)); + assertThat(field.isNullable()).isEqualTo(repetition.equals("optional")); + assertThat(field.getChildren()).hasSize(1); + assertThat(field.getChildren().get(0).getType()) + .isEqualTo(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)); + assertThat(field.getChildren().get(0).isNullable()).isFalse(); + assertThat(mapping.getParquetSchema()).isEqualTo(vector); + assertThat(converter.map(mapping.getArrowSchema(), vector).getParquetSchema()) + .isEqualTo(vector); + } private static Field field(String name, boolean nullable, ArrowType type, Field... children) { if (nullable) { diff --git a/parquet-avro/src/main/java/org/apache/parquet/avro/AvroReadSupport.java b/parquet-avro/src/main/java/org/apache/parquet/avro/AvroReadSupport.java index ee8f9ab272..9823fe4c4c 100644 --- a/parquet-avro/src/main/java/org/apache/parquet/avro/AvroReadSupport.java +++ b/parquet-avro/src/main/java/org/apache/parquet/avro/AvroReadSupport.java @@ -281,7 +281,8 @@ private static Boolean allListStructuresAreThreeLevel(Type type) { return null; } GroupType group = type.asGroupType(); - if (group.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation) { + LogicalTypeAnnotation logicalType = group.getLogicalTypeAnnotation(); + if (logicalType instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation) { if (group.isRepetition(Type.Repetition.REPEATED) || group.getFieldCount() != 1) { return false; } diff --git a/parquet-avro/src/main/java/org/apache/parquet/avro/AvroSchemaConverter.java b/parquet-avro/src/main/java/org/apache/parquet/avro/AvroSchemaConverter.java index 782c009db4..f705ee0705 100644 --- a/parquet-avro/src/main/java/org/apache/parquet/avro/AvroSchemaConverter.java +++ b/parquet-avro/src/main/java/org/apache/parquet/avro/AvroSchemaConverter.java @@ -470,7 +470,7 @@ public Optional visit( if (!repeatedType.isRepetition(REPEATED)) { throw new UnsupportedOperationException("Invalid list type " + parquetGroupType); } - if (isElementType(repeatedType, parquetGroupType.getName())) { + if (isElementType(repeatedType, parquetGroupType)) { // repeated element types are always required return of(Schema.createArray(convertField(repeatedType, names))); } else { @@ -654,7 +654,10 @@ public Optional visit(UUIDLogicalTypeAnnotation uuidLogicalType) { * the same. This ensures that the LIST rules are followed when there is no * schema and that a schema can be provided to override the default behavior. */ - private boolean isElementType(Type repeatedType, String parentName) { + private boolean isElementType(Type repeatedType, GroupType parent) { + if (parent.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation) { + return false; + } return ( // can't be a synthetic layer because it would be invalid repeatedType.isPrimitive() @@ -663,7 +666,7 @@ private boolean isElementType(Type repeatedType, String parentName) { || // known patterns without the synthetic layer repeatedType.getName().equals("array") - || repeatedType.getName().equals(parentName + "_tuple") + || repeatedType.getName().equals(parent.getName() + "_tuple") || // default assumption assumeRepeatedIsListElement); diff --git a/parquet-avro/src/test/java/org/apache/parquet/avro/TestArrayCompatibility.java b/parquet-avro/src/test/java/org/apache/parquet/avro/TestArrayCompatibility.java index 242fd2fcf0..60bf7315fa 100644 --- a/parquet-avro/src/test/java/org/apache/parquet/avro/TestArrayCompatibility.java +++ b/parquet-avro/src/test/java/org/apache/parquet/avro/TestArrayCompatibility.java @@ -39,13 +39,20 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; import org.apache.parquet.DirectWriterTest; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.io.InvalidRecordException; +import org.apache.parquet.io.LocalInputFile; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.MessageTypeParser; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; public class TestArrayCompatibility extends DirectWriterTest { @@ -1610,6 +1617,53 @@ public void testAutoDetectTwoLevelRepeatedPrimitive() throws Exception { assertReaderContains(new AvroParquetReader<>(newAutoDetectConf(), test), expectedSchema, expectedRecord); } + @Test + public void testVectorWithoutAvroSchema(@TempDir java.nio.file.Path dir) throws IOException { + MessageType schema = MessageTypeParser.parseMessageType("message VectorWithoutAvroSchema {" + + " required group v (FINITE_VECTOR(2)) { repeated group list { required float element; } }" + + " optional group o (FINITE_VECTOR(1)) { repeated group list { required float element; } } }"); + Path path = new Path(dir.resolve("vector.parquet").toUri()); + SimpleGroupFactory factory = new SimpleGroupFactory(schema); + Group first = factory.newGroup(); + first.addGroup("v").addGroup("list").append("element", 1.0f); + first.getGroup("v", 0).addGroup("list").append("element", 2.0f); + first.addGroup("o").addGroup("list").append("element", 9.0f); + Group second = factory.newGroup(); + second.addGroup("v").addGroup("list").append("element", 3.0f); + second.getGroup("v", 0).addGroup("list").append("element", 4.0f); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withType(schema) + .withWriterVersion(org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_1_0) + .withDictionaryEncoding(false) + .withCompressionCodec(CompressionCodecName.UNCOMPRESSED) + .withValidation(false) + .build()) { + writer.write(first); + writer.write(second); + } + try (ParquetFileReader reader = ParquetFileReader.open(new LocalInputFile(dir.resolve("vector.parquet")))) { + assertThat(reader.getFileMetaData().getKeyValueMetaData()) + .doesNotContainKeys( + AvroReadSupport.AVRO_SCHEMA_METADATA_KEY, AvroReadSupport.OLD_AVRO_SCHEMA_METADATA_KEY); + } + Schema expectedSchema = record( + "VectorWithoutAvroSchema", + field("v", array(primitive(Schema.Type.FLOAT))), + optionalField("o", array(primitive(Schema.Type.FLOAT)))); + assertThat(new AvroSchemaConverter().convert(schema)).isEqualTo(expectedSchema); + for (boolean assumeRepeatedIsListElement : new boolean[] {true, false}) { + assertThat(new AvroSchemaConverter(assumeRepeatedIsListElement, false).convert(schema)) + .isEqualTo(expectedSchema); + } + try (AvroParquetReader reader = autoDetectReader(path)) { + assertReaderContains( + reader, + expectedSchema, + instance(expectedSchema, "v", Arrays.asList(1.0f, 2.0f), "o", Arrays.asList(9.0f)), + instance(expectedSchema, "v", Arrays.asList(3.0f, 4.0f), "o", null)); + } + } + private static Configuration newAutoDetectConf() { Configuration conf = new Configuration(); conf.setBoolean(AvroReadSupport.AUTO_DETECT_LIST_STRUCTURE, true); diff --git a/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java b/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java index 4040b55e31..43ff641cf6 100644 --- a/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java +++ b/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java @@ -62,6 +62,13 @@ protected LogicalTypeAnnotation fromString(List params) { return listType(); } }, + FINITE_VECTOR { + @Override + protected LogicalTypeAnnotation fromString(List params) { + Preconditions.checkArgument(params.size() == 1, "Expecting 1 parameter for FINITE_VECTOR: %s", params); + return finiteVectorType(Integer.parseInt(params.get(0))); + } + }, VARIANT { @Override protected LogicalTypeAnnotation fromString(List params) { @@ -314,6 +321,10 @@ public static ListLogicalTypeAnnotation listType() { return ListLogicalTypeAnnotation.INSTANCE; } + public static FiniteVectorLogicalTypeAnnotation finiteVectorType(int numElements) { + return new FiniteVectorLogicalTypeAnnotation(numElements); + } + public static VariantLogicalTypeAnnotation variantType(byte specVersion) { return new VariantLogicalTypeAnnotation(specVersion); } @@ -503,7 +514,7 @@ LogicalTypeToken getType() { @Override public boolean equals(Object obj) { - return obj instanceof ListLogicalTypeAnnotation; + return obj instanceof ListLogicalTypeAnnotation && !(obj instanceof FiniteVectorLogicalTypeAnnotation); } @Override @@ -513,6 +524,45 @@ public int hashCode() { } } + public static class FiniteVectorLogicalTypeAnnotation extends ListLogicalTypeAnnotation { + private final int numElements; + + private FiniteVectorLogicalTypeAnnotation(int numElements) { + Preconditions.checkArgument(numElements > 0, "Invalid FINITE_VECTOR length: %s", numElements); + this.numElements = numElements; + } + + public int getNumElements() { + return numElements; + } + + @Override + public Optional accept(LogicalTypeAnnotationVisitor visitor) { + return visitor.visit(this); + } + + @Override + LogicalTypeToken getType() { + return LogicalTypeToken.FINITE_VECTOR; + } + + @Override + protected String typeParametersAsString() { + return "(" + numElements + ")"; + } + + @Override + public boolean equals(Object obj) { + return obj instanceof FiniteVectorLogicalTypeAnnotation + && numElements == ((FiniteVectorLogicalTypeAnnotation) obj).numElements; + } + + @Override + public int hashCode() { + return 31 * getClass().hashCode() + numElements; + } + } + public static class EnumLogicalTypeAnnotation extends LogicalTypeAnnotation { private static final EnumLogicalTypeAnnotation INSTANCE = new EnumLogicalTypeAnnotation(); @@ -1536,6 +1586,10 @@ default Optional visit(ListLogicalTypeAnnotation listLogicalType) { return empty(); } + default Optional visit(FiniteVectorLogicalTypeAnnotation vectorLogicalType) { + return visit((ListLogicalTypeAnnotation) vectorLogicalType); + } + default Optional visit(VariantLogicalTypeAnnotation variantLogicalType) { return empty(); } diff --git a/parquet-column/src/main/java/org/apache/parquet/schema/TypeUtil.java b/parquet-column/src/main/java/org/apache/parquet/schema/TypeUtil.java index bb1278fc57..3b6d35eb9e 100644 --- a/parquet-column/src/main/java/org/apache/parquet/schema/TypeUtil.java +++ b/parquet-column/src/main/java/org/apache/parquet/schema/TypeUtil.java @@ -18,11 +18,77 @@ */ package org.apache.parquet.schema; +import org.apache.parquet.schema.LogicalTypeAnnotation.DecimalLogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.Float16LogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.IntLogicalTypeAnnotation; + public class TypeUtil { + public static boolean hasFiniteVector(Type type) { + if (type.getLogicalTypeAnnotation() instanceof FiniteVectorLogicalTypeAnnotation) return true; + if (!type.isPrimitive()) { + for (Type child : type.asGroupType().getFields()) { + if (hasFiniteVector(child)) return true; + } + } + return false; + } + + public static FiniteVectorLogicalTypeAnnotation getFiniteVectorType(GroupType schema, String[] path) { + if (path.length < 3) { + return null; + } + GroupType group = schema; + for (int i = 0; i < path.length - 2; i++) { + group = group.getType(path[i]).asGroupType(); + } + LogicalTypeAnnotation logical = group.getLogicalTypeAnnotation(); + return logical instanceof FiniteVectorLogicalTypeAnnotation + ? (FiniteVectorLogicalTypeAnnotation) logical + : null; + } + + static void checkFiniteVector(GroupType vector) { + if (vector.isRepetition(Type.Repetition.REPEATED) || vector.getFieldCount() != 1) { + throw new InvalidSchemaException("Invalid FINITE_VECTOR schema: " + vector); + } + Type list = vector.getType(0); + if (list.isPrimitive() + || !list.getName().equals("list") + || !list.isRepetition(Type.Repetition.REPEATED) + || list.getLogicalTypeAnnotation() != null + || list.asGroupType().getFieldCount() != 1) { + throw new InvalidSchemaException("Invalid FINITE_VECTOR schema: " + vector); + } + Type element = list.asGroupType().getType(0); + if (!element.isPrimitive() + || !element.getName().equals("element") + || !element.isRepetition(Type.Repetition.REQUIRED)) { + throw new InvalidSchemaException("Invalid FINITE_VECTOR schema: " + vector); + } + LogicalTypeAnnotation logical = element.getLogicalTypeAnnotation(); + boolean allowed = switch (element.asPrimitiveType().getPrimitiveTypeName()) { + case BOOLEAN, FLOAT, DOUBLE -> logical == null; + case INT32, INT64 -> + logical == null + || logical instanceof IntLogicalTypeAnnotation + || logical instanceof DecimalLogicalTypeAnnotation; + case FIXED_LEN_BYTE_ARRAY -> + logical instanceof Float16LogicalTypeAnnotation || logical instanceof DecimalLogicalTypeAnnotation; + default -> false; + }; + if (!allowed) { + throw new InvalidSchemaException("Invalid FINITE_VECTOR schema: " + vector); + } + } + public static void checkValidWriteSchema(GroupType schema) { schema.accept(new TypeVisitor() { @Override public void visit(GroupType groupType) { + if (groupType.getLogicalTypeAnnotation() instanceof FiniteVectorLogicalTypeAnnotation) { + checkFiniteVector(groupType); + } if (groupType.getFieldCount() <= 0) { throw new InvalidSchemaException("Cannot write a schema with an empty group: " + groupType); } diff --git a/parquet-column/src/main/java/org/apache/parquet/schema/Types.java b/parquet-column/src/main/java/org/apache/parquet/schema/Types.java index 2d6f0cbf8f..08cc51cc24 100644 --- a/parquet-column/src/main/java/org/apache/parquet/schema/Types.java +++ b/parquet-column/src/main/java/org/apache/parquet/schema/Types.java @@ -29,6 +29,7 @@ import java.util.stream.Collectors; import org.apache.parquet.Preconditions; import org.apache.parquet.schema.ColumnOrder.ColumnOrderName; +import org.apache.parquet.schema.LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation; import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; import org.apache.parquet.schema.Type.ID; import org.slf4j.Logger; @@ -846,11 +847,13 @@ public THIS addFields(Type... types) { @Override protected GroupType build(String name) { - if (newLogicalTypeSet) { - return new GroupType(repetition, name, logicalTypeAnnotation, fields, id); - } else { - return new GroupType(repetition, name, getOriginalType(), fields, id); + GroupType group = newLogicalTypeSet + ? new GroupType(repetition, name, logicalTypeAnnotation, fields, id) + : new GroupType(repetition, name, getOriginalType(), fields, id); + if (logicalTypeAnnotation instanceof FiniteVectorLogicalTypeAnnotation) { + TypeUtil.checkFiniteVector(group); } + return group; } public MapBuilder map(Type.Repetition repetition) { diff --git a/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java b/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java index ae2f32477c..63f91f0f45 100644 --- a/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java +++ b/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java @@ -24,9 +24,11 @@ import static org.apache.parquet.schema.LogicalTypeAnnotation.bsonType; import static org.apache.parquet.schema.LogicalTypeAnnotation.dateType; import static org.apache.parquet.schema.LogicalTypeAnnotation.decimalType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.finiteVectorType; import static org.apache.parquet.schema.LogicalTypeAnnotation.float16Type; import static org.apache.parquet.schema.LogicalTypeAnnotation.intType; import static org.apache.parquet.schema.LogicalTypeAnnotation.jsonType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.listType; import static org.apache.parquet.schema.LogicalTypeAnnotation.stringType; import static org.apache.parquet.schema.LogicalTypeAnnotation.timeType; import static org.apache.parquet.schema.LogicalTypeAnnotation.timestampType; @@ -46,13 +48,180 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.Optional; +import java.util.stream.Stream; import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; import org.apache.parquet.schema.Type.Repetition; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; public class TestTypeBuildersWithLogicalTypes { + private static GroupType vector(Type element) { + return new GroupType(REQUIRED, "v", finiteVectorType(5), new GroupType(Repetition.REPEATED, "list", element)); + } + + @ParameterizedTest + @ValueSource(ints = {1, Integer.MAX_VALUE}) + public void testVectorAnnotation(int length) { + GroupType vector = new GroupType( + REQUIRED, + "v", + finiteVectorType(length), + new GroupType(Repetition.REPEATED, "list", Types.required(FLOAT).named("element"))); + MessageType schema = Types.buildMessage().addField(vector).named("m"); + TypeUtil.checkValidWriteSchema(schema); + assertThat(MessageTypeParser.parseMessageType(schema.toString())).isEqualTo(schema); + assertThat(finiteVectorType(length)).asString().isEqualTo("FINITE_VECTOR(" + length + ")"); + assertThat(finiteVectorType(length).getNumElements()).isEqualTo(length); + assertThat(finiteVectorType(length).toOriginalType()).isEqualTo(OriginalType.LIST); + assertThat(finiteVectorType(length)).isInstanceOf(LogicalTypeAnnotation.ListLogicalTypeAnnotation.class); + assertThat(finiteVectorType(length)) + .isEqualTo(finiteVectorType(length)) + .hasSameHashCodeAs(finiteVectorType(length)); + assertThat(finiteVectorType(length)).isNotEqualTo(finiteVectorType(length == 1 ? 5 : 1)); + assertThat(finiteVectorType(length)).isNotEqualTo(LogicalTypeAnnotation.listType()); + assertThat(listType()).isNotEqualTo(finiteVectorType(length)); + } + + @Test + public void testVectorVisitorDispatch() { + LogicalTypeAnnotation.LogicalTypeAnnotationVisitor listVisitor = + new LogicalTypeAnnotation.LogicalTypeAnnotationVisitor() { + @Override + public Optional visit(LogicalTypeAnnotation.ListLogicalTypeAnnotation list) { + return Optional.of("list"); + } + }; + assertThat(listType().accept(listVisitor)).contains("list"); + assertThat(finiteVectorType(5).accept(listVisitor)).contains("list"); + + LogicalTypeAnnotation.LogicalTypeAnnotationVisitor vectorVisitor = + new LogicalTypeAnnotation.LogicalTypeAnnotationVisitor() { + @Override + public Optional visit(LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation vector) { + return Optional.of(vector.getNumElements()); + } + }; + assertThat(finiteVectorType(5).accept(vectorVisitor)).contains(5); + } + + @ParameterizedTest + @ValueSource(ints = {0, -1}) + public void testInvalidVectorLength(int length) { + assertThatThrownBy(() -> finiteVectorType(length)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Invalid FINITE_VECTOR length: " + length); + } + + private static Stream invalidVectors() { + PrimitiveType element = Types.required(FLOAT).named("element"); + GroupType list = new GroupType(Repetition.REPEATED, "list", element); + return Stream.of( + Arguments.of(new GroupType(Repetition.REPEATED, "v", finiteVectorType(5), list)), + Arguments.of(new GroupType( + REQUIRED, "v", finiteVectorType(5), new GroupType(Repetition.REPEATED, "items", element))), + Arguments.of(new GroupType(REQUIRED, "v", finiteVectorType(5), list, list)), + Arguments.of( + new GroupType(REQUIRED, "v", finiteVectorType(5), new GroupType(REQUIRED, "list", element))), + Arguments.of(new GroupType( + REQUIRED, + "v", + finiteVectorType(5), + new GroupType(Repetition.REPEATED, "list", LogicalTypeAnnotation.listType(), element))), + Arguments.of(new GroupType( + REQUIRED, + "v", + finiteVectorType(5), + new GroupType(Repetition.REPEATED, "list", element, element))), + Arguments.of(vector(Types.required(FLOAT).named("item"))), + Arguments.of(vector(Types.optional(FLOAT).named("element"))), + Arguments.of(vector(Types.repeated(FLOAT).named("element"))), + Arguments.of(vector(new GroupType(REQUIRED, "element", element))), + Arguments.of(vector(Types.required(BINARY).as(decimalType(2, 9)).named("element"))), + Arguments.of( + vector(Types.required(FIXED_LEN_BYTE_ARRAY).length(2).named("element"))), + Arguments.of(vector(Types.required(INT96).named("element"))), + Arguments.of(vector(Types.required(BINARY).as(stringType()).named("element"))), + Arguments.of(vector(Types.required(INT32).as(dateType()).named("element")))); + } + + @ParameterizedTest + @MethodSource("invalidVectors") + public void testInvalidVectorSchema(GroupType vector) { + assertThatThrownBy(() -> Types.buildGroup(vector.getRepetition()) + .as(vector.getLogicalTypeAnnotation()) + .addFields(vector.getFields().toArray(new Type[0])) + .named(vector.getName())) + .isInstanceOf(InvalidSchemaException.class) + .hasMessage("Invalid FINITE_VECTOR schema: " + vector); + } + + private static Stream vectorConstructionBoundaries() { + GroupType list = + new GroupType(Repetition.REPEATED, "list", Types.optional(FLOAT).named("element")); + return Stream.of( + () -> Types.requiredGroup() + .as(finiteVectorType(5)) + .addField(list) + .named("v"), + () -> Types.buildMessage() + .requiredGroup() + .as(finiteVectorType(5)) + .repeatedGroup() + .optional(FLOAT) + .named("element") + .named("list") + .named("v") + .named("m"), + () -> Types.buildMessage() + .requiredGroup() + .addField(vector(Types.required(FLOAT).named("element"))) + .requiredGroup() + .as(finiteVectorType(5)) + .addField(list) + .named("bad") + .named("parent") + .named("m"), + () -> MessageTypeParser.parseMessageType("message m { required group v (FINITE_VECTOR(5)) { " + + "repeated group list { optional float element; } } }")); + } + + @ParameterizedTest + @MethodSource("vectorConstructionBoundaries") + public void testVectorConstructionBoundaries(Executable construction) { + assertThatThrownBy(construction::execute) + .isInstanceOf(InvalidSchemaException.class) + .hasMessageContaining("Invalid FINITE_VECTOR schema"); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testHasFiniteVector(boolean hasVector) { + GroupType field = hasVector + ? vector(Types.required(FLOAT).named("element")) + : new GroupType( + REQUIRED, + "v", + listType(), + new GroupType( + Repetition.REPEATED, + "list", + Types.required(FLOAT).named("element"))); + MessageType schema = new MessageType("m", new GroupType(REQUIRED, "parent", field)); + assertThat(TypeUtil.hasFiniteVector(schema)).isEqualTo(hasVector); + } + + @Test + public void testVectorOnPrimitive() { + assertThatThrownBy(() -> Types.required(FLOAT).as(finiteVectorType(5)).named("v")) + .isInstanceOf(IllegalStateException.class) + .hasMessage("FINITE_VECTOR(5) can not be applied to a primitive type"); + } + @Test public void testGroupTypeConstruction() { PrimitiveType f1 = Types.required(BINARY).as(stringType()).named("f1"); diff --git a/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java b/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java index 8aa21e0ae3..953d97d2b3 100644 --- a/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java +++ b/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java @@ -38,6 +38,10 @@ public static LogicalType VARIANT(byte specificationVersion) { return LogicalType.VARIANT(type); } + public static LogicalType FINITE_VECTOR(int numElements) { + return LogicalType.FINITE_VECTOR(new FiniteVectorType(numElements)); + } + public static final LogicalType UTF8 = LogicalType.STRING(new StringType()); public static final LogicalType MAP = LogicalType.MAP(new MapType()); public static final LogicalType LIST = LogicalType.LIST(new ListType()); diff --git a/parquet-format-structures/src/main/thrift/parquet.thrift b/parquet-format-structures/src/main/thrift/parquet.thrift index bb582017bf..54759691bc 100644 --- a/parquet-format-structures/src/main/thrift/parquet.thrift +++ b/parquet-format-structures/src/main/thrift/parquet.thrift @@ -479,6 +479,11 @@ struct GeographyType { struct FileType { } +// A fixed-length LIST with exactly num_elements required, finite primitive elements. +struct FiniteVectorType { + 1: required i32 num_elements +} + /** * LogicalType annotations to replace ConvertedType. * @@ -513,6 +518,7 @@ union LogicalType { 17: GeometryType GEOMETRY // no compatible ConvertedType 18: GeographyType GEOGRAPHY // no compatible ConvertedType 19: FileType FILE // no compatible ConvertedType + 20: FiniteVectorType FINITE_VECTOR // use ConvertedType LIST } /** @@ -1129,6 +1135,7 @@ union ColumnOrder { * GEOMETRY - undefined * GEOGRAPHY - undefined * FILE - undefined + * FINITE_VECTOR - undefined * * In the absence of logical types, the sort order is determined by the physical type: * BOOLEAN - false, true diff --git a/parquet-hadoop/pom.xml b/parquet-hadoop/pom.xml index 69e2be1d33..5ec2f33786 100644 --- a/parquet-hadoop/pom.xml +++ b/parquet-hadoop/pom.xml @@ -58,6 +58,12 @@ test-jar test + + org.apache.parquet + parquet-encoding + ${project.version} + test + org.apache.hadoop hadoop-client diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java b/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java index 811e69f458..188f14a341 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java @@ -132,7 +132,9 @@ import org.apache.parquet.io.api.Binary; import org.apache.parquet.schema.ColumnOrder.ColumnOrderName; import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.InvalidSchemaException; import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation; import org.apache.parquet.schema.LogicalTypeAnnotation.UUIDLogicalTypeAnnotation; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.OriginalType; @@ -503,6 +505,11 @@ public Optional visit(LogicalTypeAnnotation.ListLogicalTypeAnnotati return of(LogicalTypes.LIST); } + @Override + public Optional visit(FiniteVectorLogicalTypeAnnotation vectorLogicalType) { + return of(LogicalTypes.FINITE_VECTOR(vectorLogicalType.getNumElements())); + } + @Override public Optional visit(LogicalTypeAnnotation.EnumLogicalTypeAnnotation enumLogicalType) { return of(LogicalTypes.ENUM); @@ -1333,6 +1340,8 @@ LogicalTypeAnnotation getLogicalTypeAnnotation(LogicalType type) { case ENUM -> LogicalTypeAnnotation.enumType(); case JSON -> LogicalTypeAnnotation.jsonType(); case LIST -> LogicalTypeAnnotation.listType(); + case FINITE__VECTOR -> + LogicalTypeAnnotation.finiteVectorType(type.getFINITE_VECTOR().getNum_elements()); case TIME -> { TimeType time = type.getTIME(); yield LogicalTypeAnnotation.timeType(time.isAdjustedToUTC, convertTimeUnit(time.unit)); @@ -2017,13 +2026,20 @@ private static ColumnPath getPath(ColumnMetaData metaData) { MessageType fromParquetSchema(List schema, List columnOrders) { Iterator iterator = schema.iterator(); SchemaElement root = iterator.next(); + if (root.isSetLogicalType() && root.logicalType.isSetFINITE_VECTOR()) { + throw new ParquetDecodingException("Invalid FINITE_VECTOR metadata for " + root.name); + } Types.MessageTypeBuilder builder = Types.buildMessage(); if (root.isSetField_id()) { builder.id(root.field_id); } Iterator columnOrderIterator = columnOrders == null ? null : columnOrders.iterator(); - buildChildren(builder, iterator, root.getNum_children(), columnOrderIterator); - return builder.named(root.name); + try { + buildChildren(builder, iterator, root.getNum_children(), columnOrderIterator); + return builder.named(root.name); + } catch (InvalidSchemaException e) { + throw new ParquetDecodingException(e.getMessage(), e); + } } private void buildChildren( @@ -2033,6 +2049,14 @@ private void buildChildren( Iterator columnOrders) { for (int i = 0; i < childrenCount; i++) { SchemaElement schemaElement = schema.next(); + if (schemaElement.isSetLogicalType() + && schemaElement.logicalType.isSetFINITE_VECTOR() + && (schemaElement.type != null + || schemaElement.converted_type != ConvertedType.LIST + || !schemaElement.logicalType.getFINITE_VECTOR().isSetNum_elements() + || schemaElement.logicalType.getFINITE_VECTOR().getNum_elements() <= 0)) { + throw new ParquetDecodingException("Invalid FINITE_VECTOR metadata for " + schemaElement.name); + } // Create Parquet Type. Types.Builder childBuilder; diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java index 5a611b41d1..dfba5cdc69 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java @@ -117,8 +117,10 @@ import org.apache.parquet.io.ParquetDecodingException; import org.apache.parquet.io.ParquetFileRange; import org.apache.parquet.io.SeekableInputStream; +import org.apache.parquet.schema.LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.TypeUtil; import org.apache.parquet.util.AutoCloseables; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -763,6 +765,7 @@ public static ParquetFileReader open( private final InputFile file; private final ParquetReadOptions options; private final Map paths = new HashMap<>(); + private final Map vectorColumns; private final FileMetaData fileMetaData; // may be null private final List blocks; private final List blockIndexStores; @@ -817,6 +820,7 @@ public ParquetFileReader( this.converter = new ParquetMetadataConverter(configuration); this.file = HadoopInputFile.fromPath(filePath, configuration); this.fileMetaData = fileMetaData; + this.vectorColumns = vectorColumns(fileMetaData.getSchema()); this.f = file.newStream(); this.fileDecryptor = fileMetaData.getFileDecryptor(); if (null == fileDecryptor) { @@ -874,6 +878,7 @@ public ParquetFileReader(Configuration conf, Path file, ParquetMetadata footer) this.file = HadoopInputFile.fromPath(file, conf); this.f = this.file.newStream(); this.fileMetaData = footer.getFileMetaData(); + this.vectorColumns = vectorColumns(fileMetaData.getSchema()); this.fileDecryptor = fileMetaData.getFileDecryptor(); if (null == fileDecryptor) { this.options = HadoopReadOptions.builder(conf).build(); @@ -931,6 +936,7 @@ public ParquetFileReader( this.file = HadoopInputFile.fromPath(file, conf); this.f = f; this.fileMetaData = footer.getFileMetaData(); + this.vectorColumns = vectorColumns(fileMetaData.getSchema()); this.fileDecryptor = fileMetaData.getFileDecryptor(); this.options = options; this.footer = footer; @@ -985,6 +991,7 @@ public ParquetFileReader(InputFile file, ParquetMetadata footer, ParquetReadOpti this.footer = footer; this.fileMetaData = footer.getFileMetaData(); + this.vectorColumns = vectorColumns(fileMetaData.getSchema()); this.fileDecryptor = fileMetaData.getFileDecryptor(); // must be called before filterRowGroups! if (null != fileDecryptor && fileDecryptor.plaintextFile()) { this.fileDecryptor = null; // Plaintext file. No need in decryptor @@ -1158,6 +1165,57 @@ public PageReadStore readNextRowGroup() throws IOException { return currentRowGroup; } + private record VectorColumn(ColumnDescriptor descriptor, int length) {} + + private static Map vectorColumns(MessageType schema) { + if (!TypeUtil.hasFiniteVector(schema)) return Collections.emptyMap(); + Map columns = new HashMap<>(); + for (ColumnDescriptor column : schema.getColumns()) { + FiniteVectorLogicalTypeAnnotation vector = TypeUtil.getFiniteVectorType(schema, column.getPath()); + if (vector != null) { + columns.put(ColumnPath.get(column.getPath()), new VectorColumn(column, vector.getNumElements())); + } + } + return columns; + } + + private static void checkVectorCounts(ColumnDescriptor path, int numElements, long rows, long values, long nulls) { + long declaredNulls = nulls; + boolean invalid = rows < 0 || values < 0 || nulls < -1 || nulls > values; + if (!invalid && path.getMaxRepetitionLevel() == 1) { + if (path.getMaxDefinitionLevel() == 1) { + invalid = nulls > 0; + nulls = 0; + } + if (nulls >= 0) { + long present = values - nulls; + invalid |= nulls > rows || present % numElements != 0 || present / numElements != rows - nulls; + } else { + long extra = values - rows; + invalid |= values < rows + || (numElements == 1 + ? values != rows + : extra % (numElements - 1) != 0 || extra / (numElements - 1) > rows); + } + } else if (!invalid && nulls >= 0) { + invalid = (values - nulls) % numElements != 0; + } + if (invalid) { + throw new ParquetDecodingException("Invalid FINITE_VECTOR counts for " + Arrays.toString(path.getPath()) + + ": rows=" + rows + ", values=" + values + ", nulls=" + declaredNulls + + ", length=" + numElements); + } + } + + private void checkVectorColumnCounts(ColumnChunkMetaData column, long rows) { + if (vectorColumns.isEmpty()) return; + VectorColumn vector = vectorColumns.get(column.getPath()); + if (vector == null) return; + org.apache.parquet.column.statistics.Statistics stats = column.getStatistics(); + long nulls = stats != null && stats.isNumNullsSet() ? stats.getNumNulls() : -1; + checkVectorCounts(vector.descriptor(), vector.length(), rows, column.getValueCount(), nulls); + } + private ColumnChunkPageReadStore internalReadRowGroup(int blockIndex) throws IOException { if (blockIndex < 0 || blockIndex >= blocks.size()) { return null; @@ -1175,6 +1233,7 @@ private ColumnChunkPageReadStore internalReadRowGroup(int blockIndex) throws IOE ColumnPath pathKey = mc.getPath(); ColumnDescriptor columnDescriptor = paths.get(pathKey); if (columnDescriptor != null) { + checkVectorColumnCounts(mc, block.getRowCount()); BenchmarkCounter.incrementTotalBytes(mc.getTotalSize()); long startingPos = mc.getStartingPos(); // first part or not consecutive => new list @@ -1573,6 +1632,7 @@ private ColumnChunkPageReadStore internalReadFilteredRowGroup( ColumnPath pathKey = mc.getPath(); ColumnDescriptor columnDescriptor = paths.get(pathKey); if (columnDescriptor != null) { + checkVectorColumnCounts(mc, block.getRowCount()); OffsetIndex offsetIndex = ciStore.getOffsetIndex(mc.getPath()); OffsetIndex filteredOffsetIndex = @@ -2135,6 +2195,8 @@ public ColumnChunkPageReader readAllPages( .getSchema() .getType(descriptor.col.getPath()) .asPrimitiveType(); + String[] path = descriptor.col.getPath(); + VectorColumn vector = vectorColumns.isEmpty() ? null : vectorColumns.get(descriptor.metadata.getPath()); long valuesCountReadSoFar = 0L; int dataPageCountReadSoFar = 0; byte[] dataPageHeaderAAD = null; @@ -2209,12 +2271,36 @@ public ColumnChunkPageReader readAllPages( if (pageHeader.isSetCrc()) { dataPageV1.setCrc(pageHeader.getCrc()); } + if (vector != null && offsetIndex != null) { + long rows = offsetIndex.getLastRowIndex(dataPageCountReadSoFar, rowCount) + - offsetIndex.getFirstRowIndex(dataPageCountReadSoFar) + + 1; + org.apache.parquet.column.statistics.Statistics stats = dataPageV1.getStatistics(); + long nulls = stats != null && stats.isNumNullsSet() ? stats.getNumNulls() : -1; + checkVectorCounts( + vector.descriptor(), vector.length(), rows, dataPageV1.getValueCount(), nulls); + } pagesInChunk.add(dataPageV1); valuesCountReadSoFar += dataHeaderV1.getNum_values(); ++dataPageCountReadSoFar; break; case DATA_PAGE_V2: DataPageHeaderV2 dataHeaderV2 = pageHeader.getData_page_header_v2(); + if (vector != null) { + if (dataHeaderV2.getNum_nulls() < 0) { + throw new ParquetDecodingException("Invalid FINITE_VECTOR counts for " + + Arrays.toString(path) + + ": rows=" + dataHeaderV2.getNum_rows() + ", values=" + + dataHeaderV2.getNum_values() + ", nulls=" + dataHeaderV2.getNum_nulls() + + ", length=" + vector.length()); + } + checkVectorCounts( + vector.descriptor(), + vector.length(), + dataHeaderV2.getNum_rows(), + dataHeaderV2.getNum_values(), + dataHeaderV2.getNum_nulls()); + } int dataSize = compressedPageSize - dataHeaderV2.getRepetition_levels_byte_length() - dataHeaderV2.getDefinition_levels_byte_length(); diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java b/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java index c7b92dbca1..70a34072ab 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java @@ -68,6 +68,7 @@ import java.util.Random; import java.util.Set; import java.util.TreeSet; +import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -134,6 +135,7 @@ import org.apache.parquet.internal.column.columnindex.ColumnIndexBuilder; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import org.apache.parquet.internal.column.columnindex.OffsetIndexBuilder; +import org.apache.parquet.io.ParquetDecodingException; import org.apache.parquet.io.api.Binary; import org.apache.parquet.schema.ColumnOrder; import org.apache.parquet.schema.LogicalTypeAnnotation; @@ -146,9 +148,66 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; public class TestParquetMetadataConverter { + private static MessageType vectorSchema() { + return parseMessageType("message m { required group v (FINITE_VECTOR(5)) { repeated group list {" + + " required float element columnorder(TYPE_DEFINED_ORDER); } } }"); + } + + private static Stream invalidVectorMetadata() { + return Stream.of( + Arguments.of( + "missing length", + (Consumer>) wire -> + wire.get(1).getLogicalType().getFINITE_VECTOR().unsetNum_elements(), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "zero length", + (Consumer>) wire -> + wire.get(1).getLogicalType().getFINITE_VECTOR().setNum_elements(0), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "negative length", + (Consumer>) wire -> + wire.get(1).getLogicalType().getFINITE_VECTOR().setNum_elements(-1), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "missing converted", + (Consumer>) wire -> wire.get(1).unsetConverted_type(), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "wrong converted", + (Consumer>) wire -> wire.get(1).setConverted_type(ConvertedType.MAP), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "primitive outer", + (Consumer>) wire -> wire.get(1).setType(Type.FLOAT), + "Invalid FINITE_VECTOR metadata for v"), + Arguments.of( + "root", + (Consumer>) wire -> + wire.get(0).setLogicalType(org.apache.parquet.format.LogicalTypes.FINITE_VECTOR(5)), + "Invalid FINITE_VECTOR metadata for m"), + Arguments.of( + "element name", + (Consumer>) wire -> wire.get(3).setName("item"), + "Invalid FINITE_VECTOR schema:")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("invalidVectorMetadata") + public void testInvalidVectorMetadata(String name, Consumer> edit, String message) { + ParquetMetadataConverter converter = new ParquetMetadataConverter(); + List wire = converter.toParquetSchema(vectorSchema()); + edit.accept(wire); + assertThatThrownBy(() -> converter.fromParquetSchema(wire, null)) + .isInstanceOf(ParquetDecodingException.class) + .hasMessageStartingWith(message); + } + private static SecureRandom random = new SecureRandom(); private static final String CHAR_LOWER = "abcdefghijklmnopqrstuvwxyz"; private static final String CHAR_UPPER = CHAR_LOWER.toUpperCase(); diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectorReadWrite.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectorReadWrite.java new file mode 100644 index 0000000000..c5b3277be0 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectorReadWrite.java @@ -0,0 +1,591 @@ +/* + * 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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.hadoop; + +import static org.apache.parquet.column.Encoding.PLAIN; +import static org.apache.parquet.column.Encoding.RLE; +import static org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_2_0; +import static org.apache.parquet.filter.ColumnPredicates.equalTo; +import static org.apache.parquet.filter2.predicate.FilterApi.and; +import static org.apache.parquet.filter2.predicate.FilterApi.gtEq; +import static org.apache.parquet.filter2.predicate.FilterApi.intColumn; +import static org.apache.parquet.filter2.predicate.FilterApi.ltEq; +import static org.apache.parquet.hadoop.ParquetFileWriter.Mode.CREATE; +import static org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE; +import static org.apache.parquet.hadoop.ParquetWriter.MAX_PADDING_SIZE_DEFAULT; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FLOAT; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.function.Consumer; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.bytes.BytesUtils; +import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.column.ParquetProperties.WriterVersion; +import org.apache.parquet.column.page.DataPage; +import org.apache.parquet.column.page.DataPageV1; +import org.apache.parquet.column.page.DataPageV2; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.column.page.PageReader; +import org.apache.parquet.column.statistics.FloatStatistics; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroup; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.example.data.simple.convert.GroupRecordConverter; +import org.apache.parquet.filter2.columnindex.RowRanges; +import org.apache.parquet.filter2.compat.FilterCompat; +import org.apache.parquet.format.ConvertedType; +import org.apache.parquet.format.FileMetaData; +import org.apache.parquet.format.Util; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.example.GroupReadSupport; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.hadoop.util.HadoopOutputFile; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.LocalInputFile; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.ParquetDecodingException; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.InvalidSchemaException; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; + +public class TestVectorReadWrite { + @TempDir + private java.nio.file.Path tempDir; + + private static MessageType schema(String outer, int length, boolean nested) { + String vector = + outer + " group v (FINITE_VECTOR(" + length + ")) { repeated group list { required float element; } }"; + return MessageTypeParser.parseMessageType( + nested ? "message m { repeated group outer { " + vector + " } }" : "message m { " + vector + " }"); + } + + private Path writeVectors(MessageType schema, WriterVersion version, boolean dictionary, List rows) + throws IOException { + Path path = new Path(tempDir.resolve(UUID.randomUUID() + ".parquet").toUri()); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withType(schema) + .withWriterVersion(version) + .withDictionaryEncoding(dictionary) + .withPageRowCountLimit(2) + .withMinRowCountForPageSizeCheck(1) + .withMaxRowCountForPageSizeCheck(1) + .build()) { + for (Group row : rows) writer.write(row); + } + return path; + } + + private static MessageType mixedSchema() { + return MessageTypeParser.parseMessageType("message m {" + + " required int32 id;" + + " required group f (FINITE_VECTOR(5)) { repeated group list { required float element; } }" + + " optional group o (FINITE_VECTOR(2)) { repeated group list { required int32 element; } }" + + " optional group a { required group h (FINITE_VECTOR(2)) { repeated group list {" + + " required fixed_len_byte_array(2) element (FLOAT16); } } }" + + " repeated group outer { required group d (FINITE_VECTOR(2)) { repeated group list {" + + " required double element; } } }" + + " required group b (FINITE_VECTOR(2)) { repeated group list { required boolean element; } }" + + " required group l (FINITE_VECTOR(1)) { repeated group list { required int64 element; } }" + + " required group s (FINITE_VECTOR(1)) { repeated group list {" + + " required int32 element (INTEGER(16,true)); } }" + + " required group u (FINITE_VECTOR(1)) { repeated group list {" + + " required int64 element (INTEGER(64,false)); } }" + + " required group dec (FINITE_VECTOR(1)) { repeated group list {" + + " required int32 element (DECIMAL(9,2)); } }" + + " required group bytes (FINITE_VECTOR(1)) { repeated group list {" + + " required fixed_len_byte_array(4) element (DECIMAL(9,2)); } } }"); + } + + private static List mixedRows() { + List rows = new ArrayList<>(); + SimpleGroupFactory factory = new SimpleGroupFactory(mixedSchema()); + for (int i = 0; i < 6; i++) { + Group row = factory.newGroup().append("id", i); + Group f = row.addGroup("f"); + for (float value : new float[] {i + 0.5f, -0.0f, Float.MIN_VALUE, -0.25f, 100.0f}) { + f.addGroup("list").append("element", value); + } + if (i % 2 != 0) { + Group o = row.addGroup("o"); + o.addGroup("list").append("element", 7); + o.addGroup("list").append("element", 8); + } + if (i != 1 && i != 4) { + Group h = row.addGroup("a").addGroup("h"); + h.addGroup("list").append("element", Binary.fromConstantByteArray(new byte[] {0, (byte) 0x80})); + h.addGroup("list").append("element", Binary.fromConstantByteArray(new byte[] {1, 0})); + } + for (int j = 0; j < i % 3; j++) { + Group d = row.addGroup("outer").addGroup("d"); + d.addGroup("list").append("element", -0.0d); + d.addGroup("list").append("element", Double.MIN_VALUE); + } + Group b = row.addGroup("b"); + b.addGroup("list").append("element", true); + b.addGroup("list").append("element", false); + row.addGroup("l").addGroup("list").append("element", 9L); + row.addGroup("s").addGroup("list").append("element", -7); + row.addGroup("u").addGroup("list").append("element", 9L); + row.addGroup("dec").addGroup("list").append("element", 125); + row.addGroup("bytes") + .addGroup("list") + .append("element", Binary.fromConstantByteArray(new byte[] {0, 0, 0x7c, 0})); + rows.add(row); + } + return rows; + } + + private static FileMetaData readFooter(Path path) throws IOException { + byte[] bytes = Files.readAllBytes(java.nio.file.Path.of(path.toUri())); + int length = BytesUtils.readIntLittleEndian(bytes, bytes.length - 8); + return Util.readFileMetaData(new ByteArrayInputStream(bytes, bytes.length - 8 - length, length)); + } + + private static void rewriteFooter(Path path, Consumer edit) throws IOException { + java.nio.file.Path local = java.nio.file.Path.of(path.toUri()); + byte[] original = Files.readAllBytes(local); + int footerLength = BytesUtils.readIntLittleEndian(original, original.length - 8); + FileMetaData footer = readFooter(path); + edit.accept(footer); + ByteArrayOutputStream encoded = new ByteArrayOutputStream(); + Util.writeFileMetaData(footer, encoded); + ByteArrayOutputStream rewritten = new ByteArrayOutputStream(); + rewritten.write(original, 0, original.length - 8 - footerLength); + encoded.writeTo(rewritten); + BytesUtils.writeIntLittleEndian(rewritten, encoded.size()); + rewritten.write(original, original.length - 4, 4); + Files.write(local, rewritten.toByteArray()); + } + + private static Path referenceFile(String name) throws java.net.URISyntaxException { + return new Path(java.nio.file.Path.of( + TestVectorReadWrite.class.getResource("/vector/" + name).toURI()) + .toUri()); + } + + private static List listRows(MessageType schema, int... lengths) { + List rows = new ArrayList<>(); + int value = 0; + for (int length : lengths) { + Group row = new SimpleGroupFactory(schema).newGroup(); + if (length >= 0) { + Group parent = schema.containsField("outer") ? row.addGroup("outer") : row; + Group vector = parent.addGroup("v"); + for (int i = 0; i < length; i++) { + vector.addGroup("list").append("element", (float) value++); + } + } + rows.add(row); + } + return rows; + } + + private static List readGroups(Path path, MessageType requested) throws IOException { + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri())))) { + reader.setRequestedSchema(requested); + try (PageReadStore pages = reader.readRowGroup(0)) { + MessageType fileSchema = reader.getFileMetaData().getSchema(); + MessageColumnIO columns = new ColumnIOFactory().getColumnIO(requested, fileSchema); + RecordReader records = columns.getRecordReader(pages, new GroupRecordConverter(requested)); + List result = new ArrayList<>(); + for (long i = 0; i < pages.getRowCount(); i++) result.add(records.read()); + return result; + } + } + } + + private static void assertVectorRows(List actual, int... lengths) { + assertThat(actual).hasSize(lengths.length); + int value = 0; + for (int row = 0; row < lengths.length; row++) { + Group vector = actual.get(row).getGroup("v", 0); + assertThat(vector.getFieldRepetitionCount("list")).isEqualTo(lengths[row]); + for (int i = 0; i < lengths[row]; i++) { + assertThat(vector.getGroup("list", i).getFloat("element", 0)).isEqualTo((float) value++); + } + } + } + + @ParameterizedTest + @CsvSource({"PARQUET_1_0,false", "PARQUET_1_0,true", "PARQUET_2_0,false", "PARQUET_2_0,true"}) + public void testVectorRoundTrip(WriterVersion version, boolean dictionary) throws IOException { + Path path = writeVectors(mixedSchema(), version, dictionary, mixedRows()); + verifyMixedFile(path, version, dictionary); + if (!dictionary) verifyFilteredRows(path); + if (version == WriterVersion.PARQUET_1_0 && !dictionary) { + Path legacy = new Path(tempDir.resolve("legacy-list.parquet").toUri()); + Files.copy(java.nio.file.Path.of(path.toUri()), java.nio.file.Path.of(legacy.toUri())); + rewriteFooter( + legacy, + footer -> footer.getSchema().stream() + .filter(element -> element.isSetLogicalType() + && element.getLogicalType().isSetFINITE_VECTOR()) + .forEach(org.apache.parquet.format.SchemaElement::unsetLogicalType)); + List vectors = readFooter(legacy).getSchema().stream() + .filter(element -> element.getConverted_type() == org.apache.parquet.format.ConvertedType.LIST) + .toList(); + assertThat(vectors).hasSize(10); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), legacy).build()) { + for (Group expected : mixedRows()) assertRow(reader.read(), expected); + assertThat(reader.read()).isNull(); + } + } + } + + private static void assertRow(Group actual, Group expected) { + assertThat(actual).isNotNull(); + assertThat(actual.getType().getFieldCount()) + .isEqualTo(expected.getType().getFieldCount()); + for (int field = 0; field < expected.getType().getFieldCount(); field++) { + assertThat(actual.getType().getFieldName(field)) + .isEqualTo(expected.getType().getFieldName(field)); + assertThat(actual.getFieldRepetitionCount(field)).isEqualTo(expected.getFieldRepetitionCount(field)); + for (int index = 0; index < expected.getFieldRepetitionCount(field); index++) { + if (!expected.getType().getType(field).isPrimitive()) { + assertRow(actual.getGroup(field, index), expected.getGroup(field, index)); + } else { + assertThat(((SimpleGroup) actual).getObject(field, index)) + .isEqualTo(((SimpleGroup) expected).getObject(field, index)); + } + } + } + } + + private static void verifyMixedFile(Path path, WriterVersion version, boolean dictionary) throws IOException { + MessageType schema = mixedSchema(); + try (ParquetReader reader = + ParquetReader.builder(new GroupReadSupport(), path).build()) { + for (Group expected : mixedRows()) assertRow(reader.read(), expected); + assertThat(reader.read()).isNull(); + } + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri())))) { + assertThat(reader.getFooter().getFileMetaData().getSchema()).isEqualTo(schema); + assertThat(reader.getFooter().getBlocks()).hasSize(1); + assertThat(reader.getFooter().getBlocks().get(0).getRowCount()).isEqualTo(6); + try (PageReadStore pages = reader.readRowGroup(0)) { + PageReader f = pages.getPageReader(schema.getColumns().get(1)); + if (dictionary) { + assertThat(f.readDictionaryPage()).isNotNull(); + } else { + assertThat(f.readDictionaryPage()).isNull(); + } + DataPage page = f.readPage(); + assertThat(page) + .isInstanceOf(version == WriterVersion.PARQUET_1_0 ? DataPageV1.class : DataPageV2.class); + } + } + List vectors = readFooter(path).getSchema().stream() + .filter(element -> + element.isSetLogicalType() && element.getLogicalType().isSetFINITE_VECTOR()) + .toList(); + assertThat(vectors) + .hasSize(10) + .extracting(element -> element.getConverted_type()) + .containsOnly(ConvertedType.LIST); + assertThat((int) vectors.get(0).getLogicalType().getSetField().getThriftFieldId()) + .isEqualTo(20); + } + + private static void verifyFilteredRows(Path path) throws IOException { + MessageType schema = mixedSchema(); + List expected = mixedRows(); + var predicate = and(gtEq(intColumn("id"), 2), ltEq(intColumn("id"), 3)); + try (ParquetReader reader = ParquetReader.builder(new GroupReadSupport(), path) + .withFilter(FilterCompat.get(predicate)) + .useColumnIndexFilter(true) + .build()) { + assertRow(reader.read(), expected.get(2)); + assertRow(reader.read(), expected.get(3)); + assertThat(reader.read()).isNull(); + } + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri()))); + PageReadStore pages = reader.readFilteredRowGroup( + 0, RowRanges.builder().addSelectedRange(1, 4).build())) { + assertThat(pages.getRowIndexes()).isPresent(); + MessageColumnIO columns = new ColumnIOFactory().getColumnIO(schema, schema); + RecordReader records = columns.getRecordReader(pages, new GroupRecordConverter(schema)); + for (int i = 1; i <= 4; i++) assertRow(records.read(), expected.get(i)); + } + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri()))); + PageReadStore pages = reader.readRowGroup(0)) { + MessageColumnIO columns = new ColumnIOFactory().getColumnIO(schema, schema); + RecordReader records = columns.getRecordReader( + pages, + new GroupRecordConverter(schema), + FilterCompat.get(org.apache.parquet.filter.ColumnRecordFilter.column("id", equalTo(3)))); + assertRow(records.read(), expected.get(3)); + assertThat(records.read()).isNull(); + } + } + + @Test + public void testInvalidVectorWriteSchema() { + GroupType vector = schema("required", 2, false).getType("v").asGroupType(); + GroupType invalid = vector.withNewFields( + Types.repeatedGroup().optional(FLOAT).named("element").named("list")); + MessageType schema = Types.buildMessage() + .optionalGroup() + .addField(invalid) + .named("parent") + .named("m"); + assertThatThrownBy(() -> writeVectors(schema, WriterVersion.PARQUET_1_0, false, List.of())) + .isInstanceOf(InvalidSchemaException.class) + .hasMessage("Invalid FINITE_VECTOR schema: " + invalid); + } + + @Test + public void testVectorReadProjection() throws Exception { + MessageType list = MessageTypeParser.parseMessageType( + schema("required", 5, false).toString().replace("FINITE_VECTOR(5)", "LIST")); + MessageType vector = schema("required", 5, false); + Path path = writeVectors(list, WriterVersion.PARQUET_1_0, false, listRows(list, 4, 6)); + assertVectorRows(readGroups(path, vector), 4, 6); + + Path vectorPath = writeVectors(vector, WriterVersion.PARQUET_1_0, false, listRows(vector, 4, 6)); + for (MessageType requested : List.of(vector, list)) { + assertVectorRows(readGroups(vectorPath, requested), 4, 6); + } + } + + @ParameterizedTest + @CsvSource({"false", "true"}) + public void testVectorValuesAreNotValidated(boolean validation) throws IOException { + MessageType schema = schema("required", 3, false); + Group row = new SimpleGroupFactory(schema).newGroup(); + Group vector = row.addGroup("v"); + for (float value : new float[] {Float.NaN, Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY}) { + vector.addGroup("list").append("element", value); + } + Path path = new Path(tempDir.resolve("non-finite.parquet").toUri()); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withType(schema) + .withValidation(validation) + .build()) { + writer.write(row); + } + assertRow(readGroups(path, schema).get(0), row); + } + + private static org.apache.parquet.format.ColumnMetaData column(org.apache.parquet.format.FileMetaData footer) { + return footer.getRow_groups().get(0).getColumns().get(0).getMeta_data(); + } + + @ParameterizedTest + @CsvSource({"vector_mixed_v1_plain.parquet,PARQUET_1_0,false", "vector_mixed_v2_dictionary.parquet,PARQUET_2_0,true" + }) + public void testVectorReferenceFile(String name, WriterVersion version, boolean dictionary) throws Exception { + verifyMixedFile(referenceFile(name), version, dictionary); + } + + @Test + public void testSplitVectorFile() throws Exception { + Path path = referenceFile("vector_split_v1.parquet"); + MessageType schema = schema("required", 5, false); + assertVectorRows(readGroups(path, schema), 5, 5); + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri()))); + PageReadStore pages = reader.readRowGroup(0)) { + PageReader column = pages.getPageReader(schema.getColumns().get(0)); + for (int count : new int[] {3, 4, 3}) { + DataPageV1 page = (DataPageV1) column.readPage(); + assertThat(page.getValueCount()).isEqualTo(count); + } + assertThat(column.readPage()).isNull(); + } + } + + @ParameterizedTest + @CsvSource({ + "required,5,false,4,18,0,false,false", + "required,5,false,4,15,0,true,true", + "optional,5,false,3,11,0,false,true", + "optional,5,false,3,10,-1,false,false", + "required,1,false,3,2,0,false,false", + "required,5,false,9223372036854775807,15,0,false,false", + "required,2,true,3,5,0,false,false" + }) + public void testInvalidVectorFooter( + String outer, + int length, + boolean nested, + long rowCount, + long values, + long nulls, + boolean projectList, + boolean filtered) + throws IOException { + MessageType fileSchema = schema(outer, length, nested); + int count = rowCount == Long.MAX_VALUE ? 3 : (int) rowCount; + int[] lengths = new int[count]; + Arrays.fill(lengths, length); + if (outer.equals("optional")) lengths[0] = -1; + Path path = writeVectors(fileSchema, PARQUET_2_0, false, listRows(fileSchema, lengths)); + rewriteFooter(path, footer -> { + footer.setNum_rows(rowCount); + footer.getRow_groups().get(0).setNum_rows(rowCount); + column(footer).setNum_values(values); + if (nulls == -1) { + column(footer).getStatistics().unsetNull_count(); + } else { + column(footer).getStatistics().setNull_count(nulls); + } + }); + String expected = (nested ? "[outer, v, list, element]" : "[v, list, element]") + ": rows=" + rowCount + + ", values=" + values + ", nulls=" + nulls + ", length=" + length; + for (boolean select : filtered ? new boolean[] {false, true} : new boolean[] {false}) { + assertThatThrownBy(() -> { + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri())))) { + if (projectList) { + reader.setRequestedSchema(MessageTypeParser.parseMessageType( + "message m { repeated group v (LIST) { repeated group list {" + + " required float element; } } }")); + } + assertThat(reader.readOffsetIndex(reader.getFooter() + .getBlocks() + .get(0) + .getColumns() + .get(0))) + .isNotNull(); + if (select) { + reader.readFilteredRowGroup(0, RowRanges.createSingle(1)); + } else { + reader.readRowGroup(0); + } + } + }) + .isInstanceOf(ParquetDecodingException.class) + .hasMessageContaining("Invalid FINITE_VECTOR counts for " + expected); + } + } + + @ParameterizedTest + @CsvSource({"false,3,10,0", "false,2,10,1", "false,-1,10,0", "false,2,10,-1", "true,2,9,0"}) + public void testInvalidVectorPageHeader(boolean nested, int headerRows, int headerValues, int headerNulls) + throws IOException { + MessageType fileSchema = schema("required", nested ? 2 : 5, nested); + Path path = new Path(tempDir.resolve(UUID.randomUUID() + ".parquet").toUri()); + try (ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, new Configuration()), + fileSchema, + CREATE, + DEFAULT_BLOCK_SIZE, + MAX_PADDING_SIZE_DEFAULT, + null, + ParquetProperties.builder().build())) { + FloatStatistics stats = new FloatStatistics(); + stats.updateStats(1.0f); + stats.setNumNulls(0); + writer.start(); + writer.startBlock(2); + writer.startColumn(fileSchema.getColumns().get(0), 10, CompressionCodecName.UNCOMPRESSED); + writer.writeDataPageV2( + headerRows, + headerNulls, + headerValues, + BytesInput.empty(), + BytesInput.empty(), + PLAIN, + BytesInput.empty(), + false, + 0, + stats); + writer.endColumn(); + writer.endBlock(); + writer.end(Collections.emptyMap()); + } + assertThatThrownBy(() -> { + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri())))) { + reader.readRowGroup(0); + } + }) + .isInstanceOf(ParquetDecodingException.class) + .hasMessage("Invalid FINITE_VECTOR counts for " + + (nested ? "[outer, v, list, element]" : "[v, list, element]") + + ": rows=" + headerRows + ", values=" + headerValues + ", nulls=" + headerNulls + + ", length=" + (nested ? 2 : 5)); + } + + @Test + public void testInvalidIndexedVectorPage() throws IOException { + MessageType fileSchema = schema("required", 5, false); + Path path = new Path(tempDir.resolve("invalid-indexed.parquet").toUri()); + FloatStatistics stats = new FloatStatistics(); + stats.updateStats(1.0f); + stats.setNumNulls(0); + try (ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, new Configuration()), + fileSchema, + CREATE, + DEFAULT_BLOCK_SIZE, + MAX_PADDING_SIZE_DEFAULT, + null, + ParquetProperties.builder().build())) { + writer.start(); + writer.startBlock(4); + writer.startColumn(fileSchema.getColumns().get(0), 20, CompressionCodecName.UNCOMPRESSED); + writer.writeDataPage(9, 0, BytesInput.empty(), stats, 2, RLE, RLE, PLAIN); + writer.writeDataPage(11, 0, BytesInput.empty(), stats, 2, RLE, RLE, PLAIN); + writer.endColumn(); + writer.endBlock(); + writer.end(Collections.emptyMap()); + } + assertThatThrownBy(() -> { + try (ParquetFileReader reader = + ParquetFileReader.open(new LocalInputFile(java.nio.file.Path.of(path.toUri())))) { + assertThat(reader.readOffsetIndex(reader.getFooter() + .getBlocks() + .get(0) + .getColumns() + .get(0))) + .isNotNull(); + reader.readFilteredRowGroup(0, RowRanges.createSingle(1)); + } + }) + .isInstanceOf(ParquetDecodingException.class) + .hasMessage( + "Invalid FINITE_VECTOR counts for [v, list, element]: rows=2, values=9, nulls=-1, length=5"); + } +} diff --git a/parquet-hadoop/src/test/resources/vector/vector_mixed_v1_plain.parquet b/parquet-hadoop/src/test/resources/vector/vector_mixed_v1_plain.parquet new file mode 100644 index 0000000000..fc2e8aa647 Binary files /dev/null and b/parquet-hadoop/src/test/resources/vector/vector_mixed_v1_plain.parquet differ diff --git a/parquet-hadoop/src/test/resources/vector/vector_mixed_v2_dictionary.parquet b/parquet-hadoop/src/test/resources/vector/vector_mixed_v2_dictionary.parquet new file mode 100644 index 0000000000..ad809ff478 Binary files /dev/null and b/parquet-hadoop/src/test/resources/vector/vector_mixed_v2_dictionary.parquet differ diff --git a/parquet-hadoop/src/test/resources/vector/vector_split_v1.parquet b/parquet-hadoop/src/test/resources/vector/vector_split_v1.parquet new file mode 100644 index 0000000000..5b00013501 Binary files /dev/null and b/parquet-hadoop/src/test/resources/vector/vector_split_v1.parquet differ diff --git a/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoMessageConverter.java b/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoMessageConverter.java index 0da0dfca96..428a78b68b 100644 --- a/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoMessageConverter.java +++ b/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoMessageConverter.java @@ -268,6 +268,12 @@ public Optional visit(LogicalTypeAnnotation.ListLogicalTypeAnnotation return of(new ListConverter(parentBuilder, fieldDescriptor, parquetType)); } + @Override + public Optional visit( + LogicalTypeAnnotation.FiniteVectorLogicalTypeAnnotation vectorLogicalType) { + return visit(LogicalTypeAnnotation.listType()); + } + @Override public Optional visit(LogicalTypeAnnotation.MapLogicalTypeAnnotation mapLogicalType) { return of(new MapConverter(parentBuilder, fieldDescriptor, parquetType)); diff --git a/parquet-protobuf/src/test/java/org/apache/parquet/proto/ProtoRecordConverterTest.java b/parquet-protobuf/src/test/java/org/apache/parquet/proto/ProtoRecordConverterTest.java index 6a3560dc61..4935d8eb52 100644 --- a/parquet-protobuf/src/test/java/org/apache/parquet/proto/ProtoRecordConverterTest.java +++ b/parquet-protobuf/src/test/java/org/apache/parquet/proto/ProtoRecordConverterTest.java @@ -24,12 +24,48 @@ import static org.assertj.core.data.Offset.offset; import com.google.protobuf.ByteString; +import java.io.IOException; import java.util.List; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.column.ParquetProperties.WriterVersion; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.proto.test.TestProto3; import org.apache.parquet.proto.test.TestProtobuf; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; public class ProtoRecordConverterTest { + @Test + public void testVectorRepeatedPrimitive(@TempDir java.nio.file.Path dir) throws IOException { + MessageType schema = MessageTypeParser.parseMessageType("message m { required group repeatedInt" + + " (FINITE_VECTOR(3)) { repeated group list { required int32 element; } } }"); + Path path = new Path(dir.resolve("vector.parquet").toUri()); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withType(schema) + .withWriterVersion(WriterVersion.PARQUET_1_0) + .withDictionaryEncoding(false) + .withCompressionCodec(CompressionCodecName.UNCOMPRESSED) + .withValidation(false) + .build()) { + for (int row = 0; row < 2; row++) { + Group group = new SimpleGroupFactory(schema).newGroup(); + Group vector = group.addGroup("repeatedInt"); + for (int i = 1; i <= 3; i++) vector.addGroup("list").append("element", row * 3 + i); + writer.write(group); + } + } + List messages = + TestUtils.readMessages(path, TestProtobuf.RepeatedIntMessage.class); + assertThat(messages).hasSize(2); + assertThat(messages.get(0).getRepeatedIntList()).containsExactly(1, 2, 3); + assertThat(messages.get(1).getRepeatedIntList()).containsExactly(4, 5, 6); + } @Test public void testAllTypes() throws Exception {