Skip to content

Commit 64febd9

Browse files
xiangfu0xiangfu0
committed
Address UUID review feedback
1 parent d2dc259 commit 64febd9

8 files changed

Lines changed: 179 additions & 7 deletions

File tree

‎pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch
104104
vector = new Float8Vector(colName, ALLOCATOR);
105105
break;
106106
case TIMESTAMP:
107+
case UUID:
107108
case STRING:
108109
case BYTES:
109110
case BIG_DECIMAL:
@@ -230,6 +231,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch
230231
((Float8Vector) vector).setSafe(rowIndex, ((Number) value).doubleValue());
231232
break;
232233
case TIMESTAMP:
234+
case UUID:
233235
case STRING:
234236
case BYTES:
235237
case BIG_DECIMAL:
@@ -406,6 +408,7 @@ public ResultTable decodeResultTable(byte[] bytes, int rowSize, DataSchema schem
406408
row[col] = ((BitVector) vector).get(i) == 1;
407409
break;
408410
case TIMESTAMP:
411+
case UUID:
409412
case STRING:
410413
case BYTES:
411414
case BIG_DECIMAL:

‎pinot-common/src/test/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoderTest.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,26 @@ public void testEncodeDecodeMultipleRows()
114114
}
115115
}
116116

117+
@Test
118+
public void testEncodeDecodeUuidColumn()
119+
throws IOException {
120+
DataSchema schema = new DataSchema(new String[]{"uuidCol"}, new ColumnDataType[]{ColumnDataType.UUID});
121+
List<Object[]> rows = Arrays.asList(
122+
new Object[]{"550e8400-e29b-41d4-a716-446655440000"},
123+
new Object[]{"f81d4fae-7dec-11d0-a765-00a0c91e6bf6"}
124+
);
125+
126+
ResultTable resultTable = new ResultTable(schema, rows);
127+
ArrowResponseEncoder encoder = new ArrowResponseEncoder();
128+
byte[] encodedBytes = encoder.encodeResultTable(resultTable, 0, rows.size());
129+
ResultTable decodedTable = encoder.decodeResultTable(encodedBytes, rows.size(), schema);
130+
131+
assertEquals(decodedTable.getRows().size(), rows.size(), "Row count should match");
132+
for (int i = 0; i < rows.size(); i++) {
133+
assertEquals(decodedTable.getRows().get(i)[0], rows.get(i)[0], "UUID row " + i + " should match");
134+
}
135+
}
136+
117137
@Test
118138
public void testEncodeDecodeAllDataTypes()
119139
throws IOException {

‎pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapper.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.sql.Timestamp;
2424
import java.util.List;
2525
import java.util.Map;
26+
import java.util.UUID;
2627
import org.apache.commons.lang3.ArrayUtils;
2728
import org.apache.pinot.common.function.FunctionInfo;
2829
import org.apache.pinot.common.function.FunctionUtils;
@@ -35,6 +36,7 @@
3536
import org.apache.pinot.spi.data.FieldSpec.DataType;
3637
import org.apache.pinot.spi.utils.ByteArray;
3738
import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
39+
import org.apache.pinot.spi.utils.UuidUtils;
3840

3941

4042
/**
@@ -424,6 +426,16 @@ private void getNonLiteralValues(ValueBlock valueBlock) {
424426
case BYTES:
425427
_nonLiteralValues[i] = transformFunction.transformToBytesValuesSV(valueBlock);
426428
break;
429+
case UUID: {
430+
byte[][] bytesValues = transformFunction.transformToBytesValuesSV(valueBlock);
431+
int numValues = bytesValues.length;
432+
UUID[] uuidValues = new UUID[numValues];
433+
for (int j = 0; j < numValues; j++) {
434+
uuidValues[j] = UuidUtils.toUUID(bytesValues[j]);
435+
}
436+
_nonLiteralValues[i] = uuidValues;
437+
break;
438+
}
427439
case PRIMITIVE_INT_ARRAY:
428440
_nonLiteralValues[i] = transformFunction.transformToIntValuesMV(valueBlock);
429441
break;

‎pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,17 +24,23 @@
2424
import java.text.Normalizer;
2525
import java.util.Arrays;
2626
import java.util.Base64;
27+
import java.util.Collections;
2728
import java.util.Random;
29+
import java.util.UUID;
2830
import org.apache.commons.codec.digest.DigestUtils;
2931
import org.apache.commons.lang3.ArrayUtils;
3032
import org.apache.commons.lang3.StringUtils;
33+
import org.apache.pinot.common.function.FunctionInfo;
3134
import org.apache.pinot.common.function.scalar.ArithmeticFunctions;
3235
import org.apache.pinot.common.request.context.ExpressionContext;
3336
import org.apache.pinot.common.request.context.RequestContextUtils;
37+
import org.apache.pinot.core.operator.blocks.ValueBlock;
38+
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
3439
import org.apache.pinot.spi.data.FieldSpec.DataType;
3540
import org.apache.pinot.spi.utils.ArrayCopyUtils;
3641
import org.apache.pinot.spi.utils.BigDecimalUtils;
3742
import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
43+
import org.apache.pinot.spi.utils.UuidUtils;
3844
import org.roaringbitmap.RoaringBitmap;
3945
import org.testng.annotations.Test;
4046

@@ -1177,6 +1183,28 @@ public void testRandWithSeedTransformFunction() {
11771183
}
11781184
}
11791185

1186+
@Test
1187+
public void testUuidColumnArgumentTransformFunction()
1188+
throws NoSuchMethodException {
1189+
byte[][] uuidValues = new byte[NUM_ROWS][];
1190+
String[] expectedValues = new String[NUM_ROWS];
1191+
for (int i = 0; i < NUM_ROWS; i++) {
1192+
UUID uuid = new UUID(i, i + 1L);
1193+
uuidValues[i] = UuidUtils.toBytes(uuid);
1194+
expectedValues[i] = uuid.toString();
1195+
}
1196+
1197+
ScalarTransformFunctionWrapper transformFunction =
1198+
new ScalarTransformFunctionWrapper(FunctionInfo.fromMethod(
1199+
ScalarTransformFunctionWrapperTest.class.getDeclaredMethod("uuidToString", UUID.class)));
1200+
StaticUuidTransformFunction uuidTransformFunction = new StaticUuidTransformFunction(uuidValues);
1201+
uuidTransformFunction.init(Collections.emptyList(), Collections.emptyMap());
1202+
transformFunction.init(Arrays.asList(uuidTransformFunction), Collections.emptyMap());
1203+
1204+
assertEquals(transformFunction.getName(), "uuidToString");
1205+
testTransformFunction(transformFunction, expectedValues);
1206+
}
1207+
11801208
@Test
11811209
public void testStringLowerTransformFunctionNullLiteral() {
11821210
ExpressionContext expression =
@@ -1208,4 +1236,33 @@ public void testStringLowerTransformFunctionNullColumn() {
12081236
}
12091237
testTransformFunctionWithNull(transformFunction, expectedValues, bitmap);
12101238
}
1239+
1240+
public static String uuidToString(UUID uuid) {
1241+
return uuid.toString();
1242+
}
1243+
1244+
private static class StaticUuidTransformFunction extends BaseTransformFunction {
1245+
private static final TransformResultMetadata RESULT_METADATA = new TransformResultMetadata(DataType.UUID, true,
1246+
false);
1247+
private final byte[][] _uuidValues;
1248+
1249+
private StaticUuidTransformFunction(byte[][] uuidValues) {
1250+
_uuidValues = uuidValues;
1251+
}
1252+
1253+
@Override
1254+
public String getName() {
1255+
return "staticUuid";
1256+
}
1257+
1258+
@Override
1259+
public TransformResultMetadata getResultMetadata() {
1260+
return RESULT_METADATA;
1261+
}
1262+
1263+
@Override
1264+
public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) {
1265+
return _uuidValues;
1266+
}
1267+
}
12111268
}

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -179,13 +179,11 @@ public static org.apache.avro.Schema getAvroSchemaFromPinotSchema(Schema pinotSc
179179

180180
for (FieldSpec fieldSpec : pinotSchema.getAllFieldSpecs()) {
181181
if (fieldSpec.getDataType() == DataType.UUID) {
182+
Preconditions.checkState(fieldSpec.isSingleValueField(),
183+
"UUID data type only supports single-value fields: %s", fieldSpec.getName());
182184
org.apache.avro.Schema uuidSchema = LogicalTypes.uuid().addToSchema(org.apache.avro.Schema.create(
183185
org.apache.avro.Schema.Type.STRING));
184-
if (fieldSpec.isSingleValueField()) {
185-
fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type(uuidSchema).noDefault();
186-
} else {
187-
fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type().array().items(uuidSchema).noDefault();
188-
}
186+
fieldAssembler = fieldAssembler.name(fieldSpec.getName()).type(uuidSchema).noDefault();
189187
continue;
190188
}
191189
DataType storedType = fieldSpec.getDataType().getStoredType();
@@ -352,8 +350,11 @@ private static void extractSchemaWithComplexTypeHandling(org.apache.avro.Schema
352350
timeUnit, collectionNotUnnestedToJson);
353351
} else if (collectionNotUnnestedToJson == ComplexTypeConfig.CollectionNotUnnestedToJson.NON_PRIMITIVE
354352
&& AvroSchemaUtil.isPrimitiveType(elementType.getType())) {
355-
addFieldToPinotSchema(pinotSchema, AvroSchemaUtil.valueOf(elementType), path, false, fieldTypeMap,
356-
timeUnit);
353+
DataType elementDataType = AvroSchemaUtil.valueOf(elementType);
354+
Preconditions.checkState(elementDataType != DataType.UUID,
355+
"Avro ARRAY<uuid> cannot be mapped to Pinot UUID because UUID data type only supports single-value "
356+
+ "fields: %s", path);
357+
addFieldToPinotSchema(pinotSchema, elementDataType, path, false, fieldTypeMap, timeUnit);
357358
} else if (shallConvertToJson(collectionNotUnnestedToJson, elementType)) {
358359
addFieldToPinotSchema(pinotSchema, DataType.STRING, path, true, fieldTypeMap, timeUnit);
359360
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,17 @@
2424
import java.util.concurrent.TimeUnit;
2525
import org.apache.avro.LogicalTypes;
2626
import org.apache.pinot.spi.config.table.ingestion.ComplexTypeConfig;
27+
import org.apache.pinot.spi.data.DimensionFieldSpec;
2728
import org.apache.pinot.spi.data.FieldSpec;
2829
import org.apache.pinot.spi.data.FieldSpec.DataType;
2930
import org.apache.pinot.spi.data.FieldSpec.FieldType;
3031
import org.apache.pinot.spi.data.Schema;
32+
import org.testng.Assert;
3133
import org.testng.annotations.Test;
3234
import org.testng.collections.Lists;
3335

3436
import static org.testng.Assert.assertEquals;
37+
import static org.testng.Assert.assertTrue;
3538

3639

3740
public class AvroUtilsTest {
@@ -159,4 +162,27 @@ public void testGetAvroSchemaFromPinotSchemaWithUuidLogicalType() {
159162
assertEquals(fieldSchema.getType(), org.apache.avro.Schema.Type.STRING);
160163
assertEquals(fieldSchema.getLogicalType().getName(), "uuid");
161164
}
165+
166+
@Test
167+
public void testGetPinotSchemaFromAvroSchemaRejectsUuidArray() {
168+
org.apache.avro.Schema uuidSchema =
169+
LogicalTypes.uuid().addToSchema(org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING));
170+
org.apache.avro.Schema avroSchema = org.apache.avro.SchemaBuilder.record("uuidArrayRecord").fields()
171+
.name("uuidArray").type().array().items(uuidSchema).noDefault().endRecord();
172+
173+
IllegalStateException exception = Assert.expectThrows(IllegalStateException.class,
174+
() -> AvroUtils.getPinotSchemaFromAvroSchemaWithComplexTypeHandling(avroSchema, null, null,
175+
new ArrayList<>(), ".", ComplexTypeConfig.CollectionNotUnnestedToJson.NON_PRIMITIVE));
176+
assertTrue(exception.getMessage().contains("ARRAY<uuid>"));
177+
}
178+
179+
@Test
180+
public void testGetAvroSchemaFromPinotSchemaRejectsMvUuid() {
181+
Schema pinotSchema = new Schema();
182+
pinotSchema.addField(new DimensionFieldSpec("uuidMv", DataType.UUID, false));
183+
184+
IllegalStateException exception =
185+
Assert.expectThrows(IllegalStateException.class, () -> AvroUtils.getAvroSchemaFromPinotSchema(pinotSchema));
186+
assertTrue(exception.getMessage().contains("single-value"));
187+
}
162188
}

‎pinot-spi/src/main/java/org/apache/pinot/spi/utils/UuidUtils.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ public static byte[] toBytes(ByteArray uuidBytes) {
6868
}
6969

7070
public static byte[] toBytes(Object value) {
71+
validateNotNull(value, "UUID bytes");
7172
if (value instanceof UUID) {
7273
return toBytes((UUID) value);
7374
}
@@ -99,6 +100,7 @@ public static UUID toUUID(String uuidString) {
99100
}
100101

101102
public static UUID toUUID(Object value) {
103+
validateNotNull(value, "UUID");
102104
if (value instanceof UUID) {
103105
return (UUID) value;
104106
}
@@ -133,4 +135,10 @@ private static void validateLength(byte[] uuidBytes) {
133135
"Invalid UUID byte length: " + uuidBytes.length + ", expected: " + UUID_NUM_BYTES);
134136
}
135137
}
138+
139+
private static void validateNotNull(Object value, String targetType) {
140+
if (value == null) {
141+
throw new IllegalArgumentException("Cannot convert null value to " + targetType);
142+
}
143+
}
136144
}

‎pinot-spi/src/test/java/org/apache/pinot/spi/utils/UuidUtilsTest.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.pinot.spi.utils;
20+
21+
import org.testng.Assert;
22+
import org.testng.annotations.Test;
23+
24+
import static org.testng.Assert.assertEquals;
25+
26+
27+
/**
28+
* Tests for {@link UuidUtils}.
29+
*/
30+
public class UuidUtilsTest {
31+
32+
@Test
33+
public void testToBytesRejectsNull() {
34+
IllegalArgumentException exception =
35+
Assert.expectThrows(IllegalArgumentException.class, () -> UuidUtils.toBytes((Object) null));
36+
assertEquals(exception.getMessage(), "Cannot convert null value to UUID bytes");
37+
}
38+
39+
@Test
40+
public void testToUuidRejectsNull() {
41+
IllegalArgumentException exception =
42+
Assert.expectThrows(IllegalArgumentException.class, () -> UuidUtils.toUUID((Object) null));
43+
assertEquals(exception.getMessage(), "Cannot convert null value to UUID");
44+
}
45+
}

0 commit comments

Comments
 (0)