Skip to content

Commit ef90d01

Browse files
xiangfu0xiangfu0
committed
[UUID 2/8] UUID ingest and segment storage
Part 2/8 of splitting #18140 (logical UUID type). Rebased onto master; stacked on uuid-split/01-spi-foundation. See #18869 for the type foundation. Downstream references use the extracted UuidKey class.
1 parent 075e9f9 commit ef90d01

36 files changed

Lines changed: 1299 additions & 61 deletions

File tree

‎pinot-core/src/main/java/org/apache/pinot/core/segment/processing/genericrow/GenericRowDeserializer.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,18 @@ public void deserialize(long offset, GenericRow buffer) {
164164
multiValue[j] = new String(stringBytes, UTF_8);
165165
}
166166
break;
167+
case BYTES:
168+
// Used by MV UUID (storedType=BYTES). Mirrors the writer in GenericRowSerializer:
169+
// each element is length-prefixed.
170+
for (int j = 0; j < numValues; j++) {
171+
int numBytes = _dataBuffer.getInt(offset);
172+
offset += Integer.BYTES;
173+
byte[] elementBytes = new byte[numBytes];
174+
_dataBuffer.copyTo(offset, elementBytes);
175+
offset += numBytes;
176+
multiValue[j] = elementBytes;
177+
}
178+
break;
167179
default:
168180
throw new IllegalStateException("Unsupported MV stored type: " + _storedTypes[i]);
169181
}
@@ -366,6 +378,26 @@ public int compare(long offset1, long offset2, int numFieldsToCompare) {
366378
offset2 += numBytes2;
367379
}
368380
break;
381+
case BYTES:
382+
// Used by MV UUID (storedType=BYTES). Element-wise unsigned byte compare matches the SV BYTES
383+
// ordering above.
384+
for (int j = 0; j < numValues; j++) {
385+
int numBytes1 = _dataBuffer.getInt(offset1);
386+
offset1 += Integer.BYTES;
387+
byte[] elementBytes1 = new byte[numBytes1];
388+
_dataBuffer.copyTo(offset1, elementBytes1);
389+
int numBytes2 = _dataBuffer.getInt(offset2);
390+
offset2 += Integer.BYTES;
391+
byte[] elementBytes2 = new byte[numBytes2];
392+
_dataBuffer.copyTo(offset2, elementBytes2);
393+
int result = ByteArray.compare(elementBytes1, elementBytes2);
394+
if (result != 0) {
395+
return result;
396+
}
397+
offset1 += numBytes1;
398+
offset2 += numBytes2;
399+
}
400+
break;
369401
default:
370402
throw new IllegalStateException("Unsupported MV stored type: " + _storedTypes[i]);
371403
}

‎pinot-core/src/main/java/org/apache/pinot/core/segment/processing/genericrow/GenericRowSerializer.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,14 @@ public byte[] serialize(GenericRow row) {
148148
}
149149
_objectBytes[i] = stringBytesArray;
150150
break;
151+
case BYTES:
152+
// Used by MV UUID (storedType=BYTES) — also supports any MV BYTES column going through the
153+
// SegmentProcessorFramework. Each element is length-prefixed like SV BYTES.
154+
numBytes += Integer.BYTES * numValues;
155+
for (Object element : multiValue) {
156+
numBytes += ((byte[]) element).length;
157+
}
158+
break;
151159
default:
152160
throw new IllegalStateException("Unsupported MV stored type: " + _storedTypes[i]);
153161
}
@@ -240,6 +248,13 @@ public byte[] serialize(GenericRow row) {
240248
byteBuffer.put(stringBytes);
241249
}
242250
break;
251+
case BYTES:
252+
for (Object element : multiValue) {
253+
byte[] elementBytes = (byte[]) element;
254+
byteBuffer.putInt(elementBytes.length);
255+
byteBuffer.put(elementBytes);
256+
}
257+
break;
243258
default:
244259
throw new IllegalStateException("Unsupported MV stored type: " + _storedTypes[i]);
245260
}

‎pinot-core/src/main/java/org/apache/pinot/core/util/SegmentProcessorAvroUtils.java‎

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,14 @@
1919
package org.apache.pinot.core.util;
2020

