[UUID 2/8] UUID ingest and segment storage

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

@@ -191,7 +191,8 @@ private void resetBuffer()
       throws IOException {
     FileUtils.deleteQuietly(_bufferFile);
     _rowCount = 0;
-    _recordWriter = new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema));
+    _recordWriter =
+        new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema, SegmentProcessorAvroUtils.getAvroDataModel()));
     _recordWriter.create(_avroSchema, _bufferFile);
   }
 

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

@@ -24,56 +24,116 @@
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import org.apache.avro.Conversion;
+import org.apache.avro.LogicalType;
+import org.apache.avro.LogicalTypes;
 import org.apache.avro.Schema;
 import org.apache.avro.SchemaBuilder;
 import org.apache.avro.generic.GenericData;
-import org.apache.pinot.core.segment.processing.framework.SegmentProcessorFramework;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.data.readers.GenericRow;
-import org.apache.pinot.spi.ingestion.segment.writer.SegmentWriter;
+import org.apache.pinot.spi.utils.UuidUtils;
 
 
-/**
- * Helper methods for avro related conversions needed, when using AVRO as intermediate format in segment processing.
- * AVRO is used as intermediate processing format in {@link SegmentProcessorFramework} and file-based impl of
- * {@link SegmentWriter}
- */
+/// Helper methods for Avro conversions used when serializing Pinot [GenericRow]s to Avro — by the file-based
+/// `SegmentWriter` implementations and the segment-to-Avro/Parquet converters. (The multi-stage
+/// `SegmentProcessorFramework` itself serializes intermediate records with a custom binary `GenericRowFile`
+/// format, not Avro.)
 public final class SegmentProcessorAvroUtils {
 
   private SegmentProcessorAvroUtils() {
   }
 
-  /**
-   * Convert a GenericRow to an avro GenericRecord
-   */
+  /// Convert a GenericRow to an avro GenericRecord
   public static GenericData.Record convertGenericRowToAvroRecord(GenericRow genericRow,
       GenericData.Record reusableRecord) {
     return convertGenericRowToAvroRecord(genericRow, reusableRecord, genericRow.getFieldToValueMap().keySet());
   }
 
-  /**
-   * Convert a GenericRow to an avro GenericRecord
-   */
+  /// Convert a GenericRow to an avro GenericRecord
   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[]) {
-          value = ByteBuffer.wrap((byte[]) 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 {
         reusableRecord.put(field, value);
       }
     }
     return reusableRecord;
   }
 
