Skip to content

Commit 917440f

Browse files
xiangfu0xiangfu0
committed
[UUID 2/8] UUID ingest and segment storage
Part 2/8 of splitting #18140 (logical UUID type). Rebased onto master (now includes the merged #18869 type foundation); stacked on 92a0d98. Downstream references use the UuidKey class merged in #18869.
1 parent f054a90 commit 917440f

38 files changed

Lines changed: 1140 additions & 79 deletions

File tree

‎pinot-connectors/pinot-flink-connector/src/main/java/org/apache/pinot/connector/flink/sink/FlinkSegmentWriter.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,8 @@ private void resetBuffer()
191191
throws IOException {
192192
FileUtils.deleteQuietly(_bufferFile);
193193
_rowCount = 0;
194-
_recordWriter = new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema));
194+
_recordWriter =
195+
new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema, SegmentProcessorAvroUtils.getAvroDataModel()));
195196
_recordWriter.create(_avroSchema, _bufferFile);
196197
}
197198

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

Lines changed: 89 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,20 +24,24 @@
2424
import java.util.List;
2525
import java.util.Set;
2626
import java.util.stream.Collectors;
27+
import org.apache.avro.Conversion;
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;
30-
import org.apache.pinot.core.segment.processing.framework.SegmentProcessorFramework;
3133
import org.apache.pinot.spi.data.FieldSpec;
3234
import org.apache.pinot.spi.data.FieldSpec.DataType;
3335
import org.apache.pinot.spi.data.readers.GenericRow;
3436
import org.apache.pinot.spi.ingestion.segment.writer.SegmentWriter;
37+
import org.apache.pinot.spi.utils.UuidUtils;
3538

3639

3740
/**
38-
* Helper methods for avro related conversions needed, when using AVRO as intermediate format in segment processing.
39-
* AVRO is used as intermediate processing format in {@link SegmentProcessorFramework} and file-based impl of
40-
* {@link SegmentWriter}
41+
* Helper methods for Avro conversions used when serializing Pinot {@link GenericRow}s to Avro — by the file-based
42+
* {@link SegmentWriter} implementations and the segment-to-Avro/Parquet converters. (The multi-stage
43+
* {@code SegmentProcessorFramework} itself serializes intermediate records with a custom binary {@code GenericRowFile}
44+
* format, not Avro.)
4145
*/
4246
public final class SegmentProcessorAvroUtils {
4347

@@ -57,20 +61,85 @@ 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[]) {
68+
// Array elements are written as-is. For MV UUID (array<string{logicalType:uuid}>) the elements are the raw
69+
// 16-byte values; the uuid Conversion registered on the writer's data model (getAvroDataModel) renders each
70+
// element to its canonical string at write time.
6371
reusableRecord.put(field, Arrays.asList((Object[]) value));
64-
} else {
65-
if (value instanceof byte[]) {
66-
value = ByteBuffer.wrap((byte[]) value);
72+
} else if (value instanceof byte[]) {
73+
// A byte[] bound for a plain BYTES field must be wrapped as ByteBuffer (GenericDatumWriter requires it for the
74+
// bytes type). A byte[] bound for a UUID field (string{logicalType:uuid}) is left raw so the uuid Conversion
75+
// registered on the writer's data model (getAvroDataModel) renders it to a canonical string at write time.
76+
Schema.Field avroField = avroSchema.getField(field);
77+
if (avroField != null && avroField.schema().getType() == Schema.Type.BYTES) {
78+
reusableRecord.put(field, ByteBuffer.wrap((byte[]) value));
79+
} else {
80+
reusableRecord.put(field, value);
6781
}
82+
} else {
6883
reusableRecord.put(field, value);
6984
}
7085
}
7186
return reusableRecord;
7287
}
7388

