From 95839458d695547ceed212765fd90c70f67d2417 Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Tue, 21 Jul 2026 16:39:21 +0800 Subject: [PATCH 1/6] [core][spark] Support selected-key pushdown for shared-shredding MAP --- .../shredding/MapSharedShreddingReadPlan.java | 342 +++++++++++++++--- .../shredding/MapSharedShreddingUtils.java | 102 +++++- .../MapSharedShreddingReadPlanTest.java | 125 +++++++ .../MapSharedShreddingUtilsTest.java | 28 ++ .../apache/paimon/schema/SchemaManager.java | 57 ++- .../paimon/utils/FormatReaderMapping.java | 65 +++- .../paimon/schema/SchemaManagerTest.java | 68 ++++ .../table/MapSharedShreddingTableTest.java | 120 +++--- .../paimon/utils/FormatReaderMappingTest.java | 140 +++++++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ ...apSelectedKeysSharedShreddingE2ETest.scala | 21 ++ .../spark/execution/PaimonStrategy.scala | 15 +- .../PaimonSparkSessionExtensions.scala | 2 - ...lectedKeysSharedShreddingE2ETestBase.scala | 145 ++++++++ 18 files changed, 1219 insertions(+), 116 deletions(-) create mode 100644 paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-3.5/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala create mode 100644 paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETestBase.scala diff --git a/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlan.java b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlan.java index 521dbd5b5552..f09c58d685d3 100644 --- a/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlan.java +++ b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlan.java @@ -30,14 +30,17 @@ import org.apache.paimon.data.columnar.VectorizedColumnBatch; import org.apache.paimon.data.columnar.heap.CastedMapColumnVector; import org.apache.paimon.data.columnar.heap.HeapMapVector; +import org.apache.paimon.data.columnar.heap.HeapRowVector; import org.apache.paimon.data.columnar.writable.WritableColumnVector; import org.apache.paimon.types.DataField; -import org.apache.paimon.types.DataType; import org.apache.paimon.types.MapType; import org.apache.paimon.types.RowType; +import javax.annotation.Nullable; + import java.util.Arrays; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import static org.apache.paimon.utils.Preconditions.checkArgument; @@ -45,6 +48,8 @@ /** Read plan that rebuilds logical MAP values from shared-shredding physical ROW values. */ public class MapSharedShreddingReadPlan implements ShreddingReadPlan { + private static final int FIELD_MAPPING_POSITION = 0; + private final RowType logicalType; private final RowType physicalType; private final Map contextByFieldIndex; @@ -53,7 +58,7 @@ public MapSharedShreddingReadPlan( RowType logicalType, Map fieldMetas) { this.logicalType = logicalType; this.physicalType = MapSharedShreddingUtils.buildPhysicalReadType(logicalType, fieldMetas); - this.contextByFieldIndex = createContexts(logicalType, fieldMetas); + this.contextByFieldIndex = createContexts(logicalType, physicalType, fieldMetas); } @Override @@ -77,13 +82,30 @@ public ShreddingBatchAssembler batchAssembler() { } private static Map createContexts( - RowType logicalType, Map fieldMetas) { + RowType logicalType, + RowType physicalType, + Map fieldMetas) { Map contexts = new LinkedHashMap<>(); for (int i = 0; i < logicalType.getFieldCount(); i++) { DataField field = logicalType.getFields().get(i); MapSharedShreddingFieldMeta fieldMeta = fieldMetas.get(field.name()); - if (fieldMeta != null && field.type() instanceof MapType) { - contexts.put(i, new SharedShreddingContext(fieldMeta, field.type())); + if (fieldMeta == null) { + continue; + } + + RowType physicalStructType = (RowType) physicalType.getTypeAt(i); + SharedPhysicalContext physicalContext = + new SharedPhysicalContext(fieldMeta, physicalStructType); + if (field.type() instanceof MapType) { + contexts.put(i, new FullMapContext(physicalContext, (MapType) field.type())); + } else if (MapSelectedKeysMetadataUtils.isMapSelectedKeysField(field)) { + contexts.put( + i, + new SelectedKeysContext( + physicalContext, + fieldMeta, + (RowType) field.type(), + field.description())); } } return contexts; @@ -100,9 +122,8 @@ public VectorizedColumnBatch assemble(VectorizedColumnBatch physicalBatch) { logicalVectors[i] = physicalBatch.columns[i]; } else { logicalVectors[i] = - materializeLogicalMapVector( + context.materialize( (RowColumnVector) physicalBatch.columns[i], - context, physicalBatch.getNumRows()); } } @@ -111,7 +132,7 @@ public VectorizedColumnBatch assemble(VectorizedColumnBatch physicalBatch) { } private static CastedMapColumnVector materializeLogicalMapVector( - RowColumnVector physicalVector, SharedShreddingContext context, int rowCount) { + RowColumnVector physicalVector, FullMapContext context, int rowCount) { ColumnVector[] physicalChildren = physicalChildren(physicalVector); int totalElements = countLogicalMapElements(physicalVector, physicalChildren, context, rowCount); @@ -147,6 +168,46 @@ private static CastedMapColumnVector materializeLogicalMapVector( new WritableColumnVector[] {keyVector, valueVector})); } + private static ColumnVector materializeSelectedKeysRowVector( + RowColumnVector physicalVector, SelectedKeysContext context, int rowCount) { + ColumnVector[] physicalChildren = physicalChildren(physicalVector); + RowType selectedKeysType = context.selectedKeysType; + WritableColumnVector[] selectedKeyVectors = + new WritableColumnVector[selectedKeysType.getFieldCount()]; + for (int i = 0; i < selectedKeyVectors.length; i++) { + selectedKeyVectors[i] = + ColumnVectorUtils.createWritableColumnVector( + rowCount, selectedKeysType.getTypeAt(i)); + } + HeapRowVector rowVector = new HeapRowVector(rowCount, selectedKeyVectors); + + for (int row = 0; row < rowCount; row++) { + if (physicalVector.isNullAt(row)) { + rowVector.appendNull(); + continue; + } + + InternalArray fieldMapping = fieldMapping(physicalChildren, row, context.physical); + InternalMap overflow = null; + for (int ordinal = 0; ordinal < selectedKeyVectors.length; ordinal++) { + if (context.selectedKeyOverflow[ordinal] && overflow == null) { + overflow = overflowMap(physicalChildren, row, context.physical); + } + appendSelectedKeyValue( + physicalChildren, + fieldMapping, + context.selectedKeyOverflow[ordinal] ? overflow : null, + row, + context, + ordinal, + selectedKeyVectors[ordinal]); + } + rowVector.appendRow(); + } + + return ColumnVectorUtils.createReadableColumnVector(selectedKeysType, rowVector); + } + private static ColumnVector[] physicalChildren(RowColumnVector physicalVector) { ColumnVector[] physicalChildren = physicalVector.getChildren(); if (physicalChildren != null) { @@ -158,7 +219,7 @@ private static ColumnVector[] physicalChildren(RowColumnVector physicalVector) { private static int countLogicalMapElements( RowColumnVector physicalVector, ColumnVector[] physicalChildren, - SharedShreddingContext context, + FullMapContext context, int rowCount) { int totalElements = 0; for (int row = 0; row < rowCount; row++) { @@ -170,20 +231,20 @@ private static int countLogicalMapElements( } private static int countLogicalMapElements( - ColumnVector[] physicalChildren, int row, SharedShreddingContext context) { + ColumnVector[] physicalChildren, int row, FullMapContext context) { int count = 0; - InternalArray fieldMapping = fieldMapping(physicalChildren, row, context); - for (int column = 0; column < context.numColumns; column++) { - if (logicalFieldName(fieldMapping, column, context) != null) { + InternalArray fieldMapping = fieldMapping(physicalChildren, row, context.physical); + for (int column = 0; column < context.physical.numColumns; column++) { + if (logicalFieldName(fieldMapping, column, context.physical) != null) { count++; } } - InternalMap overflow = overflowMap(physicalChildren, row, context); + InternalMap overflow = overflowMap(physicalChildren, row, context.physical); if (overflow != null) { InternalArray keys = overflow.keyArray(); for (int i = 0; i < overflow.size(); i++) { - if (context.nameById.containsKey(keys.getInt(i))) { + if (context.physical.nameById.containsKey(keys.getInt(i))) { count++; } } @@ -194,32 +255,32 @@ private static int countLogicalMapElements( private static int appendLogicalMapElements( ColumnVector[] physicalChildren, int row, - SharedShreddingContext context, + FullMapContext context, WritableColumnVector keyVector, WritableColumnVector valueVector) { int count = 0; - InternalArray fieldMapping = fieldMapping(physicalChildren, row, context); - for (int column = 0; column < context.numColumns; column++) { - BinaryString fieldName = logicalFieldName(fieldMapping, column, context); + InternalArray fieldMapping = fieldMapping(physicalChildren, row, context.physical); + for (int column = 0; column < context.physical.numColumns; column++) { + BinaryString fieldName = logicalFieldName(fieldMapping, column, context.physical); if (fieldName == null) { continue; } context.keyConverter.append(fieldName, keyVector); - ColumnVector valueColumn = physicalChildren[column + 1]; - appendValue(context, valueColumn, row, valueVector); + ColumnVector valueColumn = physicalColumn(physicalChildren, column, context.physical); + appendValue(context.valueConverter, valueColumn, row, valueVector); count++; } - InternalMap overflow = overflowMap(physicalChildren, row, context); + InternalMap overflow = overflowMap(physicalChildren, row, context.physical); if (overflow != null) { InternalArray keys = overflow.keyArray(); InternalArray values = overflow.valueArray(); for (int i = 0; i < overflow.size(); i++) { - BinaryString fieldName = context.nameById.get(keys.getInt(i)); + BinaryString fieldName = context.physical.nameById.get(keys.getInt(i)); if (fieldName != null) { context.keyConverter.append(fieldName, keyVector); - appendValue(context, values, i, valueVector); + appendValue(context.valueConverter, values, i, valueVector); count++; } } @@ -227,50 +288,112 @@ private static int appendLogicalMapElements( return count; } + private static void appendSelectedKeyValue( + ColumnVector[] physicalChildren, + InternalArray fieldMapping, + @Nullable InternalMap overflow, + int row, + SelectedKeysContext context, + int ordinal, + WritableColumnVector valueVector) { + int fieldId = context.selectedKeyFieldIds[ordinal]; + if (fieldId < 0) { + valueVector.appendNull(); + return; + } + + int[] candidateColumns = context.selectedKeyColumns[ordinal]; + for (int i = 0; i < candidateColumns.length; i++) { + int column = candidateColumns[i]; + if (mappedFieldId(fieldMapping, column) == fieldId) { + appendValue( + context.selectedValueConverters[ordinal], + physicalColumn(physicalChildren, column, context.physical), + row, + valueVector); + return; + } + } + + if (overflow != null) { + InternalArray keys = overflow.keyArray(); + InternalArray values = overflow.valueArray(); + for (int i = 0; i < overflow.size(); i++) { + if (!keys.isNullAt(i) && keys.getInt(i) == fieldId) { + appendValue(context.selectedValueConverters[ordinal], values, i, valueVector); + return; + } + } + } + + valueVector.appendNull(); + } + private static void appendValue( - SharedShreddingContext context, + RowToColumnConverter.ElementConverter converter, ColumnVector source, int row, WritableColumnVector valueVector) { if (source.isNullAt(row)) { valueVector.appendNull(); } else { - context.valueConverter.append(source, row, valueVector); + converter.append(source, row, valueVector); } } private static void appendValue( - SharedShreddingContext context, + RowToColumnConverter.ElementConverter converter, InternalArray source, int pos, WritableColumnVector valueVector) { if (source.isNullAt(pos)) { valueVector.appendNull(); } else { - context.valueConverter.append(source, pos, valueVector); + converter.append(source, pos, valueVector); } } private static InternalArray fieldMapping( - ColumnVector[] physicalChildren, int row, SharedShreddingContext context) { - InternalArray fieldMapping = ((ArrayColumnVector) physicalChildren[0]).getArray(row); + ColumnVector[] physicalChildren, int row, SharedPhysicalContext context) { + InternalArray fieldMapping = + ((ArrayColumnVector) physicalChildren[FIELD_MAPPING_POSITION]).getArray(row); checkArgument( fieldMapping.size() == context.numColumns, "Shared-shredding field mapping size %s does not match metadata num columns %s.", fieldMapping.size(), context.numColumns); + for (int column = 0; column < fieldMapping.size(); column++) { + checkArgument( + !fieldMapping.isNullAt(column), + "Shared-shredding field mapping must not contain null at column %s.", + column); + } return fieldMapping; } private static BinaryString logicalFieldName( - InternalArray fieldMapping, int column, SharedShreddingContext context) { - int fieldId = fieldMapping.isNullAt(column) ? -1 : fieldMapping.getInt(column); + InternalArray fieldMapping, int column, SharedPhysicalContext context) { + int fieldId = mappedFieldId(fieldMapping, column); return fieldId < 0 ? null : context.nameById.get(fieldId); } + private static int mappedFieldId(InternalArray fieldMapping, int column) { + return fieldMapping.getInt(column); + } + + private static ColumnVector physicalColumn( + ColumnVector[] physicalChildren, int column, SharedPhysicalContext context) { + int position = context.physicalColumnPositions[column]; + checkArgument( + position >= 0, + "Shared-shredding physical column %s was not included in the read schema.", + MapSharedShreddingDefine.physicalColumnName(column)); + return physicalChildren[position]; + } + private static InternalMap overflowMap( - ColumnVector[] physicalChildren, int row, SharedShreddingContext context) { - if (context.overflowPosition >= physicalChildren.length) { + ColumnVector[] physicalChildren, int row, SharedPhysicalContext context) { + if (context.overflowPosition < 0 || context.overflowPosition >= physicalChildren.length) { return null; } @@ -278,27 +401,156 @@ private static InternalMap overflowMap( return overflowVector.isNullAt(row) ? null : ((MapColumnVector) overflowVector).getMap(row); } - private static class SharedShreddingContext { + private interface SharedShreddingContext { + + ColumnVector materialize(RowColumnVector physicalVector, int rowCount); + } + + private static class SharedPhysicalContext { - private final MapType mapType; private final Map nameById; - private final RowToColumnConverter.ElementConverter keyConverter; - private final RowToColumnConverter.ElementConverter valueConverter; private final int numColumns; + private final int[] physicalColumnPositions; private final int overflowPosition; - private SharedShreddingContext(MapSharedShreddingFieldMeta fieldMeta, DataType fieldType) { - this.mapType = (MapType) fieldType; + private SharedPhysicalContext( + MapSharedShreddingFieldMeta fieldMeta, RowType physicalStructType) { this.nameById = new LinkedHashMap<>(); for (Map.Entry entry : fieldMeta.nameToId().entrySet()) { this.nameById.put(entry.getValue(), BinaryString.fromString(entry.getKey())); } - this.keyConverter = - RowToColumnConverter.createElementConverter(this.mapType.getKeyType()); - this.valueConverter = - RowToColumnConverter.createElementConverter(this.mapType.getValueType()); this.numColumns = fieldMeta.numColumns(); - this.overflowPosition = fieldMeta.numColumns() + 1; + checkArgument( + physicalStructType.getFieldCount() > 0 + && MapSharedShreddingDefine.FIELD_MAPPING.equals( + physicalStructType.getFieldNames().get(FIELD_MAPPING_POSITION)), + "Shared-shredding physical struct must start with %s.", + MapSharedShreddingDefine.FIELD_MAPPING); + this.physicalColumnPositions = physicalColumnPositions(physicalStructType, numColumns); + this.overflowPosition = overflowPosition(physicalStructType); + } + + private static int[] physicalColumnPositions(RowType physicalStructType, int numColumns) { + int[] positions = new int[numColumns]; + Arrays.fill(positions, -1); + int physicalColumnEnd = overflowPosition(physicalStructType); + if (physicalColumnEnd < 0) { + physicalColumnEnd = physicalStructType.getFieldCount(); + } + for (int i = FIELD_MAPPING_POSITION + 1; i < physicalColumnEnd; i++) { + String fieldName = physicalStructType.getFieldNames().get(i); + int column = physicalColumnIndex(fieldName); + checkArgument( + column >= 0, + "Unexpected shared-shredding physical field %s at position %s.", + fieldName, + i); + checkArgument( + column < numColumns, + "Shared-shredding physical column %s exceeds metadata num columns %s.", + fieldName, + numColumns); + positions[column] = i; + } + return positions; + } + + private static int overflowPosition(RowType physicalStructType) { + int lastPosition = physicalStructType.getFieldCount() - 1; + if (lastPosition <= FIELD_MAPPING_POSITION) { + return -1; + } + return MapSharedShreddingDefine.OVERFLOW.equals( + physicalStructType.getFieldNames().get(lastPosition)) + ? lastPosition + : -1; + } + + private static int physicalColumnIndex(String fieldName) { + String prefix = "__col_"; + if (!fieldName.startsWith(prefix)) { + return -1; + } + return Integer.parseInt(fieldName.substring(prefix.length())); + } + } + + private static class FullMapContext implements SharedShreddingContext { + + private final SharedPhysicalContext physical; + private final MapType mapType; + private final RowToColumnConverter.ElementConverter keyConverter; + private final RowToColumnConverter.ElementConverter valueConverter; + + private FullMapContext(SharedPhysicalContext physical, MapType mapType) { + this.physical = physical; + this.mapType = mapType; + this.keyConverter = RowToColumnConverter.createElementConverter(mapType.getKeyType()); + this.valueConverter = + RowToColumnConverter.createElementConverter(mapType.getValueType()); + } + + @Override + public ColumnVector materialize(RowColumnVector physicalVector, int rowCount) { + return materializeLogicalMapVector(physicalVector, this, rowCount); + } + } + + private static class SelectedKeysContext implements SharedShreddingContext { + + private final SharedPhysicalContext physical; + private final RowType selectedKeysType; + private final int[] selectedKeyFieldIds; + private final int[][] selectedKeyColumns; + private final boolean[] selectedKeyOverflow; + private final RowToColumnConverter.ElementConverter[] selectedValueConverters; + + private SelectedKeysContext( + SharedPhysicalContext physical, + MapSharedShreddingFieldMeta fieldMeta, + RowType selectedKeysType, + String selectedKeysMetadata) { + this.physical = physical; + this.selectedKeysType = selectedKeysType; + List selectedKeys = + MapSelectedKeysMetadataUtils.selectedKeys(selectedKeysMetadata); + checkArgument( + selectedKeys.size() == selectedKeysType.getFieldCount(), + "Selected-key metadata size %s does not match selected ROW field count %s.", + selectedKeys.size(), + selectedKeysType.getFieldCount()); + this.selectedKeyFieldIds = new int[selectedKeys.size()]; + this.selectedKeyColumns = new int[selectedKeys.size()][]; + this.selectedKeyOverflow = new boolean[selectedKeys.size()]; + this.selectedValueConverters = + new RowToColumnConverter.ElementConverter[selectedKeys.size()]; + for (int i = 0; i < selectedKeys.size(); i++) { + Integer fieldId = fieldMeta.nameToId().get(selectedKeys.get(i)); + this.selectedKeyFieldIds[i] = fieldId == null ? -1 : fieldId; + this.selectedKeyColumns[i] = + fieldId == null ? new int[0] : candidateColumns(fieldMeta, fieldId); + this.selectedKeyOverflow[i] = + fieldId != null && fieldMeta.overflowFieldSet().contains(fieldId); + this.selectedValueConverters[i] = + RowToColumnConverter.createElementConverter(selectedKeysType.getTypeAt(i)); + } + } + + @Override + public ColumnVector materialize(RowColumnVector physicalVector, int rowCount) { + return materializeSelectedKeysRowVector(physicalVector, this, rowCount); + } + + private static int[] candidateColumns(MapSharedShreddingFieldMeta fieldMeta, int fieldId) { + List columns = fieldMeta.fieldToColumns().get(fieldId); + if (columns == null || columns.isEmpty()) { + return new int[0]; + } + int[] result = new int[columns.size()]; + for (int i = 0; i < columns.size(); i++) { + result[i] = columns.get(i); + } + return result; } } } diff --git a/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingUtils.java b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingUtils.java index c2526b03bbde..6fc2d13b5a72 100644 --- a/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingUtils.java +++ b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingUtils.java @@ -50,6 +50,8 @@ import java.util.TreeSet; import java.util.stream.Collectors; +import static org.apache.paimon.utils.Preconditions.checkArgument; + /** * Utility functions for the shared-shredding MAP storage layout. * @@ -115,18 +117,30 @@ public static RowType buildPhysicalReadType( for (DataField logicalReadField : logicalReadType.getFields()) { MapSharedShreddingFieldMeta fieldMeta = sharedShreddingFieldMetas.get(logicalReadField.name()); - if (fieldMeta == null || !(logicalReadField.type() instanceof MapType)) { + if (fieldMeta == null) { physicalReadFields.add(logicalReadField); continue; } - MapType mapType = (MapType) logicalReadField.type(); - DataType physicalType = - buildSpecificPhysicalStructType( - mapType.getValueType(), - fieldMeta.numColumns(), - !fieldMeta.overflowFieldSet().isEmpty()) - .copy(logicalReadField.type().isNullable()); + DataType valueType; + DataType physicalType; + if (MapSelectedKeysMetadataUtils.isMapSelectedKeysField(logicalReadField)) { + // recall partial key with shared shredding pushdown + valueType = selectedKeysValueType((RowType) logicalReadField.type()); + physicalType = + buildSpecificPhysicalStructType( + valueType, + selectedPhysicalColumnIds(logicalReadField, fieldMeta), + selectedKeysIncludeOverflow(logicalReadField, fieldMeta)) + .copy(logicalReadField.type().isNullable()); + } else { + // recall whole field without shared shredding pushdown + valueType = ((MapType) logicalReadField.type()).getValueType(); + physicalType = + buildPhysicalStructType(valueType, fieldMeta.numColumns()) + .copy(logicalReadField.type().isNullable()); + } + physicalReadFields.add(logicalReadField.newType(physicalType)); converted = true; } @@ -135,6 +149,19 @@ public static RowType buildPhysicalReadType( : logicalReadType; } + private static DataType selectedKeysValueType(RowType selectedKeysType) { + checkArgument( + selectedKeysType.getFieldCount() > 0, + "Selected-key MAP read type must contain at least one field."); + DataType valueType = selectedKeysType.getTypeAt(0); + for (int i = 1; i < selectedKeysType.getFieldCount(); i++) { + checkArgument( + selectedKeysType.getTypeAt(i).equalsIgnoreNullable(valueType), + "Selected-key MAP fields must have the same value type."); + } + return valueType; + } + public static Map buildColumnToNumColumns( List shreddingFieldNames, CoreOptions options) { Map fieldToNumColumns = new HashMap<>(); @@ -245,21 +272,21 @@ public static boolean hasShreddingMetadata(@Nullable Map metadat } private static RowType buildPhysicalStructType(DataType valueType, int numColumns) { - RowType.Builder builder = RowType.builder(); - builder.field(MapSharedShreddingDefine.FIELD_MAPPING, new ArrayType(new IntType())); - for (int i = 0; i < numColumns; i++) { - builder.field(MapSharedShreddingDefine.physicalColumnName(i), valueType); - } - builder.field(MapSharedShreddingDefine.OVERFLOW, new MapType(new IntType(), valueType)); - return builder.build(); + return buildSpecificPhysicalStructType(valueType, physicalColumnIds(numColumns), true); } - private static RowType buildSpecificPhysicalStructType( - DataType valueType, int numColumns, boolean includeOverflow) { + public static RowType buildSpecificPhysicalStructType( + DataType valueType, Set physicalColumnIds, boolean includeOverflow) { + return innerBuildSpecificPhysicalStructType( + valueType, new ArrayList<>(new TreeSet<>(physicalColumnIds)), includeOverflow); + } + + private static RowType innerBuildSpecificPhysicalStructType( + DataType valueType, List sortedColumns, boolean includeOverflow) { RowType.Builder builder = RowType.builder(); builder.field(MapSharedShreddingDefine.FIELD_MAPPING, new ArrayType(new IntType())); - for (int i = 0; i < numColumns; i++) { - builder.field(MapSharedShreddingDefine.physicalColumnName(i), valueType); + for (Integer column : sortedColumns) { + builder.field(MapSharedShreddingDefine.physicalColumnName(column), valueType); } if (includeOverflow) { builder.field(MapSharedShreddingDefine.OVERFLOW, new MapType(new IntType(), valueType)); @@ -267,6 +294,43 @@ private static RowType buildSpecificPhysicalStructType( return builder.build(); } + private static Set physicalColumnIds(int numColumns) { + Set physicalColumnIds = new TreeSet<>(); + for (int i = 0; i < numColumns; i++) { + physicalColumnIds.add(i); + } + return physicalColumnIds; + } + + private static Set selectedPhysicalColumnIds( + DataField selectedKeysField, MapSharedShreddingFieldMeta fieldMeta) { + Set selectedColumns = new TreeSet<>(); + for (String selectedKey : + MapSelectedKeysMetadataUtils.selectedKeys(selectedKeysField.description())) { + Integer fieldId = fieldMeta.nameToId().get(selectedKey); + if (fieldId == null) { + continue; + } + List columns = fieldMeta.fieldToColumns().get(fieldId); + if (columns != null) { + selectedColumns.addAll(columns); + } + } + return selectedColumns; + } + + private static boolean selectedKeysIncludeOverflow( + DataField selectedKeysField, MapSharedShreddingFieldMeta fieldMeta) { + for (String selectedKey : + MapSelectedKeysMetadataUtils.selectedKeys(selectedKeysField.description())) { + Integer fieldId = fieldMeta.nameToId().get(selectedKey); + if (fieldId != null && fieldMeta.overflowFieldSet().contains(fieldId)) { + return true; + } + } + return false; + } + private static Map> sortedFieldColumns( Map> fieldToColumns) { Map> result = new TreeMap<>(); diff --git a/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlanTest.java b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlanTest.java index 943bf94dbf75..6b781be4177e 100644 --- a/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlanTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlanTest.java @@ -20,10 +20,12 @@ import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.InternalMap; +import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.columnar.BytesColumnVector; import org.apache.paimon.data.columnar.ColumnVector; import org.apache.paimon.data.columnar.LongColumnVector; import org.apache.paimon.data.columnar.MapColumnVector; +import org.apache.paimon.data.columnar.RowColumnVector; import org.apache.paimon.data.columnar.VectorizedColumnBatch; import org.apache.paimon.data.columnar.heap.HeapArrayVector; import org.apache.paimon.data.columnar.heap.HeapIntVector; @@ -35,13 +37,16 @@ import org.junit.jupiter.api.Test; +import java.util.Arrays; import java.util.Collections; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.TreeMap; import java.util.TreeSet; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link MapSharedShreddingReadPlan}. */ class MapSharedShreddingReadPlanTest { @@ -88,6 +93,19 @@ void testReadOverflowOnlyWhenOverflowColumnExists() { assertThat(restored.valueArray().getLong(0)).isEqualTo(30L); } + @Test + void testFieldMappingRejectsNullElement() { + MapSharedShreddingFieldMeta fieldMeta = + new MapSharedShreddingFieldMeta( + nameToId("a", 0), Collections.emptyMap(), new TreeSet(), 2, 1); + HeapRowVector physicalMap = + rowVector(fieldMappingWithNull(0, null), longVector(10L), longVector(null)); + + assertThatThrownBy(() -> readMap(fieldMeta, physicalMap)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("field mapping must not contain null"); + } + @Test void testAssembledMapVectorExposesKeyValueChildren() { MapSharedShreddingFieldMeta fieldMeta = @@ -114,6 +132,59 @@ void testAssembledMapVectorExposesKeyValueChildren() { assertThat(mapVector.getMap(0).size()).isEqualTo(2); } + @Test + void testReadSelectedKeysAsLogicalRowDirectly() { + MapSharedShreddingFieldMeta fieldMeta = + new MapSharedShreddingFieldMeta( + nameToId("key1", 0, "key2", 1, "cold", 2), + fieldToColumns( + 0, Collections.singletonList(0), 1, Collections.singletonList(1)), + new TreeSet(Collections.singletonList(1)), + 2, + 2); + HeapRowVector physicalMap = + rowVector( + fieldMapping(0, -1), + longVector(10L), + longVector(null), + overflowMap(1, 20L)); + + MapSharedShreddingReadPlan readPlan = selectedKeysReadPlan(fieldMeta); + RowColumnVector selectedKeysVector = assembleSelectedKeysVector(readPlan, physicalMap); + InternalRow selectedKeys = selectedKeysVector.getRow(0); + + assertThat(selectedKeys.getLong(0)).isEqualTo(10L); + assertThat(selectedKeys.getLong(1)).isEqualTo(20L); + assertThat(selectedKeys.isNullAt(2)).isTrue(); + } + + @Test + void testReadSelectedKeysFromPrunedPhysicalColumns() { + MapSharedShreddingFieldMeta fieldMeta = + new MapSharedShreddingFieldMeta( + nameToId("key1", 0, "key2", 1, "cold", 2), + fieldToColumns( + 0, Collections.singletonList(2), 2, Collections.singletonList(0)), + new TreeSet(Collections.singletonList(1)), + 4, + 2); + + MapSharedShreddingReadPlan readPlan = selectedKeysReadPlan(fieldMeta); + RowType physicalMapType = (RowType) readPlan.physicalRowType().getTypeAt(0); + assertThat(physicalMapType.getFieldNames()) + .containsExactly("__field_mapping", "__col_2", "__overflow"); + + HeapRowVector physicalMap = + rowVector(fieldMapping(-1, -1, 0, -1), longVector(10L), overflowMap(1, 20L)); + + RowColumnVector selectedKeysVector = assembleSelectedKeysVector(readPlan, physicalMap); + InternalRow selectedKeys = selectedKeysVector.getRow(0); + + assertThat(selectedKeys.getLong(0)).isEqualTo(10L); + assertThat(selectedKeys.getLong(1)).isEqualTo(20L); + assertThat(selectedKeys.isNullAt(2)).isTrue(); + } + private static InternalMap readMap( MapSharedShreddingFieldMeta fieldMeta, HeapRowVector physicalMap) { return assembleMapVector(fieldMeta, physicalMap).getMap(0); @@ -132,6 +203,24 @@ private static MapColumnVector assembleMapVector( return (MapColumnVector) logicalBatch.columns[0]; } + private static RowColumnVector assembleSelectedKeysVector( + MapSharedShreddingReadPlan readPlan, HeapRowVector physicalMap) { + assertThat(readPlan.physicalRowType().getTypeAt(0)).isInstanceOf(RowType.class); + + VectorizedColumnBatch physicalBatch = + new VectorizedColumnBatch(new ColumnVector[] {physicalMap}); + physicalBatch.setNumRows(1); + VectorizedColumnBatch logicalBatch = readPlan.batchAssembler().assemble(physicalBatch); + return (RowColumnVector) logicalBatch.columns[0]; + } + + private static MapSharedShreddingReadPlan selectedKeysReadPlan( + MapSharedShreddingFieldMeta fieldMeta) { + Map fieldMetas = new LinkedHashMap<>(); + fieldMetas.put("metrics", fieldMeta); + return new MapSharedShreddingReadPlan(selectedKeysLogicalType(), fieldMetas); + } + private static RowType logicalType() { return DataTypes.ROW( DataTypes.FIELD( @@ -140,6 +229,19 @@ private static RowType logicalType() { DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.BIGINT()))); } + private static RowType selectedKeysLogicalType() { + RowType selectedKeysType = + DataTypes.ROW( + DataTypes.FIELD(0, "0", DataTypes.BIGINT()), + DataTypes.FIELD(1, "1", DataTypes.BIGINT()), + DataTypes.FIELD(2, "2", DataTypes.BIGINT())); + return DataTypes.ROW( + MapSelectedKeysMetadataUtils.withSelectedKeys( + DataTypes.FIELD(0, "metrics", selectedKeysType), + selectedKeysType, + Arrays.asList("key1", "key2", "missing"))); + } + private static HeapRowVector rowVector(ColumnVector... children) { HeapRowVector vector = new HeapRowVector(1, children); vector.appendRow(); @@ -156,6 +258,20 @@ private static HeapArrayVector fieldMapping(int... ids) { return vector; } + private static HeapArrayVector fieldMappingWithNull(Integer... ids) { + HeapIntVector child = new HeapIntVector(ids.length); + for (Integer id : ids) { + if (id == null) { + child.appendNull(); + } else { + child.appendInt(id); + } + } + HeapArrayVector vector = new HeapArrayVector(1, child); + vector.putOffsetLength(0, 0, ids.length); + return vector; + } + private static HeapLongVector longVector(Long value) { HeapLongVector vector = new HeapLongVector(1); if (value == null) { @@ -182,4 +298,13 @@ private static Map nameToId(Object... pairs) { } return result; } + + @SuppressWarnings("unchecked") + private static Map> fieldToColumns(Object... pairs) { + Map> result = new TreeMap<>(); + for (int i = 0; i < pairs.length; i += 2) { + result.put((Integer) pairs[i], (List) pairs[i + 1]); + } + return result; + } } diff --git a/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingUtilsTest.java b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingUtilsTest.java index 9c86109efa87..dd628bc5a37f 100644 --- a/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingUtilsTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingUtilsTest.java @@ -32,6 +32,7 @@ import java.util.List; import java.util.Map; import java.util.TreeMap; +import java.util.TreeSet; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -186,6 +187,33 @@ void testLogicalToPhysicalSchemaNoShreddingColumns() { .isEqualTo(logical); } + @Test + void testBuildSpecificPhysicalStructType() { + RowType physicalType = + MapSharedShreddingUtils.buildSpecificPhysicalStructType( + DataTypes.BIGINT().notNull(), + new TreeSet(Arrays.asList(3, 1)), + true); + + assertThat(physicalType.getFieldNames()) + .containsExactly("__field_mapping", "__col_1", "__col_3", "__overflow"); + assertThat(physicalType.getFields()).extracting(DataField::id).containsExactly(0, 1, 2, 3); + assertThat(physicalType.getField("__col_1").type()).isEqualTo(DataTypes.BIGINT().notNull()); + assertThat(physicalType.getField("__overflow").type()) + .isEqualTo(DataTypes.MAP(DataTypes.INT(), DataTypes.BIGINT().notNull())); + } + + @Test + void testBuildSpecificPhysicalStructTypeWithoutOverflow() { + RowType physicalType = + MapSharedShreddingUtils.buildSpecificPhysicalStructType( + DataTypes.STRING(), new TreeSet(Arrays.asList(3)), false); + + assertThat(physicalType.getFieldNames()).containsExactly("__field_mapping", "__col_3"); + assertThat(physicalType.getFields()).extracting(DataField::id).containsExactly(0, 1); + assertThat(physicalType.getField("__col_3").type()).isEqualTo(DataTypes.STRING()); + } + @Test void testMetadataRoundtrip() { Map nameToId = new TreeMap<>(); diff --git a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java index c9c1a6b5f5a0..17c5cd790e94 100644 --- a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java +++ b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java @@ -94,6 +94,8 @@ import static org.apache.paimon.CoreOptions.IGNORE_RETRACT; import static org.apache.paimon.CoreOptions.IGNORE_UPDATE_BEFORE; import static org.apache.paimon.CoreOptions.LIST_AGG_DELIMITER; +import static org.apache.paimon.CoreOptions.MAP_SHARED_SHREDDING_MAX_COLUMNS; +import static org.apache.paimon.CoreOptions.MAP_STORAGE_LAYOUT; import static org.apache.paimon.CoreOptions.NESTED_KEY; import static org.apache.paimon.CoreOptions.PK_CLUSTERING_OVERRIDE; import static org.apache.paimon.CoreOptions.SEQUENCE_FIELD; @@ -634,10 +636,56 @@ protected void updateLastColumn( newSchema.primaryKeys(), newSchema.options(), newSchema.comment()); + checkMapStorageLayoutUnchanged(oldTableSchema, newTableSchema); SchemaValidation.validateTableSchema(newTableSchema); return newTableSchema; } + private static void checkMapStorageLayoutUnchanged( + TableSchema oldTableSchema, TableSchema newTableSchema) { + Map oldLayouts = + mapStorageLayoutByFieldId(oldTableSchema); + Map newLayouts = + mapStorageLayoutByFieldId(newTableSchema); + Map oldNames = fieldNameById(oldTableSchema); + Map newNames = fieldNameById(newTableSchema); + + for (Map.Entry oldLayout : oldLayouts.entrySet()) { + Integer fieldId = oldLayout.getKey(); + CoreOptions.MapStorageLayout newLayout = newLayouts.get(fieldId); + if (newLayout == null || oldLayout.getValue() == newLayout) { + continue; + } + + throw new UnsupportedOperationException( + String.format( + "Cannot change map storage layout for field id %s ('%s' -> '%s') from '%s' to '%s'.", + fieldId, + oldNames.get(fieldId), + newNames.get(fieldId), + oldLayout.getValue(), + newLayout)); + } + } + + private static Map mapStorageLayoutByFieldId( + TableSchema schema) { + CoreOptions options = new CoreOptions(schema.options()); + Map layouts = new HashMap<>(); + for (DataField field : schema.fields()) { + layouts.put(field.id(), options.mapStorageLayout(field.name())); + } + return layouts; + } + + private static Map fieldNameById(TableSchema schema) { + Map names = new HashMap<>(); + for (DataField field : schema.fields()) { + names.put(field.id(), field.name()); + } + return names; + } + // gets the rootType at the defined depth // ex: ARRAY>> // if we want to update ARRAY -> ARRAY @@ -854,7 +902,14 @@ private static Map applyRenameColumnsToOptions( fieldName -> FIELDS_PREFIX + "." + fieldName + "." + AGG_FUNCTION, fieldName -> FIELDS_PREFIX + "." + fieldName + "." + IGNORE_RETRACT, fieldName -> FIELDS_PREFIX + "." + fieldName + "." + DISTINCT, - fieldName -> FIELDS_PREFIX + "." + fieldName + "." + LIST_AGG_DELIMITER); + fieldName -> FIELDS_PREFIX + "." + fieldName + "." + LIST_AGG_DELIMITER, + fieldName -> FIELDS_PREFIX + "." + fieldName + "." + MAP_STORAGE_LAYOUT, + fieldName -> + FIELDS_PREFIX + + "." + + fieldName + + "." + + MAP_SHARED_SHREDDING_MAX_COLUMNS); for (RenameColumn rename : renameColumns) { String fieldName = rename.fieldNames()[0]; diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java index 77395cd63a7d..a06f4a280c2e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java @@ -18,7 +18,10 @@ package org.apache.paimon.utils; +import org.apache.paimon.CoreOptions; +import org.apache.paimon.CoreOptions.MapStorageLayout; import org.apache.paimon.casting.CastFieldGetter; +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils; import org.apache.paimon.data.variant.VariantMetadataUtils; import org.apache.paimon.format.FileFormatDiscover; import org.apache.paimon.format.FormatReaderFactory; @@ -44,10 +47,12 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.function.Function; import static org.apache.paimon.predicate.PredicateBuilder.excludePredicateWithFields; import static org.apache.paimon.table.SpecialFields.KEY_FIELD_ID_START; +import static org.apache.paimon.utils.Preconditions.checkArgument; /** Class with index mapping and format reader. */ public class FormatReaderMapping { @@ -205,7 +210,9 @@ public FormatReaderMapping build( new ArrayList<>(fieldsExtractor.apply(dataSchema)); Map systemFields = findSystemFields(expectedFields); - List readDataFields = readDataFields(allDataFieldsInFile, expectedFields); + Set selectedKeysFieldIds = selectedKeysFieldIds(tableSchema, expectedFields); + List readDataFields = + readDataFields(allDataFieldsInFile, expectedFields, selectedKeysFieldIds); IndexCastMapping indexCastMapping = SchemaEvolutionUtil.createIndexCastMapping(expectedFields, readDataFields); @@ -314,7 +321,9 @@ static Pair trimKeyFields( } private List readDataFields( - List allDataFields, List expectedFields) { + List allDataFields, + List expectedFields, + Set selectedKeysFieldIds) { List readDataFields = new ArrayList<>(); for (DataField dataField : allDataFields) { expectedFields.stream() @@ -322,6 +331,12 @@ private List readDataFields( .findFirst() .ifPresent( field -> { + if (selectedKeysFieldIds.contains(field.id())) { + checkSelectedKeysDataField(dataField); + readDataFields.add(field); + return; + } + DataType prunedType = pruneDataType(field.type(), dataField.type()); if (prunedType != null) { @@ -332,6 +347,52 @@ private List readDataFields( return readDataFields; } + private void checkSelectedKeysDataField(DataField dataField) { + checkArgument( + dataField.type() instanceof MapType, + "Selected-key MAP field %s should be MAP type in data schema.", + dataField.name()); + } + + private Set selectedKeysFieldIds( + TableSchema tableSchema, List expectedFields) { + CoreOptions options = CoreOptions.fromMap(tableSchema.options()); + Map tableFields = tableSchema.idToFieldMap(); + Set selectedKeysFieldIds = new HashSet<>(); + for (DataField expectedField : expectedFields) { + if (MapSelectedKeysMetadataUtils.isMapSelectedKeysField(expectedField)) { + DataField tableField = tableFields.get(expectedField.id()); + checkArgument( + tableField != null, + "Cannot find selected-key MAP field id %s in table schema.", + expectedField.id()); + checkArgument( + tableField.type() instanceof MapType, + "Selected-key MAP pushdown only supports top-level MAP field: %s.", + tableField.name()); + checkArgument( + options.mapStorageLayout(tableField.name()) + == MapStorageLayout.SHARED_SHREDDING, + "Selected-key MAP pushdown only supports top-level shared-shredding MAP field: %s.", + tableField.name()); + validateSelectedKeyValueTypes(expectedField, tableField); + selectedKeysFieldIds.add(expectedField.id()); + } + } + return selectedKeysFieldIds; + } + + private void validateSelectedKeyValueTypes(DataField expectedField, DataField tableField) { + RowType selectedKeysType = (RowType) expectedField.type(); + DataType mapValueType = ((MapType) tableField.type()).getValueType(); + for (DataField selectedKeyField : selectedKeysType.getFields()) { + checkArgument( + selectedKeyField.type().equalsIgnoreNullable(mapValueType), + "Selected-key MAP pushdown does not support pruning MAP value fields: %s.", + tableField.name()); + } + } + @Nullable private DataType pruneDataType(DataType readType, DataType dataType) { switch (readType.getTypeRoot()) { diff --git a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java index 1b3140a61615..3527e6e68ab4 100644 --- a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java @@ -170,6 +170,74 @@ public void testUpdateOptions() throws Exception { assertThat(latest.get().options()).containsEntry("new_k", "new_v"); } + @Test + public void testCannotChangeMapStorageLayoutForExistingField() throws Exception { + retryArtificialException(() -> manager.createTable(mapStorageLayoutSchema("default"))); + + assertThatThrownBy( + () -> + manager.commitChanges( + SchemaChange.setOption( + "fields.metrics.map.storage-layout", + "shared-shredding"))) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining( + "Cannot change map storage layout for field id 1 ('metrics' -> 'metrics') from 'default' to 'shared-shredding'."); + } + + @Test + public void testCannotChangeMapStorageLayoutByRenameColumn() throws Exception { + retryArtificialException(() -> manager.createTable(mapStorageLayoutSchema(null))); + + assertThatThrownBy( + () -> + manager.commitChanges( + Arrays.asList( + SchemaChange.renameColumn( + "metrics", "renamed_metrics"), + SchemaChange.setOption( + "fields.renamed_metrics.map.storage-layout", + "shared-shredding")))) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining( + "Cannot change map storage layout for field id 1 ('metrics' -> 'renamed_metrics') from 'default' to 'shared-shredding'."); + } + + @Test + public void testRenameColumnKeepsMapStorageLayoutOptions() throws Exception { + retryArtificialException( + () -> manager.createTable(mapStorageLayoutSchema("shared-shredding"))); + + retryArtificialException( + () -> manager.commitChanges(SchemaChange.renameColumn("metrics", "renamed"))); + + Optional latest = retryArtificialException(() -> manager.latest()); + assertThat(latest.isPresent()).isTrue(); + assertThat(latest.get().options()) + .doesNotContainKey("fields.metrics.map.storage-layout") + .containsEntry("fields.renamed.map.storage-layout", "shared-shredding") + .containsEntry("fields.renamed.map.shared-shredding.max-columns", "2"); + } + + private Schema mapStorageLayoutSchema(String layout) { + Map options = new HashMap<>(); + if (layout != null) { + options.put("fields.metrics.map.storage-layout", layout); + options.put("fields.metrics.map.shared-shredding.max-columns", "2"); + } + return new Schema( + Arrays.asList( + new DataField(0, "id", DataTypes.INT()), + new DataField( + 1, + "metrics", + DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT()))), + Collections.emptyList(), + Collections.emptyList(), + options, + ""); + } + @Test public void testRejectRenamePrimaryKeyVectorIndexColumn() throws Exception { Map options = new HashMap<>(); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java index 35e86d9a1d23..660f31110485 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java @@ -31,6 +31,7 @@ import org.apache.paimon.data.InternalMap; import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.Timestamp; +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils; import org.apache.paimon.data.shredding.MapSharedShreddingFieldMeta; import org.apache.paimon.data.shredding.MapSharedShreddingUtils; import org.apache.paimon.format.FileFormat; @@ -75,6 +76,7 @@ import static org.apache.paimon.deletionvectors.DeletionVectorsIndexFile.DELETION_VECTORS_INDEX; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Table-level tests for MAP shared-shredding. */ public class MapSharedShreddingTableTest extends TableTestBase { @@ -159,6 +161,42 @@ public void testColumnPlacementPolicies(String format, String placementPolicy) .containsEntry(4, javaMapOf("a", 70L, "b", 80L, "c", 90L, "d", 100L)); } + @ParameterizedTest + @ValueSource(strings = {"orc", "parquet"}) + public void testReadSelectedKeysAsRow(String format) throws Exception { + Table table = createTable(format, 1, "metrics"); + + write( + table, + GenericRow.of(1, mapOf("key1", 10L, "key2", 20L)), + GenericRow.of(2, mapOf("key2", 30L, "cold", 40L)), + GenericRow.of(3, null)); + + Map> actual = new LinkedHashMap<>(); + ReadBuilder readBuilder = table.newReadBuilder().withReadType(selectedKeysReadType()); + RecordReader reader = + readBuilder.newRead().createReader(readBuilder.newScan().plan()); + reader.forEachRemaining( + row -> { + if (row.isNullAt(1)) { + actual.put(row.getInt(0), null); + } else { + InternalRow selectedKeys = row.getRow(1, 3); + actual.put( + row.getInt(0), + Arrays.asList( + selectedKeys.isNullAt(0) ? null : selectedKeys.getLong(0), + selectedKeys.isNullAt(1) ? null : selectedKeys.getLong(1), + selectedKeys.isNullAt(2) ? null : selectedKeys.getLong(2))); + } + }); + + assertThat(actual) + .containsEntry(1, Arrays.asList(10L, 20L, null)) + .containsEntry(2, Arrays.asList(null, 30L, null)) + .containsEntry(3, null); + } + @ParameterizedTest @ValueSource(strings = {"orc", "parquet"}) public void testAppendOnlyTableReadWriteWithTwoMapFields(String format) throws Exception { @@ -723,53 +761,29 @@ public void testPrimaryKeyInfersColumnCountPerFile(String format, boolean thinMo @ParameterizedTest @ValueSource(strings = {"orc", "parquet"}) - public void testSwitchMapLayoutAndInferColumns(String format) throws Exception { - Table table = - createTableWithBucket( - format, - 4, - "1", - Arrays.asList("metrics", "labels"), - Arrays.asList("labels")); - - write(table, GenericRow.of(1, mapOf("a", 11L, "b", 12L), mapOf("x", 21L))); - - catalog.alterTable( - identifier(format), - Arrays.asList( - SchemaChange.setOption( - "fields.metrics.map.storage-layout", "shared-shredding"), - SchemaChange.setOption( - "fields.metrics.map.shared-shredding.max-columns", "3"), - SchemaChange.setOption("fields.labels.map.storage-layout", "default")), - false); - table = catalog.getTable(identifier(format)); - - write(table, GenericRow.of(2, mapOf("c", 31L), mapOf("y", 41L, "z", 42L))); - - FileStoreTable fileStoreTable = (FileStoreTable) table; - List files = currentDataFiles(fileStoreTable); - files.sort(Comparator.comparingLong(file -> file.dataFile.minSequenceNumber())); - assertThat(files).hasSize(2); - - MapSharedShreddingFieldMeta metricsMeta = - readSharedShreddingFieldMeta(fileStoreTable, files.get(1), "metrics"); - assertThat(metricsMeta.numColumns()).isEqualTo(1); - assertThat(metricsMeta.maxRowWidth()).isEqualTo(1); - - Map>> actual = new LinkedHashMap<>(); - for (InternalRow row : read(table)) { - actual.put( - row.getInt(0), - Arrays.asList( - row.isNullAt(1) ? null : toJavaMap(row.getMap(1)), - row.isNullAt(2) ? null : toJavaMap(row.getMap(2)))); - } + public void testCannotSwitchMapLayoutAndUseMaxColumnsWithoutMetadata(String format) + throws Exception { + createTableWithBucket( + format, 4, "1", Arrays.asList("metrics", "labels"), Arrays.asList("labels")); - assertThat(actual) - .containsEntry(1, Arrays.asList(javaMapOf("a", 11L, "b", 12L), javaMapOf("x", 21L))) - .containsEntry( - 2, Arrays.asList(javaMapOf("c", 31L), javaMapOf("y", 41L, "z", 42L))); + assertThatThrownBy( + () -> + catalog.alterTable( + identifier(format), + Arrays.asList( + SchemaChange.setOption( + "fields.metrics.map.storage-layout", + "shared-shredding"), + SchemaChange.setOption( + "fields.metrics.map.shared-shredding.max-columns", + "3"), + SchemaChange.setOption( + "fields.labels.map.storage-layout", + "default")), + false)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining( + "Cannot change map storage layout for field id 1 ('metrics' -> 'metrics') from 'default' to 'shared-shredding'."); } @ParameterizedTest @@ -1326,6 +1340,20 @@ private Schema schema(String format, int maxColumns, String... sharedShreddingFi return schemaWithBucket(format, maxColumns, "-1", sharedShreddingFields); } + private RowType selectedKeysReadType() { + RowType selectedKeysType = + DataTypes.ROW( + DataTypes.FIELD(0, "0", DataTypes.BIGINT()), + DataTypes.FIELD(1, "1", DataTypes.BIGINT()), + DataTypes.FIELD(2, "2", DataTypes.BIGINT())); + return DataTypes.ROW( + DataTypes.FIELD(0, "id", DataTypes.INT()), + MapSelectedKeysMetadataUtils.withSelectedKeys( + DataTypes.FIELD(1, "metrics", selectedKeysType), + selectedKeysType, + Arrays.asList("key1", "key2", "missing"))); + } + private Schema schemaWithBucket( String format, int maxColumns, String bucket, String... sharedShreddingFields) { return schemaWithBucket( diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java index 7b2ad9898d27..fca8e6c1cd94 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java @@ -18,8 +18,10 @@ package org.apache.paimon.utils; +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils; import org.apache.paimon.schema.IndexCastMapping; import org.apache.paimon.schema.SchemaEvolutionUtil; +import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.SpecialFields; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; @@ -28,9 +30,14 @@ import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; +import java.lang.reflect.Method; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.Set; /** Test for {@link FormatReaderMapping.Builder}. */ public class FormatReaderMappingTest { @@ -147,4 +154,137 @@ public void testTrimKeyWithIndexMapping() { Assertions.assertThat(trimmed.get(2).id()).isEqualTo(3); Assertions.assertThat(trimmed.size()).isEqualTo(3); } + + @Test + public void testMapSelectedKeysKeepsSelectedRowTypeForFormatReader() throws Exception { + RowType dataValueType = + DataTypes.ROW( + DataTypes.FIELD(10, "a", DataTypes.INT()), + DataTypes.FIELD(11, "b", DataTypes.INT()), + DataTypes.FIELD(12, "c", DataTypes.INT())); + RowType selectedKeysType = + DataTypes.ROW( + DataTypes.FIELD(0, "0", dataValueType), + DataTypes.FIELD(1, "1", dataValueType)); + + DataField dataField = + DataTypes.FIELD(2, "attrs", DataTypes.MAP(DataTypes.STRING(), dataValueType)); + DataField expectedField = + MapSelectedKeysMetadataUtils.withSelectedKeys( + DataTypes.FIELD(2, "attrs", selectedKeysType), + selectedKeysType, + Arrays.asList("key1", "key2")); + + List dataFields = Collections.singletonList(dataField); + List expectedFields = Collections.singletonList(expectedField); + + DataField readField = invokeDataFields(dataFields, expectedFields).get(0); + + Assertions.assertThat(readField.type()).isEqualTo(selectedKeysType); + Assertions.assertThat(MapSelectedKeysMetadataUtils.isMapSelectedKeysField(readField)) + .isTrue(); + } + + @Test + public void testRejectMapSelectedKeysValuePartialProjection() { + RowType tableValueType = + DataTypes.ROW( + DataTypes.FIELD(10, "a", DataTypes.INT()), + DataTypes.FIELD(11, "b", DataTypes.INT())); + RowType selectedValueType = DataTypes.ROW(DataTypes.FIELD(10, "a", DataTypes.INT())); + DataField tableField = + DataTypes.FIELD(2, "attrs", DataTypes.MAP(DataTypes.STRING(), tableValueType)); + DataField expectedField = selectedKeysField(2, "attrs", selectedValueType); + + Map options = new HashMap<>(); + options.put("fields.attrs.map.storage-layout", "shared-shredding"); + + Assertions.assertThatThrownBy( + () -> + invokeSelectedKeysFieldIds( + tableSchema(Collections.singletonList(tableField), options), + Collections.singletonList(expectedField))) + .hasRootCauseMessage( + "Selected-key MAP pushdown does not support pruning MAP value fields: attrs."); + } + + @Test + public void testMapSelectedKeysOnlySupportsTopLevelSharedShreddingMap() throws Exception { + DataField tableField = + DataTypes.FIELD(2, "attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT())); + DataField expectedField = selectedKeysField(2, "attrs", DataTypes.BIGINT()); + + Map options = new HashMap<>(); + options.put("fields.attrs.map.storage-layout", "shared-shredding"); + Assertions.assertThat( + invokeSelectedKeysFieldIds( + tableSchema(Collections.singletonList(tableField), options), + Collections.singletonList(expectedField))) + .containsExactly(2); + + Assertions.assertThatThrownBy( + () -> + invokeSelectedKeysFieldIds( + tableSchema( + Collections.singletonList(tableField), + Collections.emptyMap()), + Collections.singletonList(expectedField))) + .hasRootCauseMessage( + "Selected-key MAP pushdown only supports top-level shared-shredding MAP field: attrs."); + } + + @Test + public void testRejectSelectedKeysDataFieldWithNonMapType() { + DataField dataField = DataTypes.FIELD(2, "attrs", DataTypes.BIGINT()); + DataField expectedField = selectedKeysField(2, "attrs", DataTypes.BIGINT()); + + Assertions.assertThatThrownBy( + () -> + invokeDataFields( + Collections.singletonList(dataField), + Collections.singletonList(expectedField))) + .hasRootCauseMessage( + "Selected-key MAP field attrs should be MAP type in data schema."); + } + + @SuppressWarnings("unchecked") + private static List invokeDataFields( + List allDataFields, List expectedFields) throws Exception { + FormatReaderMapping.Builder builder = + new FormatReaderMapping.Builder( + null, Collections.emptyList(), null, null, null, null); + Method method = + FormatReaderMapping.Builder.class.getDeclaredMethod( + "readDataFields", List.class, List.class, Set.class); + method.setAccessible(true); + return (List) + method.invoke(builder, allDataFields, expectedFields, Collections.singleton(2)); + } + + @SuppressWarnings("unchecked") + private static Set invokeSelectedKeysFieldIds( + TableSchema tableSchema, List expectedFields) throws Exception { + FormatReaderMapping.Builder builder = + new FormatReaderMapping.Builder( + null, Collections.emptyList(), null, null, null, null); + Method method = + FormatReaderMapping.Builder.class.getDeclaredMethod( + "selectedKeysFieldIds", TableSchema.class, List.class); + method.setAccessible(true); + return (Set) method.invoke(builder, tableSchema, expectedFields); + } + + private static DataField selectedKeysField( + int id, String name, org.apache.paimon.types.DataType valueType) { + RowType selectedKeysType = DataTypes.ROW(DataTypes.FIELD(0, "0", valueType)); + return MapSelectedKeysMetadataUtils.withSelectedKeys( + DataTypes.FIELD(id, name, selectedKeysType), + selectedKeysType, + Collections.singletonList("key1")); + } + + private static TableSchema tableSchema(List fields, Map options) { + return new TableSchema( + 1, fields, 100, Collections.emptyList(), Collections.emptyList(), options, ""); + } } diff --git a/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-3.5/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-3.5/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-3.5/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala b/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala new file mode 100644 index 000000000000..1bdf43ce20c2 --- /dev/null +++ b/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETest.scala @@ -0,0 +1,21 @@ +/* + * 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.spark.sql + +class MapSelectedKeysSharedShreddingE2ETest extends MapSelectedKeysSharedShreddingE2ETestBase {} diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala index 74c0e6c47dbd..32e85fb1e589 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala @@ -27,6 +27,7 @@ import org.apache.paimon.predicate.{Predicate, PredicateBuilder} import org.apache.paimon.spark.{PaimonRecordReaderIterator, SparkCatalog, SparkGenericCatalog, SparkTable, SparkUtils} import org.apache.paimon.spark.catalog.{SparkBaseCatalog, SupportView} import org.apache.paimon.spark.catalyst.analysis.ResolvedPaimonView +import org.apache.paimon.spark.catalyst.optimizer.PushDownMapSelectedKeys import org.apache.paimon.spark.catalyst.plans.logical.{CopyIntoLocationCommand, CopyIntoLocationSource, CopyIntoTableCommand, CreateOrReplaceTagCommand, CreatePaimonView, DeleteTagCommand, DropPaimonView, LateralVectorSearch, PaimonCallCommand, PaimonDropPartitions, PaimonTableValuedFunctions, RenameTagCommand, ResolvedIdentifier, ShowPaimonViews, ShowTagsCommand, TruncatePaimonTableWithFilter} import org.apache.paimon.spark.data.SparkInternalRow import org.apache.paimon.spark.read.VectorSearchResultUtils @@ -61,7 +62,19 @@ case class PaimonStrategy(spark: SparkSession) import DataSourceV2Implicits._ protected lazy val catalogManager = spark.sessionState.catalogManager - override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { + override def apply(plan: LogicalPlan): Seq[SparkPlan] = { + // Spark creates DataSourceV2ScanRelation after injected optimizer rules have run. Apply this + // rewrite during physical planning, when the scan relation is available, and let the regular + // Spark strategies plan the rewritten logical subtree. + val rewritten = PushDownMapSelectedKeys(plan) + if (!rewritten.fastEquals(plan)) { + planLater(rewritten) :: Nil + } else { + applyWithoutMapSelectedKeysPushDown(plan) + } + } + + private def applyWithoutMapSelectedKeysPushDown(plan: LogicalPlan): Seq[SparkPlan] = plan match { case ctas: CreateTableAsSelect => PaimonCreateTableAsSelectStrategy(spark)(ctas) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala index 388889bbeaaa..61481e201c0e 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala @@ -101,8 +101,6 @@ class PaimonSparkSessionExtensions extends (SparkSessionExtensions => Unit) { // optimization rules extensions.injectOptimizerRule(spark => ReplacePaimonFunctions(spark)) extensions.injectOptimizerRule(_ => OptimizeMetadataOnlyDeleteFromPaimonTable) - // TODO: Enable MAP selected-key pushdown after core reader supports - // __PAIMON_MAP_SELECTED_KEYS read type. extensions.injectOptimizerRule(_ => MergePaimonScalarSubqueries) extensions.injectOptimizerRule(_ => PushDownLateralVectorSearchFilter) diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETestBase.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETestBase.scala new file mode 100644 index 000000000000..0290bd774343 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/MapSelectedKeysSharedShreddingE2ETestBase.scala @@ -0,0 +1,145 @@ +/* + * 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.spark.sql + +import org.apache.paimon.spark.{PaimonScan, PaimonSparkTestBase} + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + +abstract class MapSelectedKeysSharedShreddingE2ETestBase extends PaimonSparkTestBase { + + Seq("parquet", "orc").foreach { + format => + test(s"read selected shared-shredding map keys directly from $format") { + withTable("T") { + sql(s""" + |CREATE TABLE T (id INT, attrs MAP) + |TBLPROPERTIES ( + | 'bucket' = '-1', + | 'file.format' = '$format', + | 'fields.attrs.map.storage-layout' = 'shared-shredding', + | 'fields.attrs.map.shared-shredding.max-columns' = '1' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO T VALUES + | (1, map('key1', CAST(10 AS BIGINT), 'key2', CAST(20 AS BIGINT))), + | (2, map('key2', CAST(30 AS BIGINT), 'cold', CAST(40 AS BIGINT))) + |""".stripMargin) + + val query = + sql("SELECT id, attrs['key1'], attrs['key2'], attrs['missing'] FROM T ORDER BY id") + val sparkPlan = query.queryExecution.sparkPlan + val pushedMapSelectedKeys = sparkPlan.collectFirst { + case scan: BatchScanExec if scan.scan.isInstanceOf[PaimonScan] => + scan.scan.asInstanceOf[PaimonScan].pushedMapSelectedKeys + } + val expectedPushedMapSelectedKeys = Map("attrs" -> Seq("key1", "key2", "missing")) + assert( + pushedMapSelectedKeys.contains(expectedPushedMapSelectedKeys), + s"""Expected selected MAP keys to be pushed down. + |Physical plan: + |$sparkPlan""".stripMargin + ) + + checkAnswer(query, Row(1, 10L, 20L, null) :: Row(2, null, 30L, null) :: Nil) + } + } + + test(s"read full shared-shredding map column from $format") { + withTable("T") { + sql(s""" + |CREATE TABLE T (id INT, attrs MAP) + |TBLPROPERTIES ( + | 'bucket' = '-1', + | 'file.format' = '$format', + | 'fields.attrs.map.storage-layout' = 'shared-shredding', + | 'fields.attrs.map.shared-shredding.max-columns' = '1' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO T VALUES + | (1, map('key1', CAST(10 AS BIGINT), 'key2', CAST(20 AS BIGINT))), + | (2, map('key2', CAST(30 AS BIGINT), 'cold', CAST(40 AS BIGINT))), + | (3, NULL) + |""".stripMargin) + + checkAnswer( + sql("SELECT id, attrs FROM T ORDER BY id"), + Row(1, Map("key1" -> 10L, "key2" -> 20L)) :: + Row(2, Map("key2" -> 30L, "cold" -> 40L)) :: + Row(3, null) :: Nil) + } + } + + test(s"read selected normal map keys from $format") { + withTable("T") { + sql(s""" + |CREATE TABLE T (id INT, attrs MAP) + |TBLPROPERTIES ( + | 'bucket' = '-1', + | 'file.format' = '$format' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO T VALUES + | (1, map('key1', CAST(10 AS BIGINT), 'key2', CAST(20 AS BIGINT))), + | (2, map('key2', CAST(30 AS BIGINT), 'cold', CAST(40 AS BIGINT))), + | (3, NULL) + |""".stripMargin) + + checkAnswer( + sql("SELECT id, attrs['key1'], attrs['key2'], attrs['missing'] FROM T ORDER BY id"), + Row(1, 10L, 20L, null) :: + Row(2, null, 30L, null) :: + Row(3, null, null, null) :: Nil + ) + } + } + + test(s"read full normal map column from $format") { + withTable("T") { + sql(s""" + |CREATE TABLE T (id INT, attrs MAP) + |TBLPROPERTIES ( + | 'bucket' = '-1', + | 'file.format' = '$format' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO T VALUES + | (1, map('key1', CAST(10 AS BIGINT), 'key2', CAST(20 AS BIGINT))), + | (2, map('key2', CAST(30 AS BIGINT), 'cold', CAST(40 AS BIGINT))), + | (3, NULL) + |""".stripMargin) + + checkAnswer( + sql("SELECT id, attrs FROM T ORDER BY id"), + Row(1, Map("key1" -> 10L, "key2" -> 20L)) :: + Row(2, Map("key2" -> 30L, "cold" -> 40L)) :: + Row(3, null) :: Nil) + } + } + } +} From 4dadbfa9af655cbec0b1a16de299b918b0a0ceb5 Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Tue, 21 Jul 2026 19:02:09 +0800 Subject: [PATCH 2/6] fix ci --- .../test/java/org/apache/paimon/schema/SchemaManagerTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java index 3527e6e68ab4..ec24ed4151a2 100644 --- a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java @@ -231,7 +231,7 @@ private Schema mapStorageLayoutSchema(String layout) { new DataField( 1, "metrics", - DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT()))), + DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.BIGINT()))), Collections.emptyList(), Collections.emptyList(), options, From 3ee70c5f26755c7ce90d5720b47ea3e127225c97 Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Wed, 22 Jul 2026 17:10:25 +0800 Subject: [PATCH 3/6] fix comment --- .../paimon/utils/FormatReaderMapping.java | 14 ++++- .../table/MapSharedShreddingTableTest.java | 51 +++++++++++++++++++ .../paimon/utils/FormatReaderMappingTest.java | 28 ++++++++++ 3 files changed, 92 insertions(+), 1 deletion(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java index a06f4a280c2e..ba70953e9704 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java @@ -333,7 +333,7 @@ private List readDataFields( field -> { if (selectedKeysFieldIds.contains(field.id())) { checkSelectedKeysDataField(dataField); - readDataFields.add(field); + readDataFields.add(selectedKeysDataField(field, dataField)); return; } @@ -347,6 +347,18 @@ private List readDataFields( return readDataFields; } + private DataField selectedKeysDataField(DataField expectedField, DataField dataField) { + RowType selectedKeysType = (RowType) expectedField.type(); + DataType dataValueType = ((MapType) dataField.type()).getValueType(); + List selectedKeysDataFields = new ArrayList<>(); + for (DataField selectedKeyField : selectedKeysType.getFields()) { + selectedKeysDataFields.add( + selectedKeyField.newType( + dataValueType.copy(selectedKeyField.type().isNullable()))); + } + return expectedField.newType(selectedKeysType.copy(selectedKeysDataFields)); + } + private void checkSelectedKeysDataField(DataField dataField) { checkArgument( dataField.type() instanceof MapType, diff --git a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java index 660f31110485..6b109c90eaea 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java @@ -197,6 +197,57 @@ public void testReadSelectedKeysAsRow(String format) throws Exception { .containsEntry(3, null); } + @ParameterizedTest + @ValueSource(strings = {"orc", "parquet"}) + public void testReadSelectedKeysAfterMapValueTypeEvolution(String format) throws Exception { + catalog.createTable( + identifier(format), + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column( + "metrics", + DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.INT())) + .option("bucket", "-1") + .option("file.format", format) + .option(CoreOptions.WRITE_ONLY.key(), "true") + .option("fields.metrics.map.storage-layout", "shared-shredding") + .option("fields.metrics.map.shared-shredding.max-columns", "1") + .build(), + true); + Table table = catalog.getTable(identifier(format)); + Map values = new LinkedHashMap<>(); + values.put(BinaryString.fromString("key1"), 10); + write(table, GenericRow.of(1, new GenericMap(values))); + + catalog.alterTable( + identifier(format), + Collections.singletonList( + SchemaChange.updateColumnType( + new String[] {"metrics", "value"}, DataTypes.BIGINT(), false)), + false); + table = catalog.getTable(identifier(format)); + + ReadBuilder readBuilder = table.newReadBuilder().withReadType(selectedKeysReadType()); + List> actual = new ArrayList<>(); + readBuilder + .newRead() + .createReader(readBuilder.newScan().plan()) + .forEachRemaining( + row -> { + InternalRow selectedKeys = row.getRow(1, 3); + actual.add( + Arrays.asList( + selectedKeys.getLong(0), + selectedKeys.isNullAt(1) + ? null + : selectedKeys.getLong(1), + selectedKeys.isNullAt(2) + ? null + : selectedKeys.getLong(2))); + }); + assertThat(actual).containsExactly(Arrays.asList(10L, null, null)); + } + @ParameterizedTest @ValueSource(strings = {"orc", "parquet"}) public void testAppendOnlyTableReadWriteWithTwoMapFields(String format) throws Exception { diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java index fca8e6c1cd94..8a43de5af33b 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java @@ -18,6 +18,8 @@ package org.apache.paimon.utils; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils; import org.apache.paimon.schema.IndexCastMapping; import org.apache.paimon.schema.SchemaEvolutionUtil; @@ -185,6 +187,32 @@ public void testMapSelectedKeysKeepsSelectedRowTypeForFormatReader() throws Exce .isTrue(); } + @Test + public void testMapSelectedKeysUsesDataValueTypeForSchemaEvolution() throws Exception { + DataField dataField = + DataTypes.FIELD(2, "attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())); + DataField expectedField = selectedKeysField(2, "attrs", DataTypes.BIGINT()); + + DataField readField = + invokeDataFields( + Collections.singletonList(dataField), + Collections.singletonList(expectedField)) + .get(0); + RowType readType = (RowType) readField.type(); + Assertions.assertThat(readType.getTypeAt(0)).isEqualTo(DataTypes.INT()); + Assertions.assertThat(MapSelectedKeysMetadataUtils.isMapSelectedKeysField(readField)) + .isTrue(); + + IndexCastMapping mapping = + SchemaEvolutionUtil.createIndexCastMapping( + Collections.singletonList(expectedField), + Collections.singletonList(readField)); + Assertions.assertThat(mapping.getCastMapping()).isNotNull(); + InternalRow selectedKeys = + mapping.getCastMapping()[0].getFieldOrNull(GenericRow.of(GenericRow.of(10))); + Assertions.assertThat(selectedKeys.getLong(0)).isEqualTo(10L); + } + @Test public void testRejectMapSelectedKeysValuePartialProjection() { RowType tableValueType = From 5708a1fd5c675f6149e3488e415486367a2ae525 Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Thu, 23 Jul 2026 11:07:43 +0800 Subject: [PATCH 4/6] fix ci --- .../paimon/schema/SchemaManagerTest.java | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java index ec24ed4151a2..5876235bd802 100644 --- a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java @@ -176,10 +176,12 @@ public void testCannotChangeMapStorageLayoutForExistingField() throws Exception assertThatThrownBy( () -> - manager.commitChanges( - SchemaChange.setOption( - "fields.metrics.map.storage-layout", - "shared-shredding"))) + retryArtificialException( + () -> + manager.commitChanges( + SchemaChange.setOption( + "fields.metrics.map.storage-layout", + "shared-shredding")))) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining( "Cannot change map storage layout for field id 1 ('metrics' -> 'metrics') from 'default' to 'shared-shredding'."); @@ -191,13 +193,16 @@ public void testCannotChangeMapStorageLayoutByRenameColumn() throws Exception { assertThatThrownBy( () -> - manager.commitChanges( - Arrays.asList( - SchemaChange.renameColumn( - "metrics", "renamed_metrics"), - SchemaChange.setOption( - "fields.renamed_metrics.map.storage-layout", - "shared-shredding")))) + retryArtificialException( + () -> + manager.commitChanges( + Arrays.asList( + SchemaChange.renameColumn( + "metrics", + "renamed_metrics"), + SchemaChange.setOption( + "fields.renamed_metrics.map.storage-layout", + "shared-shredding"))))) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining( "Cannot change map storage layout for field id 1 ('metrics' -> 'renamed_metrics') from 'default' to 'shared-shredding'."); From 4e8b5a74cc773fd5e86443a5ad654ddafafa866b Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Thu, 23 Jul 2026 14:13:22 +0800 Subject: [PATCH 5/6] fix comment --- .../apache/paimon/schema/SchemaManager.java | 9 +++++- .../paimon/schema/SchemaManagerTest.java | 29 +++++++++++++++---- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java index 17c5cd790e94..5eadd040c7d7 100644 --- a/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java +++ b/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java @@ -94,6 +94,7 @@ import static org.apache.paimon.CoreOptions.IGNORE_RETRACT; import static org.apache.paimon.CoreOptions.IGNORE_UPDATE_BEFORE; import static org.apache.paimon.CoreOptions.LIST_AGG_DELIMITER; +import static org.apache.paimon.CoreOptions.MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY; import static org.apache.paimon.CoreOptions.MAP_SHARED_SHREDDING_MAX_COLUMNS; import static org.apache.paimon.CoreOptions.MAP_STORAGE_LAYOUT; import static org.apache.paimon.CoreOptions.NESTED_KEY; @@ -909,7 +910,13 @@ private static Map applyRenameColumnsToOptions( + "." + fieldName + "." - + MAP_SHARED_SHREDDING_MAX_COLUMNS); + + MAP_SHARED_SHREDDING_MAX_COLUMNS, + fieldName -> + FIELDS_PREFIX + + "." + + fieldName + + "." + + MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY); for (RenameColumn rename : renameColumns) { String fieldName = rename.fieldNames()[0]; diff --git a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java index 5876235bd802..2b0f4f18bf66 100644 --- a/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/schema/SchemaManagerTest.java @@ -54,6 +54,7 @@ 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; import java.io.File; import java.io.IOException; @@ -208,10 +209,14 @@ public void testCannotChangeMapStorageLayoutByRenameColumn() throws Exception { "Cannot change map storage layout for field id 1 ('metrics' -> 'renamed_metrics') from 'default' to 'shared-shredding'."); } - @Test - public void testRenameColumnKeepsMapStorageLayoutOptions() throws Exception { + @ParameterizedTest + @ValueSource(strings = {"plain", "sequential"}) + public void testRenameColumnKeepsMapStorageLayoutOptions(String placementPolicy) + throws Exception { retryArtificialException( - () -> manager.createTable(mapStorageLayoutSchema("shared-shredding"))); + () -> + manager.createTable( + mapStorageLayoutSchema("shared-shredding", placementPolicy))); retryArtificialException( () -> manager.commitChanges(SchemaChange.renameColumn("metrics", "renamed"))); @@ -219,17 +224,31 @@ public void testRenameColumnKeepsMapStorageLayoutOptions() throws Exception { Optional latest = retryArtificialException(() -> manager.latest()); assertThat(latest.isPresent()).isTrue(); assertThat(latest.get().options()) - .doesNotContainKey("fields.metrics.map.storage-layout") + .doesNotContainKeys( + "fields.metrics.map.storage-layout", + "fields.metrics.map.shared-shredding.max-columns", + "fields.metrics.map.shared-shredding.column-placement-policy") .containsEntry("fields.renamed.map.storage-layout", "shared-shredding") - .containsEntry("fields.renamed.map.shared-shredding.max-columns", "2"); + .containsEntry("fields.renamed.map.shared-shredding.max-columns", "2") + .containsEntry( + "fields.renamed.map.shared-shredding.column-placement-policy", + placementPolicy); } private Schema mapStorageLayoutSchema(String layout) { + return mapStorageLayoutSchema(layout, null); + } + + private Schema mapStorageLayoutSchema(String layout, String placementPolicy) { Map options = new HashMap<>(); if (layout != null) { options.put("fields.metrics.map.storage-layout", layout); options.put("fields.metrics.map.shared-shredding.max-columns", "2"); } + if (placementPolicy != null) { + options.put( + "fields.metrics.map.shared-shredding.column-placement-policy", placementPolicy); + } return new Schema( Arrays.asList( new DataField(0, "id", DataTypes.INT()), From 5d7dfe629aaebf73f50b004100edb030df8e54ed Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Fri, 24 Jul 2026 11:27:47 +0800 Subject: [PATCH 6/6] fix comment --- .../paimon/utils/FormatReaderMapping.java | 4 +- .../table/MapSharedShreddingTableTest.java | 43 ++++++++++++++++++- .../paimon/utils/FormatReaderMappingTest.java | 19 ++++++++ 3 files changed, 64 insertions(+), 2 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java index ba70953e9704..36c1d1a6028c 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/FormatReaderMapping.java @@ -356,7 +356,9 @@ private DataField selectedKeysDataField(DataField expectedField, DataField dataF selectedKeyField.newType( dataValueType.copy(selectedKeyField.type().isNullable()))); } - return expectedField.newType(selectedKeysType.copy(selectedKeysDataFields)); + return dataField + .newType(selectedKeysType.copy(selectedKeysDataFields)) + .newDescription(expectedField.description()); } private void checkSelectedKeysDataField(DataField dataField) { diff --git a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java index 6b109c90eaea..b392b927ce91 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/MapSharedShreddingTableTest.java @@ -248,6 +248,43 @@ public void testReadSelectedKeysAfterMapValueTypeEvolution(String format) throws assertThat(actual).containsExactly(Arrays.asList(10L, null, null)); } + @ParameterizedTest + @ValueSource(strings = {"orc", "parquet"}) + public void testReadSelectedKeysAfterRenameColumn(String format) throws Exception { + Table table = createTable(format, 1, "metrics"); + write( + table, + GenericRow.of(1, mapOf("key1", 10L, "key2", 20L)), + GenericRow.of(2, mapOf("key2", 30L))); + + catalog.alterTable( + identifier(format), + Collections.singletonList(SchemaChange.renameColumn("metrics", "renamed_metrics")), + false); + table = catalog.getTable(identifier(format)); + + ReadBuilder readBuilder = + table.newReadBuilder().withReadType(selectedKeysReadType("renamed_metrics")); + Map> actual = new LinkedHashMap<>(); + try (RecordReader reader = + readBuilder.newRead().createReader(readBuilder.newScan().plan())) { + reader.forEachRemaining( + row -> { + InternalRow selectedKeys = row.getRow(1, 3); + actual.put( + row.getInt(0), + Arrays.asList( + selectedKeys.isNullAt(0) ? null : selectedKeys.getLong(0), + selectedKeys.isNullAt(1) ? null : selectedKeys.getLong(1), + selectedKeys.isNullAt(2) ? null : selectedKeys.getLong(2))); + }); + } + + assertThat(actual) + .containsEntry(1, Arrays.asList(10L, 20L, null)) + .containsEntry(2, Arrays.asList(null, 30L, null)); + } + @ParameterizedTest @ValueSource(strings = {"orc", "parquet"}) public void testAppendOnlyTableReadWriteWithTwoMapFields(String format) throws Exception { @@ -1392,6 +1429,10 @@ private Schema schema(String format, int maxColumns, String... sharedShreddingFi } private RowType selectedKeysReadType() { + return selectedKeysReadType("metrics"); + } + + private RowType selectedKeysReadType(String fieldName) { RowType selectedKeysType = DataTypes.ROW( DataTypes.FIELD(0, "0", DataTypes.BIGINT()), @@ -1400,7 +1441,7 @@ private RowType selectedKeysReadType() { return DataTypes.ROW( DataTypes.FIELD(0, "id", DataTypes.INT()), MapSelectedKeysMetadataUtils.withSelectedKeys( - DataTypes.FIELD(1, "metrics", selectedKeysType), + DataTypes.FIELD(1, fieldName, selectedKeysType), selectedKeysType, Arrays.asList("key1", "key2", "missing"))); } diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java index 8a43de5af33b..b14c6b6ea35a 100644 --- a/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/utils/FormatReaderMappingTest.java @@ -213,6 +213,25 @@ public void testMapSelectedKeysUsesDataValueTypeForSchemaEvolution() throws Exce Assertions.assertThat(selectedKeys.getLong(0)).isEqualTo(10L); } + @Test + public void testMapSelectedKeysUsesDataFieldNameAfterRename() throws Exception { + DataField dataField = + DataTypes.FIELD( + 2, "old_attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT())); + DataField expectedField = selectedKeysField(2, "new_attrs", DataTypes.BIGINT()); + + DataField readField = + invokeDataFields( + Collections.singletonList(dataField), + Collections.singletonList(expectedField)) + .get(0); + + Assertions.assertThat(readField.name()).isEqualTo("old_attrs"); + Assertions.assertThat(readField.description()).isEqualTo(expectedField.description()); + Assertions.assertThat(MapSelectedKeysMetadataUtils.isMapSelectedKeysField(readField)) + .isTrue(); + } + @Test public void testRejectMapSelectedKeysValuePartialProjection() { RowType tableValueType =