bryanck commented on code in PR #18336:
URL: https://github.com/apache/iceberg/pull/18336#discussion_r4156428074
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestPruneColumns.java:
##########
@@ -305,6 +306,381 @@ public void testVariant() {
assertThat(actual).as("Pruned schema should be
matched").isEqualTo(expected);
}
+ private static MessageType deepNestedFileSchema() {
+ return Types.buildMessage()
+ .addField(
+ Types.primitive(PrimitiveTypeName.INT64,
Type.Repetition.REQUIRED).id(1).named("id"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(3)
+ .named("own"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(5)
+ .named("x"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.INT64,
Type.Repetition.REQUIRED)
+ .id(7)
+ .named("leaf"))
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(8)
+ .named("big"))
+ .id(6)
+ .named("l3"))
+ .id(4)
+ .named("l2"))
+ .id(2)
+ .named("l1"))
+ .named("table");
+ }
+
+ @Test
+ public void testDeeplyNestedStructProjection() {
+ MessageType fileSchema = deepNestedFileSchema();
+
+ // project the deepest leaf only: intermediate structs must not widen back
to their full type
+ Schema leafProjection =
+ new Schema(
+ NestedField.required(
+ 2,
+ "l1",
+ StructType.of(
+ NestedField.required(
+ 4,
+ "l2",
+ StructType.of(
+ NestedField.required(
+ 6,
+ "l3",
+ StructType.of(
+ NestedField.required(7, "leaf",
LongType.get()))))))));
+
+ MessageType leafExpected =
+ Types.buildMessage()
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.INT64,
Type.Repetition.REQUIRED)
+ .id(7)
+ .named("leaf"))
+ .id(6)
+ .named("l3"))
+ .id(4)
+ .named("l2"))
+ .id(2)
+ .named("l1"))
+ .named("table");
+
+ MessageType leafActual = ParquetSchemaUtil.pruneColumns(fileSchema,
leafProjection);
+ assertThat(leafActual)
+ .as("Deep projection should not widen intermediate structs")
+ .isEqualTo(leafExpected);
+ }
+
+ @Test
+ public void testDeeplyNestedStructWhole() {
+ MessageType fileSchema = deepNestedFileSchema();
+
+ // project a nested struct itself: the whole struct is still read
+ Schema structProjection =
+ new Schema(
+ NestedField.required(
+ 2,
+ "l1",
+ StructType.of(
+ NestedField.required(
+ 4,
+ "l2",
+ StructType.of(
+ NestedField.required(
+ 6,
+ "l3",
+ StructType.of(
+ NestedField.required(7, "leaf",
LongType.get()),
+ NestedField.required(8, "big",
StringType.get()))),
+ NestedField.required(5, "x",
StringType.get()))))));
+
+ MessageType structExpected =
+ Types.buildMessage()
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(5)
+ .named("x"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.INT64,
Type.Repetition.REQUIRED)
+ .id(7)
+ .named("leaf"))
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+
.as(LogicalTypeAnnotation.stringType())
+ .id(8)
+ .named("big"))
+ .id(6)
+ .named("l3"))
+ .id(4)
+ .named("l2"))
+ .id(2)
+ .named("l1"))
+ .named("table");
+
+ MessageType structActual = ParquetSchemaUtil.pruneColumns(fileSchema,
structProjection);
+ assertThat(structActual)
+ .as("Projecting a nested struct keeps the whole struct")
+ .isEqualTo(structExpected);
+ }
+
+ @Test
+ public void testDeeplyNestedStructMixed() {
+ MessageType fileSchema = deepNestedFileSchema();
+
+ // project two leaves from different branches: both are kept, siblings are
dropped
+ Schema mixedProjection =
+ new Schema(
+ NestedField.required(
+ 2,
+ "l1",
+ StructType.of(
+ NestedField.required(3, "own", StringType.get()),
+ NestedField.required(
+ 4,
+ "l2",
+ StructType.of(
+ NestedField.required(
+ 6,
+ "l3",
+ StructType.of(
+ NestedField.required(7, "leaf",
LongType.get()))))))));
+
+ MessageType mixedExpected =
+ Types.buildMessage()
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(3)
+ .named("own"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(
+ PrimitiveTypeName.INT64,
Type.Repetition.REQUIRED)
+ .id(7)
+ .named("leaf"))
+ .id(6)
+ .named("l3"))
+ .id(4)
+ .named("l2"))
+ .id(2)
+ .named("l1"))
+ .named("table");
+
+ MessageType mixedActual = ParquetSchemaUtil.pruneColumns(fileSchema,
mixedProjection);
+ assertThat(mixedActual)
+ .as("Mixed projection keeps only the selected leaves")
+ .isEqualTo(mixedExpected);
+ }
+
+ @Test
+ public void testDeeplyNestedStructPartiallyProjectedBeforeFullyProjected() {
+ // a fully projected sibling (name) must not undo pruning of an earlier
sibling (contact)
+ MessageType fileSchema =
+ Types.buildMessage()
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(3)
+ .named("email"))
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(4)
+ .named("phone"))
+ .id(2)
+ .named("contact"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(6)
+ .named("first"))
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(7)
+ .named("last"))
+ .id(5)
+ .named("name"))
+ .id(1)
+ .named("event"))
+ .named("table");
+
+ Schema projection =
+ new Schema(
+ NestedField.required(
+ 1,
+ "event",
+ StructType.of(
+ NestedField.required(
+ 2,
+ "contact",
+ StructType.of(NestedField.required(3, "email",
StringType.get()))),
+ NestedField.required(
+ 5,
+ "name",
+ StructType.of(
+ NestedField.required(6, "first", StringType.get()),
+ NestedField.required(7, "last",
StringType.get()))))));
+
+ MessageType expected =
+ Types.buildMessage()
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(3)
+ .named("email"))
+ .id(2)
+ .named("contact"))
+ .addField(
+ Types.buildGroup(Type.Repetition.REQUIRED)
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(6)
+ .named("first"))
+ .addField(
+ Types.primitive(PrimitiveTypeName.BINARY,
Type.Repetition.REQUIRED)
+ .as(LogicalTypeAnnotation.stringType())
+ .id(7)
+ .named("last"))
+ .id(5)
+ .named("name"))
+ .id(1)
+ .named("event"))
+ .named("table");
+
+ MessageType actual = ParquetSchemaUtil.pruneColumns(fileSchema,
projection);
+ assertThat(actual)
+ .as("Inner pruning of an earlier sibling must survive a
fully-projected later sibling")
+ .isEqualTo(expected);
+ }
+
+ @Test
+ public void testDeeplyNestedStructInsideList() {
+ Schema schema =
+ new Schema(
+ NestedField.optional(
+ 1,
+ "events",
+ ListType.ofOptional(
+ 2,
+ StructType.of(
+ NestedField.optional(3, "id", LongType.get()),
+ NestedField.optional(
+ 4,
+ "payload",
+ StructType.of(
+ NestedField.optional(5, "a", StringType.get()),
+ NestedField.optional(6, "b",
StringType.get())))))));
+ MessageType fileSchema = ParquetSchemaUtil.convert(schema, "table");
+
+ // project events.payload.a: the element's id and payload.b are dropped
+ Schema projection =
+ new Schema(
+ NestedField.optional(
+ 1,
+ "events",
+ ListType.ofOptional(
+ 2,
+ StructType.of(
+ NestedField.optional(
+ 4,
+ "payload",
+ StructType.of(NestedField.optional(5, "a",
StringType.get())))))));
+ MessageType expected = ParquetSchemaUtil.convert(projection, "table");
+
+ MessageType actual = ParquetSchemaUtil.pruneColumns(fileSchema,
projection);
+ assertThat(actual).as("List element must prune to the projected
subtree").isEqualTo(expected);
+
+ // projecting the list itself keeps the whole element
+ MessageType wholeActual = ParquetSchemaUtil.pruneColumns(fileSchema,
schema);
+ assertThat(wholeActual).as("Whole-list projection is
unchanged").isEqualTo(fileSchema);
+ }
+
+ @Test
+ public void testDeeplyNestedStructInsideMap() {
Review Comment:
nit: the name of the test is slightly confusing, the value type in the map
is not deeply nested
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