-  /**
-   * Converts a Pinot schema to an Avro schema
-   */
+  /// 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.
+  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.
+  ///
+  /// 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.
+  public static GenericData getAvroDataModel() {
+    return AVRO_DATA_MODEL;
+  }
+
+  private static GenericData createAvroDataModel() {
+    GenericData model = new GenericData();
+    model.addLogicalTypeConversion(new UuidConversion());
+    return model;
+  }
+
+  /// Avro logical-type [Conversion] for the `uuid` logical type, operating directly on Pinot's internal storage form
+  /// (a 16-byte big-endian `byte[]`). It renders that value to its canonical RFC-4122 string when writing a
+  /// `string{logicalType:uuid}` field and parses it back on read.
+  ///
+  /// The converted type is `byte[]` rather than [java.util.UUID] so no intermediate object is allocated, and rather
+  /// than [java.nio.ByteBuffer] because Avro resolves conversions by exact datum class: a `byte[]` instance always
+  /// reports `byte[].class`, so the conversion resolves for both single values and array elements, whereas a
+  /// concrete `HeapByteBuffer` would not match a `ByteBuffer`-typed conversion.
+  private static final class UuidConversion extends Conversion<byte[]> {
+    private static final String UUID_LOGICAL_TYPE_NAME = "uuid";
+
+    @Override
+    public Class<byte[]> getConvertedType() {
+      return byte[].class;
+    }
+
+    @Override
+    public String getLogicalTypeName() {
+      return UUID_LOGICAL_TYPE_NAME;
+    }
+
+    @Override
+    public CharSequence toCharSequence(byte[] value, Schema schema, LogicalType type) {
+      return UuidUtils.toString(value);
+    }
+
+    @Override
+    public byte[] fromCharSequence(CharSequence value, Schema schema, LogicalType type) {
+      return UuidUtils.toBytes(value.toString());
+    }
+  }
+
+  /// Converts a Pinot schema to an Avro schema
   public static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Schema pinotSchema) {
     SchemaBuilder.FieldAssembler<org.apache.avro.Schema> fieldAssembler = SchemaBuilder.record("record").fields();
 
@@ -82,6 +142,19 @@ public static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Sc
         .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) {

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

@@ -0,0 +1,147 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.util;
+
+import java.io.File;
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.avro.LogicalTypes;
+import org.apache.avro.Schema;
+import org.apache.avro.SchemaBuilder;
+import org.apache.avro.file.DataFileReader;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.UuidUtils;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+public class SegmentProcessorAvroUtilsTest {
+
+  /// convertGenericRowToAvroRecord keeps a UUID value as its raw 16-byte form (it does NOT render a string here); the
+  /// rendering is deferred to the uuid Conversion on the writer's data model. Plain BYTES columns are still wrapped as
+  /// ByteBuffer.
+  @Test
+  public void testConvertGenericRowToAvroRecordKeepsUuidBytesRaw() {
+    Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING));
+    Schema recordSchema = SchemaBuilder.record("record").fields()
+        .name("uuidCol").type(uuidSchema).noDefault()
+        .name("bytesCol").type().bytesType().noDefault()
+        .endRecord();
+
+    byte[] uuidBytes = UuidUtils.toBytes("12345678-1234-1234-1234-1234567890ab");
+    byte[] rawBytes = {1, 2, 3, 4};
+    GenericRow row = new GenericRow();
+    row.putValue("uuidCol", uuidBytes);
+    row.putValue("bytesCol", rawBytes);
+
+    GenericData.Record record = new GenericData.Record(recordSchema);
+    SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, record);
+
+    assertTrue(record.get("uuidCol") instanceof byte[], "UUID must be left as raw byte[] for the uuid Conversion");
+    assertEquals((byte[]) record.get("uuidCol"), uuidBytes);
+    assertEquals(record.get("bytesCol"), ByteBuffer.wrap(rawBytes), "BYTES byte[] must be wrapped as ByteBuffer");
+  }
+
+  /// End-to-end: a GenericDatumWriter built with getAvroDataModel() serializes the raw 16-byte UUID values as their
+  /// canonical string (via the registered uuid Conversion) for both SV and MV, and the on-disk value reads back as
+  /// that string with a vanilla reader. Without the registered Conversion the write would fail (a ByteBuffer/byte[]
+  /// cannot be written to a string{logicalType:uuid} field).
+  @Test
+  public void testUuidRoundTripThroughAvroDataModel()
+      throws Exception {
+    org.apache.pinot.spi.data.Schema pinotSchema = new org.apache.pinot.spi.data.Schema.SchemaBuilder()
+        .setSchemaName("uuidSchema")
+        .addSingleValueDimension("uuidSv", DataType.UUID)
+        .addMultiValueDimension("uuidMv", DataType.UUID)
+        .build();
+    Schema avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(pinotSchema);
+
+    String svCanonical = "12345678-1234-1234-1234-1234567890ab";
+    String mvCanonical = "550e8400-e29b-41d4-a716-446655440000";
+    GenericRow row = new GenericRow();
+    row.putValue("uuidSv", UuidUtils.toBytes(svCanonical));
+    row.putValue("uuidMv", new Object[]{UuidUtils.toBytes(mvCanonical)});
+    GenericData.Record record =
+        SegmentProcessorAvroUtils.convertGenericRowToAvroRecord(row, new GenericData.Record(avroSchema));
+
+    File tmp = File.createTempFile("uuidRoundTrip", ".avro");
+    try {
+      GenericData model = SegmentProcessorAvroUtils.getAvroDataModel();
+      try (DataFileWriter<GenericData.Record> writer =
+          new DataFileWriter<>(new GenericDatumWriter<>(avroSchema, model))) {
+        writer.create(avroSchema, tmp);
+        writer.append(record);
+      }
+      // Read back with a vanilla reader: the on-disk representation must be the canonical UUID string.
+      try (DataFileReader<GenericRecord> reader = new DataFileReader<>(tmp, new GenericDatumReader<>(avroSchema))) {
+        assertTrue(reader.hasNext());
+        GenericRecord read = reader.next();
+        assertEquals(read.get("uuidSv").toString(), svCanonical, "SV UUID must serialize as canonical string");
+        List<?> mv = (List<?>) read.get("uuidMv");
+        assertEquals(mv.size(), 1);
+        assertEquals(mv.get(0).toString(), mvCanonical, "MV UUID element must serialize as canonical string");
+      }
+      // Read back WITH the data model: the uuid Conversion's fromCharSequence must reconstruct the raw 16-byte value.
+      try (DataFileReader<GenericRecord> reader = new DataFileReader<>(tmp,
+          new GenericDatumReader<>(avroSchema, avroSchema, SegmentProcessorAvroUtils.getAvroDataModel()))) {
+        GenericRecord read = reader.next();
+        assertEquals((byte[]) read.get("uuidSv"), UuidUtils.toBytes(svCanonical),
+            "reading with the model must convert the uuid string back to raw bytes");
+        assertEquals((byte[]) ((List<?>) read.get("uuidMv")).get(0), UuidUtils.toBytes(mvCanonical),
+            "MV uuid element must convert back to raw bytes when reading with the model");
+      }
+    } finally {
+      FileUtils.deleteQuietly(tmp);
+    }
+  }
+
+  /// convertPinotSchemaToAvroSchema must emit SV UUID as string{logicalType:uuid} and MV UUID as
+  /// array<string{logicalType:uuid}>, which is what the uuid Conversion above pairs with.
+  @Test
+  public void testConvertPinotSchemaToAvroSchemaEmitsUuidLogicalType() {
+    org.apache.pinot.spi.data.Schema pinotSchema = new org.apache.pinot.spi.data.Schema.SchemaBuilder()
+        .setSchemaName("uuidSchema")
+        .addSingleValueDimension("uuidSv", DataType.UUID)
+        .addMultiValueDimension("uuidMv", DataType.UUID)
+        .build();
+
+    Schema avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(pinotSchema);
+
+    Schema svFieldSchema = avroSchema.getField("uuidSv").schema();
+    assertEquals(svFieldSchema.getType(), Schema.Type.STRING, "SV UUID must be string{logicalType:uuid}");
+    assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(svFieldSchema), LogicalTypes.uuid(),
+        "SV UUID must carry the uuid logical type");
+
+    Schema mvFieldSchema = avroSchema.getField("uuidMv").schema();
+    assertEquals(mvFieldSchema.getType(), Schema.Type.ARRAY, "MV UUID must be emitted as an array");
+    Schema mvElementSchema = mvFieldSchema.getElementType();
+    assertEquals(mvElementSchema.getType(), Schema.Type.STRING, "MV UUID elements must be string{logicalType:uuid}");
+    assertEquals(LogicalTypes.fromSchemaIgnoreInvalid(mvElementSchema), LogicalTypes.uuid(),
+        "MV UUID elements must carry the uuid logical type");
+  }
+}

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

