diff --git a/pinot-connectors/pinot-flink-connector/src/main/java/org/apache/pinot/connector/flink/sink/FlinkSegmentWriter.java b/pinot-connectors/pinot-flink-connector/src/main/java/org/apache/pinot/connector/flink/sink/FlinkSegmentWriter.java index ef2b684dc032..e7ddce376bd7 100644 --- a/pinot-connectors/pinot-flink-connector/src/main/java/org/apache/pinot/connector/flink/sink/FlinkSegmentWriter.java +++ b/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); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/util/SegmentProcessorAvroUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/util/SegmentProcessorAvroUtils.java index 0a47f5f20295..0c7f5b5e2a79 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/util/SegmentProcessorAvroUtils.java +++ b/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 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) 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 { + private static final String UUID_LOGICAL_TYPE_NAME = "uuid"; + + @Override + public Class 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 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) { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/util/SegmentProcessorAvroUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/util/SegmentProcessorAvroUtilsTest.java new file mode 100644 index 000000000000..2f3012a2c66d --- /dev/null +++ b/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 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 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 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, 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"); + } +} diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java index 1ec33656460d..a90cee65c6aa 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java +++ b/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 { diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java index 0b1f8ad3ab0f..be17c9ba0d3d 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java +++ b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java @@ -20,18 +20,28 @@ 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 { + private AvroSchemaUtil() { } - /// Returns the data type stored in Pinot that is associated with the given Avro type. + // 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"; + + /// Returns the Pinot data type for a bare Avro type. This does not honor logical types (e.g. a `string` or `fixed` + /// carrying `logicalType:uuid` maps to STRING/BYTES, not UUID); prefer [#valueOf(Schema)] when a full [Schema] is + /// available. public static DataType valueOf(Schema.Type avroType) { switch (avroType) { case INT: @@ -60,6 +70,25 @@ public static DataType valueOf(Schema.Type avroType) { } } + /// Returns the Pinot data type associated with the given Avro schema, including logical types. + /// + /// 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 + /// [AvroRecordExtractor] as either a [java.util.UUID] (for STRING-backed logical UUIDs) or a 16-byte `byte[]` + /// (for FIXED-backed ones), and both are accepted by [UuidUtils#toBytes]. + 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) { diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java index 8dc688adef60..b01916df360c 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java +++ b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java @@ -27,6 +27,7 @@ import java.util.concurrent.TimeUnit; import java.util.zip.GZIPInputStream; import javax.annotation.Nullable; +import org.apache.avro.LogicalTypes; import org.apache.avro.Schema.Field; import org.apache.avro.SchemaBuilder; import org.apache.avro.file.DataFileStream; @@ -160,6 +161,16 @@ public static org.apache.avro.Schema getAvroSchemaFromPinotSchema(Schema pinotSc SchemaBuilder.FieldAssembler fieldAssembler = SchemaBuilder.record("record").fields(); for (FieldSpec fieldSpec : pinotSchema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() == DataType.UUID) { + org.apache.avro.Schema uuidSchema = LogicalTypes.uuid().addToSchema(org.apache.avro.Schema.create( + org.apache.avro.Schema.Type.STRING)); + if (fieldSpec.isSingleValueField()) { + fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type(uuidSchema).noDefault(); + } else { + fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type().array().items(uuidSchema).noDefault(); + } + continue; + } DataType storedType = fieldSpec.getDataType().getStoredType(); if (fieldSpec.isSingleValueField()) { switch (storedType) { @@ -244,12 +255,13 @@ public static boolean isSingleValueField(Field field) { public static DataType extractFieldDataType(Field field) { try { org.apache.avro.Schema fieldSchema = extractSupportedSchema(field.schema()); - org.apache.avro.Schema.Type fieldType = fieldSchema.getType(); - if (fieldType == org.apache.avro.Schema.Type.ARRAY) { - return AvroSchemaUtil.valueOf(extractSupportedSchema(fieldSchema.getElementType()).getType()); + if (fieldSchema.getType() == org.apache.avro.Schema.Type.ARRAY) { + return AvroSchemaUtil.valueOf(extractSupportedSchema(fieldSchema.getElementType())); } else { - return AvroSchemaUtil.valueOf(fieldType); + return AvroSchemaUtil.valueOf(fieldSchema); } + } catch (RuntimeException e) { + throw e; } catch (Exception e) { throw new RuntimeException("Caught exception while extracting data type from field: " + field.name(), e); } @@ -323,16 +335,17 @@ private static void extractSchemaWithComplexTypeHandling(org.apache.avro.Schema extractSchemaWithComplexTypeHandling(elementType, fieldsToUnnest, delimiter, path, pinotSchema, fieldTypeMap, timeUnit, collectionNotUnnestedToJson); } else if (collectionNotUnnestedToJson == ComplexTypeConfig.CollectionNotUnnestedToJson.NON_PRIMITIVE - && AvroSchemaUtil.isPrimitiveType(elementType.getType())) { - addFieldToPinotSchema(pinotSchema, AvroSchemaUtil.valueOf(elementType.getType()), path, false, fieldTypeMap, - timeUnit); + && (AvroSchemaUtil.isPrimitiveType(elementType.getType()) + || AvroSchemaUtil.valueOf(elementType) == DataType.UUID)) { + DataType elementDataType = AvroSchemaUtil.valueOf(elementType); + addFieldToPinotSchema(pinotSchema, elementDataType, path, false, fieldTypeMap, timeUnit); } else if (shallConvertToJson(collectionNotUnnestedToJson, elementType)) { addFieldToPinotSchema(pinotSchema, DataType.STRING, path, true, fieldTypeMap, timeUnit); } // do not include the node for other cases break; default: - DataType dataType = AvroSchemaUtil.valueOf(fieldType); + DataType dataType = AvroSchemaUtil.valueOf(fieldSchema); addFieldToPinotSchema(pinotSchema, dataType, path, true, fieldTypeMap, timeUnit); break; } diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java index 09142d201f0d..436e3d8c68f7 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java +++ b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.avro.Schema; import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.testng.annotations.Test; @@ -58,16 +59,48 @@ public void testToAvroSchemaJsonObjectRejectsUnsupportedType() { () -> AvroSchemaUtil.toAvroSchemaJsonObject(new DimensionFieldSpec("col", DataType.BIG_DECIMAL, true))); } - /// A UUID column is represented faithfully as an Avro string annotated with logicalType "uuid". - /// `DataGenerator#buildSpec` always marks the recommender schema FieldSpec single-value, so this emits a scalar - /// union like every sibling type. + /// UUID is a logical type; a single-value column maps to an Avro string carrying the "uuid" logical type. @Test public void testToAvroSchemaJsonObjectForUuid() { JsonNode type = typeOf(DataType.UUID); assertEquals(type.get(0).asText(), "null"); - JsonNode uuidBranch = type.get(1); - assertEquals(uuidBranch.get("type").asText(), "string"); - assertEquals(uuidBranch.get("logicalType").asText(), "uuid"); + assertEquals(type.get(1).get("type").asText(), "string"); + assertEquals(type.get(1).get("logicalType").asText(), "uuid"); + } + + @Test + public void testValueOfUuidStringLogicalType() { + Schema schema = new Schema.Parser().parse("{\"type\":\"record\",\"name\":\"r\",\"fields\":[{\"name\":\"id\"," + + "\"type\":{\"type\":\"string\",\"logicalType\":\"uuid\"}}]}"); + assertEquals(AvroSchemaUtil.valueOf(schema.getField("id").schema()), DataType.UUID, + "STRING logicalType:uuid should map to UUID"); + } + + @Test + public void testValueOfUuidFixed16LogicalType() { + // FIXED(16) + logicalType:uuid — produced by Confluent fixed-uuid mode and Parquet uuid + Schema schema = new Schema.Parser().parse("{\"type\":\"record\",\"name\":\"r\",\"fields\":[{\"name\":\"id\"," + + "\"type\":{\"type\":\"fixed\",\"name\":\"uuid_fixed\",\"size\":16,\"logicalType\":\"uuid\"}}]}"); + assertEquals(AvroSchemaUtil.valueOf(schema.getField("id").schema()), DataType.UUID, + "FIXED(16) logicalType:uuid should map to UUID"); + } + + @Test + public void testValueOfFixed16WithoutLogicalTypeIsBytes() { + // FIXED(16) without logicalType should stay as BYTES + Schema schema = new Schema.Parser().parse("{\"type\":\"record\",\"name\":\"r\",\"fields\":[{\"name\":\"raw\"," + + "\"type\":{\"type\":\"fixed\",\"name\":\"raw16\",\"size\":16}}]}"); + assertEquals(AvroSchemaUtil.valueOf(schema.getField("raw").schema()), DataType.BYTES, + "FIXED(16) without logicalType:uuid should stay BYTES"); + } + + @Test + public void testValueOfFixedWrongSizeWithUuidLogicalTypeIsBytes() { + // FIXED of non-16 size with logicalType:uuid should not map to UUID + Schema schema = new Schema.Parser().parse("{\"type\":\"record\",\"name\":\"r\",\"fields\":[{\"name\":\"id\"," + + "\"type\":{\"type\":\"fixed\",\"name\":\"uuid32\",\"size\":32,\"logicalType\":\"uuid\"}}]}"); + assertEquals(AvroSchemaUtil.valueOf(schema.getField("id").schema()), DataType.BYTES, + "FIXED(32) with logicalType:uuid should stay BYTES"); } private static void assertPrimitiveType(DataType dataType, String expectedAvroType) { diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroUtilsTest.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroUtilsTest.java index 5af642aca6b8..704de6ca58ae 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroUtilsTest.java +++ b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroUtilsTest.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.Map; import java.util.concurrent.TimeUnit; +import org.apache.avro.LogicalTypes; import org.apache.pinot.spi.config.table.ingestion.ComplexTypeConfig; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; @@ -131,4 +132,58 @@ public void testGetPinotSchemaFromAvroSchemaWithComplexType() .addDateTime("hoursSinceEpoch", DataType.LONG, "EPOCH|HOURS", "1:HOURS").build(); assertEquals(inferredPinotSchema, expectedSchema); } + + @Test + public void testGetPinotSchemaFromAvroSchemaWithUuidLogicalType() { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false); + org.apache.avro.Schema uuidSchema = + LogicalTypes.uuid().addToSchema(org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING)); + avroSchema.setFields(Lists.newArrayList( + new org.apache.avro.Schema.Field("uuidCol", uuidSchema, null, null) + )); + + Schema inferredPinotSchema = AvroUtils.getPinotSchemaFromAvroSchema(avroSchema, null, null); + + Schema expectedSchema = + new Schema.SchemaBuilder().addSingleValueDimension("uuidCol", DataType.UUID).build(); + assertEquals(inferredPinotSchema, expectedSchema); + } + + @Test + public void testGetAvroSchemaFromPinotSchemaWithUuidLogicalType() { + Schema pinotSchema = new Schema.SchemaBuilder().addSingleValueDimension("uuidCol", DataType.UUID).build(); + + org.apache.avro.Schema avroSchema = AvroUtils.getAvroSchemaFromPinotSchema(pinotSchema); + org.apache.avro.Schema fieldSchema = avroSchema.getField("uuidCol").schema(); + + assertEquals(fieldSchema.getType(), org.apache.avro.Schema.Type.STRING); + assertEquals(fieldSchema.getLogicalType().getName(), "uuid"); + } + + @Test + public void testGetPinotSchemaFromAvroSchemaWithUuidArray() { + org.apache.avro.Schema uuidSchema = + LogicalTypes.uuid().addToSchema(org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING)); + org.apache.avro.Schema avroSchema = org.apache.avro.SchemaBuilder.record("uuidArrayRecord").fields() + .name("uuidArray").type().array().items(uuidSchema).noDefault().endRecord(); + + Schema inferredPinotSchema = AvroUtils.getPinotSchemaFromAvroSchemaWithComplexTypeHandling(avroSchema, null, null, + new ArrayList<>(), ".", ComplexTypeConfig.CollectionNotUnnestedToJson.NON_PRIMITIVE); + + Schema expectedSchema = + new Schema.SchemaBuilder().addMultiValueDimension("uuidArray", DataType.UUID).build(); + assertEquals(inferredPinotSchema, expectedSchema); + } + + @Test + public void testGetAvroSchemaFromPinotSchemaWithMvUuid() { + Schema pinotSchema = new Schema.SchemaBuilder().addMultiValueDimension("uuidMv", DataType.UUID).build(); + + org.apache.avro.Schema avroSchema = AvroUtils.getAvroSchemaFromPinotSchema(pinotSchema); + org.apache.avro.Schema fieldSchema = avroSchema.getField("uuidMv").schema(); + + assertEquals(fieldSchema.getType(), org.apache.avro.Schema.Type.ARRAY); + assertEquals(fieldSchema.getElementType().getType(), org.apache.avro.Schema.Type.STRING); + assertEquals(fieldSchema.getElementType().getLogicalType().getName(), "uuid"); + } } diff --git a/pinot-plugins/pinot-segment-writer/pinot-segment-writer-file-based/src/main/java/org/apache/pinot/plugin/segmentwriter/filebased/FileBasedSegmentWriter.java b/pinot-plugins/pinot-segment-writer/pinot-segment-writer-file-based/src/main/java/org/apache/pinot/plugin/segmentwriter/filebased/FileBasedSegmentWriter.java index 28a1f218064c..07ee87f5cbf3 100644 --- a/pinot-plugins/pinot-segment-writer/pinot-segment-writer-file-based/src/main/java/org/apache/pinot/plugin/segmentwriter/filebased/FileBasedSegmentWriter.java +++ b/pinot-plugins/pinot-segment-writer/pinot-segment-writer-file-based/src/main/java/org/apache/pinot/plugin/segmentwriter/filebased/FileBasedSegmentWriter.java @@ -132,7 +132,8 @@ public void init(TableConfig tableConfig, Schema schema, Map bat private void resetBuffer() throws IOException { FileUtils.deleteQuietly(_bufferFile); - _recordWriter = new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema)); + _recordWriter = + new DataFileWriter<>(new GenericDatumWriter<>(_avroSchema, SegmentProcessorAvroUtils.getAvroDataModel())); _recordWriter.create(_avroSchema, _bufferFile); } diff --git a/pinot-segment-local/pom.xml b/pinot-segment-local/pom.xml index 5235fd1503cf..ac18a6da9bfd 100644 --- a/pinot-segment-local/pom.xml +++ b/pinot-segment-local/pom.xml @@ -109,6 +109,11 @@ pinot-avro test + + org.apache.pinot + pinot-avro-base + test + org.apache.pinot pinot-csv diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java index a8357d5779bb..795e35126c55 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java @@ -123,6 +123,7 @@ import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.FixedIntArray; import org.apache.pinot.spi.utils.MapUtils; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.roaringbitmap.BatchIterator; import org.roaringbitmap.buffer.MutableRoaringBitmap; @@ -370,16 +371,9 @@ public MultiColumnTextMetadata getMultiColumnTextMetadata() { } else { dictionary = null; if (!fieldSpec.isSingleValueField()) { - // Raw MV columns - switch (storedType) { - case INT: - case LONG: - case FLOAT: - case DOUBLE: - break; - default: - throw new UnsupportedOperationException( - "Unsupported data type: " + dataType + " for MV no-dictionary column: " + column); + if (!dataType.isFixedWidth()) { + throw new UnsupportedOperationException( + "Unsupported data type: " + dataType + " for MV no-dictionary column: " + column); } } } @@ -760,7 +754,8 @@ private RecordInfo getRecordInfo(GenericRow row, int docId) { private Comparable getComparisonValue(GenericRow row) { int numComparisonColumns = _upsertComparisonColumns.size(); if (numComparisonColumns == 1) { - return (Comparable) row.getValue(_upsertComparisonColumns.get(0)); + String comparisonColumn = _upsertComparisonColumns.get(0); + return toComparable(row.getValue(comparisonColumn)); } Comparable[] comparisonValues = new Comparable[numComparisonColumns]; @@ -777,11 +772,7 @@ private Comparable getComparisonValue(GenericRow row) { "Documents must have exactly 1 non-null comparison column value"); comparableIndex = i; - - Object comparisonValue = row.getValue(columnName); - Preconditions.checkState(comparisonValue instanceof Comparable, - "Upsert comparison column: %s must be comparable", columnName); - comparisonValues[i] = (Comparable) comparisonValue; + comparisonValues[i] = toComparable(row.getValue(columnName)); } } Preconditions.checkState(comparableIndex != -1, "Documents must have exactly 1 non-null comparison column value"); @@ -928,14 +919,7 @@ private void addNewRow(int docId, GenericRow row) { // Update min/max value from raw value // NOTE: Skip updating min/max value for aggregated metrics because the value will change over time. if (!isAggregateMetricsEnabled() || fieldSpec.getFieldType() != FieldSpec.FieldType.METRIC) { - Comparable comparable; - if (dataType == BYTES) { - comparable = new ByteArray((byte[]) value); - } else if (dataType == MAP) { - comparable = new ByteArray(MapUtils.serializeMap((Map) value)); - } else { - comparable = (Comparable) value; - } + Comparable comparable = toComparableValue(value, dataType, column); if (indexContainer._minValue == null) { indexContainer._minValue = comparable; indexContainer._maxValue = comparable; @@ -988,6 +972,29 @@ private void addNewRow(int docId, GenericRow row) { } } + /// Wraps a raw comparison-column value as a Comparable without a per-row schema lookup: a byte[] (a BYTES or UUID + /// comparison column) becomes a ByteArray; every other type is already Comparable. Mirrors + /// UpsertUtils.SingleComparisonColumnReader so the write and read paths agree. + private static Comparable toComparable(Object value) { + if (value instanceof byte[]) { + return new ByteArray((byte[]) value); + } + Preconditions.checkState(value instanceof Comparable, "Upsert comparison column value must be comparable: %s", + value); + return (Comparable) value; + } + + private Comparable toComparableValue(Object value, DataType dataType, @Nullable String columnName) { + if (dataType == MAP) { + return new ByteArray(MapUtils.serializeMap((Map) value)); + } + if (dataType.getStoredType() == BYTES) { + return new ByteArray((byte[]) value); + } + Preconditions.checkState(value instanceof Comparable, "Column: %s must be comparable", columnName); + return (Comparable) value; + } + private void updateIndexCapacityThresholdBreached(MutableIndex mutableIndex, IndexType indexType, String column) { // Few of the Immutable version of the mutable index are bounded by size like // {@link VarByteChunkForwardIndexWriterV4#putBytes(byte[])} and {@link FixedBitMVForwardIndex} @@ -1428,7 +1435,8 @@ private int[] getSortedDocIdsWithRawForwardIndex(String column, IndexContainer i docIds[i] = i; } - DataType storedType = indexContainer._fieldSpec.getDataType().getStoredType(); + DataType dataType = indexContainer._fieldSpec.getDataType(); + DataType storedType = dataType.getStoredType(); switch (storedType) { case INT: IntArrays.quickSort(docIds, (d1, d2) -> Integer.compare(forwardIndex.getInt(d1), forwardIndex.getInt(d2))); @@ -1450,8 +1458,13 @@ private int[] getSortedDocIdsWithRawForwardIndex(String column, IndexContainer i IntArrays.quickSort(docIds, (d1, d2) -> forwardIndex.getString(d1).compareTo(forwardIndex.getString(d2))); break; case BYTES: - IntArrays.quickSort(docIds, - (d1, d2) -> ByteArray.compare(forwardIndex.getBytes(d1), forwardIndex.getBytes(d2))); + if (dataType == DataType.UUID) { + IntArrays.quickSort(docIds, + (d1, d2) -> UuidUtils.compare(forwardIndex.getBytes(d1), forwardIndex.getBytes(d2))); + } else { + IntArrays.quickSort(docIds, + (d1, d2) -> ByteArray.compare(forwardIndex.getBytes(d1), forwardIndex.getBytes(d2))); + } break; default: throw new UnsupportedOperationException( diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/converter/stats/MutableNoDictColumnStatistics.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/converter/stats/MutableNoDictColumnStatistics.java index 92d3455425f6..85c0e2423298 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/converter/stats/MutableNoDictColumnStatistics.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/converter/stats/MutableNoDictColumnStatistics.java @@ -117,7 +117,9 @@ public boolean isSorted() { int numDocs = _dataSourceMetadata.getNumDocs(); - // Verify that values are non-decreasing when iterated in the given order + // Verify that values are non-decreasing when iterated in the given order. The BYTES path uses + // ByteArray.compare (unsigned byte-wise lexicographic), which is identical to UuidUtils.compare's unsigned + // 64-bit-word ordering on canonical 16-byte big-endian UUIDs, so a single comparator handles both. DataType storedType = getStoredType(); if (_sortedDocIds != null) { switch (storedType) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOffHeapMutableDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOffHeapMutableDictionary.java index f8afb134544e..e03d20bb832d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOffHeapMutableDictionary.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOffHeapMutableDictionary.java @@ -66,7 +66,14 @@ public int index(Object value) { @Override public int[] index(Object[] values) { - throw new UnsupportedOperationException(); + int numValues = values.length; + int[] dictIds = new int[numValues]; + for (int i = 0; i < numValues; i++) { + byte[] bytesValue = (byte[]) values[i]; + updateStats(bytesValue); + dictIds[i] = indexValue(new ByteArray(bytesValue), bytesValue); + } + return dictIds; } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOnHeapMutableDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOnHeapMutableDictionary.java index 49b20161289d..fd22f15244f5 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOnHeapMutableDictionary.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/dictionary/BytesOnHeapMutableDictionary.java @@ -46,7 +46,14 @@ public int index(Object value) { @Override public int[] index(Object[] values) { - throw new UnsupportedOperationException(); + int numValues = values.length; + int[] dictIds = new int[numValues]; + for (int i = 0; i < numValues; i++) { + byte[] bytesValue = (byte[]) values[i]; + updateStats(bytesValue); + dictIds[i] = indexValue(new ByteArray(bytesValue)); + } + return dictIds; } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/FixedByteMVMutableForwardIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/FixedByteMVMutableForwardIndex.java index d42da3add11f..e51854c8a59f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/FixedByteMVMutableForwardIndex.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/FixedByteMVMutableForwardIndex.java @@ -121,6 +121,7 @@ public class FixedByteMVMutableForwardIndex implements MutableForwardIndex { private final String _context; private final boolean _isDictionaryEncoded; private final FieldSpec.DataType _storedType; + private final FieldSpec.DataType _dataType; private FixedByteSingleValueMultiColWriter _curHeaderWriter; private FixedByteSingleValueMultiColWriter _currentDataWriter; @@ -132,6 +133,13 @@ public class FixedByteMVMutableForwardIndex implements MutableForwardIndex { public FixedByteMVMutableForwardIndex(int maxNumberOfMultiValuesPerRow, int avgMultiValueCount, int rowCountPerChunk, int columnSizeInBytes, PinotDataBufferMemoryManager memoryManager, String context, boolean isDictionaryEncoded, FieldSpec.DataType storedType) { + this(maxNumberOfMultiValuesPerRow, avgMultiValueCount, rowCountPerChunk, columnSizeInBytes, memoryManager, context, + isDictionaryEncoded, storedType, storedType); + } + + public FixedByteMVMutableForwardIndex(int maxNumberOfMultiValuesPerRow, int avgMultiValueCount, int rowCountPerChunk, + int columnSizeInBytes, PinotDataBufferMemoryManager memoryManager, String context, boolean isDictionaryEncoded, + FieldSpec.DataType storedType, FieldSpec.DataType dataType) { _memoryManager = memoryManager; _context = context; int initialCapacity = Math.max(maxNumberOfMultiValuesPerRow, rowCountPerChunk * avgMultiValueCount); @@ -148,6 +156,7 @@ public FixedByteMVMutableForwardIndex(int maxNumberOfMultiValuesPerRow, int avgM //init(_rowCountPerChunk, _columnSizeInBytes, _maxNumberOfMultiValuesPerRow, initialCapacity, _incrementalCapacity); _isDictionaryEncoded = isDictionaryEncoded; _storedType = storedType; + _dataType = dataType; } private void addHeaderBuffer() { @@ -380,6 +389,37 @@ public double[] getDoubleMV(int docId) { return valueBuffer; } + @Override + public int getBytesMV(int docId, byte[][] valueBuffer) { + checkBytesMvSupported(); + FixedByteSingleValueMultiColReader headerReader = getCurrentReader(docId); + int rowInCurrentHeader = getRowInCurrentHeader(docId); + int bufferIndex = headerReader.getInt(rowInCurrentHeader, 0); + int startIndex = headerReader.getInt(rowInCurrentHeader, 1); + int length = headerReader.getInt(rowInCurrentHeader, 2); + FixedByteSingleValueMultiColReader dataReader = _dataReaders.get(bufferIndex); + for (int i = 0; i < length; i++) { + valueBuffer[i] = dataReader.getBytes(startIndex + i, 0); + } + return length; + } + + @Override + public byte[][] getBytesMV(int docId) { + checkBytesMvSupported(); + FixedByteSingleValueMultiColReader headerReader = getCurrentReader(docId); + int rowInCurrentHeader = getRowInCurrentHeader(docId); + int bufferIndex = headerReader.getInt(rowInCurrentHeader, 0); + int startIndex = headerReader.getInt(rowInCurrentHeader, 1); + int length = headerReader.getInt(rowInCurrentHeader, 2); + FixedByteSingleValueMultiColReader dataReader = _dataReaders.get(bufferIndex); + byte[][] valueBuffer = new byte[length][]; + for (int i = 0; i < length; i++) { + valueBuffer[i] = dataReader.getBytes(startIndex + i, 0); + } + return valueBuffer; + } + @Override public int getNumValuesMV(int docId) { FixedByteSingleValueMultiColReader headerReader = getCurrentReader(docId); @@ -424,6 +464,26 @@ public void setDoubleMV(int docId, double[] values) { } } + @Override + public void setBytesMV(int docId, byte[][] values) { + checkBytesMvSupported(); + int newStartIndex = updateHeader(docId, values.length); + for (int i = 0; i < values.length; i++) { + byte[] value = values[i]; + if (value.length != _columnSizeInBytes) { + throw new IllegalArgumentException( + "Expected fixed-width bytes value of length: " + _columnSizeInBytes + ", got: " + value.length); + } + _currentDataWriter.setBytes(newStartIndex + i, 0, value); + } + } + + private void checkBytesMvSupported() { + if (_dataType != DataType.UUID) { + throw new UnsupportedOperationException("Unsupported data type: " + _dataType + " for raw bytes MV index"); + } + } + @Override public boolean canAddMore() { return _numValues < DEFAULT_THRESHOLD_FOR_NUM_OF_VALUES_PER_COLUMN; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformer.java index 6deb83908557..98c0a2f91903 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformer.java @@ -19,8 +19,12 @@ package org.apache.pinot.segment.local.recordtransformer; import java.util.HashMap; +import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Set; +import java.util.UUID; +import javax.annotation.Nullable; import org.apache.pinot.common.utils.ThrottledLogger; import org.apache.pinot.segment.local.utils.DataTypeTransformerUtils; import org.apache.pinot.spi.config.table.TableConfig; @@ -52,21 +56,52 @@ public class DataTypeTransformer implements RecordTransformer { private final Map _dataTypes; private final boolean _continueOnError; private final ThrottledLogger _throttledLogger; + // UUID primary key columns for upsert/dedup tables; non-null and non-empty when applicable. + // Non-canonical (uppercase) UUID strings in these columns will be rejected at ingest time because + // Kafka partition routing is decided by the producer before Pinot normalises the value. Accepting + // uppercase UUIDs as primary keys silently causes dedup failures in multi-partition realtime + // upsert tables when the same logical UUID is routed to different partitions. + @Nullable + private final Set _upsertUuidPrimaryKeyColumns; /// Creates a [DataTypeTransformer] that converts the (non-virtual) schema columns to the data types defined in the /// [Schema]. public DataTypeTransformer(TableConfig tableConfig, Schema schema) { - this(tableConfig, extractSchemaDataTypes(schema)); + this(tableConfig, extractSchemaDataTypes(schema), schema); } /// Creates a [DataTypeTransformer] that converts the given columns to the provided [PinotDataType]s. This is useful /// for fixing the data types of source fields before other transformers (such as [ExpressionTransformer]) consume /// them. public DataTypeTransformer(TableConfig tableConfig, Map dataTypes) { + this(tableConfig, dataTypes, null); + } + + private DataTypeTransformer(TableConfig tableConfig, Map dataTypes, @Nullable Schema schema) { _dataTypes = dataTypes; IngestionConfig ingestionConfig = tableConfig.getIngestionConfig(); _continueOnError = ingestionConfig != null && ingestionConfig.isContinueOnError(); _throttledLogger = new ThrottledLogger(LOGGER, ingestionConfig); + + // Enforce canonical-form UUID PKs only when upsert/dedup is actually enabled — a present-but-disabled config + // (e.g., UpsertConfig with Mode.NONE) must not reject otherwise-valid rows. + if (schema != null && (tableConfig.isUpsertEnabled() || tableConfig.isDedupEnabled())) { + List primaryKeyColumns = schema.getPrimaryKeyColumns(); + if (primaryKeyColumns != null && !primaryKeyColumns.isEmpty()) { + Set uuidPkCols = new HashSet<>(); + for (String col : primaryKeyColumns) { + FieldSpec spec = schema.getFieldSpecFor(col); + if (spec != null && spec.getDataType() == FieldSpec.DataType.UUID) { + uuidPkCols.add(col); + } + } + _upsertUuidPrimaryKeyColumns = uuidPkCols.isEmpty() ? null : Set.copyOf(uuidPkCols); + } else { + _upsertUuidPrimaryKeyColumns = null; + } + } else { + _upsertUuidPrimaryKeyColumns = null; + } } private static Map extractSchemaDataTypes(Schema schema) { @@ -95,6 +130,10 @@ public void transform(GenericRow record) { String column = entry.getKey(); try { Object value = record.getValue(column); + if (_upsertUuidPrimaryKeyColumns != null && _upsertUuidPrimaryKeyColumns.contains(column) + && value instanceof CharSequence) { + validateCanonicalUuidPrimaryKey(column, value.toString()); + } value = DataTypeTransformerUtils.transformValue(column, value, entry.getValue()); record.putValue(column, value); } catch (Exception e) { @@ -107,4 +146,41 @@ public void transform(GenericRow record) { } } } + + /** + * Validates that a UUID primary key string value is in canonical lowercase RFC 4122 form. + * + *

UUID primary keys in upsert/dedup realtime tables must be canonical because Kafka partition routing + * is determined by the raw string value that the producer sends. If the producer sends the same logical + * UUID with different casing (e.g. uppercase vs lowercase), Kafka will hash the strings differently + * and the messages may land on different partitions. Pinot then normalises them to the same bytes + * inside each consuming segment, so within-segment dedup works, but cross-partition dedup never + * fires - producing duplicate or stale rows silently. + * + *

By rejecting any non-canonical UUID here (before bytes conversion loses the original string), + * we ensure that a producer-side inconsistency surfaces as an ingestion error rather than a silent + * correctness failure. This covers not only casing but also the dash-less 32-hex form and + * whitespace-padded values: {@code UuidUtils.toBytes(String)} (used by the downstream conversion) + * accepts those and would normalise them to the same bytes as the canonical string, so they must be + * rejected here too. Producers must canonicalise UUID primary keys to lowercase RFC 4122 form before + * publishing to Kafka. + */ + private static void validateCanonicalUuidPrimaryKey(String column, String uuidStr) { + String canonical = null; + try { + canonical = UUID.fromString(uuidStr).toString(); + } catch (IllegalArgumentException e) { + // Not parseable as a dashed RFC 4122 UUID (e.g. dash-less 32-hex or whitespace-padded). Reject below rather + // than defer to DataTypeTransformerUtils.transformValue, whose UuidUtils.toBytes would accept it as a + // non-canonical value and defeat the Kafka partition-routing guarantee. + } + if (!uuidStr.equals(canonical)) { + throw new IllegalArgumentException( + "Non-canonical UUID primary key value '" + uuidStr + "' in upsert/dedup column '" + column + "'. " + + "Expected canonical lowercase RFC 4122 form" + + (canonical != null ? ": '" + canonical + "'" : " (8-4-4-4-12 dashed lowercase)") + ". " + + "UUID primary keys must be in canonical lowercase form to ensure consistent Kafka partition " + + "routing. Non-canonical values cause silent dedup failures in multi-partition realtime tables."); + } + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPreIndexStatsCollector.java index 169d4d61f446..fc076605c9c0 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPreIndexStatsCollector.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPreIndexStatsCollector.java @@ -74,7 +74,12 @@ public void collect(Object entry) { addressSorted(value); if (_values.add(value)) { if (isPartitionEnabled()) { - updatePartition(value.toString()); + // Use DataType.toString so UUID columns get their canonical RFC 4122 form (with dashes) rather + // than the bare hex from ByteArray.toString(). The runtime ingest path (MutableSegmentImpl) and + // every partition function (Murmur, Uuid) expect this canonical form; passing bare hex here + // produces a different partition value than runtime — silently breaking partition pruning — and + // throws for the Uuid partition function (UuidUtils.toBytes rejects no-dash strings). + updatePartition(_fieldSpec.getDataType().toString(entry)); } int length = value.length(); _minLength = Math.min(_minLength, length); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/NoDictColumnStatisticsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/NoDictColumnStatisticsCollector.java index 4aa23fa791f4..d0894b292738 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/NoDictColumnStatisticsCollector.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/NoDictColumnStatisticsCollector.java @@ -155,7 +155,11 @@ public void collect(Object entry) { } } if (isPartitionEnabled()) { - updatePartition(comparable.toString()); + // See BytesColumnPreIndexStatsCollector: UUID columns must use DataType.toString to produce the + // canonical dashed form that matches the runtime MutableSegmentImpl partition path and + // UuidPartitionFunction expectation. comparable.toString() for UUID would yield the bare hex + // (ByteArray.toString) which breaks both Murmur (different partition value) and Uuid (rejected). + updatePartition(_fieldSpec.getDataType().toString(entry)); } _totalNumberOfEntries++; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProvider.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProvider.java index a9d5bc0bcc86..41adfa4aa1e4 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProvider.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProvider.java @@ -70,6 +70,7 @@ public Dictionary buildDictionary(VirtualColumnContext context) { case STRING: return new ConstantValueStringDictionary((String) fieldSpec.getDefaultNullValue()); case BYTES: + case UUID: return new ConstantValueBytesDictionary((byte[]) fieldSpec.getDefaultNullValue()); default: throw new IllegalStateException(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java index ac542042a81e..7f4ae0a210b0 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java @@ -364,7 +364,8 @@ public MutableIndex createMutableIndex(MutableIndexContext context, ForwardIndex } String column = context.getFieldSpec().getName(); String segmentName = context.getSegmentName(); - FieldSpec.DataType storedType = context.getFieldSpec().getDataType().getStoredType(); + FieldSpec.DataType dataType = context.getFieldSpec().getDataType(); + FieldSpec.DataType storedType = dataType.getStoredType(); int fixedLengthBytes = context.getFixedLengthBytes(); boolean isSingleValue = context.getFieldSpec().isSingleValueField(); if (!context.hasDictionary()) { @@ -402,13 +403,15 @@ public MutableIndex createMutableIndex(MutableIndexContext context, ForwardIndex } } else { // TODO: Add support for variable width (bytes, string, big decimal) MV RAW column types - assert storedType.isFixedWidth(); + Preconditions.checkState(dataType.isFixedWidth(), "Unsupported stored type: %s for no-dictionary MV column: %s", + storedType, column); String allocationContext = IndexUtil.buildAllocationContext(context.getSegmentName(), context.getFieldSpec().getName(), V1Constants.Indexes.RAW_MV_FORWARD_INDEX_FILE_EXTENSION); // TODO: Start with a smaller capacity on FixedByteMVForwardIndexReaderWriter and let it expand return new FixedByteMVMutableForwardIndex(MAX_MULTI_VALUES_PER_ROW, context.getAvgNumMultiValues(), - context.getCapacity(), storedType.size(), context.getMemoryManager(), allocationContext, false, storedType); + context.getCapacity(), dataType.size(), context.getMemoryManager(), allocationContext, false, storedType, + dataType); } } else { if (isSingleValue) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java index 318b7173dc8e..4e52a9245bbf 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java @@ -360,6 +360,8 @@ private void addColumnMinMaxValueWithoutDictionary(ColumnMetadata columnMetadata break; } case BYTES: { + // ByteArray.compare is unsigned byte-wise lexicographic; for canonical 16-byte big-endian UUIDs this is + // identical to UuidUtils.compare's unsigned 64-bit-word ordering, so a single comparator is sufficient. byte[] min = null; byte[] max = null; if (isSingleValue) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java index 39ee9a240cb3..12c7ec5a747e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java @@ -641,6 +641,8 @@ public int compare(int docId1, int docId2) { return _forwardIndexReader.getString(docId1, _forwardIndexReaderContext) .compareTo(_forwardIndexReader.getString(docId2, _forwardIndexReaderContext)); case BYTES: + // ByteArray.compare is unsigned byte-wise lexicographic; for canonical 16-byte big-endian UUIDs this is + // identical to UuidUtils.compare's unsigned 64-bit-word ordering, so a single comparator handles both. return ByteArray.compare(_forwardIndexReader.getBytes(docId1, _forwardIndexReaderContext), _forwardIndexReader.getBytes(docId2, _forwardIndexReaderContext)); default: diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/HashUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/HashUtils.java index 578280fb0d27..d37d8463fcb7 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/HashUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/HashUtils.java @@ -30,6 +30,7 @@ import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.PinotMd5Mode; +import org.apache.pinot.spi.utils.UuidUtils; public class HashUtils { @@ -47,6 +48,15 @@ public static byte[] hashMD5(byte[] bytes) { /** * Returns a byte array that is a concatenation of the binary representation of each of the passed UUID values. * If any of the values is not a valid UUID, then we return the result of {@link PrimaryKey#asBytes()}. + * + *

String inputs are parsed via {@link UUID#fromString(String)} (lenient — accepts non-canonical short + * forms like {@code "1-2-3-4-5"}) for backward compatibility with tables that have been hashed under the + * pre-UUID-type {@code HashFunction.UUID} contract. Binary inputs ({@code byte[]} or {@link UUID}) go + * through {@link UuidUtils#toBytes(Object)} since they carry no parse ambiguity. Un-parseable strings + * (and any other invalid value) still fall through to {@link PrimaryKey#asBytes()} via the outer + * {@code catch (RuntimeException)} below — matching master's behavior. Newly-declared + * {@code DataType.UUID} primary-key columns are independently constrained to canonical form at ingest + * time by {@code DataTypeTransformer.validateCanonicalUuidPrimaryKey}. */ public static byte[] hashUUID(PrimaryKey primaryKey) { Object[] values = primaryKey.getValues(); @@ -56,14 +66,17 @@ public static byte[] hashUUID(PrimaryKey primaryKey) { if (value == null) { throw new IllegalArgumentException("Found null value in primary key"); } - UUID uuid; try { - uuid = UUID.fromString(value.toString()); - } catch (Throwable t) { + if (value instanceof CharSequence) { + UUID uuid = UUID.fromString(value.toString()); + byteBuffer.putLong(uuid.getMostSignificantBits()); + byteBuffer.putLong(uuid.getLeastSignificantBits()); + } else { + byteBuffer.put(UuidUtils.toBytes(value)); + } + } catch (RuntimeException e) { return primaryKey.asBytes(); } - byteBuffer.putLong(uuid.getMostSignificantBits()); - byteBuffer.putLong(uuid.getLeastSignificantBits()); } return result; } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java index 7d3ba9433c8a..b3aae871b1f1 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java @@ -20,11 +20,14 @@ import java.nio.charset.StandardCharsets; import java.util.Random; +import java.util.UUID; import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.FixedByteMVMutableForwardIndex; +import org.apache.pinot.segment.spi.index.mutable.MutableDictionary; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -88,6 +91,31 @@ public void testMultiValueIndexingWithDictionary() { } } + @Test + public void testMultiValueBytesIndexingWithDictionary() + throws Exception { + try (MutableDictionary dictionary = new BytesOnHeapMutableDictionary()) { + assertMultiValueBytesIndexingWithDictionary(dictionary); + } + try (MutableDictionary dictionary = + new BytesOffHeapMutableDictionary(4, 10, _memoryManager, "bytesDictionary", 4)) { + assertMultiValueBytesIndexingWithDictionary(dictionary); + } + } + + private static void assertMultiValueBytesIndexingWithDictionary(MutableDictionary dictionary) { + byte[] first = new byte[]{1, 2}; + byte[] second = new byte[]{3, 4, 5}; + byte[] third = new byte[]{6}; + + assertEquals(dictionary.index(new Object[]{first, second, first}), new int[]{0, 1, 0}); + assertEquals(dictionary.index(new Object[]{third, second}), new int[]{2, 1}); + assertEquals(dictionary.length(), 3); + assertEquals(dictionary.getBytesValue(0), first); + assertEquals(dictionary.getBytesValue(1), second); + assertEquals(dictionary.getBytesValue(2), third); + } + @Test public void testMultiValueIndexingWithRawInt() { long seed = System.nanoTime(); @@ -286,4 +314,29 @@ public void testMultiValueIndexingWithRawByte() { fail("Failed with random seed: " + seed, t); } } + + @Test + public void testMultiValueIndexingWithRawUuidBytes() + throws Exception { + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, UuidUtils.UUID_NUM_BYTES, memManager, "indexer", false, FieldSpec.DataType.BYTES, + FieldSpec.DataType.UUID)) { + byte[][] values = new byte[][]{ + UuidUtils.toBytes(new UUID(1L, 2L)), + UuidUtils.toBytes(new UUID(3L, 4L)) + }; + indexer.setBytesMV(0, values); + + byte[][] buffer = new byte[values.length][]; + assertEquals(indexer.getBytesMV(0, buffer), values.length); + for (int i = 0; i < values.length; i++) { + assertEquals(buffer[i], values[i]); + } + byte[][] actualValues = indexer.getBytesMV(0); + for (int i = 0; i < values.length; i++) { + assertEquals(actualValues[i], values[i]); + } + } + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformerTest.java index fa98adf22186..17c1ba1874e1 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/DataTypeTransformerTest.java @@ -23,6 +23,14 @@ import java.util.List; import java.util.Map; import org.apache.pinot.segment.local.utils.DataTypeTransformerUtils; +import org.apache.pinot.spi.config.table.DedupConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.UpsertConfig; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -201,4 +209,108 @@ public void testStandardize() { } assertEqualsNoOrder((Object[]) DataTypeTransformerUtils.standardize(COLUMN, values, false), expectedValues); } + + /** + * Verifies that non-canonical (uppercase) UUID strings in upsert/dedup primary key columns are rejected, + * while canonical lowercase UUIDs are accepted, and non-primary-key UUID columns are unaffected. + */ + @Test + public void testUuidUpsertPrimaryKeyCanonicalValidation() { + String uuidCol = "uuidPk"; + String nonPkUuidCol = "uuidOther"; + String canonicalUuid = "550e8400-e29b-41d4-a716-446655440000"; + String uppercaseUuid = "550E8400-E29B-41D4-A716-446655440000"; + // Dash-less 32-hex and whitespace-padded forms are accepted by UuidUtils.toBytes(String) (and would normalise to + // the same bytes), but they are NOT canonical and route to a different Kafka partition, so they must be rejected. + String noDashUuid = "550e8400e29b41d4a716446655440000"; + String whitespaceUuid = " 550e8400-e29b-41d4-a716-446655440000 "; + + Schema schema = new Schema.SchemaBuilder() + .addSingleValueDimension(uuidCol, FieldSpec.DataType.UUID) + .addSingleValueDimension(nonPkUuidCol, FieldSpec.DataType.UUID) + .setPrimaryKeyColumns(List.of(uuidCol)) + .build(); + TableConfig upsertTableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName("testUpsertUuid") + .setUpsertConfig(new UpsertConfig(UpsertConfig.Mode.FULL)) + .build(); + + DataTypeTransformer upsertTransformer = new DataTypeTransformer(upsertTableConfig, schema); + + // Canonical lowercase UUID primary key: accepted + GenericRow canonicalRow = new GenericRow(); + canonicalRow.putValue(uuidCol, canonicalUuid); + canonicalRow.putValue(nonPkUuidCol, canonicalUuid); + upsertTransformer.transform(canonicalRow); // must not throw + + // Non-canonical uppercase UUID primary key: rejected + GenericRow uppercaseRow = new GenericRow(); + uppercaseRow.putValue(uuidCol, uppercaseUuid); + uppercaseRow.putValue(nonPkUuidCol, canonicalUuid); + try { + upsertTransformer.transform(uppercaseRow); + fail("Expected RuntimeException for non-canonical UUID primary key in upsert table"); + } catch (RuntimeException e) { + // Expected: DataTypeTransformer wraps the IllegalArgumentException in a RuntimeException + } + + // Dash-less 32-hex UUID primary key: rejected (parses via UuidUtils.toBytes hex fallback, but not canonical) + GenericRow noDashRow = new GenericRow(); + noDashRow.putValue(uuidCol, noDashUuid); + noDashRow.putValue(nonPkUuidCol, canonicalUuid); + try { + upsertTransformer.transform(noDashRow); + fail("Expected RuntimeException for dash-less UUID primary key in upsert table"); + } catch (RuntimeException e) { + // Expected + } + + // Whitespace-padded UUID primary key: rejected + GenericRow whitespaceRow = new GenericRow(); + whitespaceRow.putValue(uuidCol, whitespaceUuid); + whitespaceRow.putValue(nonPkUuidCol, canonicalUuid); + try { + upsertTransformer.transform(whitespaceRow); + fail("Expected RuntimeException for whitespace-padded UUID primary key in upsert table"); + } catch (RuntimeException e) { + // Expected + } + + // Non-canonical uppercase UUID in a NON-primary-key column: accepted (no restriction) + GenericRow nonPkUppercaseRow = new GenericRow(); + nonPkUppercaseRow.putValue(uuidCol, canonicalUuid); + nonPkUppercaseRow.putValue(nonPkUuidCol, uppercaseUuid); + upsertTransformer.transform(nonPkUppercaseRow); // must not throw + + // For a non-upsert table, non-canonical UUID in primary-key column is also accepted + TableConfig offlineTableConfig = + new TableConfigBuilder(TableType.OFFLINE).setTableName("testOfflineUuid").build(); + DataTypeTransformer offlineTransformer = new DataTypeTransformer(offlineTableConfig, schema); + GenericRow offlineUppercaseRow = new GenericRow(); + offlineUppercaseRow.putValue(uuidCol, uppercaseUuid); + offlineUppercaseRow.putValue(nonPkUuidCol, uppercaseUuid); + offlineTransformer.transform(offlineUppercaseRow); // must not throw + + // A present-but-disabled upsert config (Mode.NONE) must not enforce the canonical-PK restriction + TableConfig disabledUpsertTableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName("testDisabledUpsertUuid") + .setUpsertConfig(new UpsertConfig(UpsertConfig.Mode.NONE)) + .build(); + DataTypeTransformer disabledUpsertTransformer = new DataTypeTransformer(disabledUpsertTableConfig, schema); + GenericRow disabledUpsertUppercaseRow = new GenericRow(); + disabledUpsertUppercaseRow.putValue(uuidCol, uppercaseUuid); + disabledUpsertUppercaseRow.putValue(nonPkUuidCol, canonicalUuid); + disabledUpsertTransformer.transform(disabledUpsertUppercaseRow); // must not throw + + // A present-but-disabled dedup config must not enforce the canonical-PK restriction either + TableConfig disabledDedupTableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName("testDisabledDedupUuid") + .setDedupConfig(new DedupConfig(false, null)) + .build(); + DataTypeTransformer disabledDedupTransformer = new DataTypeTransformer(disabledDedupTableConfig, schema); + GenericRow disabledDedupUppercaseRow = new GenericRow(); + disabledDedupUppercaseRow.putValue(uuidCol, uppercaseUuid); + disabledDedupUppercaseRow.putValue(nonPkUuidCol, canonicalUuid); + disabledDedupTransformer.transform(disabledDedupUppercaseRow); // must not throw + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProviderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProviderTest.java index f1f655ae9bd1..3e7cc400dee2 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProviderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/column/DefaultNullValueVirtualColumnProviderTest.java @@ -31,6 +31,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.BytesUtils; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -45,6 +46,7 @@ public class DefaultNullValueVirtualColumnProviderTest { private static final FieldSpec SV_STRING_WITH_DEFAULT = new DimensionFieldSpec("svStringColumn", DataType.STRING, true, "default"); private static final FieldSpec SV_BYTES = new DimensionFieldSpec("svBytesColumn", DataType.BYTES, true); + private static final FieldSpec SV_UUID = new DimensionFieldSpec("svUuidColumn", DataType.UUID, true); private static final FieldSpec MV_INT = new DimensionFieldSpec("mvIntColumn", DataType.INT, false); private static final FieldSpec MV_LONG = new DimensionFieldSpec("mvLongColumn", DataType.LONG, false); private static final FieldSpec MV_FLOAT = new DimensionFieldSpec("mvFloatColumn", DataType.FLOAT, false); @@ -146,6 +148,14 @@ public void testBuildDictionary() { assertEquals(dictionary.getClass(), ConstantValueBytesDictionary.class); assertEquals(dictionary.getBytesValue(0), new byte[0]); + virtualColumnContext = new VirtualColumnContext(SV_UUID, 1); + dictionary = new DefaultNullValueVirtualColumnProvider().buildDictionary(virtualColumnContext); + assertEquals(dictionary.getClass(), ConstantValueBytesDictionary.class); + byte[] uuidDefaultNullValue = (byte[]) SV_UUID.getDefaultNullValue(); + assertEquals(dictionary.getBytesValue(0), uuidDefaultNullValue); + assertEquals(dictionary.getStringValue(0), BytesUtils.toHexString(uuidDefaultNullValue)); + assertEquals(dictionary.indexOf(BytesUtils.toHexString(uuidDefaultNullValue)), 0); + virtualColumnContext = new VirtualColumnContext(MV_INT, 1); dictionary = new DefaultNullValueVirtualColumnProvider().buildDictionary(virtualColumnContext); assertEquals(dictionary.getClass(), ConstantValueIntDictionary.class); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java index 589016d7bc86..d81ef7fe1e9b 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java @@ -29,6 +29,7 @@ import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.spi.config.table.BloomFilterConfig; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.util.TestUtils; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -75,6 +76,33 @@ public void testBloomFilterCreator() } } + @Test + public void testUuidBloomFilterCreatorWithBytesValues() + throws Exception { + int cardinality = 100; + String columnName = "uuidColumn"; + String uuid0 = "550e8400-e29b-41d4-a716-446655440000"; + String uuid1 = "550e8400-e29b-41d4-a716-446655440001"; + try (BloomFilterCreator bloomFilterCreator = new OnHeapGuavaBloomFilterCreator(TEMP_DIR, columnName, cardinality, + new BloomFilterConfig(BloomFilterConfig.DEFAULT_FPP, 0, false), FieldSpec.DataType.UUID)) { + bloomFilterCreator.add(UuidUtils.toBytes(uuid0), -1); + bloomFilterCreator.add(UuidUtils.toBytes(uuid1), -1); + bloomFilterCreator.seal(); + } + + File bloomFilterFile = new File(TEMP_DIR, columnName + V1Constants.Indexes.BLOOM_FILTER_FILE_EXTENSION); + try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(bloomFilterFile); + BloomFilterReader onHeapBloomFilter = BloomFilterReaderFactory.getBloomFilterReader(dataBuffer, true); + BloomFilterReader offHeapBloomFilter = BloomFilterReaderFactory.getBloomFilterReader(dataBuffer, false)) { + Assert.assertTrue(onHeapBloomFilter.mightContain(uuid0)); + Assert.assertTrue(onHeapBloomFilter.mightContain(uuid1)); + Assert.assertFalse(onHeapBloomFilter.mightContain("550e8400-e29b-41d4-a716-4466554400ff")); + Assert.assertTrue(offHeapBloomFilter.mightContain(uuid0)); + Assert.assertTrue(offHeapBloomFilter.mightContain(uuid1)); + Assert.assertFalse(offHeapBloomFilter.mightContain("550e8400-e29b-41d4-a716-4466554400ff")); + } + } + @AfterClass public void tearDown() throws Exception { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java index c33aa0662a7e..1e38c4dfc698 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java @@ -39,6 +39,7 @@ import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; @@ -48,30 +49,39 @@ import static org.testng.Assert.assertTrue; -/// Regression tests for [ColumnMinMaxValueGenerator] on raw BYTES columns. +/// Regression tests for [ColumnMinMaxValueGenerator] on raw (no-dictionary) BYTES and UUID columns. /// /// The generator reads the forward index when column metadata does not contain min/max values. These tests exercise -/// both single-value and multi-value forward index paths and verify their unsigned byte-wise ordering. +/// the single-value, multi-value, and UUID (storedType=BYTES) forward index paths and verify their unsigned byte-wise +/// ordering. The raw-BYTES loop once had its comparison directions inverted, silently persisting swapped min/max +/// metadata that value-based segment pruning then consumed. public class ColumnMinMaxValueGeneratorTest { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), ColumnMinMaxValueGeneratorTest.class.getSimpleName()); private static final String SEGMENT_NAME = "testSegment"; private static final String BYTES_COLUMN = "bytesCol"; private static final String BYTES_MV_COLUMN = "bytesMvCol"; + private static final String UUID_COLUMN = "uuidCol"; + // Ordered ascending by unsigned byte-wise comparison private static final byte[] BYTES_SMALL = new byte[]{0x00, 0x01}; private static final byte[] BYTES_MID = new byte[]{0x10, (byte) 0xff}; private static final byte[] BYTES_LARGE = new byte[]{(byte) 0x80, 0x00}; + private static final String UUID_SMALL = "00000000-0000-0000-0000-000000000001"; + private static final String UUID_MID = "550e8400-e29b-41d4-a716-446655440000"; + private static final String UUID_LARGE = "ffffffff-ffff-ffff-ffff-fffffffffffe"; + @Test public void testRawBytesMinMaxDirection() throws Exception { Schema schema = new Schema.SchemaBuilder().setSchemaName("testSchema") .addSingleValueDimension(BYTES_COLUMN, DataType.BYTES) .addMultiValueDimension(BYTES_MV_COLUMN, DataType.BYTES) + .addSingleValueDimension(UUID_COLUMN, DataType.UUID) .build(); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable") - .setNoDictionaryColumns(List.of(BYTES_COLUMN, BYTES_MV_COLUMN)) + .setNoDictionaryColumns(List.of(BYTES_COLUMN, BYTES_MV_COLUMN, UUID_COLUMN)) .build(); File indexDir = buildSegment(tableConfig, schema, createRows()); @@ -80,8 +90,11 @@ public void testRawBytesMinMaxDirection() generateMinMaxValues(indexDir); SegmentMetadataImpl reloaded = new SegmentMetadataImpl(indexDir); - assertMinMax(reloaded, BYTES_COLUMN); - assertMinMax(reloaded, BYTES_MV_COLUMN); + assertMinMax(reloaded, BYTES_COLUMN, BYTES_SMALL, BYTES_LARGE); + // The raw-BYTES MV loop is a distinct code path sharing the comparator with the SV loop + assertMinMax(reloaded, BYTES_MV_COLUMN, BYTES_SMALL, BYTES_LARGE); + // UUID stores as BYTES; its min/max must order by the same unsigned byte-wise comparison + assertMinMax(reloaded, UUID_COLUMN, UuidUtils.toBytes(UUID_SMALL), UuidUtils.toBytes(UUID_LARGE)); } @Test @@ -104,11 +117,12 @@ public void testEmptyRawBytesMarksMinMaxInvalid() assertTrue(reloaded.getColumnMetadataFor(BYTES_COLUMN).isMinMaxValueInvalid()); } - private static void assertMinMax(SegmentMetadataImpl segmentMetadata, String column) { + private static void assertMinMax(SegmentMetadataImpl segmentMetadata, String column, byte[] expectedMin, + byte[] expectedMax) { ByteArray min = (ByteArray) segmentMetadata.getColumnMetadataFor(column).getMinValue(); ByteArray max = (ByteArray) segmentMetadata.getColumnMetadataFor(column).getMaxValue(); - assertEquals(min, new ByteArray(BYTES_SMALL)); - assertEquals(max, new ByteArray(BYTES_LARGE)); + assertEquals(min, new ByteArray(expectedMin)); + assertEquals(max, new ByteArray(expectedMax)); assertTrue(min.compareTo(max) < 0); } @@ -125,10 +139,17 @@ private static void generateMinMaxValues(File indexDir) private static List createRows() { List rows = new ArrayList<>(); - for (byte[] value : new byte[][]{BYTES_MID, BYTES_LARGE, BYTES_SMALL}) { + // Deliberately insert in non-sorted order so the running min/max comparisons are both exercised + for (Object[] values : new Object[][]{ + {BYTES_MID, UuidUtils.toBytes(UUID_MID)}, + {BYTES_LARGE, UuidUtils.toBytes(UUID_LARGE)}, + {BYTES_SMALL, UuidUtils.toBytes(UUID_SMALL)} + }) { GenericRow row = new GenericRow(); - row.putValue(BYTES_COLUMN, value); - row.putValue(BYTES_MV_COLUMN, new Object[]{value, BYTES_MID}); + row.putValue(BYTES_COLUMN, values[0]); + // Each MV row carries two values so the inner per-value loop is exercised as well + row.putValue(BYTES_MV_COLUMN, new Object[]{values[0], BYTES_MID}); + row.putValue(UUID_COLUMN, values[1]); rows.add(row); } return rows; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/LazyRowTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/LazyRowTest.java index d941988f30cd..63f906721618 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/LazyRowTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/LazyRowTest.java @@ -22,9 +22,12 @@ import java.util.HashSet; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; import org.testng.annotations.Test; import static org.mockito.ArgumentMatchers.any; @@ -93,6 +96,15 @@ private IndexSegment getMockSegment() { when(segment.getDataSource("col1")).thenReturn(_col1Datasource); when(segment.getDataSource("col2")).thenReturn(col2Datasource); + DataSourceMetadata col1Metadata = mock(DataSourceMetadata.class); + DataSourceMetadata col2Metadata = mock(DataSourceMetadata.class); + FieldSpec col1FieldSpec = new DimensionFieldSpec("col1", FieldSpec.DataType.STRING, true); + FieldSpec col2FieldSpec = new DimensionFieldSpec("col2", FieldSpec.DataType.STRING, true); + when(col1Metadata.getFieldSpec()).thenReturn(col1FieldSpec); + when(col2Metadata.getFieldSpec()).thenReturn(col2FieldSpec); + when(_col1Datasource.getDataSourceMetadata()).thenReturn(col1Metadata); + when(col2Datasource.getDataSourceMetadata()).thenReturn(col2Metadata); + NullValueVectorReader col1NullVectorReader = mock(NullValueVectorReader.class); when(col1NullVectorReader.isNull(1)).thenReturn(true); NullValueVectorReader col2NullVectorReader = mock(NullValueVectorReader.class); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java index 7acfc0ace545..760055735a2f 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java @@ -44,12 +44,14 @@ import org.apache.pinot.segment.spi.MutableSegment; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; import org.apache.pinot.spi.config.table.HashFunction; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.UpsertConfig; +import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.PrimaryKey; @@ -105,6 +107,10 @@ private static ImmutableSegmentImpl mockImmutableSegment(int sequenceNumber, when(forwardIndex.getInt(anyInt(), any())).thenAnswer( invocation -> primaryKeys.get(invocation.getArgument(0)).getValues()[0]); when(dataSource.getForwardIndex()).thenReturn(forwardIndex); + DataSourceMetadata dataSourceMetadata = mock(DataSourceMetadata.class); + when(dataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), FieldSpec.DataType.INT, true)); + when(dataSource.getDataSourceMetadata()).thenReturn(dataSourceMetadata); SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); long creationTimeMs = System.currentTimeMillis(); when(segmentMetadata.getIndexCreationTime()).thenReturn(creationTimeMs); @@ -134,6 +140,10 @@ private static ImmutableSegmentImpl mockUploadedImmutableSegment(String suffix, when(forwardIndex.getInt(anyInt(), any())).thenAnswer( invocation -> primaryKeys.get(invocation.getArgument(0)).getValues()[0]); when(dataSource.getForwardIndex()).thenReturn(forwardIndex); + DataSourceMetadata uploadedDataSourceMetadata = mock(DataSourceMetadata.class); + when(uploadedDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), FieldSpec.DataType.INT, true)); + when(dataSource.getDataSourceMetadata()).thenReturn(uploadedDataSourceMetadata); SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); when(segmentMetadata.getIndexCreationTime()).thenReturn(creationTimeMs); when(segmentMetadata.getZkCreationTime()).thenReturn(creationTimeMs); @@ -186,6 +196,10 @@ private static MutableSegment mockMutableSegmentWithDataSource(int sequenceNumbe return docId; }); when(dataSource.getForwardIndex()).thenReturn(forwardIndex); + DataSourceMetadata mutableDataSourceMetadata = mock(DataSourceMetadata.class); + when(mutableDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), FieldSpec.DataType.INT, true)); + when(dataSource.getDataSourceMetadata()).thenReturn(mutableDataSourceMetadata); when(segment.getDataSource(anyString())).thenReturn(dataSource); when(segment.getDataSource(PRIMARY_KEY_COLUMNS.get(0))).thenReturn(dataSource); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java index dbba9b0e2eba..2e9dc43ddb8b 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java @@ -50,6 +50,7 @@ import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; @@ -924,6 +925,10 @@ private static ImmutableSegmentImpl mockImmutableSegmentWithTimestamps(int seque return MOCK_FALLBACK_BASE_OFFSET + docId; }); when(primaryKeyDataSource.getForwardIndex()).thenReturn(primaryKeyForwardIndex); + DataSourceMetadata primaryKeyDataSourceMetadata = mock(DataSourceMetadata.class); + when(primaryKeyDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), DataType.INT, true)); + when(primaryKeyDataSource.getDataSourceMetadata()).thenReturn(primaryKeyDataSourceMetadata); // Mock comparison column data source DataSource comparisonDataSource = mock(DataSource.class); @@ -940,6 +945,10 @@ private static ImmutableSegmentImpl mockImmutableSegmentWithTimestamps(int seque return MOCK_FALLBACK_BASE_OFFSET + (docId * 100); }); when(comparisonDataSource.getForwardIndex()).thenReturn(comparisonForwardIndex); + DataSourceMetadata comparisonDataSourceMetadata = mock(DataSourceMetadata.class); + when(comparisonDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(COMPARISON_COLUMNS.get(0), DataType.INT, true)); + when(comparisonDataSource.getDataSourceMetadata()).thenReturn(comparisonDataSourceMetadata); // Set up data source mapping - IMPORTANT: anyString() must be registered FIRST, // then specific matchers override it (Mockito uses last matching stub) @@ -1037,6 +1046,10 @@ private static ImmutableSegmentImpl mockUploadedImmutableSegment(String suffix, when(forwardIndex.getInt(anyInt(), any())).thenAnswer( invocation -> primaryKeys.get(invocation.getArgument(0)).getValues()[0]); when(dataSource.getForwardIndex()).thenReturn(forwardIndex); + DataSourceMetadata uploadedDataSourceMetadata = mock(DataSourceMetadata.class); + when(uploadedDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), DataType.INT, true)); + when(dataSource.getDataSourceMetadata()).thenReturn(uploadedDataSourceMetadata); SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); when(segmentMetadata.getIndexCreationTime()).thenReturn(creationTimeMs); when(segmentMetadata.getZkCreationTime()).thenReturn(creationTimeMs); @@ -1108,6 +1121,10 @@ private static MutableSegment mockMutableSegmentWithDataSource(int sequenceNumbe return docId; }); when(dataSource.getForwardIndex()).thenReturn(forwardIndex); + DataSourceMetadata mutableDataSourceMetadata = mock(DataSourceMetadata.class); + when(mutableDataSourceMetadata.getFieldSpec()).thenReturn( + new DimensionFieldSpec(PRIMARY_KEY_COLUMNS.get(0), DataType.INT, true)); + when(dataSource.getDataSourceMetadata()).thenReturn(mutableDataSourceMetadata); when(segment.getDataSource(anyString())).thenReturn(dataSource); when(segment.getDataSource(PRIMARY_KEY_COLUMNS.get(0))).thenReturn(dataSource); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/HashUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/HashUtilsTest.java index 8c60ad6f374d..c573265424bc 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/HashUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/HashUtilsTest.java @@ -21,9 +21,11 @@ import java.util.UUID; import org.apache.pinot.spi.config.table.HashFunction; import org.apache.pinot.spi.data.readers.PrimaryKey; +import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.PinotMd5Mode; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.Test; import static org.testng.Assert.*; @@ -64,6 +66,46 @@ public void testHashUUID() { } } + @Test + public void testHashUUIDNormalizedPrimaryKeyValues() { + UUID firstUuid = UUID.randomUUID(); + UUID secondUuid = UUID.randomUUID(); + + byte[] expectedHash = HashUtils.hashUUID(new PrimaryKey(new Object[]{firstUuid, secondUuid})); + assertEquals(HashUtils.hashUUID(new PrimaryKey( + new Object[]{new ByteArray(UuidUtils.toBytes(firstUuid)), new ByteArray(UuidUtils.toBytes(secondUuid))})), + expectedHash); + assertEquals(HashUtils.hashUUID( + new PrimaryKey(new Object[]{UuidUtils.toBytes(firstUuid), UuidUtils.toBytes(secondUuid)})), expectedHash); + } + + /** + * Regression: {@link HashUtils#hashUUID} must accept non-canonical-but-Java-parseable UUID strings + * (e.g. {@code "1-2-3-4-5"}) so existing upsert tables using {@code HashFunction.UUID} keep producing + * the same hash before and after this PR. {@link UuidUtils#toBytes(String)} is strict and would reject + * such inputs — for the legacy hash path we route String inputs through {@code UUID.fromString} directly + * to preserve the lenient pre-PR behavior. + */ + @Test + public void testHashUUIDLenientForNonCanonicalStrings() { + String nonCanonical = "1-2-3-4-5"; + UUID expected = UUID.fromString(nonCanonical); + + byte[] hashResult = HashUtils.hashUUID(new PrimaryKey(new Object[]{nonCanonical})); + assertEquals(hashResult.length, 16); + + long msb = 0; + long lsb = 0; + for (int i = 0; i < 8; i++) { + msb = (msb << 8) | (hashResult[i] & 0xFF); + } + for (int i = 8; i < 16; i++) { + lsb = (lsb << 8) | (hashResult[i] & 0xFF); + } + assertEquals(new UUID(msb, lsb), expected, + "hashUUID(\"1-2-3-4-5\") must equal UUID.fromString output for backward compatibility"); + } + @Test public void testHashPrimaryKeyWithMd5Disabled() { PrimaryKey primaryKey = new PrimaryKey(new Object[]{"hello world"}); diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/BloomFilterCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/BloomFilterCreator.java index 760a2310ece6..271188b08314 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/BloomFilterCreator.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/BloomFilterCreator.java @@ -23,6 +23,7 @@ import org.apache.pinot.segment.spi.index.IndexCreator; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.UuidUtils; public interface BloomFilterCreator extends IndexCreator { @@ -33,6 +34,8 @@ public interface BloomFilterCreator extends IndexCreator { default void add(Object value, int dictId) { if (getDataType() == FieldSpec.DataType.BYTES) { add(BytesUtils.toHexString((byte[]) value)); + } else if (getDataType() == FieldSpec.DataType.UUID) { + add(uuidToCanonicalString(value)); } else { add(value.toString()); } @@ -44,12 +47,25 @@ default void add(Object[] values, @Nullable int[] dictIds) { for (Object value : values) { add(BytesUtils.toHexString((byte[]) value)); } + } else if (getDataType() == FieldSpec.DataType.UUID) { + for (Object value : values) { + add(uuidToCanonicalString(value)); + } } else { for (Object value : values) { add(value.toString()); } } } + + /// Renders a UUID value (typically a 16-byte big-endian `byte[]` from segment ingest) as its canonical + /// lowercase RFC 4122 string. The `byte[]` fast path just skips the type dispatch of `UuidUtils.toBytes(Object)`; + /// neither path copies the buffer (`UuidUtils.toBytes(byte[])` validates the width and returns it as-is). + private static String uuidToCanonicalString(Object value) { + return value instanceof byte[] + ? UuidUtils.toString((byte[]) value) + : UuidUtils.toString(UuidUtils.toBytes(value)); + } /** * Adds a value to the bloom filter. */ diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToAvroConverter.java b/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToAvroConverter.java index dd0d56d2cf43..a7a6d41dcc52 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToAvroConverter.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToAvroConverter.java @@ -56,7 +56,8 @@ public void convert() PinotSegmentRecordReader pinotSegmentRecordReader = new PinotSegmentRecordReader(); pinotSegmentRecordReader.init(indexDir, null, null, false, _forwardIndexOnly); try (pinotSegmentRecordReader) { - try (DataFileWriter recordWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) { + try (DataFileWriter recordWriter = + new DataFileWriter<>(new GenericDatumWriter<>(avroSchema, SegmentProcessorAvroUtils.getAvroDataModel()))) { recordWriter.create(avroSchema, new File(_outputFile)); GenericRow row = new GenericRow(); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToParquetConverter.java b/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToParquetConverter.java index 8a901d567790..f50ff8e1dd43 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToParquetConverter.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/segment/converter/PinotSegmentToParquetConverter.java @@ -69,6 +69,7 @@ public void convert() try (pinotSegmentRecordReader) { try (ParquetWriter parquetWriter = AvroParquetWriter.builder(outputFile).withSchema(avroSchema) + .withDataModel(SegmentProcessorAvroUtils.getAvroDataModel()) .withCompressionCodec(_compressionCodec) .withConf(hadoopConf).build()) { GenericRow row = new GenericRow(); diff --git a/pinot-tools/src/test/java/org/apache/pinot/tools/segment/converter/PinotSegmentConverterTest.java b/pinot-tools/src/test/java/org/apache/pinot/tools/segment/converter/PinotSegmentConverterTest.java index bc5244b17f67..666c06a4ba61 100644 --- a/pinot-tools/src/test/java/org/apache/pinot/tools/segment/converter/PinotSegmentConverterTest.java +++ b/pinot-tools/src/test/java/org/apache/pinot/tools/segment/converter/PinotSegmentConverterTest.java @@ -36,6 +36,7 @@ import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -216,6 +217,62 @@ public void testParquetConverter() } } + private static final String UUID_SV_COLUMN = "uuidSVColumn"; + private static final String UUID_MV_COLUMN = "uuidMVColumn"; + private static final String UUID_SV_VALUE = "12345678-1234-1234-1234-1234567890ab"; + private static final String UUID_MV_VALUE_1 = "550e8400-e29b-41d4-a716-446655440000"; + private static final String UUID_MV_VALUE_2 = "550e8400-e29b-41d4-a716-446655440001"; + + /// Builds a segment with SV + MV UUID columns and converts it through both the Avro and Parquet converters. UUID is + /// stored as a 16-byte value but exported as a string{logicalType:uuid} field, so this exercises the uuid + /// Conversion registered on each writer's data model (getAvroDataModel) — including the distinct Parquet write path + /// (AvroParquetWriter.withDataModel), which a plain writer without the Conversion could not serialize. + @Test + public void testUuidConverters() + throws Exception { + Schema uuidSchema = new Schema.SchemaBuilder().addSingleValueDimension(UUID_SV_COLUMN, DataType.UUID) + .addMultiValueDimension(UUID_MV_COLUMN, DataType.UUID).build(); + TableConfig uuidTableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("uuidTable").build(); + + GenericRow record = new GenericRow(); + record.putValue(UUID_SV_COLUMN, UUID_SV_VALUE); + record.putValue(UUID_MV_COLUMN, new Object[]{UUID_MV_VALUE_1, UUID_MV_VALUE_2}); + + SegmentGeneratorConfig config = new SegmentGeneratorConfig(uuidTableConfig, uuidSchema); + config.setTableName("uuidTable"); + config.setSegmentName("uuidSegment"); + config.setOutDir(new File(TEMP_DIR, "uuidSegment").getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(List.of(record))); + driver.build(); + String segmentDir = driver.getOutputDirectory().getPath(); + + File avroOut = new File(TEMP_DIR, "uuidSegment.avro"); + new PinotSegmentToAvroConverter(segmentDir, avroOut.getPath()).convert(); + try (AvroRecordReader reader = new AvroRecordReader()) { + reader.init(avroOut, uuidSchema.getFieldSpecMap().keySet(), null); + assertUuidRecord(reader.next()); + assertFalse(reader.hasNext()); + } + + File parquetOut = new File(TEMP_DIR, "uuidSegment.parquet"); + new PinotSegmentToParquetConverter(segmentDir, parquetOut.getPath()).convert(); + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(parquetOut, uuidSchema.getFieldSpecMap().keySet(), null); + assertUuidRecord(reader.next()); + assertFalse(reader.hasNext()); + } + } + + private static void assertUuidRecord(GenericRow record) { + // The reader may surface the uuid value as a String/UUID/byte[]; UuidUtils.toBytes normalizes all of them. + assertEquals(UuidUtils.toBytes(record.getValue(UUID_SV_COLUMN)), UuidUtils.toBytes(UUID_SV_VALUE)); + Object[] mv = (Object[]) record.getValue(UUID_MV_COLUMN); + assertEquals(mv.length, 2); + assertEquals(UuidUtils.toBytes(mv[0]), UuidUtils.toBytes(UUID_MV_VALUE_1)); + assertEquals(UuidUtils.toBytes(mv[1]), UuidUtils.toBytes(UUID_MV_VALUE_2)); + } + @AfterClass public void tearDown() throws IOException {