89+
/// Shared Avro data model with [UuidConversion] registered. Populated once at class initialization and never
90+
/// mutated afterward (effectively immutable), so it is safe to share across writers.
91+
private static final GenericData AVRO_DATA_MODEL = createAvroDataModel();
92+
93+
/// Returns the shared Avro data model that a `GenericDatumWriter` (or `AvroParquetWriter`) must be constructed with
94+
/// to serialize UUID columns produced by [#convertGenericRowToAvroRecord]: it registers [UuidConversion] so the
95+
/// internal 16-byte UUID form is rendered as the canonical string required by `string{logicalType:uuid}` fields.
96+
/// The UUID column's field schema must be `string{logicalType:uuid}` — as emitted by
97+
/// [#convertPinotSchemaToAvroSchema] and `AvroUtils.getAvroSchemaFromPinotSchema` — for the conversion to apply.
98+
///
99+
/// The returned instance is shared and must be treated as read-only: do not call its mutators
100+
/// (`addLogicalTypeConversion`, `setStringType`, ...), which are not thread-safe against concurrent writer reads.
101+
public static GenericData getAvroDataModel() {
102+
return AVRO_DATA_MODEL;
103+
}
104+
105+
private static GenericData createAvroDataModel() {
106+
GenericData model = new GenericData();
107+
model.addLogicalTypeConversion(new UuidConversion());
108+
return model;
109+
}
110+
111+
/// Avro logical-type [Conversion] for the `uuid` logical type, operating directly on Pinot's internal storage form
112+
/// (a 16-byte big-endian `byte[]`). It renders that value to its canonical RFC-4122 string when writing a
113+
/// `string{logicalType:uuid}` field and parses it back on read.
114+
///
115+
/// The converted type is `byte[]` rather than [java.util.UUID] so no intermediate object is allocated, and rather
116+
/// than [java.nio.ByteBuffer] because Avro resolves conversions by exact datum class: a `byte[]` instance always
117+
/// reports `byte[].class`, so the conversion resolves for both single values and array elements, whereas a
118+
/// concrete `HeapByteBuffer` would not match a `ByteBuffer`-typed conversion.
119+
private static final class UuidConversion extends Conversion<byte[]> {
120+
private static final String UUID_LOGICAL_TYPE_NAME = "uuid";
121+
122+
@Override
123+
public Class<byte[]> getConvertedType() {
124+
return byte[].class;
125+
}
126+
127+
@Override
128+
public String getLogicalTypeName() {
129+
return UUID_LOGICAL_TYPE_NAME;
130+
}
131+
132+
@Override
133+
public CharSequence toCharSequence(byte[] value, Schema schema, LogicalType type) {
134+
return UuidUtils.toString(value);
135+
}
136+
137+
@Override
138+
public byte[] fromCharSequence(CharSequence value, Schema schema, LogicalType type) {
139+
return UuidUtils.toBytes(value.toString());
140+
}
141+
}
142+
74143
/**
75144
* Converts a Pinot schema to an Avro schema
76145
*/
@@ -82,6 +151,19 @@ public static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Sc
82151
.collect(Collectors.toList());
83152
for (FieldSpec fieldSpec : orderedFieldSpecs) {
84153
String name = fieldSpec.getName();
154+
// Emit UUID columns as Avro string{logicalType:uuid} (matching AvroUtils.getAvroSchemaFromPinotSchema)
155+
// so the runtime byte[] → canonical-string conversion in convertGenericRowToAvroRecord lines up with
156+
// the field schema. Without this branch SV UUID would fall through to BYTES (losing UUID semantics) and
157+
// MV UUID would throw at this point (MV switch below has no BYTES case).
158+
if (fieldSpec.getDataType() == DataType.UUID) {
159+
Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
160+
if (fieldSpec.isSingleValueField()) {
161+
fieldAssembler = fieldAssembler.name(name).type(uuidSchema).noDefault();
162+
} else {
163+
fieldAssembler = fieldAssembler.name(name).type().array().items(uuidSchema).noDefault();
164+
}
165+
continue;
166+
}
85167
DataType storedType = fieldSpec.getDataType().getStoredType();
86168
if (fieldSpec.isSingleValueField()) {
87169
switch (storedType) {

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

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
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.io.File;
22+
import java.nio.ByteBuffer;
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.file.DataFileReader;
28+
import org.apache.avro.file.DataFileWriter;
29+
import org.apache.avro.generic.GenericData;
30+
import org.apache.avro.generic.GenericDatumReader;
31+
import org.apache.avro.generic.GenericDatumWriter;
32+
import org.apache.avro.generic.GenericRecord;
33+
import org.apache.commons.io.FileUtils;
34+
import org.apache.pinot.spi.data.FieldSpec.DataType;
35+
import org.apache.pinot.spi.data.readers.GenericRow;
36+
import org.apache.pinot.spi.utils.UuidUtils;
37+
import org.testng.annotations.Test;
38+
39+
import static org.testng.Assert.assertEquals;
40+
import static org.testng.Assert.assertTrue;
41+
42+
43+
public class SegmentProcessorAvroUtilsTest {
44+
45+
/// convertGenericRowToAvroRecord keeps a UUID value as its raw 16-byte form (it does NOT render a string here); the
46+
/// rendering is deferred to the uuid Conversion on the writer's data model. Plain BYTES columns are still wrapped as
47+
/// ByteBuffer.
48+
@Test
49+
public void testConvertGenericRowToAvroRecordKeepsUuidBytesRaw() {
50+
Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
51+
Schema recordSchema = SchemaBuilder.record("record").fields()
52+
.name("uuidCol").type(uuidSchema).noDefault()
53+
.name("bytesCol").type().bytesType().noDefault()
54+
.endRecord();
55+
56+
byte[] uuidBytes = UuidUtils.toBytes("12345678-1234-1234-1234-1234567890ab");
57+
byte[] rawBytes = {1, 2, 3, 4};
58+
GenericRow row = new GenericRow();
59+
row.putValue("uuidCol", uuidBytes);
60+
row.putValue("bytesCol", rawBytes);
61+
62+
GenericData.Record record = new GenericData.Record(recordSchema);
63+
SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, record);
64+
65+
assertTrue(record.get("uuidCol") instanceof byte[], "UUID must be left as raw byte[] for the uuid Conversion");
66+
assertEquals((byte[]) record.get("uuidCol"), uuidBytes);
67+
assertEquals(record.get("bytesCol"), ByteBuffer.wrap(rawBytes), "BYTES byte[] must be wrapped as ByteBuffer");
68+
}
69+
70+
/// End-to-end: a GenericDatumWriter built with getAvroDataModel() serializes the raw 16-byte UUID values as their
71+
/// canonical string (via the registered uuid Conversion) for both SV and MV, and the on-disk value reads back as
72+
/// that string with a vanilla reader. Without the registered Conversion the write would fail (a ByteBuffer/byte[]
73+
/// cannot be written to a string{logicalType:uuid} field).
74+
@Test
75+
public void testUuidRoundTripThroughAvroDataModel()
76+
throws Exception {
77+
org.apache.pinot.spi.data.Schema pinotSchema = new org.apache.pinot.spi.data.Schema.SchemaBuilder()
78+
.setSchemaName("uuidSchema")
79+
.addSingleValueDimension("uuidSv", DataType.UUID)
80+
.addMultiValueDimension("uuidMv", DataType.UUID)
81+
.build();
82+
Schema avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(pinotSchema);
83+
84+
String svCanonical = "12345678-1234-1234-1234-1234567890ab";
85+
String mvCanonical = "550e8400-e29b-41d4-a716-446655440000";
86+
GenericRow row = new GenericRow();
87+
row.putValue("uuidSv", UuidUtils.toBytes(svCanonical));
88+
row.putValue("uuidMv", new Object[]{UuidUtils.toBytes(mvCanonical)});
89+
GenericData.Record record =
90+
SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, new GenericData.Record(avroSchema));
91+
92+
File tmp = File.createTempFile("uuidRoundTrip", ".avro");
93+
try {
94+
GenericData model = SegmentProcessorAvroUtils.getAvroDataModel();
95+
try (DataFileWriter<GenericData.Record> writer =
96+
new DataFileWriter<>(new GenericDatumWriter<>(avroSchema, model))) {
97+
writer.create(avroSchema, tmp);
98+
writer.append(record);
99+
}
100+
// Read back with a vanilla reader: the on-disk representation must be the canonical UUID string.
101+
try (DataFileReader<GenericRecord> reader = new DataFileReader<>(tmp, new GenericDatumReader<>(avroSchema))) {
102+
assertTrue(reader.hasNext());
103+
GenericRecord read = reader.next();
104+
assertEquals(read.get("uuidSv").toString(), svCanonical, "SV UUID must serialize as canonical string");
105+
List<?> mv = (List<?>) read.get("uuidMv");
106+
assertEquals(mv.size(), 1);
107+
assertEquals(mv.get(0).toString(), mvCanonical, "MV UUID element must serialize as canonical string");
108+
}
109+
// Read back WITH the data model: the uuid Conversion's fromCharSequence must reconstruct the raw 16-byte value.
110+
try (DataFileReader<GenericRecord> reader = new DataFileReader<>(tmp,
111+
new GenericDatumReader<>(avroSchema, avroSchema, SegmentProcessorAvroUtils.getAvroDataModel()))) {
112+
GenericRecord read = reader.next();
113+
assertEquals((byte[]) read.get("uuidSv"), UuidUtils.toBytes(svCanonical),
114+
"reading with the model must convert the uuid string back to raw bytes");
115+
assertEquals((byte[]) ((List<?>) read.get("uuidMv")).get(0), UuidUtils.toBytes(mvCanonical),
116+
"MV uuid element must convert back to raw bytes when reading with the model");
117+
}
118+
} finally {
119+
FileUtils.deleteQuietly(tmp);
120+
}
121+
}
122+
123+
/// convertPinotSchemaToAvroSchema must emit SV UUID as string{logicalType:uuid} and MV UUID as
124+
/// array<string{logicalType:uuid}>, which is what the uuid Conversion above pairs with.
125+
@Test
126+
public void testConvertPinotSchemaToAvroSchemaEmitsUuidLogicalType() {
127+
org.apache.pinot.spi.data.Schema pinotSchema = new org.apache.pinot.spi.data.Schema.SchemaBuilder()
128+
.setSchemaName("uuidSchema")
129+
.addSingleValueDimension("uuidSv", DataType.UUID)
130+
.addMultiValueDimension("uuidMv", DataType.UUID)
131+
.build();
132+
133+
Schema avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(pinotSchema);
134+
135+
Schema svFieldSchema = avroSchema.getField("uuidSv").schema();
136+
assertEquals(svFieldSchema.getType(), Schema.Type.STRING, "SV UUID must be string{logicalType:uuid}");
137+
assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(svFieldSchema), LogicalTypes.uuid(),
138+
"SV UUID must carry the uuid logical type");
139+
140+
Schema mvFieldSchema = avroSchema.getField("uuidMv").schema();
141+
assertEquals(mvFieldSchema.getType(), Schema.Type.ARRAY, "MV UUID must be emitted as an array");
142+
Schema mvElementSchema = mvFieldSchema.getElementType();
143+
assertEquals(mvElementSchema.getType(), Schema.Type.STRING, "MV UUID elements must be string{logicalType:uuid}");
144+
assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(mvElementSchema), LogicalTypes.uuid(),
145+
"MV UUID elements must carry the uuid logical type");
146+
}
147+
}