2121
import java.nio.ByteBuffer;
22+
import java.util.ArrayList;
2223
import java.util.Arrays;
2324
import java.util.Comparator;
2425
import java.util.List;
2526
import java.util.Set;
2627
import java.util.stream.Collectors;
28+
import org.apache.avro.LogicalType;
29+
import org.apache.avro.LogicalTypes;
2730
import org.apache.avro.Schema;
2831
import org.apache.avro.SchemaBuilder;
2932
import org.apache.avro.generic.GenericData;
@@ -32,6 +35,7 @@
3235
import org.apache.pinot.spi.data.FieldSpec.DataType;
3336
import org.apache.pinot.spi.data.readers.GenericRow;
3437
import org.apache.pinot.spi.ingestion.segment.writer.SegmentWriter;
38+
import org.apache.pinot.spi.utils.UuidUtils;
3539

3640

3741
/**
@@ -57,20 +61,50 @@ public static GenericData.Record convertGenericRowToAvroRecord(GenericRow generi
5761
*/
5862
public static GenericData.Record convertGenericRowToAvroRecord(GenericRow genericRow,
5963
GenericData.Record reusableRecord, Set<String> fields) {
64+
Schema avroSchema = reusableRecord.getSchema();
6065
for (String field : fields) {
6166
Object value = genericRow.getValue(field);
6267
if (value instanceof Object[]) {
63-
reusableRecord.put(field, Arrays.asList((Object[]) value));
68+
Schema.Field avroField = avroSchema.getField(field);
69+
if (avroField != null && isUuidArrayLogicalType(avroField.schema())) {
70+
// MV UUID columns are emitted with an Avro array<string{logicalType:uuid}> schema; convert each
71+
// 16-byte element to its canonical UUID string so GenericDatumWriter accepts the record.
72+
Object[] elements = (Object[]) value;
73+
List<Object> converted = new ArrayList<>(elements.length);
74+
for (Object element : elements) {
75+
converted.add(element instanceof byte[] ? UuidUtils.toString((byte[]) element) : element);
76+
}
77+
reusableRecord.put(field, converted);
78+
} else {
79+
reusableRecord.put(field, Arrays.asList((Object[]) value));
80+
}
6481
} else {
6582
if (value instanceof byte[]) {
66-
value = ByteBuffer.wrap((byte[]) value);
83+
// UUID columns are emitted with an Avro string{logicalType:uuid} schema (Avro 1.x only allows the
84+
// uuid logical type on string), so the 16-byte canonical form must be rendered as a canonical
85+
// UUID string here. Plain BYTES columns continue to be wrapped as ByteBuffer.
86+
Schema.Field avroField = avroSchema.getField(field);
87+
if (avroField != null && isUuidLogicalType(avroField.schema())) {
88+
value = UuidUtils.toString((byte[]) value);
89+
} else {
90+
value = ByteBuffer.wrap((byte[]) value);
91+
}
6792
}
6893
reusableRecord.put(field, value);
6994
}
7095
}
7196
return reusableRecord;
7297
}
7398

