Skip to content

Generate Avro schemas from Pinot logical data types (#19071) #19073

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
xiangfu0 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Generate Avro schemas from Pinot logical data types (#19071) #19073

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
211 changes: 127 additions & 84 deletions pinot-core/src/main/java/org/apache/pinot/core/util/SegmentProcessorAvroUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,15 @@
package org.apache.pinot.core.util;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.avro.Conversion;
import org.apache.avro.Conversions;
import org.apache.avro.LogicalType;
import org.apache.avro.LogicalTypes;
import org.apache.avro.Schema;
Expand All @@ -51,43 +54,99 @@ public static GenericData.Record convertGenericRowToAvroRecord(GenericRow generi
return convertGenericRowToAvroRecord(genericRow, reusableRecord, genericRow.getFieldToValueMap().keySet());
}

/// Convert a GenericRow to an avro GenericRecord
/// Convert a GenericRow to an avro GenericRecord.
///
/// Values arrive in Pinot's internal (stored) representation and are coordinated with the Avro field type produced
/// by `AvroSchemaUtil.toAvroSchema`: whatever a registered logical-type [Conversion] can handle is left untouched
/// for the writer, and only the two cases Avro cannot resolve on its own are fixed up here (see
/// [#convertValue(Schema, Object)]).
public static GenericData.Record convertGenericRowToAvroRecord(GenericRow genericRow,
GenericData.Record reusableRecord, Set<String> fields) {
Schema avroSchema = reusableRecord.getSchema();
for (String field : fields) {
Object value = genericRow.getValue(field);
if (value instanceof Object[]) {
// Array elements are written as-is. For MV UUID (array<string{logicalType:uuid}>) the elements are the raw
// 16-byte values; the uuid Conversion registered on the writer's data model (getAvroDataModel) renders each
// element to its canonical string at write time.
reusableRecord.put(field, Arrays.asList((Object[]) value));
} else if (value instanceof byte[]) {
// A byte[] bound for a plain BYTES field must be wrapped as ByteBuffer (GenericDatumWriter requires it for the
// bytes type). A byte[] bound for a UUID field (string{logicalType:uuid}) is left raw so the uuid Conversion
// registered on the writer's data model (getAvroDataModel) renders it to a canonical string at write time.
Schema.Field avroField = avroSchema.getField(field);
if (avroField != null && avroField.schema().getType() == Schema.Type.BYTES) {
reusableRecord.put(field, ByteBuffer.wrap((byte[]) value));
} else {
reusableRecord.put(field, value);
}
} else {
Schema.Field avroField = avroSchema.getField(field);
if (avroField == null) {
// Let Avro raise its own "Not a valid schema field" error for a column missing from the Avro schema.
reusableRecord.put(field, value);
} else {
reusableRecord.put(avroField.pos(), convertValue(avroField.schema(), value));
}
}
return reusableRecord;
}

/// Shared Avro data model with [UuidConversion] registered. Populated once at class initialization and never
/// mutated afterward (effectively immutable), so it is safe to share across writers.
/// Adapts a Pinot value to the representation the Avro writer expects for the given field schema, recursing into
/// array elements for multi-value columns.
@Nullable
private static Object convertValue(Schema fieldSchema, @Nullable Object value) {
if (value == null) {
return null;
}
if (value instanceof Object[]) {
Object[] values = (Object[]) value;
Schema elementSchema =
fieldSchema.getType() == Schema.Type.ARRAY ? fieldSchema.getElementType() : fieldSchema;
// Only BOOLEAN (stored int -> Boolean) and BYTES (byte[] -> ByteBuffer) element schemas can require a
// per-element transform. Every other MV element type is written as-is — INT/LONG/FLOAT/DOUBLE/STRING directly,
// and UUID (string element, raw byte[] rendered by the registered Conversion) — so hand the writer a zero-copy
// view over the existing array; allocating and copying a fresh list per row would be pure overhead on the
// segment-write hot path. (BIG_DECIMAL has a BYTES element schema and so takes the copy path below, but its
// BigDecimal values still pass through convertSingleValue unchanged for the registered Conversion.)
Schema.Type elementType = elementSchema.getType();
if (elementType != Schema.Type.BOOLEAN && elementType != Schema.Type.BYTES) {
return Arrays.asList(values);
}
List<Object> converted = new ArrayList<>(values.length);
for (Object singleValue : values) {
converted.add(convertSingleValue(elementSchema, singleValue));
}
return converted;
}
return convertSingleValue(fieldSchema, value);
}

/// Adapts a single (non-array) Pinot value to what `GenericDatumWriter` expects for `valueSchema`.
///
/// Only two cases need fixing up; everything else is written as-is, either because the Pinot representation already
/// *is* the Avro representation (`Integer` for `int`, `Long` for `long{timestamp-millis}`, `String` for `string`)
/// or because a [Conversion] registered on [#getAvroDataModel] handles it (`byte[]` for `string{uuid}`,
/// [java.math.BigDecimal] for `bytes{big-decimal}`).
@Nullable
private static Object convertSingleValue(Schema valueSchema, @Nullable Object value) {
if (value == null) {
return null;
}
switch (valueSchema.getType()) {
case BOOLEAN:
// BOOLEAN is the one Pinot logical type with no Avro logical type to carry a Conversion, so its stored int
// 0/1 has to be coerced here. Values that are already Boolean (e.g. from a source that never went through
// Pinot's stored form) pass through.
return value instanceof Boolean ? value : ((Number) value).intValue() != 0;
case BYTES:
// A byte[] bound for a plain BYTES field must be wrapped as ByteBuffer (GenericDatumWriter requires it for
// the bytes type). A BigDecimal bound for bytes{big-decimal} is left alone for the registered Conversion.
return value instanceof byte[] ? ByteBuffer.wrap((byte[]) value) : value;
default:
return value;
}
}

/// Shared Avro data model with the logical-type conversions registered. Populated once at class initialization and
/// never mutated afterward (effectively immutable), so it is safe to share across writers.
private static final GenericData AVRO_DATA_MODEL = createAvroDataModel();

/// Returns the shared Avro data model that a `GenericDatumWriter` (or `AvroParquetWriter`) must be constructed with
/// to serialize UUID columns produced by [#convertGenericRowToAvroRecord]: it registers [UuidConversion] so the
/// internal 16-byte UUID form is rendered as the canonical string required by `string{logicalType:uuid}` fields.
/// The UUID column's field schema must be `string{logicalType:uuid}` — as emitted by
/// [#convertPinotSchemaToAvroSchema] and `AvroUtils.getAvroSchemaFromPinotSchema` — for the conversion to apply.
/// to serialize the logical-type columns produced by [#convertGenericRowToAvroRecord]:
/// - [UuidConversion] renders Pinot's internal 16-byte UUID form as the canonical string required by
/// `string{logicalType:uuid}` fields.
/// - Avro's `BigDecimalConversion` encodes a [java.math.BigDecimal] — unscaled value plus its own scale — into the
/// `bytes{logicalType:big-decimal}` fields Pinot emits for BIG_DECIMAL columns.
///
/// The column's field schema must be the one emitted by [#convertPinotSchemaToAvroSchema] /
/// `AvroUtils.getAvroSchemaFromPinotSchema` for the conversions to apply. BOOLEAN and TIMESTAMP need no entry
/// here: Avro has no `boolean` logical type (the coercion happens in [#convertGenericRowToAvroRecord]), and a
/// `Long` already *is* the base representation of `long{logicalType:timestamp-millis}`.
///
/// The returned instance is shared and must be treated as read-only: do not call its mutators
/// (`addLogicalTypeConversion`, `setStringType`, ...), which are not thread-safe against concurrent writer reads.
Expand All @@ -98,6 +157,7 @@ public static GenericData getAvroDataModel() {
private static GenericData createAvroDataModel() {
GenericData model = new GenericData();
model.addLogicalTypeConversion(new UuidConversion());
model.addLogicalTypeConversion(new Conversions.BigDecimalConversion());
return model;
}

Expand Down Expand Up @@ -133,74 +193,57 @@ public byte[] fromCharSequence(CharSequence value, Schema schema, LogicalType ty
}
}

/// Converts a Pinot schema to an Avro schema
/// Converts a Pinot schema to an Avro schema, with the fields ordered by column name.
///
/// Field types are derived from the **original (logical)** Pinot data type, so BOOLEAN becomes Avro `boolean`,
/// TIMESTAMP a `timestamp-millis` long, BIG_DECIMAL a `big-decimal` bytes and UUID a `uuid` string, instead of all
/// four collapsing to their physical storage type. This must stay identical to
/// `AvroSchemaUtil.toAvroSchema(FieldSpec)` in `pinot-avro-base` — the two live in different modules (neither can
/// depend on the other) but feed the same writers, so `SegmentProcessorAvroUtilsTest` pins them together. See that
/// method for the full mapping table and the value representation each Avro type expects.
public static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Schema pinotSchema) {
SchemaBuilder.FieldAssembler<org.apache.avro.Schema> fieldAssembler = SchemaBuilder.record("record").fields();

List<FieldSpec> orderedFieldSpecs = pinotSchema.getAllFieldSpecs().stream()
.sorted(Comparator.comparing(FieldSpec::getName))
.collect(Collectors.toList());
for (FieldSpec fieldSpec : orderedFieldSpecs) {
String name = fieldSpec.getName();
// Emit UUID columns as Avro string{logicalType:uuid} (matching AvroUtils.getAvroSchemaFromPinotSchema)
// so the runtime byte[] → canonical-string conversion in convertGenericRowToAvroRecord lines up with
// the field schema. Without this branch SV UUID would fall through to BYTES (losing UUID semantics) and
// MV UUID would throw at this point (MV switch below has no BYTES case).
if (fieldSpec.getDataType() == DataType.UUID) {
Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
if (fieldSpec.isSingleValueField()) {
fieldAssembler = fieldAssembler.name(name).type(uuidSchema).noDefault();
} else {
fieldAssembler = fieldAssembler.name(name).type().array().items(uuidSchema).noDefault();
}
continue;
}
DataType storedType = fieldSpec.getDataType().getStoredType();
if (fieldSpec.isSingleValueField()) {
switch (storedType) {
case INT:
fieldAssembler = fieldAssembler.name(name).type().intType().noDefault();
break;
case LONG:
fieldAssembler = fieldAssembler.name(name).type().longType().noDefault();
break;
case FLOAT:
fieldAssembler = fieldAssembler.name(name).type().floatType().noDefault();
break;
case DOUBLE:
fieldAssembler = fieldAssembler.name(name).type().doubleType().noDefault();
break;
case STRING:
fieldAssembler = fieldAssembler.name(name).type().stringType().noDefault();
break;
case BYTES:
fieldAssembler = fieldAssembler.name(name).type().bytesType().noDefault();
break;
default:
throw new RuntimeException("Unsupported data type: " + storedType);
}
} else {
switch (storedType) {
case INT:
fieldAssembler = fieldAssembler.name(name).type().array().items().intType().noDefault();
break;
case LONG:
fieldAssembler = fieldAssembler.name(name).type().array().items().longType().noDefault();
break;
case FLOAT:
fieldAssembler = fieldAssembler.name(name).type().array().items().floatType().noDefault();
break;
case DOUBLE:
fieldAssembler = fieldAssembler.name(name).type().array().items().doubleType().noDefault();
break;
case STRING:
fieldAssembler = fieldAssembler.name(name).type().array().items().stringType().noDefault();
break;
default:
throw new RuntimeException("Unsupported data type: " + storedType);
}
}
fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type(toAvroSchema(fieldSpec)).noDefault();
}
return fieldAssembler.endRecord();
}

/// Returns the Avro schema for a whole Pinot column: the single-value type from [#toAvroSchema(DataType)], or an
/// array of it for a multi-value column.
private static Schema toAvroSchema(FieldSpec fieldSpec) {
Schema valueSchema = toAvroSchema(fieldSpec.getDataType());
return fieldSpec.isSingleValueField() ? valueSchema : Schema.createArray(valueSchema);
}

private static Schema toAvroSchema(DataType dataType) {
switch (dataType) {
case INT:
return Schema.create(Schema.Type.INT);
case LONG:
return Schema.create(Schema.Type.LONG);
case FLOAT:
return Schema.create(Schema.Type.FLOAT);
case DOUBLE:
return Schema.create(Schema.Type.DOUBLE);
case BOOLEAN:
return Schema.create(Schema.Type.BOOLEAN);
case TIMESTAMP:
return LogicalTypes.timestampMillis().addToSchema(Schema.create(Schema.Type.LONG));
case BIG_DECIMAL:
return LogicalTypes.bigDecimal().addToSchema(Schema.create(Schema.Type.BYTES));
case STRING:
case JSON:
return Schema.create(Schema.Type.STRING);
case BYTES:
return Schema.create(Schema.Type.BYTES);
case UUID:
return LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
default:
throw new UnsupportedOperationException("Unsupported data type: " + dataType);
}
}
}
Loading
Loading