diff --git a/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingReadPlanFactory.java b/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingReadPlanFactory.java index eaeeca1d2dd9..0e902c56e307 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingReadPlanFactory.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingReadPlanFactory.java @@ -33,7 +33,6 @@ import org.apache.paimon.data.variant.PaimonShreddingUtils; import org.apache.paimon.data.variant.PaimonShreddingUtils.FieldToExtract; import org.apache.paimon.data.variant.VariantMetadataUtils; -import org.apache.paimon.data.variant.VariantPathSegment; import org.apache.paimon.format.shredding.ShreddingReadPlanFactory; import org.apache.paimon.types.ArrayType; import org.apache.paimon.types.DataField; @@ -52,14 +51,11 @@ import javax.annotation.Nullable; import java.util.ArrayList; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; import static org.apache.paimon.data.variant.Variant.METADATA; import static org.apache.paimon.data.variant.Variant.VALUE; -import static org.apache.paimon.data.variant.VariantMetadataUtils.path; import static org.apache.paimon.format.parquet.ParquetSchemaConverter.parquetListElementType; import static org.apache.paimon.format.parquet.ParquetSchemaConverter.parquetMapKeyValueType; @@ -124,53 +120,7 @@ public static Type clipParquetType(DataType logicalType, Type parquetType) { /** Clips a Variant Parquet field according to the logical Variant row read type. */ public static Type clipVariantType(RowType variantRowType, GroupType parquetType) { - if (!parquetType.containsField(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME)) { - return parquetType; - } - - boolean canClip = true; - Set fieldsToRead = new HashSet<>(); - for (DataField field : variantRowType.getFields()) { - String path = path(field.description()); - VariantPathSegment[] pathSegments = VariantPathSegment.parse(path); - if (pathSegments.length < 1) { - canClip = false; - break; - } - - // TODO: support nested column pruning. - VariantPathSegment pathSegment = pathSegments[0]; - if (pathSegment instanceof VariantPathSegment.ObjectExtraction) { - fieldsToRead.add(((VariantPathSegment.ObjectExtraction) pathSegment).getKey()); - } else { - canClip = false; - break; - } - } - - if (!canClip) { - return parquetType; - } - - List typedFieldsToRead = new ArrayList<>(); - GroupType typedValue = - parquetType.getType(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME).asGroupType(); - for (Type field : typedValue.getFields()) { - if (fieldsToRead.contains(field.getName())) { - typedFieldsToRead.add(field); - fieldsToRead.remove(field.getName()); - } - } - - List rowGroupFields = new ArrayList<>(); - rowGroupFields.add(parquetType.getType(PaimonShreddingUtils.METADATA_FIELD_NAME)); - if (!fieldsToRead.isEmpty()) { - rowGroupFields.add(parquetType.getType(PaimonShreddingUtils.VARIANT_VALUE_FIELD_NAME)); - } - if (!typedFieldsToRead.isEmpty()) { - rowGroupFields.add(typedValue.withNewFields(typedFieldsToRead)); - } - return parquetType.withNewFields(rowGroupFields); + return VariantShreddingTypePruner.clip(variantRowType, parquetType); } private static boolean containsVariantFields(DataType dataType) { diff --git a/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingTypePruner.java b/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingTypePruner.java new file mode 100644 index 000000000000..5561983bd828 --- /dev/null +++ b/paimon-format/src/main/java/org/apache/paimon/format/parquet/VariantShreddingTypePruner.java @@ -0,0 +1,251 @@ +/* + * 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.paimon.format.parquet; + +import org.apache.paimon.data.variant.PaimonShreddingUtils; +import org.apache.paimon.data.variant.VariantMetadataUtils; +import org.apache.paimon.data.variant.VariantPathSegment; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.RowType; + +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.Type; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.paimon.format.parquet.ParquetSchemaConverter.parquetListElementType; +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** + * Prunes a shredded Variant Parquet type according to a logical Variant projection. + * + *

It builds a trie from the requested object/array paths and recursively removes unneeded fields + * from the {@code typed_value} (and list element) groups while preserving {@code value} fallbacks + * when a requested path cannot be satisfied from typed columns. + * + *

Variant object keys are matched case-sensitively, independently of Parquet column-name + * resolution, because Variant path extraction downstream resolves keys exactly through {@code + * objectSchemaMap}. + */ +public class VariantShreddingTypePruner { + private static final String LIST_WRAPPER_NAME = "list"; + private static final String LIST_ELEMENT_NAME = "element"; + + @Nullable private final PathNode root; + + VariantShreddingTypePruner(RowType variantRowType) { + this.root = buildPathTree(variantRowType); + } + + /** + * Clips the given Parquet Variant type to only include fields needed for {@code + * variantRowType}. + * + * @param variantRowType the logical Variant projection row type + * @param parquetType the physical Parquet Variant type + * @return a clipped Parquet type + */ + public static Type clip(RowType variantRowType, GroupType parquetType) { + return new VariantShreddingTypePruner(variantRowType).clip(parquetType); + } + + private Type clip(GroupType parquetType) { + return clipShreddingRow(parquetType, root); + } + + /** A projection trie for Variant object paths and array element paths. */ + private static class PathNode { + private final Map children = new HashMap<>(); + private PathNode arrayElement; + private boolean keepAll; + + private PathNode getOrCreateChild(String key) { + return children.computeIfAbsent(key, k -> new PathNode()); + } + } + + @Nullable + private PathNode buildPathTree(RowType variantRowType) { + PathNode root = new PathNode(); + for (DataField field : variantRowType.getFields()) { + String path = VariantMetadataUtils.path(field.description()); + VariantPathSegment[] segments = VariantPathSegment.parse(path); + if (segments.length == 0) { + return null; + } + + PathNode node = root; + for (VariantPathSegment segment : segments) { + if (segment instanceof VariantPathSegment.ArrayExtraction) { + // Array indices cannot prune individual elements at the Parquet level, + // but we can still prune nested fields inside each array element. + if (node.arrayElement == null) { + node.arrayElement = new PathNode(); + } + node = node.arrayElement; + } else if (segment instanceof VariantPathSegment.ObjectExtraction) { + String key = ((VariantPathSegment.ObjectExtraction) segment).getKey(); + node = node.getOrCreateChild(key); + } else { + return null; + } + } + node.keepAll = true; + } + return root; + } + + private Type clipShreddingRow(Type type, PathNode node) { + if (type.isPrimitive() || node == null) { + return type; + } + + GroupType group = type.asGroupType(); + if (node.keepAll || !group.containsField(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME)) { + return group; + } + + List newFields = new ArrayList<>(); + if (group.containsField(PaimonShreddingUtils.METADATA_FIELD_NAME)) { + newFields.add(group.getType(PaimonShreddingUtils.METADATA_FIELD_NAME)); + } + + Type typedValue = group.getType(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME); + if (isCanonicalList(typedValue) && node.arrayElement != null) { + return clipListShreddingRow(group, node, newFields); + } else if (isObjectGroup(typedValue) && node.arrayElement == null) { + return clipObjectShreddingRow(group, node, newFields); + } else { + return group; + } + } + + private GroupType clipObjectShreddingRow(GroupType group, PathNode node, List newFields) { + Type typedValueType = group.getType(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME); + GroupType typedValue = typedValueType.asGroupType(); + // typed_value is an object group: prune by object key. + boolean needValue = false; + List clippedTypedFields = new ArrayList<>(); + Set requestedKeys = new HashSet<>(node.children.keySet()); + + for (Type field : typedValue.getFields()) { + String fieldName = field.getName(); + PathNode child = node.children.get(fieldName); + if (child == null) { + continue; + } + requestedKeys.remove(fieldName); + + checkArgument(!field.isPrimitive()); + Type clippedChild = clipShreddingRow(field, child); + clippedTypedFields.add(clippedChild); + } + + if (!requestedKeys.isEmpty()) { + needValue = true; + } + + if (needValue) { + checkArgument(group.containsField(PaimonShreddingUtils.VARIANT_VALUE_FIELD_NAME)); + newFields.add(group.getType(PaimonShreddingUtils.VARIANT_VALUE_FIELD_NAME)); + } + + if (!clippedTypedFields.isEmpty()) { + newFields.add(typedValue.withNewFields(clippedTypedFields)); + } + return group.withNewFields(newFields); + } + + private GroupType clipListShreddingRow(GroupType group, PathNode node, List newFields) { + Type type = group.getType(PaimonShreddingUtils.TYPED_VALUE_FIELD_NAME); + GroupType listGroup = type.asGroupType(); + if (node.arrayElement.keepAll) { + // The projection reads the whole array element (e.g. $.arr[0] read as VARIANT); + return group; + } + + Type elementType = parquetListElementType(listGroup); + Type clippedElement = clipShreddingRow(elementType.asGroupType(), node.arrayElement); + GroupType repeated = listGroup.getType(0).asGroupType(); + GroupType clippedRepeated = + repeated.withNewFields(Collections.singletonList(clippedElement)); + newFields.add(listGroup.withNewFields(Collections.singletonList(clippedRepeated))); + return group.withNewFields(newFields); + } + + /** + * Returns true if the given group follows the canonical three-level Parquet list layout. + * + *

The canonical layout is described in the Parquet spec: LogicalTypes#Lists + */ + private static boolean isCanonicalList(Type type) { + if (type.isPrimitive()) { + return false; + } + + GroupType listGroup = type.asGroupType(); + // 1. Must be a LIST logical type. + if (!(listGroup.getLogicalTypeAnnotation() + instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation)) { + return false; + } + + // 2. LIST group must have exactly one child named "list". + if (listGroup.getFieldCount() != 1) { + return false; + } + Type middle = listGroup.getType(0); + if (!LIST_WRAPPER_NAME.equals(middle.getName())) { + return false; + } + + // 3. The child must be a repeated group. + if (middle.isPrimitive() || middle.getRepetition() != Type.Repetition.REPEATED) { + return false; + } + GroupType repeatedWrapper = middle.asGroupType(); + + // 4. The repeated wrapper must contain exactly one child named "element". + if (repeatedWrapper.getFieldCount() != 1) { + return false; + } + + Type element = repeatedWrapper.getType(0); + return LIST_ELEMENT_NAME.equals(element.getName()); + } + + /** Returns true if the given group is a plain struct (not a Parquet list or map). */ + private static boolean isObjectGroup(Type type) { + if (type.isPrimitive()) { + return false; + } + GroupType groupType = type.asGroupType(); + return groupType.getLogicalTypeAnnotation() == null; + } +} diff --git a/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingReadTest.java b/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingReadTest.java index 7bad5f390990..dffb879fdfd7 100644 --- a/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingReadTest.java +++ b/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingReadTest.java @@ -281,6 +281,217 @@ public void testReadNestedVariantInArray(String shreddingSchema) throws Exceptio assertThat(result2.get(1).getArray(0).getRow(0, 1).getInt(0)).isEqualTo(5); } + @ParameterizedTest + @ValueSource( + strings = { + "null", + "{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"v\",\"type\":{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"a\",\"type\":{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"a\",\"type\":\"INT\"}," + + " {\"name\":\"b\",\"type\":\"INT\"}," + + " {\"name\":\"c\",\"type\":{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"d\",\"type\":\"STRING\"}," + + " {\"name\":\"e\",\"type\":\"INT\"}" + + " ]}}" + + " ]}}," + + " {\"name\":\"b\",\"type\":{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"a\",\"type\":\"INT\"}," + + " {\"name\":\"c\",\"type\":\"INT\"}" + + " ]}}," + + " {\"name\":\"d\",\"type\":\"INT\"}," + + " {\"name\":\"arr\",\"type\":{\"type\":\"ARRAY\",\"element\":{\"type\":\"ROW\",\"fields\":[" + + " {\"name\":\"x\",\"type\":\"INT\"}," + + " {\"name\":\"y\",\"type\":\"INT\"}" + + " ]}}}" + + " ]}}" + + "]}" + }) + public void testReadNestedVariantWithPruning(String shreddingSchema) throws Exception { + Options options = new Options(); + if (!shreddingSchema.equals("null")) { + options.set("parquet.variant.shreddingSchema", shreddingSchema); + } + ParquetFileFormat format = + new ParquetFileFormat(new FileFormatFactory.FormatContext(options, 1024, 1024)); + + RowType writeType = DataTypes.ROW(DataTypes.FIELD(0, "v", DataTypes.VARIANT())); + + FormatWriterFactory factory = format.createWriterFactory(writeType); + writeRows( + factory, + GenericRow.of( + GenericVariant.fromJson( + " {" + + " \"a\": {" + + " \"a\": 0," + + " \"b\": 1," + + " \"c\": {" + + " \"d\": \"hello\"," + + " \"e\": 2" + + " }" + + " }," + + " \"b\": {" + + " \"a\": 3," + + " \"c\": 4" + + " }," + + " \"c\": {" + + " \"a\": 5," + + " \"c\": 6" + + " }," + + " \"d\": 7," + + " \"arr\": [" + + " {\"x\":10,\"y\":11,\"z\":12}," + + " {\"x\":12,\"y\":13,\"z\":14}" + + " ]" + + " }"))); + + // case1: multiple nested projections with a missing top-level typed key ($.c.a). + RowType readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a.b") + .field(DataTypes.STRING(), "$.a.c.d") + .field(DataTypes.INT(), "$.b.a") + .field(DataTypes.INT(), "$.c.a") + .build())); + List result = readRows(format, readType); + assertThat(result).hasSize(1); + InternalRow projected = result.get(0).getRow(0, 4); + assertThat(projected.getInt(0)).isEqualTo(1); + assertThat(projected.getString(1).toString()).isEqualTo("hello"); + assertThat(projected.getInt(2)).isEqualTo(3); + assertThat(projected.getInt(3)).isEqualTo(5); + + // case2: read a whole shredded object ($.b). + readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field( + DataTypes.ROW( + DataTypes.FIELD(0, "a", DataTypes.INT()), + DataTypes.FIELD(1, "c", DataTypes.INT())), + "$.b") + .build())); + result = readRows(format, readType); + assertThat(result.get(0).getRow(0, 1).getRow(0, 2).getInt(1)).isEqualTo(4); + + // case3: read an object ($.c) that is not described by the shredding schema. + readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field( + DataTypes.ROW( + DataTypes.FIELD(0, "a", DataTypes.INT()), + DataTypes.FIELD(1, "c", DataTypes.INT())), + "$.c") + .build())); + result = readRows(format, readType); + assertThat(result.get(0).getRow(0, 1).getRow(0, 2).getInt(1)).isEqualTo(6); + + // case4: array index path: nested field inside each element can be pruned. + readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0].x") + .build())); + result = readRows(format, readType); + assertThat(result.get(0).getRow(0, 1).getInt(0)).isEqualTo(10); + + // case5: missing field inside array element falls back to binary value. + readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[1].z") + .build())); + result = readRows(format, readType); + assertThat(result.get(0).getRow(0, 1).getInt(0)).isEqualTo(14); + + // case6: read the whole array element as a VARIANT. + readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.VARIANT(), "$.arr[1]") + .build())); + result = readRows(format, readType); + assertThat(result.get(0).getRow(0, 1).getVariant(0).toJson()) + .isEqualTo("{\"x\":12,\"y\":13,\"z\":14}"); + } + + @Test + public void testReadHeterogeneousVariantObjectTypedAsArray() throws Exception { + // The shredding schema defines "v.a" as an object, but the JSON value is an array. + // Reading $.a[0].x must fall back to the binary value column. + Options options = new Options(); + options.set( + "parquet.variant.shreddingSchema", + "{\"type\":\"ROW\",\"fields\":[{\"name\":\"v\",\"type\":{\"type\":\"ROW\",\"fields\":[{\"name\":\"a\",\"type\":{\"type\":\"ROW\",\"fields\":[{\"name\":\"b\",\"type\":\"INT\"}]}}]}}]}"); + ParquetFileFormat format = + new ParquetFileFormat(new FileFormatFactory.FormatContext(options, 1024, 1024)); + + RowType writeType = DataTypes.ROW(DataTypes.FIELD(0, "v", DataTypes.VARIANT())); + FormatWriterFactory factory = format.createWriterFactory(writeType); + writeRows(factory, GenericRow.of(GenericVariant.fromJson("{\"a\":[{\"x\":1}]}"))); + + RowType readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a[0].x") + .build())); + List result = readRows(format, readType); + assertThat(result).hasSize(1); + assertThat(result.get(0).getRow(0, 1).getInt(0)).isEqualTo(1); + } + + @Test + public void testReadHeterogeneousVariantListTypedAsObject() throws Exception { + // The shredding schema defines "v.a" as a list, but the JSON value is an object. + // Reading $.a.x must fall back to the binary value column. + Options options = new Options(); + options.set( + "parquet.variant.shreddingSchema", + "{\"type\":\"ROW\",\"fields\":[{\"name\":\"v\",\"type\":{\"type\":\"ROW\",\"fields\":[{\"name\":\"a\",\"type\":{\"type\":\"ARRAY\",\"element\":\"INT\"}}]}}]}"); + ParquetFileFormat format = + new ParquetFileFormat(new FileFormatFactory.FormatContext(options, 1024, 1024)); + + RowType writeType = DataTypes.ROW(DataTypes.FIELD(0, "v", DataTypes.VARIANT())); + FormatWriterFactory factory = format.createWriterFactory(writeType); + writeRows(factory, GenericRow.of(GenericVariant.fromJson("{\"a\":{\"x\":2}}"))); + + RowType readType = + DataTypes.ROW( + DataTypes.FIELD( + 0, + "v", + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a.x") + .build())); + List result = readRows(format, readType); + assertThat(result).hasSize(1); + assertThat(result.get(0).getRow(0, 1).getInt(0)).isEqualTo(2); + } + protected List readRows(ParquetFileFormat format, RowType rowType) throws IOException { List result = new ArrayList<>(); diff --git a/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingTypePrunerTest.java b/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingTypePrunerTest.java new file mode 100644 index 000000000000..fa40f3650c10 --- /dev/null +++ b/paimon-format/src/test/java/org/apache/paimon/format/parquet/reader/VariantShreddingTypePrunerTest.java @@ -0,0 +1,542 @@ +/* + * 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.paimon.format.parquet.reader; + +import org.apache.paimon.data.variant.PaimonShreddingUtils; +import org.apache.paimon.data.variant.VariantMetadataUtils; +import org.apache.paimon.format.parquet.ParquetSchemaConverter; +import org.apache.paimon.format.parquet.VariantShreddingTypePruner; +import org.apache.paimon.types.ArrayType; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; + +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.BINARY; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32; +import static org.apache.parquet.schema.Type.Repetition.OPTIONAL; +import static org.apache.parquet.schema.Type.Repetition.REPEATED; +import static org.assertj.core.api.Assertions.assertThat; + +/** Unit tests for {@link VariantShreddingTypePruner}. */ +public class VariantShreddingTypePrunerTest { + + @Test + public void testObjectNestedFieldsPrunesTypedValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a.b") + .field(DataTypes.INT(), "$.a.c") + .field(DataTypes.INT(), "$.d") + .build(); + + GroupType parquetType = createParquetVariantType(nestedObjectShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("a", "d"); + + RowType physicalA = getRowType(physicalTypedValue, "a"); + RowType physicalATypedValue = getRowType(physicalA, "typed_value"); + assertThat(physicalATypedValue.getFieldNames()).containsExactly("b", "c"); + } + + @Test + public void testObjectMissingTopLevelFieldFallsBackToValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.z") + .build(); + + GroupType parquetType = createParquetVariantType(nestedObjectShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).contains("value"); + assertThat(physicalV.getFieldNames()).doesNotContain("typed_value"); + } + + @Test + public void testObjectMixedExistingAndMissingFieldsKeepsTypedValueAndValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a.b") + .field(DataTypes.INT(), "$.z") + .build(); + + GroupType parquetType = createParquetVariantType(nestedObjectShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).contains("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("a"); + + RowType physicalA = getRowType(physicalTypedValue, "a"); + RowType physicalATypedValue = getRowType(physicalA, "typed_value"); + assertThat(physicalATypedValue.getFieldNames()).containsExactly("b"); + } + + @Test + public void testObjectKeepAllKeepsWholeField() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field( + DataTypes.ROW( + DataTypes.FIELD(0, "b", DataTypes.INT()), + DataTypes.FIELD(1, "c", DataTypes.INT())), + "$.a") + .build(); + + GroupType parquetType = createParquetVariantType(nestedObjectShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("a"); + RowType physicalA = getRowType(physicalTypedValue, "a"); + RowType physicalATypedValue = getRowType(physicalA, "typed_value"); + assertThat(physicalATypedValue.getFieldNames()).containsExactly("b", "c"); + } + + private static RowType nestedObjectShreddingSchema() { + return PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ROW( + DataTypes.FIELD( + 0, + "a", + DataTypes.ROW( + DataTypes.FIELD(0, "b", DataTypes.INT()), + DataTypes.FIELD(1, "c", DataTypes.INT()))), + DataTypes.FIELD(1, "d", DataTypes.INT()), + DataTypes.FIELD(2, "e", DataTypes.INT()))); + } + + @Test + public void testCaseSensitive() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.a") + .build(); + + GroupType parquetType = + createParquetVariantType( + PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ROW( + DataTypes.FIELD(0, "a", DataTypes.INT()), + DataTypes.FIELD(1, "A", DataTypes.INT())))); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("a"); + } + + @Test + public void testArrayKeepAllArrayKeepsElementValueAndTypedValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.ARRAY(DataTypes.VARIANT()), "$.arr") + .build(); + + GroupType parquetType = createParquetVariantType(arrayShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("arr"); + + RowType physicalArr = getRowType(physicalTypedValue, "arr"); + ArrayType physicalArrList = + (ArrayType) physicalArr.getTypeAt(physicalArr.getFieldIndex("typed_value")); + RowType physicalArrElement = (RowType) physicalArrList.getElementType(); + assertThat(physicalArrElement.getFieldNames()).contains("value", "typed_value"); + RowType physicalArrElementTypedValue = getRowType(physicalArrElement, "typed_value"); + assertThat(physicalArrElementTypedValue.getFieldNames()).containsExactly("x", "y"); + } + + @Test + public void testArrayKeepAllElementKeepsElementValueAndTypedValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.VARIANT(), "$.arr[0]") + .build(); + + GroupType parquetType = createParquetVariantType(arrayShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("arr"); + + RowType physicalArr = getRowType(physicalTypedValue, "arr"); + ArrayType physicalArrList = + (ArrayType) physicalArr.getTypeAt(physicalArr.getFieldIndex("typed_value")); + RowType physicalArrElement = (RowType) physicalArrList.getElementType(); + assertThat(physicalArrElement.getFieldNames()).contains("value", "typed_value"); + RowType physicalArrElementTypedValue = getRowType(physicalArrElement, "typed_value"); + assertThat(physicalArrElementTypedValue.getFieldNames()).containsExactly("x", "y"); + } + + @Test + public void testArrayElementFieldPrunesElementTypedValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0].x") + .build(); + + GroupType parquetType = createParquetVariantType(arrayShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("arr"); + + RowType physicalArr = getRowType(physicalTypedValue, "arr"); + ArrayType physicalArrList = + (ArrayType) physicalArr.getTypeAt(physicalArr.getFieldIndex("typed_value")); + RowType physicalArrElement = (RowType) physicalArrList.getElementType(); + RowType physicalArrElementTypedValue = getRowType(physicalArrElement, "typed_value"); + assertThat(physicalArrElementTypedValue.getFieldNames()).containsExactly("x"); + } + + @Test + public void testArrayElementFieldMissingFallsBackToElementValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0].z") + .build(); + + GroupType parquetType = createParquetVariantType(arrayShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("arr"); + + RowType physicalArr = getRowType(physicalTypedValue, "arr"); + ArrayType physicalArrList = + (ArrayType) physicalArr.getTypeAt(physicalArr.getFieldIndex("typed_value")); + RowType physicalArrElement = (RowType) physicalArrList.getElementType(); + assertThat(physicalArrElement.getFieldNames()).contains("value"); + assertThat(physicalArrElement.getFieldNames()).doesNotContain("typed_value"); + } + + private static RowType arrayShreddingSchema() { + return PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ROW( + DataTypes.FIELD( + 0, + "arr", + DataTypes.ARRAY( + DataTypes.ROW( + DataTypes.FIELD(0, "x", DataTypes.INT()), + DataTypes.FIELD(1, "y", DataTypes.INT())))))); + } + + @Test + public void testArrayPrimitiveElementKeepAllKeepsElementShreddingRow() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0]") + .build(); + + GroupType parquetType = createParquetVariantType(primitiveArrayShreddingSchema()); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + RowType physicalTypedValue = getRowType(physicalV, "typed_value"); + assertThat(physicalTypedValue.getFieldNames()).containsExactly("arr"); + + RowType physicalArr = getRowType(physicalTypedValue, "arr"); + ArrayType physicalArrList = + (ArrayType) physicalArr.getTypeAt(physicalArr.getFieldIndex("typed_value")); + RowType physicalArrElement = (RowType) physicalArrList.getElementType(); + assertThat(physicalArrElement.getFieldNames()).contains("value", "typed_value"); + assertThat(physicalArrElement.getTypeAt(physicalArrElement.getFieldIndex("typed_value"))) + .isEqualTo(DataTypes.INT()); + } + + @Test + public void testTopLevelArrayElementFieldPrunesElementTypedValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$[0].x") + .build(); + + GroupType parquetType = + createParquetVariantType( + PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ARRAY( + DataTypes.ROW( + DataTypes.FIELD(0, "x", DataTypes.INT()), + DataTypes.FIELD(1, "y", DataTypes.INT()))))); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + ArrayType physicalList = + (ArrayType) physicalV.getTypeAt(physicalV.getFieldIndex("typed_value")); + RowType physicalElement = (RowType) physicalList.getElementType(); + RowType physicalElementTypedValue = getRowType(physicalElement, "typed_value"); + assertThat(physicalElementTypedValue.getFieldNames()).containsExactly("x"); + } + + @Test + public void testTopLevelArrayKeepAllVariantKeepsElementValue() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.VARIANT(), "$[0]") + .build(); + + GroupType parquetType = + createParquetVariantType( + PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ARRAY(DataTypes.VARIANT()))); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + assertThat(physicalV.getFieldNames()).contains("value"); + ArrayType physicalList = + (ArrayType) physicalV.getTypeAt(physicalV.getFieldIndex("typed_value")); + RowType physicalElement = (RowType) physicalList.getElementType(); + assertThat(physicalElement.getFieldNames()).containsExactly("value"); + } + + @Test + public void testTopLevelArrayNestedProjectionKeepsWholeList() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$[0].x") + .build(); + + GroupType parquetType = + createParquetVariantType( + PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ARRAY(DataTypes.VARIANT()))); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + RowType physicalV = toRowType(clipped); + + // The element has no typed_value column, so we cannot prune inside it. + // The whole list must be retained and the assembler reads x from the binary value. + assertThat(physicalV.getFieldNames()).doesNotContain("value"); + ArrayType physicalList = + (ArrayType) physicalV.getTypeAt(physicalV.getFieldIndex("typed_value")); + RowType physicalElement = (RowType) physicalList.getElementType(); + assertThat(physicalElement.getFieldNames()).containsExactly("value"); + } + + private static RowType primitiveArrayShreddingSchema() { + return PaimonShreddingUtils.variantShreddingSchema( + DataTypes.ROW(DataTypes.FIELD(0, "arr", DataTypes.ARRAY(DataTypes.INT())))); + } + + @Test + public void testTwoLevelArrayGroupElementKeepsWholeList() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0].x") + .build(); + + GroupType parquetType = createTwoLevelArrayVariantType(); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + GroupType actual = + clipped.asGroupType() + .getType("typed_value") + .asGroupType() + .getType(0) + .asGroupType() + .getType(1) + .asGroupType(); + GroupType expected = + parquetType + .asGroupType() + .getType("typed_value") + .asGroupType() + .getType(0) + .asGroupType() + .getType(1) + .asGroupType(); + assertThat(actual).isEqualTo(expected); + } + + private static GroupType createTwoLevelArrayVariantType() { + // x shredding row: value + typed_value + GroupType xShreddingRow = + Types.buildGroup(Type.Repetition.REQUIRED) + .optional(BINARY) + .named("value") + .optional(INT32) + .named("typed_value") + .named("x"); + + // y shredding row: value + typed_value + GroupType yShreddingRow = + Types.buildGroup(Type.Repetition.REQUIRED) + .optional(BINARY) + .named("value") + .optional(INT32) + .named("typed_value") + .named("y"); + + // element shredding row: value + typed_value(x, y) + GroupType elementTypedValue = + Types.buildGroup(OPTIONAL) + .addField(xShreddingRow) + .addField(yShreddingRow) + .named("typed_value"); + + GroupType elementRow = + Types.buildGroup(REPEATED) + .optional(BINARY) + .named("value") + .addField(elementTypedValue) + .named("element"); + + // Two-level list: the list group's immediate child is the repeated element row. + GroupType arrList = + Types.buildGroup(OPTIONAL) + .as(LogicalTypeAnnotation.listType()) + .addField(elementRow) + .named("typed_value"); + + // arr shredding row: value + typed_value(list) + GroupType arrShreddingRow = + Types.buildGroup(Type.Repetition.REQUIRED) + .optional(BINARY) + .named("value") + .addField(arrList) + .named("arr"); + + GroupType typedValue = + Types.buildGroup(OPTIONAL).addField(arrShreddingRow).named("typed_value"); + + return Types.buildGroup(OPTIONAL) + .required(BINARY) + .named("metadata") + .required(BINARY) + .named("value") + .addField(typedValue) + .named("v"); + } + + @Test + public void testTwoLevelPrimitiveArrayKeepAllKeepsWholeList() { + RowType variantRowType = + VariantMetadataUtils.VariantRowTypeBuilder.builder() + .field(DataTypes.INT(), "$.arr[0]") + .build(); + + GroupType parquetType = createTwoLevelPrimitiveArrayVariantType(); + Type clipped = VariantShreddingTypePruner.clip(variantRowType, parquetType); + GroupType actual = + clipped.asGroupType() + .getType("typed_value") + .asGroupType() + .getType(0) + .asGroupType() + .getType(1) + .asGroupType(); + GroupType expected = + parquetType + .asGroupType() + .getType("typed_value") + .asGroupType() + .getType(0) + .asGroupType() + .getType(1) + .asGroupType(); + assertThat(actual).isEqualTo(expected); + } + + private static GroupType createTwoLevelPrimitiveArrayVariantType() { + // Primitive array elements are still wrapped in a shredding row. + GroupType elementRow = + Types.buildGroup(Type.Repetition.REPEATED) + .optional(BINARY) + .id(0) + .named("value") + .optional(INT32) + .id(1) + .named("typed_value") + .named("element") + .withId(123); + + // Two-level list: the list group's immediate child is the repeated element row. + GroupType arrList = + Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.listType()) + .addField(elementRow) + .named("typed_value") + .withId(1); + + // arr shredding row: value + typed_value(list) + GroupType arrShreddingRow = + Types.buildGroup(Type.Repetition.REQUIRED) + .optional(BINARY) + .id(0) + .named("value") + .addField(arrList) + .named("arr") + .withId(0); + + GroupType typedValue = + Types.buildGroup(Type.Repetition.OPTIONAL) + .addField(arrShreddingRow) + .named("typed_value") + .withId(2); + + return Types.buildGroup(Type.Repetition.OPTIONAL) + .required(PrimitiveTypeName.BINARY) + .id(0) + .named("metadata") + .required(PrimitiveTypeName.BINARY) + .id(1) + .named("value") + .addField(typedValue) + .named("v") + .withId(0); + } + + private static GroupType createParquetVariantType(RowType shreddingSchema) { + DataField field = new DataField(0, "v", shreddingSchema); + return ParquetSchemaConverter.convertToParquetType(field).asGroupType(); + } + + private static RowType toRowType(Type type) { + return (RowType) ParquetSchemaConverter.convertToPaimonField(type).type(); + } + + private static RowType getRowType(RowType rowType, String fieldName) { + return (RowType) rowType.getTypeAt(rowType.getFieldIndex(fieldName)); + } +}