99+
private static boolean isUuidLogicalType(Schema schema) {
100+
LogicalType logicalType = LogicalTypes.fromSchemaIgnoreInvalid(schema);
101+
return logicalType != null && "uuid".equals(logicalType.getName());
102+
}
103+
104+
private static boolean isUuidArrayLogicalType(Schema schema) {
105+
return schema.getType() == Schema.Type.ARRAY && isUuidLogicalType(schema.getElementType());
106+
}
107+
74108
/**
75109
* Converts a Pinot schema to an Avro schema
76110
*/
@@ -82,6 +116,19 @@ public static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Sc
82116
.collect(Collectors.toList());
83117
for (FieldSpec fieldSpec : orderedFieldSpecs) {
84118
String name = fieldSpec.getName();
119+
// Emit UUID columns as Avro string{logicalType:uuid} (matching AvroUtils.getAvroSchemaFromPinotSchema)
120+
// so the runtime byte[] → canonical-string conversion in convertGenericRowToAvroRecord lines up with
121+
// the field schema. Without this branch SV UUID would fall through to BYTES (losing UUID semantics) and
122+
// MV UUID would throw at this point (MV switch below has no BYTES case).
123+
if (fieldSpec.getDataType() == DataType.UUID) {
124+
Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
125+
if (fieldSpec.isSingleValueField()) {
126+
fieldAssembler = fieldAssembler.name(name).type(uuidSchema).noDefault();
127+
} else {
128+
fieldAssembler = fieldAssembler.name(name).type().array().items(uuidSchema).noDefault();
129+
}
130+
continue;
131+
}
85132
DataType storedType = fieldSpec.getDataType().getStoredType();
86133
if (fieldSpec.isSingleValueField()) {
87134
switch (storedType) {

‎pinot-core/src/test/java/org/apache/pinot/core/segment/processing/genericrow/GenericRowSerDeTest.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,10 @@ public void setUp() {
5555
new DimensionFieldSpec("floatMV", DataType.FLOAT, false),
5656
new DimensionFieldSpec("doubleMV", DataType.DOUBLE, false),
5757
new DimensionFieldSpec("stringMV", DataType.STRING, false),
58+
// MV BYTES is the storage shape for MV UUID columns; the GenericRow serializer/deserializer must
59+
// round-trip these (used by every SegmentProcessorFramework minion task — MergeRollup,
60+
// RealtimeToOffline, UpsertCompactMerge).
61+
new DimensionFieldSpec("bytesMV", DataType.BYTES, false),
5862
new DimensionFieldSpec("nullMV", DataType.LONG, false));
5963

6064
_row = new GenericRow();
@@ -78,6 +82,7 @@ public void setUp() {
7882
_row.putValue("floatMV", new Object[]{123.0f, 456.0f});
7983
_row.putValue("doubleMV", new Object[]{123.0, 456.0});
8084
_row.putValue("stringMV", new Object[]{"123", "456"});
85+
_row.putValue("bytesMV", new Object[]{new byte[]{1, 2, 3}, new byte[]{4, 5, 6, 7}, new byte[]{}});
8186
_row.putDefaultNullValue("nullMV", new Object[]{Long.MIN_VALUE});
8287
}
8388

‎pinot-core/src/test/java/org/apache/pinot/core/util/SegmentProcessorAvroUtilsTest.java‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.pinot.core.util;
20+
21+
import java.nio.ByteBuffer;
22+
import java.util.Arrays;
23+
import java.util.List;
24+
import org.apache.avro.LogicalTypes;
25+
import org.apache.avro.Schema;
26+
import org.apache.avro.SchemaBuilder;
27+
import org.apache.avro.generic.GenericData;
28+
import org.apache.pinot.spi.data.readers.GenericRow;
29+
import org.apache.pinot.spi.utils.UuidUtils;
30+
import org.testng.annotations.Test;
31+
32+
import static org.testng.Assert.assertEquals;
33+
34+
35+
public class SegmentProcessorAvroUtilsTest {
36+
37+
/**
38+
* Regression test: UUID columns are exported with an Avro {@code string{logicalType:uuid}} schema, so the
39+
* 16-byte canonical form held internally must be converted to a canonical UUID string at write time. Wrapping
40+
* it as a ByteBuffer (the default for plain BYTES columns) would be rejected by GenericDatumWriter against a
41+
* STRING-typed schema.
42+
*/
43+
@Test
44+
public void testConvertGenericRowToAvroRecordRendersUuidBytesAsCanonicalString() {
45+
Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
46+
Schema bytesSchema = Schema.create(Schema.Type.BYTES);
47+
Schema recordSchema = SchemaBuilder.record("record").fields()
48+
.name("uuidCol").type(uuidSchema).noDefault()
49+
.name("bytesCol").type(bytesSchema).noDefault()
50+
.endRecord();
51+
52+
String canonical = "12345678-1234-1234-1234-1234567890ab";
53+
byte[] uuidBytes = UuidUtils.toBytes(canonical);
54+
byte[] rawBytes = new byte[]{1, 2, 3, 4};
55+
56+
GenericRow row = new GenericRow();
57+
row.putValue("uuidCol", uuidBytes);
58+
row.putValue("bytesCol", rawBytes);
59+
60+
GenericData.Record reusableRecord = new GenericData.Record(recordSchema);
61+
SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, reusableRecord);
62+
63+
assertEquals(reusableRecord.get("uuidCol"), canonical,
64+
"UUID byte[] must be converted to canonical UUID string for string{logicalType:uuid} fields");
65+
assertEquals(reusableRecord.get("bytesCol"), ByteBuffer.wrap(rawBytes),
66+
"Plain BYTES columns must continue to be wrapped as ByteBuffer");
67+
}
68+
69+
/**
70+
* Regression: MV UUID columns are emitted as Avro {@code array<string{logicalType:uuid}>}. The 16-byte
71+
* canonical-form elements must be converted to canonical UUID strings, otherwise GenericDatumWriter would
72+
* reject the byte[] elements against the string-typed array schema.
73+
*/
74+
@Test
75+
public void testConvertGenericRowToAvroRecordRendersMvUuidBytesAsCanonicalStrings() {
76+
Schema uuidElementSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
77+
Schema uuidArraySchema = Schema.createArray(uuidElementSchema);
78+
Schema recordSchema = SchemaBuilder.record("record").fields()
79+
.name("uuidArrayCol").type(uuidArraySchema).noDefault()
80+
.endRecord();
81+
82+
String canonicalA = "12345678-1234-1234-1234-1234567890ab";
83+
String canonicalB = "550e8400-e29b-41d4-a716-446655440000";
84+
Object[] uuidMv = new Object[]{UuidUtils.toBytes(canonicalA), UuidUtils.toBytes(canonicalB)};
85+
86+
GenericRow row = new GenericRow();
87+
row.putValue("uuidArrayCol", uuidMv);
88+
89+
GenericData.Record reusableRecord = new GenericData.Record(recordSchema);
90+
SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, reusableRecord);
91+
92+
Object emitted = reusableRecord.get("uuidArrayCol");
93+
assertEquals(emitted instanceof List, true, "MV UUID column must be emitted as a List for Avro array schema");
94+
assertEquals(emitted, Arrays.asList(canonicalA, canonicalB),
95+
"MV UUID byte[] elements must each be converted to canonical UUID strings");
96+
}
97+
98+
/**
99+
* Regression: convertPinotSchemaToAvroSchema must emit SV UUID as {@code string{logicalType:uuid}} (not plain
100+
* BYTES) and MV UUID as {@code array<string{logicalType:uuid}>}. Without the UUID branch, SV UUID would
101+
* silently fall through to plain BYTES (losing UUID semantics in the output Avro file) and MV UUID would
102+
* throw at schema-construction time because the MV switch has no BYTES case.
103+
*/
104+
@Test
105+
public void testConvertPinotSchemaToAvroSchemaEmitsUuidLogicalType() {
106+
org.apache.pinot.spi.data.Schema pinotSchema = new org.apache.pinot.spi.data.Schema.SchemaBuilder()
107+
.setSchemaName("uuidSchema")
108+
.addSingleValueDimension("uuidSv", org.apache.pinot.spi.data.FieldSpec.DataType.UUID)
109+
.addMultiValueDimension("uuidMv", org.apache.pinot.spi.data.FieldSpec.DataType.UUID)
110+
.build();
111+
112+
Schema avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(pinotSchema);
113+
114+
Schema svFieldSchema = avroSchema.getField("uuidSv").schema();
115+
assertEquals(svFieldSchema.getType(), Schema.Type.STRING, "SV UUID must be string{logicalType:uuid}");
116+
assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(svFieldSchema), LogicalTypes.uuid(),
117+
"SV UUID must carry the uuid logical type");
118+
119+
Schema mvFieldSchema = avroSchema.getField("uuidMv").schema();
120+
assertEquals(mvFieldSchema.getType(), Schema.Type.ARRAY, "MV UUID must be emitted as an array");
121+
Schema mvElementSchema = mvFieldSchema.getElementType();
122+
assertEquals(mvElementSchema.getType(), Schema.Type.STRING,
123+
"MV UUID elements must be string{logicalType:uuid}");
124+
assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(mvElementSchema), LogicalTypes.uuid(),
125+
"MV UUID elements must carry the uuid logical type");
126+
}
127+
}

‎pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ private void validateSchemas() {
139139
if (fieldSpec.getDataType() != dataTypeForSVColumn) {
140140
_dataTypeMismatch.addMismatchReason(String
141141
.format("The Pinot column: (%s: %s) doesn't match with the column (%s: %s) in input %s schema.",
142-
columnName, fieldSpec.getDataType().name(), avroColumnName, avroColumnType.name(),
142+
columnName, fieldSpec.getDataType().name(), avroColumnName, dataTypeForSVColumn.name(),
143143
getInputSchemaType()));
144144
}
145145
} else {

0 commit comments

Comments
 (0)