jzhuge commented on code in PR #18336:
URL: https://github.com/apache/iceberg/pull/18336#discussion_r4152038070
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestPruneColumns.java:
##########
@@ -305,6 +306,384 @@ 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() {
+ // event struct<contact struct<email, phone>, name struct<first, last>>:
+ // contact is partially projected (email only) and declared BEFORE name,
which is fully
+ // projected. The partially-projected sibling's inner pruning must survive
the later
+ // fully-projected sibling (its visit returns the original field object).
Review Comment:
Done, simplified to one line.
--
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]