@@ -139,7 +139,7 @@ private void validateSchemas() {
         if (fieldSpec.getDataType() != dataTypeForSVColumn) {
           _dataTypeMismatch.addMismatchReason(String
               .format("The Pinot column: (%s: %s) doesn't match with the column (%s: %s) in input %s schema.",
-                  columnName, fieldSpec.getDataType().name(), avroColumnName, avroColumnType.name(),
+                  columnName, fieldSpec.getDataType().name(), avroColumnName, dataTypeForSVColumn.name(),
                   getInputSchemaType()));
         }
       } else {

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

@@ -20,14 +20,21 @@
 
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.avro.LogicalType;
+import org.apache.avro.LogicalTypes;
 import org.apache.avro.Schema;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.UuidUtils;
 
 
 /// Stateless helpers for mapping between Avro schema shapes and Pinot's [DataType] / Avro JSON schema representations.
 public class AvroSchemaUtil {
+  // Avro logical-type name for UUID (see org.apache.avro.LogicalTypes). Value-level logical-type conversion lives
+  // in AvroRecordExtractor; this class only deals with schema-shape mapping.
+  private static final String UUID = "uuid";
+
   private AvroSchemaUtil() {
   }
 
@@ -60,6 +67,28 @@ public static DataType valueOf(Schema.Type avroType) {
     }
   }
 
+  /**
+   * Returns the Pinot data type associated with the given Avro schema, including logical types.
+   *
+   * <p>Recognizes the UUID logical type on both STRING-backed schemas (Avro spec §logical-types.uuid) and FIXED(16)
+   * schemas (used by some producers including Confluent's fixed-uuid mode). Both forms arrive at
+   * {@link org.apache.pinot.plugin.inputformat.avro.AvroRecordExtractor} as either a {@link java.util.UUID} (for
+   * STRING-backed logical UUIDs) or a 16-byte {@code byte[]} (for FIXED-backed ones), and both are accepted by
+   * {@link org.apache.pinot.spi.utils.UuidUtils#toBytes(Object)}.
+   */
+  public static DataType valueOf(Schema schema) {
+    LogicalType logicalType = LogicalTypes.fromSchemaIgnoreInvalid(schema);
+    if (logicalType != null && UUID.equals(logicalType.getName())) {
+      if (schema.getType() == Schema.Type.STRING) {
+        return DataType.UUID;
+      }
+      if (schema.getType() == Schema.Type.FIXED && schema.getFixedSize() == UuidUtils.UUID_NUM_BYTES) {
+        return DataType.UUID;
+      }
+    }
+    return valueOf(schema.getType());
+  }
+
   /// Returns whether the given Avro type is a primitive type.
   public static boolean isPrimitiveType(Schema.Type avroType) {
     switch (avroType) {