‎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 {

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,21 @@
2020

2121
import com.fasterxml.jackson.databind.node.ArrayNode;
2222
import com.fasterxml.jackson.databind.node.ObjectNode;
23+
import org.apache.avro.LogicalType;
24+
import org.apache.avro.LogicalTypes;
2325
import org.apache.avro.Schema;
2426
import org.apache.pinot.spi.data.FieldSpec;
2527
import org.apache.pinot.spi.data.FieldSpec.DataType;
2628
import org.apache.pinot.spi.utils.JsonUtils;
29+
import org.apache.pinot.spi.utils.UuidUtils;
2730

2831

2932
/// Stateless helpers for mapping between Avro schema shapes and Pinot's [DataType] / Avro JSON schema representations.
3033
public class AvroSchemaUtil {
34+
// Avro logical-type name for UUID (see org.apache.avro.LogicalTypes). Value-level logical-type conversion lives
35+
// in AvroRecordExtractor; this class only deals with schema-shape mapping.
36+
private static final String UUID = "uuid";
37+
3138
private AvroSchemaUtil() {
3239
}
3340

@@ -60,7 +67,31 @@ public static DataType valueOf(Schema.Type avroType) {
6067
}
6168
}
6269

63-
/// Returns whether the given Avro type is a primitive type.
70+
/**
71+
* Returns the Pinot data type associated with the given Avro schema, including logical types.
72+
*
73+
* <p>Recognizes the UUID logical type on both STRING-backed schemas (Avro spec §logical-types.uuid) and FIXED(16)
74+
* schemas (used by some producers including Confluent's fixed-uuid mode). Both forms arrive at
75+
* {@link org.apache.pinot.plugin.inputformat.avro.AvroRecordExtractor} as either a {@link java.util.UUID} (for
76+
* STRING-backed logical UUIDs) or a 16-byte {@code byte[]} (for FIXED-backed ones), and both are accepted by
77+
* {@link org.apache.pinot.spi.utils.UuidUtils#toBytes(Object)}.
78+
*/
79+
public static DataType valueOf(Schema schema) {
80+
LogicalType logicalType = LogicalTypes.fromSchemaIgnoreInvalid(schema);
81+
if (logicalType != null && UUID.equals(logicalType.getName())) {
82+
if (schema.getType() == Schema.Type.STRING) {
83+
return DataType.UUID;
84+
}
85+
if (schema.getType() == Schema.Type.FIXED && schema.getFixedSize() == UuidUtils.UUID_NUM_BYTES) {
86+
return DataType.UUID;
87+
}
88+
}
89+
return valueOf(schema.getType());
90+
}
91+
92+
/**
93+
* @return if the given avro type is a primitive type.
94+
*/
6495
public static boolean isPrimitiveType(Schema.Type avroType) {
6596
switch (avroType) {
6697
case INT:

0 commit comments

Comments
 (0)