Skip to content

Commit 3351e51

Browse files
committed
[core] Support append read/write for MAP shared-shredding
1 parent d9383ae commit 3351e51

15 files changed

Lines changed: 2704 additions & 37 deletions

File tree

Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,304 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.data.shredding;
20+
21+
import org.apache.paimon.data.BinaryString;
22+
import org.apache.paimon.data.InternalArray;
23+
import org.apache.paimon.data.InternalMap;
24+
import org.apache.paimon.data.columnar.ArrayColumnVector;
25+
import org.apache.paimon.data.columnar.ColumnVector;
26+
import org.apache.paimon.data.columnar.ColumnVectorUtils;
27+
import org.apache.paimon.data.columnar.MapColumnVector;
28+
import org.apache.paimon.data.columnar.RowColumnVector;
29+
import org.apache.paimon.data.columnar.RowToColumnConverter;
30+
import org.apache.paimon.data.columnar.VectorizedColumnBatch;
31+
import org.apache.paimon.data.columnar.heap.CastedMapColumnVector;
32+
import org.apache.paimon.data.columnar.heap.HeapMapVector;
33+
import org.apache.paimon.data.columnar.writable.WritableColumnVector;
34+
import org.apache.paimon.types.DataField;
35+
import org.apache.paimon.types.DataType;
36+
import org.apache.paimon.types.MapType;
37+
import org.apache.paimon.types.RowType;
38+
39+
import java.util.Arrays;
40+
import java.util.LinkedHashMap;
41+
import java.util.Map;
42+
43+
import static org.apache.paimon.utils.Preconditions.checkArgument;
44+
45+
/** Read plan that rebuilds logical MAP values from shared-shredding physical ROW values. */
46+
public class MapSharedShreddingReadPlan implements ShreddingReadPlan {
47+
48+
private final RowType logicalType;
49+
private final RowType physicalType;
50+
private final Map<Integer, SharedShreddingContext> contextByFieldIndex;
51+
52+
public MapSharedShreddingReadPlan(
53+
RowType logicalType, Map<String, MapSharedShreddingFieldMeta> fieldMetas) {
54+
this.logicalType = logicalType;
55+
this.physicalType = MapSharedShreddingUtils.buildPhysicalReadType(logicalType, fieldMetas);
56+
this.contextByFieldIndex = createContexts(logicalType, fieldMetas);
57+
}
58+
59+
@Override
60+
public RowType logicalRowType() {
61+
return logicalType;
62+
}
63+
64+
@Override
65+
public RowType physicalRowType() {
66+
return physicalType;
67+
}
68+
69+
@Override
70+
public boolean isIdentity() {
71+
return contextByFieldIndex.isEmpty();
72+
}
73+
74+
@Override
75+
public ShreddingBatchAssembler batchAssembler() {
76+
return new MapSharedShreddingBatchAssembler();
77+
}
78+
79+
private static Map<Integer, SharedShreddingContext> createContexts(
80+
RowType logicalType, Map<String, MapSharedShreddingFieldMeta> fieldMetas) {
81+
Map<Integer, SharedShreddingContext> contexts = new LinkedHashMap<>();
82+
for (int i = 0; i < logicalType.getFieldCount(); i++) {
83+
DataField field = logicalType.getFields().get(i);
84+
MapSharedShreddingFieldMeta fieldMeta = fieldMetas.get(field.name());
85+
if (fieldMeta != null && field.type() instanceof MapType) {
86+
contexts.put(i, new SharedShreddingContext(fieldMeta, field.type()));
87+
}
88+
}
89+
return contexts;
90+
}
91+
92+
private class MapSharedShreddingBatchAssembler implements ShreddingBatchAssembler {
93+
94+
@Override
95+
public VectorizedColumnBatch assemble(VectorizedColumnBatch physicalBatch) {
96+
ColumnVector[] logicalVectors = new ColumnVector[logicalType.getFieldCount()];
97+
for (int i = 0; i < logicalVectors.length; i++) {
98+
SharedShreddingContext context = contextByFieldIndex.get(i);
99+
if (context == null) {
100+
logicalVectors[i] = physicalBatch.columns[i];
101+
} else {
102+
logicalVectors[i] =
103+
materializeLogicalMapVector(
104+
(RowColumnVector) physicalBatch.columns[i],
105+
context,
106+
physicalBatch.getNumRows());
107+
}
108+
}
109+
return physicalBatch.copy(logicalVectors);
110+
}
111+
}
112+
113+
private static CastedMapColumnVector materializeLogicalMapVector(
114+
RowColumnVector physicalVector, SharedShreddingContext context, int rowCount) {
115+
ColumnVector[] physicalChildren = physicalChildren(physicalVector);
116+
int totalElements =
117+
countLogicalMapElements(physicalVector, physicalChildren, context, rowCount);
118+
119+
WritableColumnVector keyVector =
120+
ColumnVectorUtils.createWritableColumnVector(
121+
totalElements, context.mapType.getKeyType());
122+
WritableColumnVector valueVector =
123+
ColumnVectorUtils.createWritableColumnVector(
124+
totalElements, context.mapType.getValueType());
125+
HeapMapVector mapVector = new HeapMapVector(rowCount, keyVector, valueVector);
126+
127+
int offset = 0;
128+
for (int row = 0; row < rowCount; row++) {
129+
if (physicalVector.isNullAt(row)) {
130+
mapVector.setNullAt(row);
131+
mapVector.putOffsetLength(row, offset, 0);
132+
continue;
133+
}
134+
135+
int elementCount =
136+
appendLogicalMapElements(
137+
physicalChildren, row, context, keyVector, valueVector);
138+
mapVector.putOffsetLength(row, offset, elementCount);
139+
offset += elementCount;
140+
}
141+
mapVector.addElementsAppended(rowCount);
142+
143+
return new CastedMapColumnVector(
144+
mapVector,
145+
ColumnVectorUtils.createReadableColumnVectors(
146+
Arrays.asList(context.mapType.getKeyType(), context.mapType.getValueType()),
147+
new WritableColumnVector[] {keyVector, valueVector}));
148+
}
149+
150+
private static ColumnVector[] physicalChildren(RowColumnVector physicalVector) {
151+
ColumnVector[] physicalChildren = physicalVector.getChildren();
152+
if (physicalChildren != null) {
153+
return physicalChildren;
154+
}
155+
return physicalVector.getBatch().columns;
156+
}
157+
158+
private static int countLogicalMapElements(
159+
RowColumnVector physicalVector,
160+
ColumnVector[] physicalChildren,
161+
SharedShreddingContext context,
162+
int rowCount) {
163+
int totalElements = 0;
164+
for (int row = 0; row < rowCount; row++) {
165+
if (!physicalVector.isNullAt(row)) {
166+
totalElements += countLogicalMapElements(physicalChildren, row, context);
167+
}
168+
}
169+
return totalElements;
170+
}
171+
172+
private static int countLogicalMapElements(
173+
ColumnVector[] physicalChildren, int row, SharedShreddingContext context) {
174+
int count = 0;
175+
InternalArray fieldMapping = fieldMapping(physicalChildren, row, context);
176+
for (int column = 0; column < context.numColumns; column++) {
177+
if (logicalFieldName(fieldMapping, column, context) != null) {
178+
count++;
179+
}
180+
}
181+
182+
InternalMap overflow = overflowMap(physicalChildren, row, context);
183+
if (overflow != null) {
184+
InternalArray keys = overflow.keyArray();
185+
for (int i = 0; i < overflow.size(); i++) {
186+
if (context.nameById.containsKey(keys.getInt(i))) {
187+
count++;
188+
}
189+
}
190+
}
191+
return count;
192+
}
193+
194+
private static int appendLogicalMapElements(
195+
ColumnVector[] physicalChildren,
196+
int row,
197+
SharedShreddingContext context,
198+
WritableColumnVector keyVector,
199+
WritableColumnVector valueVector) {
200+
int count = 0;
201+
InternalArray fieldMapping = fieldMapping(physicalChildren, row, context);
202+
for (int column = 0; column < context.numColumns; column++) {
203+
BinaryString fieldName = logicalFieldName(fieldMapping, column, context);
204+
if (fieldName == null) {
205+
continue;
206+
}
207+
208+
context.keyConverter.append(fieldName, keyVector);
209+
ColumnVector valueColumn = physicalChildren[column + 1];
210+
appendValue(context, valueColumn, row, valueVector);
211+
count++;
212+
}
213+
214+
InternalMap overflow = overflowMap(physicalChildren, row, context);
215+
if (overflow != null) {
216+
InternalArray keys = overflow.keyArray();
217+
InternalArray values = overflow.valueArray();
218+
for (int i = 0; i < overflow.size(); i++) {
219+
BinaryString fieldName = context.nameById.get(keys.getInt(i));
220+
if (fieldName != null) {
221+
context.keyConverter.append(fieldName, keyVector);
222+
appendValue(context, values, i, valueVector);
223+
count++;
224+
}
225+
}
226+
}
227+
return count;
228+
}
229+
230+
private static void appendValue(
231+
SharedShreddingContext context,
232+
ColumnVector source,
233+
int row,
234+
WritableColumnVector valueVector) {
235+
if (source.isNullAt(row)) {
236+
valueVector.appendNull();
237+
} else {
238+
context.valueConverter.append(source, row, valueVector);
239+
}
240+
}
241+
242+
private static void appendValue(
243+
SharedShreddingContext context,
244+
InternalArray source,
245+
int pos,
246+
WritableColumnVector valueVector) {
247+
if (source.isNullAt(pos)) {
248+
valueVector.appendNull();
249+
} else {
250+
context.valueConverter.append(source, pos, valueVector);
251+
}
252+
}
253+
254+
private static InternalArray fieldMapping(
255+
ColumnVector[] physicalChildren, int row, SharedShreddingContext context) {
256+
InternalArray fieldMapping = ((ArrayColumnVector) physicalChildren[0]).getArray(row);
257+
checkArgument(
258+
fieldMapping.size() == context.numColumns,
259+
"Shared-shredding field mapping size %s does not match metadata num columns %s.",
260+
fieldMapping.size(),
261+
context.numColumns);
262+
return fieldMapping;
263+
}
264+
265+
private static BinaryString logicalFieldName(
266+
InternalArray fieldMapping, int column, SharedShreddingContext context) {
267+
int fieldId = fieldMapping.isNullAt(column) ? -1 : fieldMapping.getInt(column);
268+
return fieldId < 0 ? null : context.nameById.get(fieldId);
269+
}
270+
271+
private static InternalMap overflowMap(
272+
ColumnVector[] physicalChildren, int row, SharedShreddingContext context) {
273+
if (context.overflowPosition >= physicalChildren.length) {
274+
return null;
275+
}
276+
277+
ColumnVector overflowVector = physicalChildren[context.overflowPosition];
278+
return overflowVector.isNullAt(row) ? null : ((MapColumnVector) overflowVector).getMap(row);
279+
}
280+
281+
private static class SharedShreddingContext {
282+
283+
private final MapType mapType;
284+
private final Map<Integer, BinaryString> nameById;
285+
private final RowToColumnConverter.ElementConverter keyConverter;
286+
private final RowToColumnConverter.ElementConverter valueConverter;
287+
private final int numColumns;
288+
private final int overflowPosition;
289+
290+
private SharedShreddingContext(MapSharedShreddingFieldMeta fieldMeta, DataType fieldType) {
291+
this.mapType = (MapType) fieldType;
292+
this.nameById = new LinkedHashMap<>();
293+
for (Map.Entry<String, Integer> entry : fieldMeta.nameToId().entrySet()) {
294+
this.nameById.put(entry.getValue(), BinaryString.fromString(entry.getKey()));
295+
}
296+
this.keyConverter =
297+
RowToColumnConverter.createElementConverter(this.mapType.getKeyType());
298+
this.valueConverter =
299+
RowToColumnConverter.createElementConverter(this.mapType.getValueType());
300+
this.numColumns = fieldMeta.numColumns();
301+
this.overflowPosition = fieldMeta.numColumns() + 1;
302+
}
303+
}
304+
}
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.data.shredding;
20+
21+
import org.apache.paimon.format.shredding.ShreddingReadPlanFactory;
22+
import org.apache.paimon.types.DataField;
23+
import org.apache.paimon.types.RowType;
24+
25+
import javax.annotation.Nullable;
26+
27+
import java.util.LinkedHashMap;
28+
import java.util.Map;
29+
30+
/** Creates per-file shared-shredding MAP read plans. */
31+
public class MapSharedShreddingReadPlanFactory implements ShreddingReadPlanFactory {
32+
33+
private final RowType logicalRowType;
34+
35+
public MapSharedShreddingReadPlanFactory(RowType logicalRowType) {
36+
this.logicalRowType = logicalRowType;
37+
}
38+
39+
@Override
40+
public RowType logicalRowType() {
41+
return logicalRowType;
42+
}
43+
44+
@Override
45+
public boolean shouldCreateReadPlan(
46+
Map<String, Map<String, String>> fieldMetadata, @Nullable Object fileSchema) {
47+
for (DataField field : logicalRowType.getFields()) {
48+
Map<String, String> metadata = fieldMetadata.get(field.name());
49+
if (MapSharedShreddingUtils.hasShreddingMetadata(metadata)) {
50+
return true;
51+
}
52+
}
53+
return false;
54+
}
55+
56+
@Override
57+
public ShreddingReadPlan createReadPlan(
58+
Map<String, Map<String, String>> fieldMetadata, @Nullable Object fileSchema) {
59+
return new MapSharedShreddingReadPlan(
60+
logicalRowType, readSharedShreddingMetas(logicalRowType, fieldMetadata));
61+
}
62+
63+
private static Map<String, MapSharedShreddingFieldMeta> readSharedShreddingMetas(
64+
RowType logicalRowType, Map<String, Map<String, String>> fieldMetadata) {
65+
Map<String, MapSharedShreddingFieldMeta> metas = new LinkedHashMap<>();
66+
for (Map.Entry<String, Map<String, String>> entry : fieldMetadata.entrySet()) {
67+
if (logicalRowType != null && !logicalRowType.containsField(entry.getKey())) {
68+
continue;
69+
}
70+
if (MapSharedShreddingUtils.hasShreddingMetadata(entry.getValue())) {
71+
metas.put(
72+
entry.getKey(),
73+
MapSharedShreddingUtils.deserializeMetadata(entry.getValue()));
74+
}
75+
}
76+
return metas;
77+
}
78+
}

0 commit comments

Comments
 (0)