Skip to content

[core][spark] Support selected-key pushdown for shared-shredding MAP #8782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lszskye wants to merge 6 commits into
base: master
Choose a base branch
from
Open

[core][spark] Support selected-key pushdown for shared-shredding MAP #8782

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9583945
[core][spark] Support selected-key pushdown for shared-shredding MAP
lszskye Jul 21, 2026
4dadbfa
fix ci
lszskye Jul 21, 2026
3ee70c5
fix comment
lszskye Jul 22, 2026
5708a1f
fix ci
lszskye Jul 23, 2026
4e8b5a7
fix comment
lszskye Jul 23, 2026
5d7dfe6
fix comment
lszskye Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
342 changes: 297 additions & 45 deletions paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlan.java
View file
Open in desktop

Large diffs are not rendered by default.

102 changes: 83 additions & 19 deletions paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSharedShreddingUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down Expand Up @@ -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;
}
Expand All @@ -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<String, Integer> buildColumnToNumColumns(
List<String> shreddingFieldNames, CoreOptions options) {
Map<String, Integer> fieldToNumColumns = new HashMap<>();
Expand Down Expand Up @@ -245,28 +272,65 @@ public static boolean hasShreddingMetadata(@Nullable Map<String, String> 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<Integer> physicalColumnIds, boolean includeOverflow) {
return innerBuildSpecificPhysicalStructType(
valueType, new ArrayList<>(new TreeSet<>(physicalColumnIds)), includeOverflow);
}

private static RowType innerBuildSpecificPhysicalStructType(
DataType valueType, List<Integer> 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));
}
return builder.build();
}

private static Set<Integer> physicalColumnIds(int numColumns) {
Set<Integer> physicalColumnIds = new TreeSet<>();
for (int i = 0; i < numColumns; i++) {
physicalColumnIds.add(i);
}
return physicalColumnIds;
}

private static Set<Integer> selectedPhysicalColumnIds(
DataField selectedKeysField, MapSharedShreddingFieldMeta fieldMeta) {
Set<Integer> selectedColumns = new TreeSet<>();
for (String selectedKey :
MapSelectedKeysMetadataUtils.selectedKeys(selectedKeysField.description())) {
Integer fieldId = fieldMeta.nameToId().get(selectedKey);
if (fieldId == null) {
continue;
}
List<Integer> 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<Integer, List<Integer>> sortedFieldColumns(
Map<Integer, List<Integer>> fieldToColumns) {
Map<Integer, List<Integer>> result = new TreeMap<>();
Expand Down
125 changes: 125 additions & 0 deletions ...common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingReadPlanTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -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<Integer>(), 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 =
Expand All @@ -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<Integer>(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<Integer>(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);
Expand All @@ -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<String, MapSharedShreddingFieldMeta> fieldMetas = new LinkedHashMap<>();
fieldMetas.put("metrics", fieldMeta);
return new MapSharedShreddingReadPlan(selectedKeysLogicalType(), fieldMetas);
}

private static RowType logicalType() {
return DataTypes.ROW(
DataTypes.FIELD(
Expand All @@ -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();
Expand All @@ -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) {
Expand All @@ -182,4 +298,13 @@ private static Map<String, Integer> nameToId(Object... pairs) {
}
return result;
}

@SuppressWarnings("unchecked")
private static Map<Integer, List<Integer>> fieldToColumns(Object... pairs) {
Map<Integer, List<Integer>> result = new TreeMap<>();
for (int i = 0; i < pairs.length; i += 2) {
result.put((Integer) pairs[i], (List<Integer>) pairs[i + 1]);
}
return result;
}
}
28 changes: 28 additions & 0 deletions ...on-common/src/test/java/org/apache/paimon/data/shredding/MapSharedShreddingUtilsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -186,6 +187,33 @@ void testLogicalToPhysicalSchemaNoShreddingColumns() {
.isEqualTo(logical);
}

@Test
void testBuildSpecificPhysicalStructType() {
RowType physicalType =
MapSharedShreddingUtils.buildSpecificPhysicalStructType(
DataTypes.BIGINT().notNull(),
new TreeSet<Integer>(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<Integer>(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<String, Integer> nameToId = new TreeMap<>();
Expand Down
Loading
Loading