diff --git a/.gitignore b/.gitignore index bbd3d34656..3d0535d0aa 100644 --- a/.gitignore +++ b/.gitignore @@ -58,3 +58,4 @@ kubernetes/helm/**/Chart.lock #Develocity .mvn/.gradle-enterprise/ .mvn/.develocity/ +.claude/ diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java index 2f8a2e9054..f022c8b7eb 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java @@ -295,6 +295,7 @@ public boolean equals(DataBlock left, DataBlock right) { break; case STRING: case JSON: + case URN: for (int did = 0; did < numRows; did++) { if (!left.getString(did, colId).equals(right.getString(did, colId))) { if (_failOnFalse) { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionUtils.java index c1445e98bd..81fcdbf105 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionUtils.java @@ -193,6 +193,7 @@ public static RelDataType getRelDataType(RelDataTypeFactory typeFactory, ClassAll functions accept a URN string and return a component of it. They are null-safe: a + * {@code null} input always produces a {@code null} output. An unparseable (malformed) URN + * produces a {@code null} rather than throwing so that queries over mixed-quality data do not fail + * entirely. + * + *

Example usage in SQL: + * 

+ *   SELECT urnEntity(memberId), COUNT(*) FROM myTable GROUP BY 1
+ *   SELECT urnKey(memberId) FROM myTable WHERE urnEntity(memberId) = 'urn:li:memberId'
+ *   SELECT * FROM myTable WHERE urnEntityType(memberId) = 'memberId'
+ *   SELECT urnIsComposite(datasetId), urnKeyPart(datasetId, 1) FROM myTable
+ *
+ */ +public class UrnFunctions { + private UrnFunctions() { + } + + /** + * Returns the fully-qualified entity type prefix of a URN, e.g. {@code "urn:li:memberId"} for + * {@code "urn:li:memberId:123456"}. + * + *

This is the canonical string to use when grouping or filtering on entity type — it includes + * the full namespace, so {@code "urn:li:memberId"} is unambiguous across different namespaces. + */ + @Nullable + @ScalarFunction(names = {"urnEntity", "urn_entity"}, nullableParameters = true) + public static String urnEntity(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + return urn != null ? urn.entityPrefix() : null; + } + + /** + * Returns the entity namespace component, e.g. {@code "li"} for {@code "urn:li:memberId:123"}. + */ + @Nullable + @ScalarFunction(names = {"urnNamespace", "urn_namespace"}, nullableParameters = true) + public static String urnNamespace(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + return urn != null ? urn.namespace() : null; + } + + /** + * Returns the entity type component (without namespace), e.g. {@code "memberId"} for + * {@code "urn:li:memberId:123456"}. + * + *

Useful for grouping by entity type across namespaces or when the namespace is constant. + */ + @Nullable + @ScalarFunction(names = {"urnEntityType", "urn_entity_type"}, nullableParameters = true) + public static String urnEntityType(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + return urn != null ? urn.entityType() : null; + } + + /** + * Returns the key component of a simple (non-composite) URN as a string, e.g. {@code "123456"} + * for {@code "urn:li:memberId:123456"}. + * + *

For composite URNs the full tuple string including parentheses is returned, e.g. + * {@code "(urn:li:dataPlatform:pinot,memberEducation,EI)"} for a dataset URN. + */ + @Nullable + @ScalarFunction(names = {"urnKey", "urn_key"}, nullableParameters = true) + public static String urnKey(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + if (urn == null) { + return null; + } + if (urn.isSimple()) { + return urn.simpleKey(); + } + // Composite: reconstruct the key tuple string + StringBuilder sb = new StringBuilder("("); + for (int i = 0; i < urn.parts().size(); i++) { + if (i > 0) { + sb.append(','); + } + sb.append(urn.parts().get(i).toString()); + } + sb.append(')'); + return sb.toString(); + } + + /** + * Returns the key of a simple URN as a long, or {@code Long.MIN_VALUE} if the key is not + * numeric. Returns {@code null} on null input or unparseable URN. + * + *

This is the most efficient function for numeric-key URN columns — it allows the query + * engine to operate directly on long values without repeated string-to-long conversion. + */ + @Nullable + @ScalarFunction(names = {"urnKeyLong", "urn_key_long"}, nullableParameters = true) + public static Long urnKeyLong(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + if (urn == null || !urn.isSimple()) { + return null; + } + try { + return Long.parseLong(urn.simpleKey()); + } catch (NumberFormatException e) { + return null; + } + } + + /** + * Returns {@code true} if the URN is composite (its key is a tuple), {@code false} if it is + * simple. Returns {@code null} on null input or unparseable URN. + */ + @Nullable + @ScalarFunction(names = {"urnIsComposite", "urn_is_composite"}, nullableParameters = true) + public static Boolean urnIsComposite(@Nullable String urnStr) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + return urn != null ? !urn.isSimple() : null; + } + + /** + * Returns the {@code index}-th part of a composite URN's tuple key as a string (0-based). + * Returns {@code null} for simple URNs, out-of-range indices, null input, or unparseable URN. + * + *

Example: {@code urnKeyPart('urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)', 1)} + * returns {@code "foo"}. + */ + @Nullable + @ScalarFunction(names = {"urnKeyPart", "urn_key_part"}, nullableParameters = true) + public static String urnKeyPart(@Nullable String urnStr, int index) { + if (urnStr == null) { + return null; + } + Urn urn = Urn.tryParse(urnStr); + if (urn == null || urn.isSimple() || index < 0 || index >= urn.parts().size()) { + return null; + } + return urn.parts().get(index).toString(); + } + + /** + * Returns {@code true} if the string is a valid, parseable URN; {@code false} otherwise. + * Never returns {@code null}; a null input returns {@code false}. + */ + @ScalarFunction(names = {"isUrn", "is_urn"}, nullableParameters = true) + public static boolean isUrn(@Nullable String value) { + return Urn.isUrn(value); + } +} diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java index a17eeb7053..09ac5c6fa1 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java @@ -108,6 +108,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch case BYTES: case BIG_DECIMAL: case JSON: + case URN: case OBJECT: field = new Field(colName, FieldType.nullable(new ArrowType.Utf8()), null); vector = new VarCharVector(colName, ALLOCATOR); @@ -234,6 +235,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch case BYTES: case BIG_DECIMAL: case JSON: + case URN: case OBJECT: byte[] bytes = ((String) value).getBytes(StandardCharsets.UTF_8); ((VarCharVector) vector).setSafe(rowIndex, bytes); @@ -410,6 +412,7 @@ public ResultTable decodeResultTable(byte[] bytes, int rowSize, DataSchema schem case BYTES: case BIG_DECIMAL: case JSON: + case URN: case OBJECT: row[col] = new String(((VarCharVector) vector).get(i), StandardCharsets.UTF_8); break; diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java index 8cac71ea40..70e0d689a4 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java @@ -223,6 +223,7 @@ private static Object extractValue(DataSchema.ColumnDataType columnDataType, Jso case BYTES: case TIMESTAMP: case JSON: + case URN: case BIG_DECIMAL: case OBJECT: return jsonValue.textValue(); diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java index 405ab99378..b75abd7921 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java @@ -290,6 +290,12 @@ public RelDataType toType(RelDataTypeFactory typeFactory) { return typeFactory.createSqlType(SqlTypeName.VARCHAR); } }, + URN(STRING, NullValuePlaceHolder.STRING) { + @Override + public RelDataType toType(RelDataTypeFactory typeFactory) { + return typeFactory.createSqlType(SqlTypeName.VARCHAR); + } + }, MAP(NullValuePlaceHolder.MAP) { @Override public RelDataType toType(RelDataTypeFactory typeFactory) { @@ -452,6 +458,8 @@ public DataType toDataType() { return DataType.STRING; case JSON: return DataType.JSON; + case URN: + return DataType.URN; case BYTES: case BYTES_ARRAY: return DataType.BYTES; @@ -571,6 +579,7 @@ public Serializable convert(Object value) { return new Timestamp((long) value); case STRING: case JSON: + case URN: return value.toString(); case BYTES: return ((ByteArray) value).getBytes(); @@ -641,6 +650,7 @@ public Serializable convertAndFormat(Object value) { return new Timestamp((long) value).toString(); case STRING: case JSON: + case URN: return value.toString(); case BYTES: return ((ByteArray) value).toHexString(); @@ -844,6 +854,8 @@ public static ColumnDataType fromDataTypeSV(DataType dataType) { return STRING; case JSON: return JSON; + case URN: + return URN; case BYTES: return BYTES; case MAP: diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotDataType.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotDataType.java index 030c5a9df9..d3d915483d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotDataType.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotDataType.java @@ -1493,6 +1493,11 @@ public static PinotDataType getPinotDataTypeForIngestion(FieldSpec fieldSpec) { return JSON; } throw new IllegalStateException("There is no multi-value type for JSON"); + case URN: + if (fieldSpec.isSingleValueField()) { + return STRING; + } + throw new IllegalStateException("There is no multi-value type for URN"); case STRING: return fieldSpec.isSingleValueField() ? STRING : STRING_ARRAY; case BYTES: @@ -1532,6 +1537,8 @@ public static PinotDataType getPinotDataTypeForExecution(ColumnDataType columnDa return STRING; case JSON: return JSON; + case URN: + return STRING; case BYTES: return BYTES; case OBJECT: diff --git a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/UrnFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/UrnFunctionsTest.java new file mode 100644 index 0000000000..7b002f320b --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/UrnFunctionsTest.java @@ -0,0 +1,202 @@ +/** + * 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.common.function.scalar; + +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +public class UrnFunctionsTest { + + // --- urnEntity --- + + @Test + public void urnEntitySimple() { + assertEquals(UrnFunctions.urnEntity("urn:li:memberId:123456"), "urn:li:memberId"); + } + + @Test + public void urnEntityComposite() { + assertEquals(UrnFunctions.urnEntity("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)"), + "urn:li:dataset"); + } + + @Test + public void urnEntityNullInput() { + assertNull(UrnFunctions.urnEntity(null)); + } + + @Test + public void urnEntityMalformed() { + assertNull(UrnFunctions.urnEntity("not-a-urn")); + } + + // --- urnNamespace --- + + @Test + public void urnNamespaceReturnsLi() { + assertEquals(UrnFunctions.urnNamespace("urn:li:memberId:1"), "li"); + } + + @Test + public void urnNamespaceNull() { + assertNull(UrnFunctions.urnNamespace(null)); + } + + // --- urnEntityType --- + + @Test + public void urnEntityTypeSimple() { + assertEquals(UrnFunctions.urnEntityType("urn:li:memberId:123"), "memberId"); + } + + @Test + public void urnEntityTypeMalformed() { + assertNull(UrnFunctions.urnEntityType("bad-urn")); + } + + // --- urnKey --- + + @Test + public void urnKeySimple() { + assertEquals(UrnFunctions.urnKey("urn:li:memberId:123456"), "123456"); + } + + @Test + public void urnKeyComposite() { + String result = UrnFunctions.urnKey("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)"); + assertEquals(result, "(urn:li:dataPlatform:pinot,foo,PROD)"); + } + + @Test + public void urnKeyNull() { + assertNull(UrnFunctions.urnKey(null)); + } + + @Test + public void urnKeyMalformed() { + assertNull(UrnFunctions.urnKey("bad")); + } + + // --- urnKeyLong --- + + @Test + public void urnKeyLongNumericKey() { + assertEquals(UrnFunctions.urnKeyLong("urn:li:memberId:123456"), Long.valueOf(123456L)); + } + + @Test + public void urnKeyLongZero() { + assertEquals(UrnFunctions.urnKeyLong("urn:li:memberId:0"), Long.valueOf(0L)); + } + + @Test + public void urnKeyLongNonNumericKey() { + assertNull(UrnFunctions.urnKeyLong("urn:li:corpUser:alice")); + } + + @Test + public void urnKeyLongCompositeReturnsNull() { + assertNull(UrnFunctions.urnKeyLong("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)")); + } + + @Test + public void urnKeyLongNull() { + assertNull(UrnFunctions.urnKeyLong(null)); + } + + @Test + public void urnKeyLongMaxValue() { + assertEquals(UrnFunctions.urnKeyLong("urn:li:memberId:" + Long.MAX_VALUE), Long.valueOf(Long.MAX_VALUE)); + } + + // --- urnIsComposite --- + + @Test + public void urnIsCompositeSimpleUrn() { + assertFalse(UrnFunctions.urnIsComposite("urn:li:memberId:1")); + } + + @Test + public void urnIsCompositeCompositeUrn() { + assertTrue(UrnFunctions.urnIsComposite("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)")); + } + + @Test + public void urnIsCompositeNull() { + assertNull(UrnFunctions.urnIsComposite(null)); + } + + // --- urnKeyPart --- + + @Test + public void urnKeyPartFirstPart() { + assertEquals( + UrnFunctions.urnKeyPart("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)", 0), + "urn:li:dataPlatform:pinot"); + } + + @Test + public void urnKeyPartMiddlePart() { + assertEquals( + UrnFunctions.urnKeyPart("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)", 1), + "foo"); + } + + @Test + public void urnKeyPartLastPart() { + assertEquals( + UrnFunctions.urnKeyPart("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)", 2), + "PROD"); + } + + @Test + public void urnKeyPartOutOfRange() { + assertNull(UrnFunctions.urnKeyPart("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)", 5)); + } + + @Test + public void urnKeyPartSimpleUrnReturnsNull() { + assertNull(UrnFunctions.urnKeyPart("urn:li:memberId:123", 0)); + } + + @Test + public void urnKeyPartNullInput() { + assertNull(UrnFunctions.urnKeyPart(null, 0)); + } + + // --- isUrn --- + + @Test + public void isUrnValid() { + assertTrue(UrnFunctions.isUrn("urn:li:memberId:1")); + assertTrue(UrnFunctions.isUrn("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)")); + } + + @Test + public void isUrnInvalid() { + assertFalse(UrnFunctions.isUrn("not-a-urn")); + assertFalse(UrnFunctions.isUrn(null)); + assertFalse(UrnFunctions.isUrn("")); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/EqualsPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/EqualsPredicateEvaluatorFactory.java index d99b8249d4..b9f6470095 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/EqualsPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/EqualsPredicateEvaluatorFactory.java @@ -81,6 +81,7 @@ public static EqRawPredicateEvaluator newRawValueBasedEvaluator(EqPredicate eqPr return new LongRawValueBasedEqPredicateEvaluator(eqPredicate, TimestampUtils.toMillisSinceEpoch(value)); case STRING: case JSON: + case URN: return new StringRawValueBasedEqPredicateEvaluator(eqPredicate, value); case BYTES: return new BytesRawValueBasedEqPredicateEvaluator(eqPredicate, BytesUtils.toBytes(value)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/InPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/InPredicateEvaluatorFactory.java index ed7fb5273f..1bc51b2c8f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/InPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/InPredicateEvaluatorFactory.java @@ -129,7 +129,8 @@ public static InRawPredicateEvaluator newRawValueBasedEvaluator(InPredicate inPr return new LongRawValueBasedInPredicateEvaluator(inPredicate, matchingValues); } case STRING: - case JSON: { + case JSON: + case URN: { List stringValues = inPredicate.getValues(); Set matchingValues = new ObjectOpenHashSet<>(HashUtil.getMinHashSetSize(stringValues.size())); // NOTE: Add value-by-value to avoid overhead diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotEqualsPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotEqualsPredicateEvaluatorFactory.java index 742b69c1b1..c70e948d68 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotEqualsPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotEqualsPredicateEvaluatorFactory.java @@ -77,6 +77,7 @@ public static NeqRawPredicateEvaluator newRawValueBasedEvaluator(NotEqPredicate return new LongRawValueBasedNeqPredicateEvaluator(notEqPredicate, TimestampUtils.toMillisSinceEpoch(value)); case STRING: case JSON: + case URN: return new StringRawValueBasedNeqPredicateEvaluator(notEqPredicate, value); case BYTES: return new BytesRawValueBasedNeqPredicateEvaluator(notEqPredicate, BytesUtils.toBytes(value)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotInPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotInPredicateEvaluatorFactory.java index daabde1069..1bc216a646 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotInPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/NotInPredicateEvaluatorFactory.java @@ -129,7 +129,8 @@ public static NotInRawPredicateEvaluator newRawValueBasedEvaluator(NotInPredicat return new LongRawValueBasedNotInPredicateEvaluator(notInPredicate, nonMatchingValues); } case STRING: - case JSON: { + case JSON: + case URN: { List stringValues = notInPredicate.getValues(); Set nonMatchingValues = new ObjectOpenHashSet<>(HashUtil.getMinHashSetSize(stringValues.size())); // NOTE: Add value-by-value to avoid overhead diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateUtils.java index c7b93cf086..41ccc0c397 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateUtils.java @@ -144,6 +144,7 @@ public static IntSet getDictIdSet(BaseInPredicate inPredicate, Dictionary dictio } } break; + case URN: case STRING: if (queryContext == null || values.size() <= 1) { dictionary.getDictIds(values, dictIdSet); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RangePredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RangePredicateEvaluatorFactory.java index e9bd3b4b0a..ae9e1e9b89 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RangePredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RangePredicateEvaluatorFactory.java @@ -104,6 +104,7 @@ public static BaseRawValueBasedPredicateEvaluator newRawValueBasedEvaluator(Rang lowerUnbounded ? Long.MIN_VALUE : TimestampUtils.toMillisSinceEpoch(lowerBound), upperUnbounded ? Long.MAX_VALUE : TimestampUtils.toMillisSinceEpoch(upperBound), lowerInclusive, upperInclusive); + case URN: case STRING: return new StringRawValueBasedRangePredicateEvaluator(rangePredicate, lowerUnbounded ? null : lowerBound, upperUnbounded ? null : upperBound, lowerInclusive, upperInclusive); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java index e9c2bb73fd..6e49b27f26 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java @@ -42,7 +42,7 @@ private RegexpLikePredicateEvaluatorFactory() { */ public static BaseDictionaryBasedPredicateEvaluator newDictionaryBasedEvaluator( RegexpLikePredicate regexpLikePredicate, Dictionary dictionary, DataType dataType) { - boolean condition = (dataType == DataType.STRING || dataType == DataType.JSON); + boolean condition = (dataType == DataType.STRING || dataType == DataType.JSON || dataType == DataType.URN); Preconditions.checkArgument(condition, "Unsupported data type: " + dataType); return new DictionaryBasedRegexpLikePredicateEvaluator(regexpLikePredicate, dictionary); } @@ -56,7 +56,8 @@ public static BaseDictionaryBasedPredicateEvaluator newDictionaryBasedEvaluator( */ public static BaseRawValueBasedPredicateEvaluator newRawValueBasedEvaluator(RegexpLikePredicate regexpLikePredicate, DataType dataType) { - Preconditions.checkArgument(dataType == DataType.STRING, "Unsupported data type: " + dataType); + Preconditions.checkArgument(dataType == DataType.STRING || dataType == DataType.URN, + "Unsupported data type: " + dataType); return new RawValueBasedRegexpLikePredicateEvaluator(regexpLikePredicate); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunction.java index 019040a273..f6f58ffa93 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CaseTransformFunction.java @@ -208,6 +208,7 @@ private void checkLiteral(DataType dataType, String literal) { break; case STRING: case JSON: + case URN: break; case BYTES: try { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java index 86c2e0a6e4..3120951a16 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java @@ -269,6 +269,11 @@ private static Map> createRegistry() registry.put(canonicalize(name), entry.getValue()); } } + // URN-aware optimized implementations — override the scalar function fallback. + registry.put(canonicalize(UrnEntityTransformFunction.FUNCTION_NAME), UrnEntityTransformFunction.class); + registry.put(canonicalize(UrnEntityTransformFunction.FUNCTION_NAME_ALT), UrnEntityTransformFunction.class); + registry.put(canonicalize(UrnKeyLongTransformFunction.FUNCTION_NAME), UrnKeyLongTransformFunction.class); + registry.put(canonicalize(UrnKeyLongTransformFunction.FUNCTION_NAME_ALT), UrnKeyLongTransformFunction.class); return registry; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnEntityTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnEntityTransformFunction.java new file mode 100644 index 0000000000..9e5d65571e --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnEntityTransformFunction.java @@ -0,0 +1,181 @@ +/** + * 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.operator.transform.function; + +import com.google.common.base.Preconditions; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.DimensionFieldSpec.EntityDictRange; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.utils.Urn; + + +/** + * {@code urnEntity(col)} transform function — returns the entity-type prefix of a URN column. + * + *

Performance tiers, picked at {@link #init} time. None of them parse URN strings per row: + *

    + *
  1. Constant — single-entity URN column (declared or auto-detected at segment build): + * every row gets the same literal string, so the block is filled with one + * {@link Arrays#fill} call.
    2. + *
  2. Range-lookup — multi-entity URN column whose segment metadata records the + * per-entity dict-ID ranges. Reads main-column dict IDs and finds the owning entity range + * with a small linear scan (typically ≤ 16 entries).
    3. + *
  3. Dict-lookup — multi-entity URN column stored as STRING (mixed-type entities, e.g. + * LONG memberId + STRING corpuser). At init time we walk the dictionary once and build a + * {@code dictId → entityPrefix} array; per row we read the dict ID and index into the + * array. {@link org.apache.pinot.segment.local.recordtransformer.UrnValidationTransformer} + * guarantees every value in the dictionary is a simple URN, so the init walk never sees + * malformed data.
    4. + *
+ */ +public class UrnEntityTransformFunction extends BaseTransformFunction { + + public static final String FUNCTION_NAME = "urnEntity"; + public static final String FUNCTION_NAME_ALT = "urn_entity"; + + private enum Strategy { + CONSTANT, RANGE_LOOKUP, DICT_LOOKUP + } + + private Strategy _strategy; + private String _constantEntityType; + private String _mainColumnName; + private int[] _rangeEnds; + private String[] _entityTypesByRange; + private String[] _entityTypeByDictId; + + @Override + public String getName() { + return FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap) { + super.init(arguments, columnContextMap); + Preconditions.checkArgument(arguments.size() == 1, "urnEntity requires exactly 1 argument"); + + Preconditions.checkArgument(arguments.get(0) instanceof IdentifierTransformFunction, + "urnEntity argument must be a column identifier"); + String colName = ((IdentifierTransformFunction) arguments.get(0)).getColumnName(); + ColumnContext colCtx = columnContextMap.get(colName); + Preconditions.checkArgument(colCtx != null && colCtx.getDataType() == FieldSpec.DataType.URN, + "urnEntity argument must reference a URN column: %s", colName); + + DataSource dataSource = colCtx.getDataSource(); + Preconditions.checkNotNull(dataSource, "urnEntity requires a backing data source for column: %s", colName); + FieldSpec fieldSpec = dataSource.getDataSourceMetadata().getFieldSpec(); + Preconditions.checkArgument(fieldSpec instanceof DimensionFieldSpec, + "urnEntity expects a dimension URN column: %s", colName); + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + + // Single-entity: same prefix on every row. + if (dim.isSingleUrnEntityType()) { + _strategy = Strategy.CONSTANT; + _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); + return; + } + + // Multi-entity sorted-LONG layout: routing via per-entity dict-ID ranges. + List ranges = dim.getEntityDictRanges(); + if (ranges != null && !ranges.isEmpty()) { + _mainColumnName = colName; + _rangeEnds = new int[ranges.size()]; + _entityTypesByRange = new String[ranges.size()]; + for (int i = 0; i < ranges.size(); i++) { + _rangeEnds[i] = ranges.get(i).getEndDictIdExclusive(); + _entityTypesByRange[i] = ranges.get(i).getEntityType(); + } + _strategy = Strategy.RANGE_LOOKUP; + return; + } + + // Multi-entity STRING-stored column: precompute entity-per-dict-id once. + Dictionary dictionary = dataSource.getDictionary(); + Preconditions.checkNotNull(dictionary, + "urnEntity requires a dictionary for multi-entity STRING-stored URN column: %s", colName); + int dictLength = dictionary.length(); + _entityTypeByDictId = new String[dictLength]; + for (int dictId = 0; dictId < dictLength; dictId++) { + String value = dictionary.getStringValue(dictId); + if (value == null) { + continue; + } + Urn urn = Urn.tryParse(value); + // UrnValidationTransformer guarantees parseability for ingested data; a parse failure + // here would indicate a corrupted segment or a non-URN value smuggled in by an upstream + // bypass — surface the column name to help diagnose. + Preconditions.checkState(urn != null, + "Unparseable URN value at dictId %s in column %s: %s", dictId, colName, value); + _entityTypeByDictId[dictId] = urn.entityPrefix(); + } + _mainColumnName = colName; + _strategy = Strategy.DICT_LOOKUP; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return STRING_SV_NO_DICTIONARY_METADATA; + } + + @Override + public String[] transformToStringValuesSV(ValueBlock valueBlock) { + int length = valueBlock.getNumDocs(); + initStringValuesSV(length); + switch (_strategy) { + case CONSTANT: + Arrays.fill(_stringValuesSV, 0, length, _constantEntityType); + break; + case RANGE_LOOKUP: + BlockValSet mainBlock = valueBlock.getBlockValueSet(_mainColumnName); + int[] dictIds = mainBlock.getDictionaryIdsSV(); + int numRanges = _rangeEnds.length; + for (int i = 0; i < length; i++) { + int dictId = dictIds[i]; + String entityType = null; + for (int r = 0; r < numRanges; r++) { + if (dictId < _rangeEnds[r]) { + entityType = _entityTypesByRange[r]; + break; + } + } + _stringValuesSV[i] = entityType; + } + break; + case DICT_LOOKUP: + default: + BlockValSet dictBlock = valueBlock.getBlockValueSet(_mainColumnName); + int[] dictLookupIds = dictBlock.getDictionaryIdsSV(); + for (int i = 0; i < length; i++) { + int id = dictLookupIds[i]; + _stringValuesSV[i] = id >= 0 ? _entityTypeByDictId[id] : null; + } + break; + } + return _stringValuesSV; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnKeyLongTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnKeyLongTransformFunction.java new file mode 100644 index 0000000000..06b237bbbe --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnKeyLongTransformFunction.java @@ -0,0 +1,116 @@ +/** + * 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.operator.transform.function; + +import com.google.common.base.Preconditions; +import java.util.List; +import java.util.Map; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.utils.Urn; + + +/** + * {@code urnKeyLong(col)} transform function — returns the numeric key of a URN as a LONG. + * + *

Performance tiers, evaluated at {@link #init} time: + *

    + *
  1. Direct LONG read — typed URN column stored as LONG: reads the raw LONG key + * directly from the column without any URN string construction or parsing. This is the + * fast path for queries like {@code WHERE urnKeyLong(memberId) > 1000}.
    2. + *
  2. Parse-string — untyped URN column stored as STRING: reads the string value and + * parses it to extract the numeric key; returns {@link Long#MIN_VALUE} for non-numeric + * keys (matches null-value sentinel).
    3. + *
+ */ +public class UrnKeyLongTransformFunction extends BaseTransformFunction { + + public static final String FUNCTION_NAME = "urnKeyLong"; + public static final String FUNCTION_NAME_ALT = "urn_key_long"; + + private boolean _isTypedLongUrn; + + @Override + public String getName() { + return FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap) { + super.init(arguments, columnContextMap); + Preconditions.checkArgument(arguments.size() == 1, "urnKeyLong requires exactly 1 argument"); + + TransformFunction arg = arguments.get(0); + if (arg instanceof IdentifierTransformFunction) { + String colName = ((IdentifierTransformFunction) arg).getColumnName(); + ColumnContext colCtx = columnContextMap.get(colName); + if (colCtx != null && colCtx.getDataType() == FieldSpec.DataType.URN) { + DataSource dataSource = colCtx.getDataSource(); + if (dataSource != null) { + FieldSpec fieldSpec = dataSource.getDataSourceMetadata().getFieldSpec(); + if (fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.hasTypedUrnStorage()) { + _isTypedLongUrn = true; + } + } + } + } + } + } + + @Override + public TransformResultMetadata getResultMetadata() { + return LONG_SV_NO_DICTIONARY_METADATA; + } + + @Override + public long[] transformToLongValuesSV(ValueBlock valueBlock) { + int length = valueBlock.getNumDocs(); + if (_isTypedLongUrn) { + // Fast path: the column physically stores LONG keys — read them directly. + return _arguments.get(0).transformToLongValuesSV(valueBlock); + } + // Slow path: parse the URN string and extract the numeric key. + initLongValuesSV(length); + String[] strValues = _arguments.get(0).transformToStringValuesSV(valueBlock); + for (int i = 0; i < length; i++) { + String urnStr = strValues[i]; + if (urnStr == null) { + _longValuesSV[i] = Long.MIN_VALUE; + continue; + } + Urn urn = Urn.tryParse(urnStr); + if (urn == null || !urn.isSimple()) { + _longValuesSV[i] = Long.MIN_VALUE; + continue; + } + try { + _longValuesSV[i] = Long.parseLong(urn.simpleKey()); + } catch (NumberFormatException e) { + _longValuesSV[i] = Long.MIN_VALUE; + } + } + return _longValuesSV; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java index 805c70a1c2..7123239ae2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java @@ -28,8 +28,13 @@ import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FilterContext; import org.apache.pinot.common.request.context.FunctionContext; +import org.apache.pinot.common.request.context.predicate.EqPredicate; +import org.apache.pinot.common.request.context.predicate.InPredicate; import org.apache.pinot.common.request.context.predicate.JsonMatchPredicate; +import org.apache.pinot.common.request.context.predicate.NotEqPredicate; +import org.apache.pinot.common.request.context.predicate.NotInPredicate; import org.apache.pinot.common.request.context.predicate.Predicate; +import org.apache.pinot.common.request.context.predicate.RangePredicate; import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate; import org.apache.pinot.common.request.context.predicate.TextContainsPredicate; import org.apache.pinot.common.request.context.predicate.TextMatchPredicate; @@ -52,6 +57,7 @@ import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator; import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluatorProvider; import org.apache.pinot.core.operator.transform.function.ItemTransformFunction; +import org.apache.pinot.core.operator.transform.function.UrnKeyLongTransformFunction; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.local.realtime.impl.invertedindex.NativeMutableTextIndex; import org.apache.pinot.segment.local.segment.index.readers.text.NativeTextIndexReader; @@ -65,6 +71,9 @@ import org.apache.pinot.segment.spi.index.reader.TextIndexReader; import org.apache.pinot.segment.spi.index.reader.VectorIndexReader; import org.apache.pinot.spi.config.table.FieldConfig; +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.exception.BadQueryRequestException; import org.roaringbitmap.buffer.MutableRoaringBitmap; @@ -247,7 +256,7 @@ private BaseFilterOperator constructPhysicalOperator(FilterContext filter, int n BaseFilterOperator childFilterOperator = constructPhysicalOperator(childFilters.get(0), numDocs); return FilterOperatorUtils.getNotFilterOperator(_queryContext, childFilterOperator, numDocs); case PREDICATE: - Predicate predicate = filter.getPredicate(); + Predicate predicate = rewriteUrnKeyLongPredicate(filter.getPredicate()); ExpressionContext lhs = predicate.getLhs(); if (lhs.getType() == ExpressionContext.Type.FUNCTION) { if (canApplyH3IndexForDistanceCheck(predicate, lhs.getFunction())) { @@ -347,4 +356,87 @@ private BaseFilterOperator constructPhysicalOperator(FilterContext filter, int n throw new IllegalStateException(); } } + + /** + * If {@code predicate} has the form {@code urnKeyLong(col) OP literal} where {@code col} is a + * single-entity typed URN column with LONG/INT storage, rewrites it to {@code col OP + * entityPrefix:literal} so the existing identifier-based predicate path (with prefix-stripping + * inside {@code UrnReconstructingLongDictionary}) handles it. Otherwise returns the predicate + * unchanged. + * + *

Supports EQ, NOT_EQ, IN, NOT_IN, RANGE. Multi-entity URN columns are skipped because the + * literal alone cannot identify which entity type's prefix to apply. + */ + // Package-private for testing. + Predicate rewriteUrnKeyLongPredicate(Predicate predicate) { + ExpressionContext lhs = predicate.getLhs(); + if (lhs.getType() != ExpressionContext.Type.FUNCTION) { + return predicate; + } + FunctionContext function = lhs.getFunction(); + String name = function.getFunctionName(); + if (!UrnKeyLongTransformFunction.FUNCTION_NAME.equalsIgnoreCase(name) + && !UrnKeyLongTransformFunction.FUNCTION_NAME_ALT.equalsIgnoreCase(name)) { + return predicate; + } + if (function.getArguments().size() != 1) { + return predicate; + } + ExpressionContext arg = function.getArguments().get(0); + if (arg.getType() != ExpressionContext.Type.IDENTIFIER) { + return predicate; + } + Schema schema = _queryContext.getSchema(); + if (schema == null) { + return predicate; + } + String colName = arg.getIdentifier(); + FieldSpec fs = schema.getFieldSpecFor(colName); + if (!(fs instanceof DimensionFieldSpec)) { + return predicate; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fs; + if (!dim.hasTypedUrnStorage() || dim.usesSortedUrnDictionary()) { + return predicate; + } + String prefix = dim.getUrnTypes().get(0).getEntityType() + ":"; + ExpressionContext newLhs = ExpressionContext.forIdentifier(colName); + switch (predicate.getType()) { + case EQ: { + EqPredicate eq = (EqPredicate) predicate; + return new EqPredicate(newLhs, prefix + eq.getValue()); + } + case NOT_EQ: { + NotEqPredicate neq = (NotEqPredicate) predicate; + return new NotEqPredicate(newLhs, prefix + neq.getValue()); + } + case IN: { + InPredicate in = (InPredicate) predicate; + List rewritten = new ArrayList<>(in.getValues().size()); + for (String v : in.getValues()) { + rewritten.add(prefix + v); + } + return new InPredicate(newLhs, rewritten); + } + case NOT_IN: { + NotInPredicate nin = (NotInPredicate) predicate; + List rewritten = new ArrayList<>(nin.getValues().size()); + for (String v : nin.getValues()) { + rewritten.add(prefix + v); + } + return new NotInPredicate(newLhs, rewritten); + } + case RANGE: { + RangePredicate r = (RangePredicate) predicate; + String newLower = RangePredicate.UNBOUNDED.equals(r.getLowerBound()) ? r.getLowerBound() + : prefix + r.getLowerBound(); + String newUpper = RangePredicate.UNBOUNDED.equals(r.getUpperBound()) ? r.getUpperBound() + : prefix + r.getUpperBound(); + return new RangePredicate(newLhs, r.isLowerInclusive(), newLower, r.isUpperInclusive(), newUpper, + FieldSpec.DataType.URN); + } + default: + return predicate; + } + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 79913301a3..acee10b036 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -225,6 +225,7 @@ public static Object getConvertedFinalResult(DataTable dataTable, ColumnDataType return new Timestamp(dataTable.getLong(rowId, colId)); case STRING: case JSON: + case URN: return dataTable.getString(rowId, colId); case BYTES: return dataTable.getBytes(rowId, colId).getBytes(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ParentExprMinMaxAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ParentExprMinMaxAggregationFunction.java index 15daa2f2a4..f61cf7dd49 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ParentExprMinMaxAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/ParentExprMinMaxAggregationFunction.java @@ -216,6 +216,11 @@ private void initializeProjectionColumnValSet(Map measuringColTypes[i] = DataSchema.ColumnDataType.DOUBLE; break; case STRING: + case URN: exprMinMaxWrapperMeasuringColumnSets.add( new ExprMinMaxMeasuringValSetWrapper(true, DataSchema.ColumnDataType.STRING, blockValSet)); measuringColTypes[i] = DataSchema.ColumnDataType.STRING; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxObject.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxObject.java index 437e6fc304..3459e58b2f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxObject.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxObject.java @@ -239,6 +239,7 @@ public Comparable[] getExtremumKey() { extremumKeys[i] = _immutableMeasuringKeys.getDouble(0, i); break; case STRING: + case URN: extremumKeys[i] = _immutableMeasuringKeys.getString(0, i); break; case BIG_DECIMAL: @@ -273,6 +274,7 @@ public Object getField(int rowId, int colId) { return _immutableProjectionVals.getDouble(rowId, colId); case JSON: case STRING: + case URN: return _immutableProjectionVals.getString(rowId, colId); case BYTES: return _immutableProjectionVals.getBytes(rowId, colId); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxProjectionValSetWrapper.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxProjectionValSetWrapper.java index e855c55c9d..1e0281e57f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxProjectionValSetWrapper.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxProjectionValSetWrapper.java @@ -50,6 +50,7 @@ public Object getValue(int i) { case BIG_DECIMAL: case BYTES: case JSON: + case URN: return _objectsValues[i]; case INT_ARRAY: return _intValuesMV[i].length == 0 ? null : _intValuesMV[i]; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxWrapperValSet.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxWrapperValSet.java index 7090e4e9e0..8c73233043 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxWrapperValSet.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/utils/exprminmax/ExprMinMaxWrapperValSet.java @@ -65,6 +65,7 @@ public void setNewBlock(BlockValSet blockValSet) { break; case STRING: case JSON: + case URN: _objectsValues = blockValSet.getStringValuesSV(); break; case BIG_DECIMAL: diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index 09553bf9e5..f14aa5820a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -465,6 +465,7 @@ private Object getConvertedKey(DataTable dataTable, ColumnDataType columnDataTyp return new Timestamp(dataTable.getLong(rowId, colId)); case STRING: case JSON: + case URN: return dataTable.getString(rowId, colId); case BYTES: return dataTable.getBytes(rowId, colId).getBytes(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/util/GapfillUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/util/GapfillUtils.java index d5e75de3f3..8782154bbd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/util/GapfillUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/util/GapfillUtils.java @@ -89,6 +89,7 @@ static public Serializable getDefaultValue(DataSchema.ColumnDataType dataType) { return dataType.convertAndFormat(0); case STRING: case JSON: + case URN: case BYTES: return ""; case INT_ARRAY: diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityDictLookupTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityDictLookupTest.java new file mode 100644 index 0000000000..849912bd83 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityDictLookupTest.java @@ -0,0 +1,155 @@ +/** + * 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.operator.transform.function; + +import com.google.common.io.Files; +import java.io.File; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.core.operator.DocIdSetOperator; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.filter.MatchAllFilterOperator; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentUtil; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +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.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + + +/** + * End-to-end coverage for mixed-type multi-entity URN columns (LONG memberId + STRING corpuser). + * With the polymorphic wire format, each entity gets its own typed sub-dictionary in one + * dictionary file, the per-entity dict-ID ranges are persisted in segment metadata, and + * {@code urnEntity()} reads the entity prefix per row through the RANGE_LOOKUP fast path — no + * dictionary-string lookup, no URN parsing. + */ +public class UrnEntityDictLookupTest { + + private static final String SEGMENT_NAME = "urnEntityMixedTypeSegment"; + private static final String URN_COL = "subject"; + private static final String TIME_COL = "t"; + private static final int NUM_ROWS = 60; + + private File _segmentOutputDir; + private ImmutableSegment _segment; + private Map _dataSourceMap; + private ProjectionBlock _projectionBlock; + private final String[] _expectedEntities = new String[NUM_ROWS]; + + @BeforeClass + public void setUp() + throws Exception { + _segmentOutputDir = Files.createTempDir(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); + + // Mixed-type multi-entity URN: LONG memberId + STRING corpuser. The polymorphic creator + // writes a LongDictionary region for memberId and a StringDictionary region for corpuser + // back-to-back in one dictionary file, with per-entity dict-ID ranges recorded in metadata. + DimensionFieldSpec multiDim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + multiDim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:corpuser", FieldSpec.DataType.STRING))); + + Schema schema = new Schema.SchemaBuilder() + .addField(multiDim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + List rows = new ArrayList<>(); + for (int i = 0; i < NUM_ROWS; i++) { + String entity = (i % 2 == 0) ? "urn:li:memberId" : "urn:li:corpuser"; + String key = (i % 2 == 0) ? String.valueOf(1000 + i) : ("user" + i); + _expectedEntities[i] = entity; + GenericRow row = new GenericRow(); + row.putValue(URN_COL, entity + ":" + key); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + } + + File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + _segment = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); + + _dataSourceMap = new HashMap<>(); + _dataSourceMap.put(URN_COL, _segment.getDataSource(URN_COL)); + _dataSourceMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); + + int numDocs = _segment.getSegmentMetadata().getTotalDocs(); + _projectionBlock = new ProjectionOperator(_dataSourceMap, + new DocIdSetOperator(new MatchAllFilterOperator(numDocs), + DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock(); + } + + @AfterClass + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + FileUtils.deleteQuietly(_segmentOutputDir); + } + + @Test + public void mixedTypeColumnPersistsTypedDictRanges() { + DimensionFieldSpec dim = (DimensionFieldSpec) _segment.getDataSource(URN_COL) + .getDataSourceMetadata().getFieldSpec(); + List ranges = dim.getEntityDictRanges(); + assertNotNull(ranges); + assertEquals(ranges.size(), 2); + assertEquals(ranges.get(0).getEntityType(), "urn:li:memberId"); + assertEquals(ranges.get(0).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(ranges.get(0).getNumBytesPerValue(), Long.BYTES); + assertEquals(ranges.get(1).getEntityType(), "urn:li:corpuser"); + assertEquals(ranges.get(1).getDictClass(), FieldSpec.DataType.STRING); + } + + @Test + public void mixedTypeColumnUrnEntityReturnsPerRowEntityType() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + String[] result = fn.transformToStringValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + assertEquals(result[i], _expectedEntities[i], "Row " + i); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityMultiEntityTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityMultiEntityTest.java new file mode 100644 index 0000000000..5bda752857 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityMultiEntityTest.java @@ -0,0 +1,159 @@ +/** + * 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.operator.transform.function; + +import com.google.common.io.Files; +import java.io.File; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.core.operator.DocIdSetOperator; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.filter.MatchAllFilterOperator; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentUtil; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +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.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + + +/** + * Tests the multi-entity {@code urnEntity()} fast path: COMPANION_LOOKUP strategy that reads + * the entity dict ID ranges loaded from segment metadata to determine each row's owning entity + * type from its main-column dict ID, without reading any per-row companion data. + */ +public class UrnEntityMultiEntityTest { + + private static final String SEGMENT_NAME = "urnEntityMultiSegment"; + private static final String URN_COL = "actorId"; + private static final String TIME_COL = "t"; + private static final int NUM_ROWS = 60; + + private File _segmentOutputDir; + private ImmutableSegment _segment; + private Map _dataSourceMap; + private ProjectionBlock _projectionBlock; + private final String[] _expectedEntities = new String[NUM_ROWS]; + + @BeforeClass + public void setUp() + throws Exception { + _segmentOutputDir = Files.createTempDir(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); + + DimensionFieldSpec multiDim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + multiDim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG))); + + Schema schema = new Schema.SchemaBuilder() + .addField(multiDim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + List rows = new ArrayList<>(); + for (int i = 0; i < NUM_ROWS; i++) { + String entity = (i % 2 == 0) ? "urn:li:memberId" : "urn:li:groupId"; + String urn = entity + ":" + (1000 + i); + _expectedEntities[i] = entity; + GenericRow row = new GenericRow(); + row.putValue(URN_COL, urn); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + } + + File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + _segment = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); + + _dataSourceMap = new HashMap<>(); + _dataSourceMap.put(URN_COL, _segment.getDataSource(URN_COL)); + _dataSourceMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); + + int numDocs = _segment.getSegmentMetadata().getTotalDocs(); + _projectionBlock = new ProjectionOperator(_dataSourceMap, + new DocIdSetOperator(new MatchAllFilterOperator(numDocs), + DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock(); + } + + @AfterClass + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + FileUtils.deleteQuietly(_segmentOutputDir); + } + + @Test + public void entityDictRangesPersistedInMetadata() { + org.apache.pinot.spi.data.DimensionFieldSpec dim = + (org.apache.pinot.spi.data.DimensionFieldSpec) _segment.getDataSource(URN_COL) + .getDataSourceMetadata().getFieldSpec(); + assertNotNull(dim.getEntityDictRanges()); + assertEquals(dim.getEntityDictRanges().size(), 2); + assertEquals(dim.getEntityDictRanges().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(dim.getEntityDictRanges().get(1).getEntityType(), "urn:li:groupId"); + } + + @Test + public void urnEntityUsesRangeLookupFastPath() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + + String[] result = fn.transformToStringValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + + // Each row's main-column dict ID falls into one of the per-entity ranges; the transform + // returns the owning entity type for that dict ID. + int memberIds = 0; + int groupIds = 0; + for (String r : result) { + if ("urn:li:memberId".equals(r)) { + memberIds++; + } else if ("urn:li:groupId".equals(r)) { + groupIds++; + } + } + assertEquals(memberIds, NUM_ROWS / 2); + assertEquals(groupIds, NUM_ROWS / 2); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java new file mode 100644 index 0000000000..5a00496d8a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java @@ -0,0 +1,332 @@ +/** + * 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.operator.transform.function; + +import com.google.common.collect.Lists; +import com.google.common.io.Files; +import it.unimi.dsi.fastutil.ints.IntOpenHashSet; +import it.unimi.dsi.fastutil.ints.IntSet; +import java.io.File; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.common.request.context.predicate.EqPredicate; +import org.apache.pinot.common.request.context.predicate.InPredicate; +import org.apache.pinot.common.request.context.predicate.RangePredicate; +import org.apache.pinot.core.operator.DocIdSetOperator; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.filter.MatchAllFilterOperator; +import org.apache.pinot.core.operator.filter.predicate.BaseDictionaryBasedPredicateEvaluator; +import org.apache.pinot.core.operator.filter.predicate.EqualsPredicateEvaluatorFactory; +import org.apache.pinot.core.operator.filter.predicate.InPredicateEvaluatorFactory; +import org.apache.pinot.core.operator.filter.predicate.RangePredicateEvaluatorFactory; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentUtil; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +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.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/** + * Correctness tests for {@link UrnEntityTransformFunction} and {@link UrnKeyLongTransformFunction} + * on typed URN LONG columns (fast path) and untyped URN STRING columns (fallback path). + */ +public class UrnTransformFunctionTest { + + private static final String SEGMENT_NAME = "urnTransformFunctionTestSegment"; + private static final String URN_COL = "memberId"; + private static final String UNTYPED_URN_COL = "untypedMemberId"; + private static final String TIME_COL = "t"; + private static final int NUM_ROWS = 100; + + private File _segmentOutputDir; + private ImmutableSegment _segment; + private Map _dataSourceMap; + private ProjectionBlock _projectionBlock; + + // Track expected values for assertions + private final long[] _expectedKeys = new long[NUM_ROWS]; + + @BeforeClass + public void setUp() + throws Exception { + _segmentOutputDir = Files.createTempDir(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); + + DimensionFieldSpec typedDim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + typedDim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + Schema schema = new Schema.SchemaBuilder() + .addField(typedDim) + .addSingleValueDimension(UNTYPED_URN_COL, FieldSpec.DataType.URN) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + List rows = new ArrayList<>(); + for (int i = 0; i < NUM_ROWS; i++) { + long key = 100L + i; + _expectedKeys[i] = key; + GenericRow row = new GenericRow(); + row.putValue(URN_COL, "urn:li:memberId:" + key); + row.putValue(UNTYPED_URN_COL, "urn:li:memberId:" + key); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + } + + File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + _segment = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); + + _dataSourceMap = new HashMap<>(); + // Include the typed URN column and its companion column (if any) + _dataSourceMap.put(URN_COL, _segment.getDataSource(URN_COL)); + _dataSourceMap.put(UNTYPED_URN_COL, _segment.getDataSource(UNTYPED_URN_COL)); + _dataSourceMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); + + int numDocs = _segment.getSegmentMetadata().getTotalDocs(); + _projectionBlock = new ProjectionOperator(_dataSourceMap, + new DocIdSetOperator(new MatchAllFilterOperator(numDocs), + DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock(); + } + + @AfterClass + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + FileUtils.deleteQuietly(_segmentOutputDir); + } + + // --- urnEntity on typed LONG column (CONSTANT fast path) --- + + @Test + public void testUrnEntityOnTypedLongColumn() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + + String[] result = fn.transformToStringValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + assertEquals(result[i], "urn:li:memberId", "Row " + i); + } + } + + @Test + public void testUrnEntityUnderscoreAliasOnTypedLongColumn() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urn_entity(" + URN_COL + ")"), _dataSourceMap); + String[] result = fn.transformToStringValuesSV(_projectionBlock); + for (int i = 0; i < NUM_ROWS; i++) { + assertEquals(result[i], "urn:li:memberId"); + } + } + + // --- urnEntity on untyped URN column (auto-detected to CONSTANT here: single-prefix data) --- + + @Test + public void testUrnEntityOnUntypedUrnColumn() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + UNTYPED_URN_COL + ")"), _dataSourceMap); + + String[] result = fn.transformToStringValuesSV(_projectionBlock); + assertNotNull(result); + for (int i = 0; i < NUM_ROWS; i++) { + assertEquals(result[i], "urn:li:memberId", "Row " + i); + } + } + + // --- urnKeyLong on typed LONG column (DIRECT LONG fast path) --- + + @Test + public void testUrnKeyLongOnTypedLongColumn() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnKeyLong(" + URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + + long[] result = fn.transformToLongValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + + // Values should match the sorted order of keys stored in the dictionary. + // Build expected: keys sorted ascending (100..199 are already sorted) + long[] allKeys = Arrays.copyOf(_expectedKeys, NUM_ROWS); + Arrays.sort(allKeys); + + // The segment stores docs in ingestion order and the dict is sorted, so we need to verify + // via the actual doc order. Build a set of expected keys. + java.util.Set expectedSet = new java.util.HashSet<>(); + for (long k : _expectedKeys) { + expectedSet.add(k); + } + for (int i = 0; i < NUM_ROWS; i++) { + expectedSet.remove(result[i]); + } + assertEquals(expectedSet.isEmpty(), true, "Some expected keys were not returned: " + expectedSet); + } + + @Test + public void testUrnKeyLongResultIsLongType() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnKeyLong(" + URN_COL + ")"), _dataSourceMap); + assertEquals(fn.getResultMetadata().getDataType(), FieldSpec.DataType.LONG); + } + + // --- urnKeyLong on untyped URN column --- + + @Test + public void testUrnKeyLongOnUntypedUrnColumn() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnKeyLong(" + UNTYPED_URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + + long[] result = fn.transformToLongValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + + java.util.Set expectedSet = new java.util.HashSet<>(); + for (long k : _expectedKeys) { + expectedSet.add(k); + } + for (int i = 0; i < NUM_ROWS; i++) { + expectedSet.remove(result[i]); + } + assertEquals(expectedSet.isEmpty(), true, "Untyped urnKeyLong missing keys: " + expectedSet); + } + + // --- Consistency: typed and untyped columns return same entity --- + + @Test + public void testTypedAndUntypedReturnSameEntity() { + TransformFunction typedFn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); + TransformFunction untypedFn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + UNTYPED_URN_COL + ")"), _dataSourceMap); + + String[] typed = typedFn.transformToStringValuesSV(_projectionBlock); + String[] untyped = untypedFn.transformToStringValuesSV(_projectionBlock); + + for (int i = 0; i < NUM_ROWS; i++) { + assertEquals(typed[i], untyped[i], "Entity mismatch at row " + i); + } + } + + // --- Dictionary reconstruction: SELECT on typed column returns full URN strings --- + + @Test + public void testTypedUrnDictionaryReconstructsUrn() { + DataSource ds = _dataSourceMap.get(URN_COL); + assertNotNull(ds); + Dictionary dict = ds.getDictionary(); + assertNotNull(dict, "Dictionary should exist for typed URN column"); + + for (int dictId = 0; dictId < dict.length(); dictId++) { + String strVal = dict.getStringValue(dictId); + assertNotNull(strVal); + assertTrue(strVal.startsWith("urn:li:memberId:"), "Expected URN string but got: " + strVal); + } + } + + // --- Predicate evaluator: EQ and IN on typed LONG URN column --- + + @Test + public void testEqualityPredicateOnTypedUrn() { + DataSource ds = _dataSourceMap.get(URN_COL); + Dictionary dict = ds.getDictionary(); + + // "urn:li:memberId:100" is the first ingested row, so its key exists in the dict + ExpressionContext col = ExpressionContext.forIdentifier(URN_COL); + EqPredicate eq = new EqPredicate(col, "urn:li:memberId:100"); + BaseDictionaryBasedPredicateEvaluator eval = + EqualsPredicateEvaluatorFactory.newDictionaryBasedEvaluator(eq, dict, FieldSpec.DataType.URN); + + int[] matching = eval.getMatchingDictIds(); + assertEquals(matching.length, 1); + assertEquals(dict.getStringValue(matching[0]), "urn:li:memberId:100"); + } + + @Test + public void testInPredicateOnTypedUrn() { + DataSource ds = _dataSourceMap.get(URN_COL); + Dictionary dict = ds.getDictionary(); + + ExpressionContext col = ExpressionContext.forIdentifier(URN_COL); + InPredicate in = new InPredicate(col, Lists.newArrayList("urn:li:memberId:100", "urn:li:memberId:101")); + BaseDictionaryBasedPredicateEvaluator eval = + InPredicateEvaluatorFactory.newDictionaryBasedEvaluator(in, dict, FieldSpec.DataType.URN, null); + + int[] matching = eval.getMatchingDictIds(); + assertEquals(matching.length, 2); + IntSet matchingStrings = new IntOpenHashSet(); + for (int dictId : matching) { + String val = dict.getStringValue(dictId); + assertTrue(val.equals("urn:li:memberId:100") || val.equals("urn:li:memberId:101"), + "Unexpected value: " + val); + matchingStrings.add(dictId); + } + assertEquals(matchingStrings.size(), 2); + } + + @Test + public void testRangePredicateOnTypedUrn() { + DataSource ds = _dataSourceMap.get(URN_COL); + Dictionary dict = ds.getDictionary(); + + // Range: memberId BETWEEN 'urn:li:memberId:100' AND 'urn:li:memberId:102' (inclusive) + ExpressionContext col = ExpressionContext.forIdentifier(URN_COL); + RangePredicate range = new RangePredicate(col, true, "urn:li:memberId:100", true, "urn:li:memberId:102", + FieldSpec.DataType.URN); + BaseDictionaryBasedPredicateEvaluator eval = + RangePredicateEvaluatorFactory.newDictionaryBasedEvaluator(range, dict, FieldSpec.DataType.URN); + + // Should match keys 100, 101, 102 + int[] matching = eval.getMatchingDictIds(); + assertEquals(matching.length, 3, "Expected 3 matching dict IDs for range [100, 102]"); + for (int dictId : matching) { + String val = dict.getStringValue(dictId); + long key = Long.parseLong(val.substring("urn:li:memberId:".length())); + assertTrue(key >= 100 && key <= 102, "Key out of range: " + key); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeUrnRewriteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeUrnRewriteTest.java new file mode 100644 index 0000000000..6490d14e58 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeUrnRewriteTest.java @@ -0,0 +1,185 @@ +/** + * 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.plan; + +import com.google.common.collect.Lists; +import java.util.Arrays; +import java.util.List; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.FunctionContext; +import org.apache.pinot.common.request.context.predicate.EqPredicate; +import org.apache.pinot.common.request.context.predicate.InPredicate; +import org.apache.pinot.common.request.context.predicate.NotEqPredicate; +import org.apache.pinot.common.request.context.predicate.NotInPredicate; +import org.apache.pinot.common.request.context.predicate.Predicate; +import org.apache.pinot.common.request.context.predicate.RangePredicate; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.mockito.Mockito; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; + + +/** + * Unit tests for {@link FilterPlanNode#rewriteUrnKeyLongPredicate} — the rewrite that turns + * {@code urnKeyLong(typed_long_urn_col) OP literal} into {@code col OP entityPrefix:literal} so + * the predicate hits the dictionary-based fast path instead of going through + * {@code ExpressionFilterOperator}. + */ +public class FilterPlanNodeUrnRewriteTest { + + private static final String TYPED_COL = "memberId"; + private static final String MULTI_COL = "actorId"; + private static final String UNTYPED_COL = "rawUrn"; + private static final String STR_COL = "name"; + + private FilterPlanNode _planNode; + + @BeforeClass + public void setUp() { + DimensionFieldSpec typedDim = new DimensionFieldSpec(TYPED_COL, FieldSpec.DataType.URN, true); + typedDim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + DimensionFieldSpec multiDim = new DimensionFieldSpec(MULTI_COL, FieldSpec.DataType.URN, true); + multiDim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG))); + + Schema schema = new Schema.SchemaBuilder() + .addField(typedDim) + .addField(multiDim) + .addSingleValueDimension(UNTYPED_COL, FieldSpec.DataType.URN) + .addSingleValueDimension(STR_COL, FieldSpec.DataType.STRING) + .build(); + + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT 1 FROM t"); + queryContext.setSchema(schema); + + IndexSegment indexSegment = Mockito.mock(IndexSegment.class); + SegmentContext segmentContext = Mockito.mock(SegmentContext.class); + Mockito.when(segmentContext.getIndexSegment()).thenReturn(indexSegment); + _planNode = new FilterPlanNode(segmentContext, queryContext); + } + + @Test + public void rewritesEqOnTypedLongUrn() { + Predicate in = new EqPredicate(urnKeyLongOf(TYPED_COL), "1000"); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertTrue(out instanceof EqPredicate); + EqPredicate eq = (EqPredicate) out; + assertEquals(eq.getLhs().getIdentifier(), TYPED_COL); + assertEquals(eq.getValue(), "urn:li:memberId:1000"); + } + + @Test + public void rewritesNotEqOnTypedLongUrn() { + Predicate out = _planNode.rewriteUrnKeyLongPredicate(new NotEqPredicate(urnKeyLongOf(TYPED_COL), "42")); + assertTrue(out instanceof NotEqPredicate); + NotEqPredicate neq = (NotEqPredicate) out; + assertEquals(neq.getLhs().getIdentifier(), TYPED_COL); + assertEquals(neq.getValue(), "urn:li:memberId:42"); + } + + @Test + public void rewritesInOnTypedLongUrn() { + Predicate out = _planNode.rewriteUrnKeyLongPredicate( + new InPredicate(urnKeyLongOf(TYPED_COL), Lists.newArrayList("1", "2", "3"))); + assertTrue(out instanceof InPredicate); + InPredicate in = (InPredicate) out; + assertEquals(in.getLhs().getIdentifier(), TYPED_COL); + assertEquals(in.getValues(), + Arrays.asList("urn:li:memberId:1", "urn:li:memberId:2", "urn:li:memberId:3")); + } + + @Test + public void rewritesNotInOnTypedLongUrn() { + Predicate out = _planNode.rewriteUrnKeyLongPredicate( + new NotInPredicate(urnKeyLongOf(TYPED_COL), Lists.newArrayList("9", "10"))); + assertTrue(out instanceof NotInPredicate); + NotInPredicate nin = (NotInPredicate) out; + assertEquals(nin.getValues(), Arrays.asList("urn:li:memberId:9", "urn:li:memberId:10")); + } + + @Test + public void rewritesBoundedRangeOnTypedLongUrn() { + RangePredicate in = new RangePredicate(urnKeyLongOf(TYPED_COL), true, "100", false, "200", + FieldSpec.DataType.LONG); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertTrue(out instanceof RangePredicate); + RangePredicate r = (RangePredicate) out; + assertEquals(r.getLhs().getIdentifier(), TYPED_COL); + assertEquals(r.getLowerBound(), "urn:li:memberId:100"); + assertEquals(r.getUpperBound(), "urn:li:memberId:200"); + assertEquals(r.isLowerInclusive(), true); + assertEquals(r.isUpperInclusive(), false); + } + + @Test + public void rangeWithUnboundedLowerKeepsSentinel() { + RangePredicate in = new RangePredicate(urnKeyLongOf(TYPED_COL), false, RangePredicate.UNBOUNDED, true, "500", + FieldSpec.DataType.LONG); + RangePredicate out = (RangePredicate) _planNode.rewriteUrnKeyLongPredicate(in); + assertEquals(out.getLowerBound(), RangePredicate.UNBOUNDED); + assertEquals(out.getUpperBound(), "urn:li:memberId:500"); + } + + @Test + public void skipsMultiEntityUrnColumn() { + Predicate in = new EqPredicate(urnKeyLongOf(MULTI_COL), "1000"); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertSame(out, in, "Multi-entity URN columns are ambiguous; rewrite must be skipped"); + } + + @Test + public void skipsUntypedUrnColumn() { + Predicate in = new EqPredicate(urnKeyLongOf(UNTYPED_COL), "1000"); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertSame(out, in); + } + + @Test + public void skipsNonUrnFunctionCall() { + ExpressionContext expr = ExpressionContext.forFunction(new FunctionContext(FunctionContext.Type.TRANSFORM, + "length", List.of(ExpressionContext.forIdentifier(STR_COL)))); + Predicate in = new EqPredicate(expr, "5"); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertSame(out, in); + } + + @Test + public void skipsIdentifierPredicate() { + Predicate in = new EqPredicate(ExpressionContext.forIdentifier(STR_COL), "x"); + Predicate out = _planNode.rewriteUrnKeyLongPredicate(in); + assertSame(out, in); + } + + private static ExpressionContext urnKeyLongOf(String column) { + return ExpressionContext.forFunction(new FunctionContext(FunctionContext.Type.TRANSFORM, "urnKeyLong", + List.of(ExpressionContext.forIdentifier(column)))); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java new file mode 100644 index 0000000000..a3251e0fed --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java @@ -0,0 +1,508 @@ +/** + * 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.perf; + +import com.google.common.io.Files; +import java.io.File; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.core.common.BlockDocIdIterator; +import org.apache.pinot.core.operator.DocIdSetOperator; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.blocks.ProjectionBlock; +import org.apache.pinot.core.operator.filter.BaseFilterOperator; +import org.apache.pinot.core.operator.filter.MatchAllFilterOperator; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.plan.FilterPlanNode; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentUtil; +import org.apache.pinot.segment.spi.Constants; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +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.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/** + * End-to-end JMH benchmark comparing typed and untyped URN columns: + *

    + *
  • Transform function microbenchmarks (urnEntity, urnKeyLong) on single-entity columns.
    • + *
  • Dictionary microbenchmarks (getStringValue, indexOf).
    • + *
  • Full filter pipeline (FilterPlanNode -> execution) for "col = literal" and + * "urnEntity(col) = literal" queries, on both single-entity and multi-entity columns.
    • + *
  • On-disk storage sizes for single-entity and multi-entity segments, printed at setup.
    • + *
+ */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Fork(1) +@Warmup(iterations = 2, time = 2) +@Measurement(iterations = 3, time = 2) +@State(Scope.Benchmark) +public class BenchmarkUrnTransform { + + private static final int NUM_ROWS = 100_000; + private static final int CARDINALITY = 10_000; + + private static final String SINGLE_ENTITY = "urn:li:memberId"; + private static final String[] MULTI_ENTITIES = + {"urn:li:corpuser", "urn:li:corpGroup", "urn:li:applicationName"}; + + // Existing transform-benchmark segment with both typed and untyped columns. + private static final String TYPED_COL = "typedId"; + private static final String UNTYPED_COL = "untypedId"; + private static final String TIME_COL = "t"; + + private File _segmentDir; + private ImmutableSegment _segment; + private Map _dataSourceMap; + private ProjectionBlock _projectionBlock; + private Dictionary _typedDict; + private Dictionary _untypedDict; + private String[] _lookupValues; + + // Dedicated segments per layout (typed vs untyped, single vs multi) for storage and filter + // benchmarks. Each segment contains exactly one URN column plus the time column so the + // on-disk size is directly comparable. + private SegmentBundle _singleTyped; + private SegmentBundle _singleUntyped; + private SegmentBundle _multiTyped; + private SegmentBundle _multiUntyped; + + private QueryContext _qEqSingleTyped; + private QueryContext _qEqSingleUntyped; + private QueryContext _qEqMultiTyped; + private QueryContext _qEqMultiUntyped; + + private QueryContext _qUrnEntitySingleTyped; + private QueryContext _qUrnEntitySingleUntyped; + private QueryContext _qUrnEntityMultiTyped; + private QueryContext _qUrnEntityMultiUntyped; + + public static void main(String[] args) throws Exception { + new Runner(new OptionsBuilder() + .include(BenchmarkUrnTransform.class.getSimpleName()) + .build()).run(); + } + + @Setup + public void setUp() throws Exception { + setUpTransformSegment(); + setUpFilterSegments(); + setUpQueryContexts(); + printStorageSizes(); + } + + private void setUpTransformSegment() throws Exception { + _segmentDir = Files.createTempDir(); + + DimensionFieldSpec typed = new DimensionFieldSpec(TYPED_COL, FieldSpec.DataType.URN, true); + typed.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of(SINGLE_ENTITY, FieldSpec.DataType.LONG))); + + Schema schema = new Schema.SchemaBuilder() + .addField(typed) + .addSingleValueDimension(UNTYPED_COL, FieldSpec.DataType.URN) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("t").build(); + + List rows = new ArrayList<>(NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + long key = i % CARDINALITY; + String urn = SINGLE_ENTITY + ":" + key; + GenericRow row = new GenericRow(); + row.putValue(TYPED_COL, urn); + row.putValue(UNTYPED_COL, urn); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + } + + File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, "benchmarkUrnSegment", + _segmentDir.toString(), new GenericRowRecordReader(rows)); + _segment = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); + + _dataSourceMap = new HashMap<>(); + _dataSourceMap.put(TYPED_COL, _segment.getDataSource(TYPED_COL)); + _dataSourceMap.put(UNTYPED_COL, _segment.getDataSource(UNTYPED_COL)); + _dataSourceMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); + + int numDocs = _segment.getSegmentMetadata().getTotalDocs(); + _projectionBlock = new ProjectionOperator(_dataSourceMap, + new DocIdSetOperator(new MatchAllFilterOperator(numDocs), + DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock(); + + _typedDict = _dataSourceMap.get(TYPED_COL).getDictionary(); + _untypedDict = _dataSourceMap.get(UNTYPED_COL).getDictionary(); + + _lookupValues = new String[CARDINALITY]; + for (int i = 0; i < CARDINALITY; i++) { + _lookupValues[i] = SINGLE_ENTITY + ":" + i; + } + } + + private void setUpFilterSegments() throws Exception { + List singleRows = generateSingleEntityRows(); + List multiRows = generateMultiEntityRows(); + + _singleTyped = buildBundle("singleTyped", singleEntitySchema(true), singleRows, "memberId"); + _singleUntyped = buildBundle("singleUntyped", singleEntitySchema(false), singleRows, "memberId"); + _multiTyped = buildBundle("multiTyped", multiEntitySchema(true), multiRows, "actorId"); + _multiUntyped = buildBundle("multiUntyped", multiEntitySchema(false), multiRows, "actorId"); + } + + private void setUpQueryContexts() { + String singleLiteral = SINGLE_ENTITY + ":1234"; + String multiLiteral = MULTI_ENTITIES[0] + ":1234"; + + _qEqSingleTyped = compile(_singleTyped, "SELECT t FROM t WHERE memberId = '" + singleLiteral + "'"); + _qEqSingleUntyped = compile(_singleUntyped, "SELECT t FROM t WHERE memberId = '" + singleLiteral + "'"); + _qEqMultiTyped = compile(_multiTyped, "SELECT t FROM t WHERE actorId = '" + multiLiteral + "'"); + _qEqMultiUntyped = compile(_multiUntyped, "SELECT t FROM t WHERE actorId = '" + multiLiteral + "'"); + + _qUrnEntitySingleTyped = compile(_singleTyped, + "SELECT t FROM t WHERE urnEntity(memberId) = '" + SINGLE_ENTITY + "'"); + _qUrnEntitySingleUntyped = compile(_singleUntyped, + "SELECT t FROM t WHERE urnEntity(memberId) = '" + SINGLE_ENTITY + "'"); + _qUrnEntityMultiTyped = compile(_multiTyped, + "SELECT t FROM t WHERE urnEntity(actorId) = '" + MULTI_ENTITIES[0] + "'"); + _qUrnEntityMultiUntyped = compile(_multiUntyped, + "SELECT t FROM t WHERE urnEntity(actorId) = '" + MULTI_ENTITIES[0] + "'"); + } + + private void printStorageSizes() { + System.out.println(); + System.out.println("===== Storage sizes (NUM_ROWS=" + NUM_ROWS + ", CARDINALITY=" + CARDINALITY + ") ====="); + long st = FileUtils.sizeOfDirectory(_singleTyped._segDir); + long su = FileUtils.sizeOfDirectory(_singleUntyped._segDir); + long mt = FileUtils.sizeOfDirectory(_multiTyped._segDir); + long mu = FileUtils.sizeOfDirectory(_multiUntyped._segDir); + System.out.println("Total segment size:"); + System.out.printf(" single-entity typed: %,11d bytes%n", st); + System.out.printf(" single-entity untyped: %,11d bytes (%.2fx vs typed)%n", su, ratio(su, st)); + System.out.printf(" multi-entity typed: %,11d bytes%n", mt); + System.out.printf(" multi-entity untyped: %,11d bytes (%.2fx vs typed)%n", mu, ratio(mu, mt)); + System.out.println(); + long std = dictSize(_singleTyped); + long sud = dictSize(_singleUntyped); + long mtd = dictSize(_multiTyped); + long mud = dictSize(_multiUntyped); + System.out.println("Dictionary file size (.dict only):"); + System.out.printf(" single-entity typed: %,11d bytes%n", std); + System.out.printf(" single-entity untyped: %,11d bytes (%.2fx vs typed)%n", sud, ratio(sud, std)); + System.out.printf(" multi-entity typed: %,11d bytes%n", mtd); + System.out.printf(" multi-entity untyped: %,11d bytes (%.2fx vs typed)%n", mud, ratio(mud, mtd)); + System.out.println("================================================================"); + System.out.println(); + } + + /** + * Returns the on-disk size of the URN main column's dictionary in this segment bundle. + * Reads v3-layout {@code index_map} for the {@code .dictionary.size} entry; falls back + * to a v1-layout {@code .dict} sidecar file when present. + */ + private static long dictSize(SegmentBundle bundle) { + File segDir = bundle._segDir; + File v3IndexMap = new File(new File(segDir, "v3"), "index_map"); + if (v3IndexMap.isFile()) { + try { + String key = bundle._urnColName + ".dictionary.size"; + for (String line : java.nio.file.Files.readAllLines(v3IndexMap.toPath())) { + String trimmed = line.trim(); + if (trimmed.startsWith(key)) { + int eq = trimmed.indexOf('='); + if (eq > 0) { + return Long.parseLong(trimmed.substring(eq + 1).trim()); + } + } + } + } catch (java.io.IOException e) { + // fall through to v1 sidecar + } + } + File v1Dict = new File(segDir, bundle._urnColName + ".dict"); + return v1Dict.isFile() ? v1Dict.length() : 0; + } + + private static double ratio(long a, long b) { + return b == 0 ? 0.0 : (double) a / (double) b; + } + + // ---------------- Schemas / data ---------------- + + private static Schema singleEntitySchema(boolean typed) { + Schema.SchemaBuilder builder = new Schema.SchemaBuilder(); + if (typed) { + DimensionFieldSpec dim = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of(SINGLE_ENTITY, FieldSpec.DataType.LONG))); + builder.addField(dim); + } else { + builder.addSingleValueDimension("memberId", FieldSpec.DataType.URN); + } + return builder.addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS").build(); + } + + private static Schema multiEntitySchema(boolean typed) { + Schema.SchemaBuilder builder = new Schema.SchemaBuilder(); + if (typed) { + DimensionFieldSpec dim = new DimensionFieldSpec("actorId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of(MULTI_ENTITIES[0], FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of(MULTI_ENTITIES[1], FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of(MULTI_ENTITIES[2], FieldSpec.DataType.LONG))); + builder.addField(dim); + } else { + builder.addSingleValueDimension("actorId", FieldSpec.DataType.URN); + } + return builder.addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS").build(); + } + + private List generateSingleEntityRows() { + long now = System.currentTimeMillis(); + List rows = new ArrayList<>(NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + long key = i % CARDINALITY; + GenericRow row = new GenericRow(); + row.putValue("memberId", SINGLE_ENTITY + ":" + key); + row.putValue(TIME_COL, now); + rows.add(row); + } + return rows; + } + + private List generateMultiEntityRows() { + long now = System.currentTimeMillis(); + List rows = new ArrayList<>(NUM_ROWS); + int perEntity = CARDINALITY / MULTI_ENTITIES.length; + for (int i = 0; i < NUM_ROWS; i++) { + String entity = MULTI_ENTITIES[i % MULTI_ENTITIES.length]; + long key = (i / MULTI_ENTITIES.length) % perEntity; + GenericRow row = new GenericRow(); + row.putValue("actorId", entity + ":" + key); + row.putValue(TIME_COL, now); + rows.add(row); + } + return rows; + } + + private SegmentBundle buildBundle(String name, Schema schema, List rows, String urnColName) + throws Exception { + File parent = Files.createTempDir(); + TableConfig tc = new TableConfigBuilder(TableType.OFFLINE).setTableName("t").build(); + File segDir = + PinotSegmentUtil.createSegment(tc, schema, name, parent.toString(), new GenericRowRecordReader(rows)); + ImmutableSegment seg = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); + return new SegmentBundle(parent, segDir, seg, schema, urnColName); + } + + private QueryContext compile(SegmentBundle bundle, String sql) { + QueryContext qc = QueryContextConverterUtils.getQueryContext(sql); + qc.setSchema(bundle._schema); + return qc; + } + + @TearDown + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + FileUtils.deleteQuietly(_segmentDir); + closeBundle(_singleTyped); + closeBundle(_singleUntyped); + closeBundle(_multiTyped); + closeBundle(_multiUntyped); + } + + private static void closeBundle(SegmentBundle bundle) { + if (bundle == null) { + return; + } + if (bundle._segment != null) { + bundle._segment.destroy(); + } + FileUtils.deleteQuietly(bundle._parentDir); + } + + // ---------------- Transform microbenchmarks (existing) ---------------- + + @Benchmark + public void transformTypedEntity(Blackhole bh) { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + TYPED_COL + ")"), _dataSourceMap); + bh.consume(fn.transformToStringValuesSV(_projectionBlock)); + } + + @Benchmark + public void transformUntypedEntity(Blackhole bh) { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + UNTYPED_COL + ")"), _dataSourceMap); + bh.consume(fn.transformToStringValuesSV(_projectionBlock)); + } + + @Benchmark + public void transformTypedKey(Blackhole bh) { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnKeyLong(" + TYPED_COL + ")"), _dataSourceMap); + bh.consume(fn.transformToLongValuesSV(_projectionBlock)); + } + + @Benchmark + public void transformUntypedKey(Blackhole bh) { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnKeyLong(" + UNTYPED_COL + ")"), _dataSourceMap); + bh.consume(fn.transformToLongValuesSV(_projectionBlock)); + } + + @Benchmark + public void dictGetStringTyped(Blackhole bh) { + int n = _typedDict.length(); + for (int i = 0; i < n; i++) { + bh.consume(_typedDict.getStringValue(i)); + } + } + + @Benchmark + public void dictGetStringUntyped(Blackhole bh) { + int n = _untypedDict.length(); + for (int i = 0; i < n; i++) { + bh.consume(_untypedDict.getStringValue(i)); + } + } + + @Benchmark + public void dictIndexOfTyped(Blackhole bh) { + for (String v : _lookupValues) { + bh.consume(_typedDict.indexOf(v)); + } + } + + @Benchmark + public void dictIndexOfUntyped(Blackhole bh) { + for (String v : _lookupValues) { + bh.consume(_untypedDict.indexOf(v)); + } + } + + // ---------------- End-to-end "col = literal" filter ---------------- + + @Benchmark + public void filterEqSingleTyped(Blackhole bh) { + runFilter(_singleTyped, _qEqSingleTyped, bh); + } + + @Benchmark + public void filterEqSingleUntyped(Blackhole bh) { + runFilter(_singleUntyped, _qEqSingleUntyped, bh); + } + + @Benchmark + public void filterEqMultiTyped(Blackhole bh) { + runFilter(_multiTyped, _qEqMultiTyped, bh); + } + + @Benchmark + public void filterEqMultiUntyped(Blackhole bh) { + runFilter(_multiUntyped, _qEqMultiUntyped, bh); + } + + // ---------------- End-to-end "urnEntity(col) = literal" filter ---------------- + + @Benchmark + public void filterUrnEntitySingleTyped(Blackhole bh) { + runFilter(_singleTyped, _qUrnEntitySingleTyped, bh); + } + + @Benchmark + public void filterUrnEntitySingleUntyped(Blackhole bh) { + runFilter(_singleUntyped, _qUrnEntitySingleUntyped, bh); + } + + @Benchmark + public void filterUrnEntityMultiTyped(Blackhole bh) { + runFilter(_multiTyped, _qUrnEntityMultiTyped, bh); + } + + @Benchmark + public void filterUrnEntityMultiUntyped(Blackhole bh) { + runFilter(_multiUntyped, _qUrnEntityMultiUntyped, bh); + } + + // ---------------- Helpers ---------------- + + private static void runFilter(SegmentBundle bundle, QueryContext qc, Blackhole bh) { + SegmentContext sc = new SegmentContext(bundle._segment); + BaseFilterOperator op = new FilterPlanNode(sc, qc).run(); + BlockDocIdIterator it = op.nextBlock().getBlockDocIdSet().iterator(); + int count = 0; + int docId; + while ((docId = it.next()) != Constants.EOF) { + count++; + } + bh.consume(count); + } + + private static final class SegmentBundle { + private final File _parentDir; + private final File _segDir; + private final ImmutableSegment _segment; + private final Schema _schema; + private final String _urnColName; + + SegmentBundle(File parentDir, File segDir, ImmutableSegment segment, Schema schema, String urnColName) { + _parentDir = parentDir; + _segDir = segDir; + _segment = segment; + _schema = schema; + _urnColName = urnColName; + } + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java index afef239ff9..d8bad3eb8e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java @@ -137,7 +137,8 @@ public static RexLiteral toRexLiteral(RelBuilder builder, RexExpression.Literal return rexBuilder.makeTimestampLiteral(tsString, 1); } case JSON: - case STRING: { + case STRING: + case URN: { assert value != null; return rexBuilder.makeLiteral((String) value); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java index 072256fe01..edd3f5c068 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java @@ -150,6 +150,8 @@ public static Expressions.ColumnDataType convertColumnDataType(ColumnDataType da return Expressions.ColumnDataType.STRING; case JSON: return Expressions.ColumnDataType.JSON; + case URN: + return Expressions.ColumnDataType.STRING; case BYTES: return Expressions.ColumnDataType.BYTES; case MAP: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java index 8a1450b584..82fb2800d3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java @@ -93,6 +93,7 @@ private static SqlTypeName getSqlTypeName(FieldSpec fieldSpec) { return SqlTypeName.TIMESTAMP; case STRING: case JSON: + case URN: return SqlTypeName.VARCHAR; case BYTES: return SqlTypeName.VARBINARY; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/CompositeTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/CompositeTransformer.java index 8643dbd43c..c2d2794cdf 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/CompositeTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/CompositeTransformer.java @@ -110,6 +110,8 @@ public static List getDefaultTransformers(TableConfig tableCo addIfNotNoOp(transformers, new FilterTransformer(tableConfig)); addIfNotNoOp(transformers, new SchemaConformingTransformer(tableConfig, schema)); addIfNotNoOp(transformers, new DataTypeTransformer(tableConfig, schema)); + addIfNotNoOp(transformers, new UrnValidationTransformer(tableConfig, schema)); + addIfNotNoOp(transformers, new UrnKeyEncodingTransformer(schema)); addIfNotNoOp(transformers, new TimeValidationTransformer(tableConfig, schema)); addIfNotNoOp(transformers, new SpecialValueTransformer(schema)); addIfNotNoOp(transformers, new NullValueTransformer(tableConfig, schema)); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/SchemaConformingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/SchemaConformingTransformer.java index 43a08a2f62..640c0d6f72 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/SchemaConformingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/SchemaConformingTransformer.java @@ -591,6 +591,7 @@ private void putExtrasField(String fieldName, DataType fieldType, MapThree handling modes: + *
    + *
  • Single-entity typed + * ({@link DimensionFieldSpec#hasTypedUrnStorage()} but not polymorphic): the URN string + * is replaced with its numeric key. No transient field — one entity means no ambiguity.
    • + *
  • Multi-entity polymorphic, all-numeric + * ({@link DimensionFieldSpec#usesSortedUrnDictionary()}): URN string replaced with the LONG + * (or INT) numeric key, entity index stashed in {@code __urnEntityIdx}. Downstream + * the stats collector pairs (entityIdx, longKey) and the polymorphic creator lays out + * per-entity dict-ID ranges.
    • + *
  • Multi-entity polymorphic, mixed-type + * ({@link DimensionFieldSpec#usesPolymorphicUrnDictionary()} but not all-numeric): URN + * string replaced with the per-entity-declared typed key — {@link Long} for {@code LONG}, + * {@link Integer} for {@code INT}, {@link String} for {@code STRING}. Entity index stashed + * as above. The row's value type is now heterogeneous across rows of the same column, but + * the downstream consumers read it through the entity index so the type stays paired with + * its entity.
    • + *
+ * + *

If the URN is {@code null} or unparseable, both the main and transient fields are nulled; + * downstream {@link NullValueTransformer} fills defaults. Must run after + * {@link UrnValidationTransformer}, which guarantees URN-shaped input. + */ +public class UrnKeyEncodingTransformer implements RecordTransformer { + + private final List _encodings; + + public UrnKeyEncodingTransformer(Schema schema) { + List encodings = new ArrayList<>(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() != FieldSpec.DataType.URN + || !(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.usesPolymorphicUrnDictionary()) { + // Multi-entity polymorphic (numeric or mixed). Per-row encoding is type-driven by the + // matched entity's declared valueType. + encodings.add(new ColumnEncoding(dim.getName(), /*singleEntityStoredType=*/null, + dim.getUrnTypes(), dim.transientUrnEntityIdxFieldName())); + } else if (dim.hasTypedUrnStorage()) { + // Single-entity typed LONG/INT — replace URN string with numeric key, no transient field. + encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), + dim.getUrnTypes(), null)); + } + } + _encodings = Collections.unmodifiableList(encodings); + } + + @Override + public boolean isNoOp() { + return _encodings.isEmpty(); + } + + @Nullable + @Override + public GenericRow transform(GenericRow record) { + for (ColumnEncoding enc : _encodings) { + Object raw = record.getValue(enc._columnName); + if (raw == null) { + if (enc._entityIdxFieldName != null) { + record.putValue(enc._entityIdxFieldName, null); + } + continue; + } + + Urn urn = Urn.tryParse(raw.toString()); + if (urn == null || !urn.isSimple()) { + nullOut(record, enc); + continue; + } + + if (enc._entityIdxFieldName != null) { + // Polymorphic multi-entity path. + int entityIdx = -1; + String entityType = urn.entityPrefix(); + for (int i = 0; i < enc._urnTypes.size(); i++) { + if (enc._urnTypes.get(i).getEntityType().equals(entityType)) { + entityIdx = i; + break; + } + } + if (entityIdx < 0) { + nullOut(record, enc); + continue; + } + DimensionFieldSpec.UrnType matched = enc._urnTypes.get(entityIdx); + Object typedKey = parseKey(urn.simpleKey(), matched.getValueType()); + if (typedKey == null) { + nullOut(record, enc); + continue; + } + record.putValue(enc._entityIdxFieldName, entityIdx); + record.putValue(enc._columnName, typedKey); + } else if (enc._singleEntityStoredType != null) { + // Single-entity typed path. + Object typedKey = parseKey(urn.simpleKey(), enc._singleEntityStoredType); + if (typedKey == null) { + record.putValue(enc._columnName, null); + } else { + record.putValue(enc._columnName, typedKey); + } + } + } + return record; + } + + @Nullable + private static Object parseKey(String keyStr, @Nullable FieldSpec.DataType type) { + if (type == null) { + // Untyped entity in a polymorphic column — should not happen because + // usesPolymorphicUrnDictionary() requires every entity to declare a valueType. + return keyStr; + } + try { + switch (type) { + case LONG: + return Long.parseLong(keyStr); + case INT: + return Integer.parseInt(keyStr); + case STRING: + return keyStr; + default: + return null; + } + } catch (NumberFormatException e) { + return null; + } + } + + private static void nullOut(GenericRow record, ColumnEncoding enc) { + record.putValue(enc._columnName, null); + if (enc._entityIdxFieldName != null) { + record.putValue(enc._entityIdxFieldName, null); + } + } + + private static final class ColumnEncoding { + final String _columnName; + /** Single-entity typed path only. {@code null} for polymorphic columns. */ + @Nullable + final FieldSpec.DataType _singleEntityStoredType; + final List _urnTypes; + /** Transient entity-index field name for polymorphic columns. {@code null} for single-entity. */ + @Nullable + final String _entityIdxFieldName; + + ColumnEncoding(String columnName, @Nullable FieldSpec.DataType singleEntityStoredType, + List urnTypes, @Nullable String entityIdxFieldName) { + _columnName = columnName; + _singleEntityStoredType = singleEntityStoredType; + _urnTypes = urnTypes; + _entityIdxFieldName = entityIdxFieldName; + } + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformer.java new file mode 100644 index 0000000000..41f9ef090e --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformer.java @@ -0,0 +1,158 @@ +/** + * 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.segment.local.recordtransformer; + +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; +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.GenericRow; +import org.apache.pinot.spi.recordtransformer.RecordTransformer; +import org.apache.pinot.spi.utils.Urn; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Validates values in URN-typed columns and rejects non-URN-shaped input. + * + *

Two validation modes, selected per column based on its schema: + *

    + *
  • Declared urnTypes: the value must parse as a simple URN and its entity prefix + * must appear in the column's {@link DimensionFieldSpec#getUrnTypes()} list.
    • + *
  • Untyped URN (no {@code urnTypes}): the value must parse as a simple URN. This is + * the load-bearing constraint that lets {@link + * org.apache.pinot.segment.local.segment.creator.impl.SegmentColumnarIndexCreator} + * always auto-detect URN structure at segment build. Together with the per-segment + * structure metadata, it removes any need for the {@code urnEntity()} transform to parse + * URN strings at query time.
    • + *
+ * + *

Behavior on validation failure follows the table's {@code continueOnError} flag, mirroring + * Pinot's JSON ingest mode pattern: + *

    + *
  • {@code continueOnError = false} (default): throws {@link IllegalStateException}; the row + * is rejected and ingestion fails (the equivalent of JSON's {@code FAIL} mode).
    • + *
  • {@code continueOnError = true}: the value is replaced with {@code null} + * ({@link NullValueTransformer} fills the column default), the row is marked + * {@link GenericRow#INCOMPLETE_RECORD_KEY incomplete}, and ingestion continues — the + * equivalent of JSON's {@code SKIP} mode.
    • + *
+ * + *

NOTE: must run after {@link DataTypeTransformer} so that all values are already + * converted to their schema types, and before {@link NullValueTransformer} so that + * invalid values set to {@code null} are handled correctly. + */ +public class UrnValidationTransformer implements RecordTransformer { + private static final Logger LOGGER = LoggerFactory.getLogger(UrnValidationTransformer.class); + + /** Maps column name → set of allowed entity prefixes for typed columns (empty set means untyped). */ + private final Map> _columnAllowedPrefixes; + /** Untyped URN columns -- value must parse as a simple URN but any entity is accepted. */ + private final Set _untypedUrnColumns; + private final boolean _continueOnError; + + public UrnValidationTransformer(TableConfig tableConfig, Schema schema) { + Map> typed = new HashMap<>(); + Set untyped = new LinkedHashSet<>(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() != FieldSpec.DataType.URN + || !(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + List urnTypes = dim.getUrnTypes(); + if (urnTypes == null || urnTypes.isEmpty()) { + untyped.add(fieldSpec.getName()); + continue; + } + Set prefixes = new LinkedHashSet<>(); + for (DimensionFieldSpec.UrnType t : urnTypes) { + prefixes.add(t.getEntityType()); + } + typed.put(fieldSpec.getName(), Collections.unmodifiableSet(prefixes)); + } + _columnAllowedPrefixes = Collections.unmodifiableMap(typed); + _untypedUrnColumns = Collections.unmodifiableSet(untyped); + + IngestionConfig ingestionConfig = tableConfig.getIngestionConfig(); + _continueOnError = ingestionConfig != null && ingestionConfig.isContinueOnError(); + } + + @Override + public boolean isNoOp() { + return _columnAllowedPrefixes.isEmpty() && _untypedUrnColumns.isEmpty(); + } + + @Nullable + @Override + public GenericRow transform(GenericRow record) { + for (Map.Entry> entry : _columnAllowedPrefixes.entrySet()) { + validateAgainstAllowedPrefixes(record, entry.getKey(), entry.getValue()); + } + for (String column : _untypedUrnColumns) { + validateUntypedUrnShape(record, column); + } + return record; + } + + private void validateAgainstAllowedPrefixes(GenericRow record, String column, Set allowed) { + Object rawValue = record.getValue(column); + if (rawValue == null) { + return; + } + String valueStr = rawValue.toString(); + Urn urn = Urn.tryParse(valueStr); + if (urn == null || !allowed.contains(urn.entityPrefix())) { + reject(record, column, String.format( + "URN value '%s' in column '%s' does not match any declared urnTypes %s", valueStr, column, allowed)); + } + } + + private void validateUntypedUrnShape(GenericRow record, String column) { + Object rawValue = record.getValue(column); + if (rawValue == null) { + return; + } + String valueStr = rawValue.toString(); + Urn urn = Urn.tryParse(valueStr); + if (urn == null || !urn.isSimple()) { + reject(record, column, + String.format("URN value '%s' in untyped URN column '%s' is not a simple URN", valueStr, column)); + } + } + + private void reject(GenericRow record, String column, String errorMessage) { + if (_continueOnError) { + LOGGER.debug(errorMessage); + record.putValue(column, null); + record.putValue(GenericRow.INCOMPLETE_RECORD_KEY, true); + } else { + throw new IllegalStateException(errorMessage); + } + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java index a55ef41dfd..a1db012465 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java @@ -25,6 +25,7 @@ import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -37,6 +38,7 @@ import org.apache.pinot.common.utils.FileUtils; import org.apache.pinot.segment.local.io.util.PinotDataBitSet; import org.apache.pinot.segment.local.segment.creator.impl.nullvalue.NullValueVectorCreator; +import org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexPlugin; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; @@ -64,12 +66,14 @@ import org.apache.pinot.spi.data.ComplexFieldSpec; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DateTimeFormatSpec; +import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.FieldSpec.FieldType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.env.CommonsConfigurationUtils; +import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.TimeUtils; import org.joda.time.DateTimeZone; import org.joda.time.Interval; @@ -104,6 +108,8 @@ import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.SCHEMA_MAX_LENGTH_EXCEED_STRATEGY; import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.TOTAL_DOCS; import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.TOTAL_NUMBER_OF_ENTRIES; +import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.URN_ENTITY_DICT_RANGES; +import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.URN_TYPES; import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Column.getKeyFor; import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Segment.COMPLEX_COLUMNS; import static org.apache.pinot.segment.spi.V1Constants.MetadataKeys.Segment.DATETIME_COLUMNS; @@ -140,6 +146,17 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { */ private Map, IndexCreator>> _creatorsByColAndIndex = new HashMap<>(); private final Map _nullValueVectorCreatorMap = new HashMap<>(); + /** + * For untyped URN columns, records the entity types and per-entity value types discovered from + * the data at build time. Each entity gets {@code LONG} when every value for that entity parses + * as a long; otherwise {@code STRING}. The discovered list surfaces through + * {@link #resolveFieldSpec} so the column is treated as polymorphic for the rest of the build + * and load paths. + */ + private final Map> _discoveredMultiEntityUrnTypes = new HashMap<>(); + /** Cache for synthetic FieldSpecs produced by {@link #resolveFieldSpec} so per-build mutations + * (like {@link DimensionFieldSpec#setEntityDictRanges}) persist across resolve calls. */ + private final Map _syntheticFieldSpecCache = new HashMap<>(); private String _segmentName; private Schema _schema; private File _indexDir; @@ -230,17 +247,40 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio dictConfig = new DictionaryIndexConfig(dictConfig, columnIndexCreationInfo.isUseVarLengthDictionary()); } - SegmentDictionaryCreator creator = - new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); + // URN column preprocessing: detect single-prefix or multi-prefix patterns and choose the + // right dict creator. The standard path handles everything else. + Object dictValues = maybeStripUrnPrefix(columnName, fieldSpec, context.getSortedUniqueElementsArray()); + SegmentDictionaryCreator creator; + boolean usesPolymorphicUrnDict = dictValues instanceof org.apache.pinot.segment.local.segment.creator.impl.urn + .EntityKeyPair[] + || dictValues instanceof org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair[] + || (fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).usesPolymorphicUrnDictionary()); + if (usesPolymorphicUrnDict) { + List urnTypesForCreator = + _discoveredMultiEntityUrnTypes.containsKey(columnName) + ? _discoveredMultiEntityUrnTypes.get(columnName) + : ((DimensionFieldSpec) fieldSpec).getUrnTypes(); + creator = new UrnPolymorphicDictionaryCreator(fieldSpec, _indexDir, urnTypesForCreator); + } else { + creator = new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); + } try { - creator.build(context.getSortedUniqueElementsArray()); + creator.build(dictValues); } catch (Exception e) { LOGGER.error("Error building dictionary for field: {}, cardinality: {}, number of bytes per entry: {}", context.getFieldSpec().getName(), context.getCardinality(), creator.getNumBytesPerEntry()); throw e; } + // For multi-entity URN columns, surface the computed per-entity dict ID ranges so the + // metadata write can persist them under urn.entityDictRanges. + if (creator instanceof UrnPolymorphicDictionaryCreator) { + DimensionFieldSpec dim = (DimensionFieldSpec) resolveFieldSpec(columnName); + dim.setEntityDictRanges(((UrnPolymorphicDictionaryCreator) creator).getRanges()); + } + _dictionaryCreatorMap.put(columnName, creator); } @@ -283,6 +323,184 @@ private boolean isNullable(FieldSpec fieldSpec) { return _schema.isEnableColumnBasedNullHandling() ? fieldSpec.isNullable() : _config.isDefaultNullHandlingEnabled(); } + private FieldSpec resolveFieldSpec(String columnName) { + FieldSpec fs = _schema.getFieldSpecFor(columnName); + if (fs == null) { + return null; + } + if (fs instanceof DimensionFieldSpec && fs.getDataType() == DataType.URN + && !((DimensionFieldSpec) fs).hasUrnTypes()) { + DimensionFieldSpec cached = _syntheticFieldSpecCache.get(columnName); + if (cached != null) { + return cached; + } + DimensionFieldSpec original = (DimensionFieldSpec) fs; + // For untyped URN columns where multiple prefixes were discovered with numeric keys, + // surface synthetic multi-entity urnTypes so the column behaves like declared multi-entity + // typed when the segment is loaded. + List discovered = _discoveredMultiEntityUrnTypes.get(columnName); + if (discovered != null) { + DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, + original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), + original.getMaxLengthExceedStrategy()); + synthetic.setUrnTypes(discovered); + _syntheticFieldSpecCache.put(columnName, synthetic); + return synthetic; + } + } + return fs; + } + + /** + * Preprocesses URN-column dictionary values. Returns either: + *

    + *
  • {@code EntityKeyPair[]} or {@code EntityValuePair[]} when the column is declared + * polymorphic (the stats collector already produced these) — pass-through to {@link + * org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator}.
    • + *
  • {@code EntityValuePair[]} from auto-detection when the column is untyped but every + * value parses as a simple URN — feeds the polymorphic creator, with per-entity + * {@code valueType} chosen from the data (LONG when all keys for an entity parse as + * long, otherwise STRING).
    • + *
  • {@code values} unchanged for non-URN columns and untyped URN columns whose data isn't + * URN-shaped.
    • + *
+ * Side effect: populates {@link #_discoveredMultiEntityUrnTypes} so {@link #resolveFieldSpec} + * surfaces a synthetic polymorphic FieldSpec at metadata-write time. + */ + private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Object values) { + if (!(fieldSpec instanceof DimensionFieldSpec) || fieldSpec.getDataType() != DataType.URN) { + return values; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + // Declared polymorphic URN (single- or multi-entity): the stats collector already produced + // EntityKeyPair[] (all-numeric) or EntityValuePair[] (mixed-type). Pass straight through. + if (dim.usesPolymorphicUrnDictionary()) { + return values; + } + if (dim.hasUrnTypes() || !(values instanceof Object[])) { + return values; + } + Object[] arr = (Object[]) values; + if (arr.length == 0) { + return values; + } + // Untyped URN column: discover entities + per-entity value types from the data and build a + // polymorphic EntityValuePair[] for the polymorphic dictionary creator. Per entity: LONG + // when every value's key parses as a long, otherwise STRING. + return tryUrnAutoDetection(columnName, arr); + } + + /** + * Builds a sorted {@link org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair}{@code []} + * for an untyped URN column by parsing every value as a URN, grouping by entity prefix, and + * choosing a per-entity sub-dictionary type (LONG when all keys for that entity parse as long, + * otherwise STRING). Records the discovered urnTypes so {@link #resolveFieldSpec} surfaces a + * synthetic FieldSpec at metadata-write time. Returns the original {@code values} when any + * value isn't a simple URN — the column then falls back to plain STRING storage. + */ + private Object tryUrnAutoDetection(String columnName, Object[] values) { + java.util.LinkedHashMap entityToIdx = new java.util.LinkedHashMap<>(); + boolean[] perEntityAllLong = new boolean[8]; // grows as needed + Object[] parsedSuffixes = new Object[values.length]; // String suffix per row, null for skipped + int[] perRowEntityIdx = new int[values.length]; + java.util.Arrays.fill(perRowEntityIdx, -1); + + String nullSentinel = FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_URN; + int skipped = 0; + for (int rowIdx = 0; rowIdx < values.length; rowIdx++) { + String s = values[rowIdx].toString(); + // Filter the URN null sentinel before entity discovery — NullValueTransformer fills null + // URN rows with this value, and we don't want it to be promoted to a real entity. + if (nullSentinel.equals(s)) { + skipped++; + continue; + } + org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(s); + if (urn == null || !urn.isSimple()) { + return values; + } + String entityType = urn.entityPrefix(); + Integer existing = entityToIdx.get(entityType); + int idx; + if (existing != null) { + idx = existing; + } else { + idx = entityToIdx.size(); + entityToIdx.put(entityType, idx); + if (idx >= perEntityAllLong.length) { + perEntityAllLong = java.util.Arrays.copyOf(perEntityAllLong, perEntityAllLong.length * 2); + } + perEntityAllLong[idx] = true; + } + perRowEntityIdx[rowIdx] = idx; + parsedSuffixes[rowIdx] = urn.simpleKey(); + if (perEntityAllLong[idx]) { + try { + Long.parseLong(urn.simpleKey()); + } catch (NumberFormatException e) { + perEntityAllLong[idx] = false; + } + } + } + + int numEntities = entityToIdx.size(); + if (numEntities == 0) { + // Every value was the URN null sentinel; nothing to auto-detect. + return values; + } + List urnTypes = new ArrayList<>(numEntities); + DataType[] perEntityType = new DataType[numEntities]; + int e = 0; + for (String entityType : entityToIdx.keySet()) { + DataType vt = perEntityAllLong[e] ? DataType.LONG : DataType.STRING; + perEntityType[e] = vt; + urnTypes.add(DimensionFieldSpec.UrnType.of(entityType, vt)); + e++; + } + _discoveredMultiEntityUrnTypes.put(columnName, urnTypes); + + // Skipped null-sentinel rows still need a dict-ID slot in the forward index. Reserve one + // entry in entity 0's sub-dictionary with a type-appropriate placeholder so the per-row + // indexer can resolve them; NullValueVector marks the actual null status so the placeholder + // value is never user-visible. nullPlaceholder() picks Long.MIN_VALUE for LONG entity 0, + // Integer.MIN_VALUE for INT, the URN sentinel string itself for STRING. + Comparable nullPlaceholder = skipped > 0 ? nullPlaceholderFor(perEntityType[0], nullSentinel) : null; + + int pairCount = skipped > 0 ? values.length - skipped + 1 : values.length; + org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair[] pairs = + new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair[pairCount]; + int pi = 0; + for (int rowIdx = 0; rowIdx < values.length; rowIdx++) { + int idx = perRowEntityIdx[rowIdx]; + if (idx < 0) { + continue; // skipped null sentinel + } + String suffix = (String) parsedSuffixes[rowIdx]; + Comparable typedValue = perEntityType[idx] == DataType.LONG ? Long.valueOf(suffix) : suffix; + pairs[pi++] = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair(idx, typedValue); + } + if (nullPlaceholder != null) { + pairs[pi] = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair(0, nullPlaceholder); + } + java.util.Arrays.sort(pairs); + return pairs; + } + + private static Comparable nullPlaceholderFor(DataType entityZeroType, String nullSentinel) { + switch (entityZeroType) { + case LONG: + return Long.MIN_VALUE; + case INT: + return Integer.MIN_VALUE; + default: + return nullSentinel; + } + } + + /** + * Strips the discovered URN entity prefix from a row's value at index time so the dictionary + * lookup hits the suffix-only entry. + */ private FieldIndexConfigs adaptConfig(String columnName, FieldIndexConfigs config, ColumnIndexCreationInfo columnIndexCreationInfo, SegmentGeneratorConfig segmentCreationSpec) { FieldIndexConfigs.Builder builder = new FieldIndexConfigs.Builder(config); @@ -366,11 +584,73 @@ public void indexRow(GenericRow row) if (columnValueToIndex == null) { throw new RuntimeException("Null value for column:" + columnName); } - Map, IndexCreator> creatorsByIndex = byColEntry.getValue(); - FieldSpec fieldSpec = _schema.getFieldSpecFor(columnName); + FieldSpec fieldSpec = resolveFieldSpec(columnName); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); + // For multi-entity URN columns, pair the row's main column LONG with the entity index the + // URN encoding transformer stashed in a transient field so the UrnPolymorphicDictionaryCreator + // can map (entityIdx, longKey) → dict ID. The transient field is dropped after build. + // + // For *auto-detected* multi-entity URN columns (untyped schema, multi-prefix discovered), + // the URN encoding transformer didn't run -- the row's value is still a URN string. We + // parse it inline here. + if (dictionaryCreator instanceof UrnPolymorphicDictionaryCreator) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + Object entityIdxVal = row.getValue(dim.transientUrnEntityIdxFieldName()); + if (entityIdxVal != null) { + // Declared multi-entity path: UrnKeyEncodingTransformer already stashed entityIdx and + // converted the column value to the per-entity-typed key (Long for LONG, Integer for + // INT, String for STRING). Wrap it as an EntityValuePair so the polymorphic creator + // can look it up against its per-entity sub-dictionary. + int entityIdx = ((Number) entityIdxVal).intValue(); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( + entityIdx, (Comparable) columnValueToIndex); + } else if (!(columnValueToIndex instanceof String) && columnValueToIndex instanceof Comparable) { + // Null URN path: the row's URN was null, UrnKeyEncodingTransformer cleared the + // transient entity index, and NullValueTransformer filled the column with the default + // null sentinel. Stats collector associated the sentinel with entity 0; do the same + // here so indexOfSV finds it. + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( + 0, (Comparable) columnValueToIndex); + } else if (columnValueToIndex instanceof String) { + // Auto-detected polymorphic path: UrnKeyEncodingTransformer didn't run (column was + // untyped at schema time) — parse the URN here against the discovered urnTypes and + // emit an EntityValuePair carrying the entity's declared typed key (Long for LONG, + // String for STRING). Null URN rows arrive here as the URN sentinel string; map them + // to entity 0's null-placeholder reserved by tryUrnAutoDetection. + String s = (String) columnValueToIndex; + if (FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_URN.equals(s)) { + List urnTypes = dim.getUrnTypes(); + if (urnTypes != null && !urnTypes.isEmpty()) { + DataType entityZeroType = urnTypes.get(0).getValueType(); + Comparable placeholder = nullPlaceholderFor(entityZeroType, s); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( + 0, placeholder); + } + } else { + org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(s); + if (urn != null && urn.isSimple()) { + List urnTypes = dim.getUrnTypes(); + int entityIdx = -1; + String entityType = urn.entityPrefix(); + for (int i = 0; i < urnTypes.size(); i++) { + if (urnTypes.get(i).getEntityType().equals(entityType)) { + entityIdx = i; + break; + } + } + if (entityIdx >= 0) { + DataType valueType = urnTypes.get(entityIdx).getValueType(); + Comparable typedKey = + valueType == DataType.LONG ? (Comparable) Long.valueOf(urn.simpleKey()) : urn.simpleKey(); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( + entityIdx, typedKey); + } + } + } + } + } try { if (fieldSpec.isSingleValueField()) { indexSingleValueRow(dictionaryCreator, columnValueToIndex, creatorsByIndex); @@ -404,7 +684,7 @@ public void indexColumn(String columnName, @Nullable int[] sortedDocIds, IndexSe try (PinotSegmentColumnReader colReader = new PinotSegmentColumnReader(segment, columnName)) { Map, IndexCreator> creatorsByIndex = _creatorsByColAndIndex.get(columnName); NullValueVectorCreator nullVec = _nullValueVectorCreatorMap.get(columnName); - FieldSpec fieldSpec = _schema.getFieldSpecFor(columnName); + FieldSpec fieldSpec = resolveFieldSpec(columnName); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); if (sortedDocIds != null) { int onDiskDocId = 0; @@ -604,7 +884,7 @@ private void writeMetadata() ColumnIndexCreationInfo columnIndexCreationInfo = entry.getValue(); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(column); int dictionaryElementSize = (dictionaryCreator != null) ? dictionaryCreator.getNumBytesPerEntry() : 0; - addColumnMetadataInfo(properties, column, columnIndexCreationInfo, _totalDocs, _schema.getFieldSpecFor(column), + addColumnMetadataInfo(properties, column, columnIndexCreationInfo, _totalDocs, resolveFieldSpec(column), dictionaryCreator != null, dictionaryElementSize); } @@ -666,6 +946,26 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str properties.setProperty(getKeyFor(column, DATETIME_GRANULARITY), dateTimeFieldSpec.getGranularity()); } + // URN field: persist declared entity types so they survive a segment reload + if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dimSpec = (DimensionFieldSpec) fieldSpec; + if (dimSpec.hasUrnTypes()) { + try { + properties.setProperty(getKeyFor(column, URN_TYPES), JsonUtils.objectToString(dimSpec.getUrnTypes())); + } catch (Exception e) { + LOGGER.warn("Failed to serialize urnTypes for column {}: {}", column, e.getMessage()); + } + } + // Multi-entity sorted-LONG UrnDictionary layout: persist per-entity dict ID ranges. + List entityDictRanges = dimSpec.getEntityDictRanges(); + if (entityDictRanges != null && !entityDictRanges.isEmpty()) { + String encoded = DimensionFieldSpec.serializeEntityDictRanges(entityDictRanges); + if (encoded != null) { + properties.setProperty(getKeyFor(column, URN_ENTITY_DICT_RANGES), encoded); + } + } + } + // complex field if (fieldSpec.getFieldType() == FieldType.COMPLEX) { ComplexFieldSpec complexFieldSpec = (ComplexFieldSpec) fieldSpec; @@ -677,16 +977,21 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str } // NOTE: Min/max could be null for real-time aggregate metrics. - if ((fieldSpec.getFieldType() != FieldType.COMPLEX) && (totalDocs > 0)) { + // Skip min/max for multi-entity URN columns: cross-entity comparisons aren't semantically + // meaningful, and the stats values (collected as strings for auto-detected columns) won't + // round-trip through the LONG effective-stored-type parsing at load time. + boolean isMultiEntityUrn = fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).usesPolymorphicUrnDictionary(); + if ((fieldSpec.getFieldType() != FieldType.COMPLEX) && (totalDocs > 0) && !isMultiEntityUrn) { Object min = columnIndexCreationInfo.getMin(); Object max = columnIndexCreationInfo.getMax(); if (min != null && max != null) { - addColumnMinMaxValueInfo(properties, column, min, max, dataType.getStoredType()); + addColumnMinMaxValueInfo(properties, column, min, max, fieldSpec.getEffectiveStoredType()); } } String defaultNullValue = columnIndexCreationInfo.getDefaultNullValue().toString(); - if (dataType.getStoredType() == DataType.STRING) { + if (fieldSpec.getEffectiveStoredType() == DataType.STRING) { // NOTE: Do not limit length of default null value because we need exact value to determine whether the default // null value changes defaultNullValue = CommonsConfigurationUtils.replaceSpecialCharacterInPropertyValue(defaultNullValue); @@ -734,6 +1039,26 @@ public static void addFieldSpec(PropertiesConfiguration properties, String colum properties.setProperty(getKeyFor(column, DATETIME_GRANULARITY), dateTimeFieldSpec.getGranularity()); } + // URN field: persist declared entity types so they survive a segment reload + if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dimSpec = (DimensionFieldSpec) fieldSpec; + if (dimSpec.hasUrnTypes()) { + try { + properties.setProperty(getKeyFor(column, URN_TYPES), JsonUtils.objectToString(dimSpec.getUrnTypes())); + } catch (Exception e) { + LOGGER.warn("Failed to serialize urnTypes for column {}: {}", column, e.getMessage()); + } + } + // Multi-entity sorted-LONG UrnDictionary layout: persist per-entity dict ID ranges. + List entityDictRanges = dimSpec.getEntityDictRanges(); + if (entityDictRanges != null && !entityDictRanges.isEmpty()) { + String encoded = DimensionFieldSpec.serializeEntityDictRanges(entityDictRanges); + if (encoded != null) { + properties.setProperty(getKeyFor(column, URN_ENTITY_DICT_RANGES), encoded); + } + } + } + // complex field if (fieldSpec.getFieldType() == FieldType.COMPLEX) { ComplexFieldSpec complexFieldSpec = (ComplexFieldSpec) fieldSpec; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java index 2e4db2a6e6..db7f7f065a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java @@ -81,7 +81,7 @@ public SegmentDictionaryCreator(String columnName, DataType storedType, File ind public SegmentDictionaryCreator(FieldSpec fieldSpec, File indexDir, boolean useVarLengthDictionary) { _columnName = fieldSpec.getName(); - _storedType = fieldSpec.getDataType().getStoredType(); + _storedType = fieldSpec.getEffectiveStoredType(); _dictionaryFile = new File(indexDir, _columnName + DictionaryIndexType.getFileExtension()); _useVarLengthDictionary = useVarLengthDictionary; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java index 612b696e27..4958812836 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentIndexCreationDriverImpl.java @@ -528,7 +528,7 @@ void collectStatsAndIndexCreationInfo() } String column = fieldSpec.getName(); - DataType storedType = fieldSpec.getDataType().getStoredType(); + DataType storedType = fieldSpec.getEffectiveStoredType(); ColumnStatistics columnProfile = _segmentStats.getColumnProfileFor(column); DictionaryIndexConfig dictionaryIndexConfig = indexConfigsMap.get(column).getConfig(StandardIndexes.dictionary()); boolean createDictionary = dictionaryIndexConfig.isDisabled(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java index 45c75e0c49..acf3dfb09a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/SegmentPreIndexStatsCollectorImpl.java @@ -23,6 +23,7 @@ import org.apache.pinot.segment.spi.creator.ColumnStatistics; import org.apache.pinot.segment.spi.creator.SegmentPreIndexStatsCollector; import org.apache.pinot.segment.spi.creator.StatsCollectorConfig; +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.GenericRow; @@ -48,7 +49,20 @@ public void init() { Schema dataSchema = _statsCollectorConfig.getSchema(); for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) { String column = fieldSpec.getName(); - switch (fieldSpec.getDataType().getStoredType()) { + // Polymorphic multi-entity URN columns route to the URN collector regardless of stored + // type — they track (entityIdx, value) pairs so the polymorphic dictionary creator can + // lay out per-entity dict ID ranges with the right sub-dict types. Covers both the + // all-numeric (effectiveStoredType=LONG) and mixed-type (effectiveStoredType=STRING) + // configurations. + if (fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).usesPolymorphicUrnDictionary()) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + _columnStatsCollectorMap.put(column, new UrnMultiEntityPreIndexStatsCollector(column, + dim.transientUrnEntityIdxFieldName(), _statsCollectorConfig)); + continue; + } + // For typed URN columns the physical stored type may differ from the logical type + switch (fieldSpec.getEffectiveStoredType()) { case INT: _columnStatsCollectorMap.put(column, new IntColumnPreIndexStatsCollector(column, _statsCollectorConfig)); break; @@ -98,13 +112,21 @@ public void collectRow(GenericRow row) { final String columnName = columnNameAndValue.getKey(); final Object value = columnNameAndValue.getValue(); - if (_columnStatsCollectorMap.containsKey(columnName)) { - try { - _columnStatsCollectorMap.get(columnName).collect(value); - } catch (Exception e) { - LOGGER.error("Exception while collecting stats for column:{} in row:{}", columnName, row); - throw e; + AbstractColumnStatisticsCollector collector = _columnStatsCollectorMap.get(columnName); + if (collector == null) { + continue; + } + try { + // Multi-entity URN collectors need the row's companion column value alongside the main + // column's LONG key to pair (entityIdx, longKey) for the sorted-LONG dictionary. + if (collector instanceof UrnMultiEntityPreIndexStatsCollector) { + ((UrnMultiEntityPreIndexStatsCollector) collector).collect(value, row); + } else { + collector.collect(value); } + } catch (Exception e) { + LOGGER.error("Exception while collecting stats for column:{} in row:{}", columnName, row); + throw e; } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/UrnMultiEntityPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/UrnMultiEntityPreIndexStatsCollector.java new file mode 100644 index 0000000000..9a4c1716ca --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/UrnMultiEntityPreIndexStatsCollector.java @@ -0,0 +1,149 @@ +/** + * 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.segment.local.segment.creator.impl.stats; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; +import org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair; +import org.apache.pinot.segment.spi.creator.StatsCollectorConfig; +import org.apache.pinot.spi.data.readers.GenericRow; + + +/** + * Stats collector for multi-entity URN columns using the polymorphic sub-dictionary layout. + * + *

Per row, the column carries the per-entity-typed key value ({@link Long} for LONG entities, + * {@link Integer} for INT, {@link String} for STRING — set by + * {@link org.apache.pinot.segment.local.recordtransformer.UrnKeyEncodingTransformer}) and the + * companion transient field carries the 0-based entity index. This collector dedupes by + * {@code (entityIdx, value)} so two URNs that share a key but belong to different entity types + * remain distinct dictionary entries. + * + *

At {@link #seal()} the unique pairs are sorted by {@code (entityIdx, value-natural-order)}, + * which is the order + * {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator} + * expects: entities in declared order so each entity's sub-dictionary occupies a contiguous + * dict-ID and byte range, and within each entity, values in natural order so the per-entity + * sub-dictionary (LongDictionary, StringDictionary, ...) can use its native ordered lookup. + */ +public class UrnMultiEntityPreIndexStatsCollector extends AbstractColumnStatisticsCollector { + private final String _transientEntityIdxFieldName; + private final Set _pairs = new HashSet<>(); + private EntityValuePair[] _sortedPairs; + private int _totalNumberOfEntries; + private boolean _sealed = false; + + public UrnMultiEntityPreIndexStatsCollector(String column, String transientEntityIdxFieldName, + StatsCollectorConfig statsCollectorConfig) { + super(column, statsCollectorConfig); + _transientEntityIdxFieldName = transientEntityIdxFieldName; + } + + /** + * Collects a row's main-column typed key value paired with the entity index stashed by + * {@link org.apache.pinot.segment.local.recordtransformer.UrnKeyEncodingTransformer} in the + * transient row field. + */ + public void collect(Object value, GenericRow row) { + assert !_sealed; + if (value == null) { + return; + } + Object entityIdxVal = row.getValue(_transientEntityIdxFieldName); + // A null entity index means the row's URN was null and NullValueTransformer filled the + // column with the default null sentinel. Associate that sentinel with entity 0 so it gets + // a dict ID -- the row is also tracked in the column's NullValueVector, but the forward + // index still needs a dict-ID slot per doc. + int entityIdx = entityIdxVal == null ? 0 : ((Number) entityIdxVal).intValue(); + _pairs.add(new EntityValuePair(entityIdx, (Comparable) value)); + _totalNumberOfEntries++; + } + + /** + * Default per-column collector path used by code that doesn't have row context (e.g. realtime + * to-immutable conversion). The polymorphic URN column is normally routed through + * {@link #collect(Object, GenericRow)} via {@code SegmentPreIndexStatsCollectorImpl.collectRow} + * — falling back here would lose the {@code (entityIdx, value)} pairing, so we no-op rather + * than silently corrupting the dictionary. + */ + @Override + public void collect(Object entry) { + assert !_sealed; + } + + /** + * Returns the sorted unique {@code (entityIdx, value)} pairs after {@link #seal()}. Consumed + * by {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator}. + */ + @Override + public Object getUniqueValuesSet() { + if (!_sealed) { + throw new IllegalStateException("collector must be sealed before unique values are requested"); + } + return _sortedPairs; + } + + @Override + public int getCardinality() { + return _sealed ? _sortedPairs.length : _pairs.size(); + } + + @Override + public int getLengthOfLargestElement() { + return 0; + } + + @Override + public int getTotalNumberOfEntries() { + return _totalNumberOfEntries; + } + + @Override + public Object getMinValue() { + return null; + } + + @Override + public Object getMaxValue() { + return null; + } + + @Override + public void seal() { + if (_sealed) { + return; + } + _sortedPairs = _pairs.toArray(new EntityValuePair[0]); + Arrays.sort(_sortedPairs); + _pairs.clear(); + _sealed = true; + } + + /** + * Polymorphic multi-entity URN columns are never considered sorted for forward-index purposes: + * the value space is partitioned by entity into typed regions, so consecutive rows may jump + * across entity ranges in dict-ID space even when per-entity values would themselves be + * monotone. + */ + @Override + public boolean isSorted() { + return false; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityKeyPair.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityKeyPair.java new file mode 100644 index 0000000000..3612f83393 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityKeyPair.java @@ -0,0 +1,71 @@ +/** + * 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.segment.local.segment.creator.impl.urn; + +/** + * A single dictionary entry for a multi-entity URN column: the entity index (into the declared + * {@code urnTypes} list) paired with the numeric key. Naturally ordered by {@code (entityIdx, + * longKey)}, which is also the order the on-disk dictionary buffer is written in. This ordering + * lets {@code UrnSortedDictionary} use a range-restricted binary search at read time. + */ +public final class EntityKeyPair implements Comparable { + private final int _entityIdx; + private final long _longKey; + + public EntityKeyPair(int entityIdx, long longKey) { + _entityIdx = entityIdx; + _longKey = longKey; + } + + public int getEntityIdx() { + return _entityIdx; + } + + public long getLongKey() { + return _longKey; + } + + @Override + public int compareTo(EntityKeyPair other) { + int byEntity = Integer.compare(_entityIdx, other._entityIdx); + return byEntity != 0 ? byEntity : Long.compare(_longKey, other._longKey); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof EntityKeyPair)) { + return false; + } + EntityKeyPair other = (EntityKeyPair) o; + return _entityIdx == other._entityIdx && _longKey == other._longKey; + } + + @Override + public int hashCode() { + return _entityIdx * 31 + Long.hashCode(_longKey); + } + + @Override + public String toString() { + return "(" + _entityIdx + "," + _longKey + ")"; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityValuePair.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityValuePair.java new file mode 100644 index 0000000000..443e8e9637 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityValuePair.java @@ -0,0 +1,79 @@ +/** + * 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.segment.local.segment.creator.impl.urn; + +/** + * A single dictionary entry for a polymorphic multi-entity URN column: the entity index (into + * the declared {@code urnTypes} list) paired with the per-entity typed key value. Unlike + * {@link EntityKeyPair} (always a {@code long}), the value here may be {@link Long}, {@link + * Integer}, {@link String}, etc. — whichever type matches the entity's declared {@code + * valueType}. + * + *

Naturally ordered by {@code (entityIdx, value-in-natural-order)}, which is also the order + * the on-disk dictionary buffer is written in: entities appear in declared order, and within an + * entity, values are sorted by their type-appropriate {@link Comparable} ordering. This lets + * each per-entity sub-dictionary use its native lookup (binary search on LONG for numeric + * entities, byte-comparison on STRING for textual entities) without cross-entity confusion. + */ +public final class EntityValuePair implements Comparable { + private final int _entityIdx; + private final Comparable _value; + + @SuppressWarnings("unchecked") + public EntityValuePair(int entityIdx, Comparable value) { + _entityIdx = entityIdx; + _value = (Comparable) value; + } + + public int getEntityIdx() { + return _entityIdx; + } + + public Comparable getValue() { + return _value; + } + + @Override + public int compareTo(EntityValuePair other) { + int byEntity = Integer.compare(_entityIdx, other._entityIdx); + return byEntity != 0 ? byEntity : _value.compareTo(other._value); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof EntityValuePair)) { + return false; + } + EntityValuePair other = (EntityValuePair) o; + return _entityIdx == other._entityIdx && _value.equals(other._value); + } + + @Override + public int hashCode() { + return _entityIdx * 31 + _value.hashCode(); + } + + @Override + public String toString() { + return "(" + _entityIdx + "," + _value + ")"; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreator.java new file mode 100644 index 0000000000..ae82e5e150 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreator.java @@ -0,0 +1,335 @@ +/** + * 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.segment.local.segment.creator.impl.urn; + +import com.google.common.base.Preconditions; +import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap; +import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap; +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.io.File; +import java.io.IOException; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; +import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.DimensionFieldSpec.EntityDictRange; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Dictionary creator for multi-entity URN columns. Writes a polymorphic dictionary file: each + * declared entity gets its own typed sub-dictionary region (LongDictionary for {@code LONG} + * entities, IntDictionary for {@code INT}, StringDictionary for {@code STRING}), concatenated + * back-to-back in declared-entity order. The companion {@link + * org.apache.pinot.segment.local.segment.index.readers.UrnPolymorphicDictionary} reader slices + * the file at per-entity byte boundaries and wraps each region in the matching scalar + * dictionary class. + * + *

Wire-format compatibility: when every declared entity uses {@code LONG} (the only + * configuration supported before the polymorphic wire format existed), this creator produces a + * byte-identical dictionary file and metadata to the original sorted-LONG layout — the + * concatenated per-entity LONG regions form a single sorted-LONG buffer, and {@link + * DimensionFieldSpec#serializeEntityDictRanges} writes the legacy 2-field metadata form. + * + *

Build input is a sorted {@link EntityValuePair}{@code []}, sorted by {@code (entityIdx, + * value)} using the value's natural ordering — analogous to the sorted {@link EntityKeyPair}{@code + * []} the legacy creator accepted, but with mixed-type values. For backwards compatibility, the + * legacy {@code EntityKeyPair[]} input is also accepted and treated as all-LONG. + */ +public class UrnPolymorphicDictionaryCreator extends SegmentDictionaryCreator { + + private static final Logger LOGGER = LoggerFactory.getLogger(UrnPolymorphicDictionaryCreator.class); + + private final String _columnName; + private final File _dictionaryFile; + private final List _urnTypes; + + // Computed at build(): + private List _ranges; + private int[] _perEntityStartDictId; + // Per-entity value→relative-dictId map for indexOfSV. Type varies by dictClass: Long2IntOpenHashMap + // for LONG, Int2IntOpenHashMap for INT, Object2IntOpenHashMap for STRING. + private Object[] _perEntityValueMaps; + private int _numEntries; + private int _maxNumBytesPerEntry; + + public UrnPolymorphicDictionaryCreator(FieldSpec fieldSpec, File indexDir, + List urnTypes) { + super(fieldSpec, indexDir, false); + _columnName = fieldSpec.getName(); + _dictionaryFile = new File(indexDir, _columnName + DictionaryIndexType.getFileExtension()); + _urnTypes = urnTypes; + } + + @Override + public void build(Object sortedValues) + throws IOException { + EntityValuePair[] pairs; + if (sortedValues instanceof EntityKeyPair[]) { + // Legacy LONG-only input: lift each pair to an EntityValuePair with a Long value. + EntityKeyPair[] legacy = (EntityKeyPair[]) sortedValues; + pairs = new EntityValuePair[legacy.length]; + for (int i = 0; i < legacy.length; i++) { + pairs[i] = new EntityValuePair(legacy[i].getEntityIdx(), legacy[i].getLongKey()); + } + } else { + Preconditions.checkArgument(sortedValues instanceof EntityValuePair[], + "UrnPolymorphicDictionaryCreator expects EntityValuePair[] or EntityKeyPair[]"); + pairs = (EntityValuePair[]) sortedValues; + } + // Empty input is valid -- a column whose only rows are null URN values has no non-null + // dictionary entries. We still need to emit an empty file plus a zero-width range per entity + // so the loader sees the column as a well-formed (empty) polymorphic dictionary. + if (pairs.length == 0) { + buildEmpty(); + return; + } + + int numEntities = _urnTypes.size(); + DataType[] dictClasses = new DataType[numEntities]; + int[] numBytesPerValue = new int[numEntities]; + for (int e = 0; e < numEntities; e++) { + // Sub-dictionary class is the entity's declared valueType. URN UrnType.validate() restricts + // valueType to INT/LONG/STRING (null only happens during build-time auto-detection, and the + // caller resolves it before reaching this creator). + DataType vt = _urnTypes.get(e).getValueType(); + dictClasses[e] = vt != null ? vt : DataType.STRING; + switch (dictClasses[e]) { + case INT: + numBytesPerValue[e] = Integer.BYTES; + break; + case LONG: + numBytesPerValue[e] = Long.BYTES; + break; + case STRING: + // Computed below from actual data; min 1 byte. + numBytesPerValue[e] = 1; + break; + default: + throw new IllegalStateException( + "Unsupported URN sub-dictionary type: " + dictClasses[e] + " for entity " + _urnTypes.get(e)); + } + } + + // Pass 1: count per-entity entries; compute STRING padding width from longest value. + int[] perEntityCount = new int[numEntities]; + for (EntityValuePair p : pairs) { + int e = p.getEntityIdx(); + perEntityCount[e]++; + if (dictClasses[e] == DataType.STRING) { + int len = ((String) p.getValue()).getBytes(StandardCharsets.UTF_8).length; + if (len > numBytesPerValue[e]) { + numBytesPerValue[e] = len; + } + } + } + + // Pass 2: per-entity start dict IDs and byte offsets (cumulative). + int[] perEntityStartDictId = new int[numEntities + 1]; + long[] perEntityStartByte = new long[numEntities + 1]; + for (int e = 0; e < numEntities; e++) { + perEntityStartDictId[e + 1] = perEntityStartDictId[e] + perEntityCount[e]; + perEntityStartByte[e + 1] = perEntityStartByte[e] + (long) perEntityCount[e] * numBytesPerValue[e]; + } + _perEntityStartDictId = perEntityStartDictId; + _numEntries = perEntityStartDictId[numEntities]; + for (int n : numBytesPerValue) { + if (n > _maxNumBytesPerEntry) { + _maxNumBytesPerEntry = n; + } + } + + // Pass 3: allocate per-entity value maps for indexOfSV lookups. + _perEntityValueMaps = new Object[numEntities]; + for (int e = 0; e < numEntities; e++) { + switch (dictClasses[e]) { + case LONG: + Long2IntOpenHashMap longMap = new Long2IntOpenHashMap(perEntityCount[e]); + longMap.defaultReturnValue(-1); + _perEntityValueMaps[e] = longMap; + break; + case INT: + Int2IntOpenHashMap intMap = new Int2IntOpenHashMap(perEntityCount[e]); + intMap.defaultReturnValue(-1); + _perEntityValueMaps[e] = intMap; + break; + case STRING: + Object2IntOpenHashMap strMap = new Object2IntOpenHashMap<>(perEntityCount[e]); + strMap.defaultReturnValue(-1); + _perEntityValueMaps[e] = strMap; + break; + default: + throw new IllegalStateException("Unreachable: dictClass already validated above"); + } + } + + // Pass 4: write the file. The buffer is a single big-endian file containing each entity's + // typed sub-dictionary back-to-back. Per-entity layout is identical to what the corresponding + // standalone dictionary creator would write (LONG/INT: fixed-width values; STRING: null-padded + // fixed-width UTF-8 bytes), which is what lets the reader wrap each region in a stock + // LongDictionary / IntDictionary / StringDictionary without any URN-specific decoding. + FileUtils.touch(_dictionaryFile); + long totalBytes = perEntityStartByte[numEntities]; + int[] perEntityWriteCounter = new int[numEntities]; + int currentEntity = -1; + try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapFile(_dictionaryFile, false, 0, totalBytes, + ByteOrder.BIG_ENDIAN, getClass().getSimpleName())) { + for (EntityValuePair p : pairs) { + int e = p.getEntityIdx(); + Preconditions.checkState(e >= currentEntity, + "UrnPolymorphicDictionaryCreator input must be sorted by entityIdx"); + currentEntity = e; + int relDictId = perEntityWriteCounter[e]++; + long byteOffset = perEntityStartByte[e] + (long) relDictId * numBytesPerValue[e]; + switch (dictClasses[e]) { + case LONG: { + long longVal = ((Number) p.getValue()).longValue(); + dataBuffer.putLong(byteOffset, longVal); + ((Long2IntOpenHashMap) _perEntityValueMaps[e]).put(longVal, relDictId); + break; + } + case INT: { + int intVal = ((Number) p.getValue()).intValue(); + dataBuffer.putInt(byteOffset, intVal); + ((Int2IntOpenHashMap) _perEntityValueMaps[e]).put(intVal, relDictId); + break; + } + case STRING: { + String strVal = (String) p.getValue(); + byte[] utf8 = strVal.getBytes(StandardCharsets.UTF_8); + for (int i = 0; i < utf8.length; i++) { + dataBuffer.putByte(byteOffset + i, utf8[i]); + } + // Null-pad to numBytesPerValue (the convention StringDictionary expects). + for (int i = utf8.length; i < numBytesPerValue[e]; i++) { + dataBuffer.putByte(byteOffset + i, (byte) 0); + } + @SuppressWarnings("unchecked") + Object2IntOpenHashMap strMap = (Object2IntOpenHashMap) _perEntityValueMaps[e]; + strMap.put(strVal, relDictId); + break; + } + default: + throw new IllegalStateException("Unreachable"); + } + } + } + + // Pass 5: build the EntityDictRange list with cumulative end dict IDs plus per-entity + // dictClass and numBytesPerValue. Even entities with zero rows get a range so the reader + // sees a uniform lookup table. + _ranges = new ArrayList<>(numEntities); + for (int e = 0; e < numEntities; e++) { + _ranges.add(new EntityDictRange( + _urnTypes.get(e).getEntityType(), + perEntityStartDictId[e + 1], + dictClasses[e], + numBytesPerValue[e])); + } + + LOGGER.info("Built UrnPolymorphicDictionary for column: {} with cardinality: {} across {} entities", + _columnName, _numEntries, numEntities); + } + + private void buildEmpty() + throws IOException { + FileUtils.touch(_dictionaryFile); + int numEntities = _urnTypes.size(); + _perEntityStartDictId = new int[numEntities + 1]; + _perEntityValueMaps = new Object[numEntities]; + _ranges = new ArrayList<>(numEntities); + for (int e = 0; e < numEntities; e++) { + DataType vt = _urnTypes.get(e).getValueType(); + DataType dictClass = vt != null ? vt : DataType.STRING; + int width; + switch (dictClass) { + case INT: + width = Integer.BYTES; + break; + case LONG: + width = Long.BYTES; + break; + default: + width = 1; + break; + } + _ranges.add(new EntityDictRange(_urnTypes.get(e).getEntityType(), 0, dictClass, width)); + } + _numEntries = 0; + _maxNumBytesPerEntry = 0; + } + + /** {@inheritDoc} Accepts {@link EntityValuePair} or — for the all-LONG legacy path — {@link EntityKeyPair}. */ + @Override + public int indexOfSV(Object value) { + Preconditions.checkState(_perEntityValueMaps != null, "build() must be called before indexOfSV()"); + int entityIdx; + Object key; + if (value instanceof EntityValuePair) { + EntityValuePair p = (EntityValuePair) value; + entityIdx = p.getEntityIdx(); + key = p.getValue(); + } else if (value instanceof EntityKeyPair) { + EntityKeyPair p = (EntityKeyPair) value; + entityIdx = p.getEntityIdx(); + key = p.getLongKey(); + } else { + throw new IllegalArgumentException("Expected EntityValuePair or EntityKeyPair, got: " + value); + } + int relDictId; + Object map = _perEntityValueMaps[entityIdx]; + if (map instanceof Long2IntOpenHashMap) { + relDictId = ((Long2IntOpenHashMap) map).get(((Number) key).longValue()); + } else if (map instanceof Int2IntOpenHashMap) { + relDictId = ((Int2IntOpenHashMap) map).get(((Number) key).intValue()); + } else { + @SuppressWarnings("unchecked") + Object2IntOpenHashMap strMap = (Object2IntOpenHashMap) map; + relDictId = strMap.getInt((String) key); + } + Preconditions.checkState(relDictId >= 0, + "Row references (entity=%s, key=%s) that wasn't seen during stats collection", + _urnTypes.get(entityIdx).getEntityType(), key); + return _perEntityStartDictId[entityIdx] + relDictId; + } + + /** + * Returns the maximum per-entry byte width across all sub-dictionaries. Used by callers that + * size buffers based on the worst-case entry width — they treat the dictionary as if it were + * a single uniform-width buffer. + */ + @Override + public int getNumBytesPerEntry() { + return _maxNumBytesPerEntry; + } + + /** Returns the per-entity dictionary-ID ranges computed by {@link #build}. */ + public List getRanges() { + Preconditions.checkState(_ranges != null, "build() must be called before getRanges()"); + return _ranges; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java index ed04808e14..b87385b2af 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java @@ -49,6 +49,7 @@ import org.apache.pinot.segment.local.segment.index.readers.OnHeapLongDictionary; import org.apache.pinot.segment.local.segment.index.readers.OnHeapStringDictionary; import org.apache.pinot.segment.local.segment.index.readers.StringDictionary; +import org.apache.pinot.segment.local.segment.index.readers.UrnPolymorphicDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -75,6 +76,7 @@ import org.apache.pinot.spi.config.table.IndexingConfig; import org.apache.pinot.spi.config.table.Intern; import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; @@ -107,7 +109,7 @@ public DictionaryIndexConfig getDefaultConfig() { public void validate(FieldIndexConfigs indexConfigs, FieldSpec fieldSpec, TableConfig tableConfig) { DictionaryIndexConfig dictionaryConfig = indexConfigs.getConfig(StandardIndexes.dictionary()); if (dictionaryConfig.isEnabled() && dictionaryConfig.getUseVarLengthDictionary()) { - DataType storedType = fieldSpec.getDataType().getStoredType(); + DataType storedType = fieldSpec.getEffectiveStoredType(); Preconditions.checkState(storedType == DataType.STRING || storedType == DataType.BYTES, "Cannot create var-length dictionary on column: %s of stored type other than STRING or BYTES", fieldSpec.getName()); @@ -319,8 +321,24 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat } } + // Multi-entity URN columns with per-entity dict ranges always load polymorphically, even + // when effectiveStoredType reports STRING — the wire-format slices the buffer into + // typed per-entity regions and each region is wrapped in its own typed sub-dictionary + // (LongDictionary for LONG entities, StringDictionary for STRING entities, etc.). + if (dataType == DataType.URN && metadata.getFieldSpec() instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) metadata.getFieldSpec(); + List ranges = dim.getEntityDictRanges(); + if (ranges != null) { + Dictionary polymorphic = loadPolymorphicSubDicts(dataBuffer, ranges, loadOnHeap); + if (polymorphic != null) { + return polymorphic; + } + } + } + // Use effective stored type: for typed URN LONG/INT columns the physical data is numeric. + DataType effectiveStoredType = metadata.getFieldSpec().getEffectiveStoredType(); int length = metadata.getCardinality(); - switch (dataType.getStoredType()) { + switch (effectiveStoredType) { case INT: return loadOnHeap ? new OnHeapIntDictionary(dataBuffer, length) : new IntDictionary(dataBuffer, length); @@ -337,10 +355,12 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat int numBytesPerValue = metadata.getColumnMaxLength(); return loadOnHeap ? new OnHeapBigDecimalDictionary(dataBuffer, length, numBytesPerValue) : new BigDecimalDictionary(dataBuffer, length, numBytesPerValue); - case STRING: + case STRING: { numBytesPerValue = metadata.getColumnMaxLength(); - return loadOnHeap ? new OnHeapStringDictionary(dataBuffer, length, numBytesPerValue, strInterner, byteInterner) + return loadOnHeap + ? new OnHeapStringDictionary(dataBuffer, length, numBytesPerValue, strInterner, byteInterner) : new StringDictionary(dataBuffer, length, numBytesPerValue); + } case BYTES: numBytesPerValue = metadata.getColumnMaxLength(); return loadOnHeap ? new OnHeapBytesDictionary(dataBuffer, length, numBytesPerValue, byteInterner) @@ -350,6 +370,56 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat } } + /** + * Slices {@code dataBuffer} into per-entity sub-regions using cumulative byte offsets derived + * from {@code (endDictIdExclusive - prevEnd) * numBytesPerValue}, and wraps each region in the + * sub-dictionary class declared by the range's {@code dictClass}. Returns {@code null} when no + * entity has any rows (i.e. every range is empty), which lets the caller fall back to + * non-polymorphic load paths. + */ + @Nullable + private static Dictionary loadPolymorphicSubDicts(PinotDataBuffer dataBuffer, + List ranges, boolean loadOnHeap) { + List subDicts = new ArrayList<>(ranges.size()); + long prevByteEnd = 0; + int prevDictEnd = 0; + for (DimensionFieldSpec.EntityDictRange r : ranges) { + int subLen = r.getEndDictIdExclusive() - prevDictEnd; + long byteEnd = prevByteEnd + (long) subLen * r.getNumBytesPerValue(); + if (subLen > 0) { + PinotDataBuffer sub = dataBuffer.view(prevByteEnd, byteEnd); + Dictionary subDict = createTypedSubDict(sub, subLen, r.getDictClass(), r.getNumBytesPerValue(), loadOnHeap); + subDicts.add(new UrnPolymorphicDictionary.SubDict(r.getEntityType(), subDict)); + } + prevByteEnd = byteEnd; + prevDictEnd = r.getEndDictIdExclusive(); + } + return subDicts.isEmpty() ? null : new UrnPolymorphicDictionary(subDicts); + } + + private static Dictionary createTypedSubDict(PinotDataBuffer buf, int length, DataType dictClass, + int numBytesPerValue, boolean loadOnHeap) { + switch (dictClass) { + case INT: + return loadOnHeap ? new OnHeapIntDictionary(buf, length) : new IntDictionary(buf, length); + case LONG: + return loadOnHeap ? new OnHeapLongDictionary(buf, length) : new LongDictionary(buf, length); + case FLOAT: + return loadOnHeap ? new OnHeapFloatDictionary(buf, length) : new FloatDictionary(buf, length); + case DOUBLE: + return loadOnHeap ? new OnHeapDoubleDictionary(buf, length) : new DoubleDictionary(buf, length); + case STRING: + return loadOnHeap + ? new OnHeapStringDictionary(buf, length, numBytesPerValue, null, null) + : new StringDictionary(buf, length, numBytesPerValue); + case BYTES: + return loadOnHeap ? new OnHeapBytesDictionary(buf, length, numBytesPerValue, null) + : new BytesDictionary(buf, length, numBytesPerValue); + default: + throw new IllegalStateException("Unsupported sub-dictionary type for URN polymorphic layout: " + dictClass); + } + } + @Override protected IndexReaderFactory createReaderFactory() { return ReaderFactory.INSTANCE; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexCreatorFactory.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexCreatorFactory.java index 6084c77b4e..080e94b55e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexCreatorFactory.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexCreatorFactory.java @@ -71,7 +71,7 @@ public static ForwardIndexCreator createIndexCreator(IndexCreationContext contex } } else { // Dictionary disabled columns - DataType storedType = fieldSpec.getDataType().getStoredType(); + DataType storedType = fieldSpec.getEffectiveStoredType(); if (indexConfig.getCompressionCodec() == FieldConfig.CompressionCodec.CLP) { // CLP (V1) uses hard-coded chunk compressor which is set to `PassThrough` return new CLPForwardIndexCreatorV1(indexDir, columnName, numTotalDocs, context.getColumnStatistics()); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java index 59a69047c0..ca92bb69a2 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java @@ -104,7 +104,8 @@ public static ForwardIndexReader createIndexReader(PinotDataBuffer dataBuffer, C return new CLPForwardIndexReaderV2(dataBuffer, metadata.getTotalDocs()); } } - return createRawIndexReader(dataBuffer, metadata.getDataType().getStoredType(), metadata.isSingleValue()); + DataType storedType = metadata.getFieldSpec().getEffectiveStoredType(); + return createRawIndexReader(dataBuffer, storedType, metadata.isSingleValue()); } } 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 618bfafed9..060f794599 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 @@ -120,7 +120,7 @@ private void validateForwardIndexEnabled(ForwardIndexConfig forwardIndexConfig, boolean isCLPCodec = compressionCodec == CompressionCodec.CLP || compressionCodec == CompressionCodec.CLPV2 || compressionCodec == CompressionCodec.CLPV2_ZSTD || compressionCodec == CompressionCodec.CLPV2_LZ4; if (isCLPCodec) { - Preconditions.checkState(fieldSpec.getDataType().getStoredType() == FieldSpec.DataType.STRING, + Preconditions.checkState(fieldSpec.getEffectiveStoredType() == FieldSpec.DataType.STRING, "Cannot apply CLP compression codec to column: %s of stored type other than STRING", column); } else { Preconditions.checkState(compressionCodec == null || compressionCodec.isApplicableToRawIndex(), diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionary.java new file mode 100644 index 0000000000..6a2df92caf --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionary.java @@ -0,0 +1,357 @@ +/** + * 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.segment.local.segment.index.readers; + +import it.unimi.dsi.fastutil.ints.IntOpenHashSet; +import it.unimi.dsi.fastutil.ints.IntSet; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/** + * Composite dictionary for multi-entity URN columns. Each entity declared on the column has its + * own delegate {@link Dictionary} — a standard Pinot dictionary of whatever type is natural for + * that entity's key (LONG, INT, STRING, …). The dictionary ID space is the concatenation of + * those sub-dictionaries' ID spaces; this wrapper routes every operation to the appropriate + * sub-dictionary and reattaches the entity prefix on read. + * + *

Picturing the layout: + * 

+ *   actorId
+ *   ├── entity: urn:li:memberId         → LongDictionary    (global dict IDs [0, N0))
+ *   ├── entity: urn:li:corpUser         → StringDictionary  (global dict IDs [N0, N0+N1))
+ *   └── entity: urn:li:applicationName  → IntDictionary     (global dict IDs [N0+N1, N0+N1+N2))
+ *
+ * + *

For each call: + *

    + *
  • {@link #getStringValue(int) getStringValue(globalId)} → find owning sub-dict via a small + * linear scan over offsets, compute the local ID, ask the sub-dict for its string value + * and prepend the entity prefix.
    • + *
  • {@link #indexOf(String) indexOf("urn:li:corpUser:foo")} → parse the prefix to find the + * owning entity, call that sub-dict's {@code indexOf("foo")}, add the entity's global + * offset.
    • + *
  • {@link #getLongValue(int) getLongValue(globalId)} / {@link #getIntValue(int)} → forward + * to the sub-dict; works whenever the sub-dict's value type can answer the call.
    • + *
+ * + *

Same-entity range queries delegate to the sub-dict's range API and adjust IDs by the + * offset. Cross-entity ranges return empty (semantically undefined). + */ +public class UrnPolymorphicDictionary implements Dictionary { + + private final Dictionary[] _subDicts; + private final String[] _entityTypes; + private final String[] _entityPrefixes; + /** Cumulative global offset that each sub-dict starts at. {@code _offsets[0] == 0}. */ + private final int[] _offsets; + private final int _totalLength; + + public UrnPolymorphicDictionary(List subDicts) { + if (subDicts.isEmpty()) { + throw new IllegalArgumentException("UrnPolymorphicDictionary requires at least one sub-dictionary"); + } + int n = subDicts.size(); + _subDicts = new Dictionary[n]; + _entityTypes = new String[n]; + _entityPrefixes = new String[n]; + _offsets = new int[n]; + int cumulative = 0; + for (int i = 0; i < n; i++) { + SubDict sd = subDicts.get(i); + _subDicts[i] = sd._dictionary; + _entityTypes[i] = sd._entityType; + _entityPrefixes[i] = sd._entityType + ":"; + _offsets[i] = cumulative; + cumulative += sd._dictionary.length(); + } + _totalLength = cumulative; + } + + /** Sub-dictionary descriptor: an entity type and the Pinot {@link Dictionary} backing it. */ + public static final class SubDict { + private final String _entityType; + private final Dictionary _dictionary; + + public SubDict(String entityType, Dictionary dictionary) { + _entityType = entityType; + _dictionary = dictionary; + } + + public String getEntityType() { + return _entityType; + } + + public Dictionary getDictionary() { + return _dictionary; + } + } + + /** + * Index of the sub-dict that owns {@code globalDictId}, or -1 when out of bounds. Linear scan; + * the entity count on multi-entity URN columns is small (typically ≤ 16). + */ + int findSubDictIdx(int globalDictId) { + if (globalDictId < 0 || globalDictId >= _totalLength) { + return -1; + } + for (int i = 0; i < _offsets.length; i++) { + int end = (i + 1 < _offsets.length) ? _offsets[i + 1] : _totalLength; + if (globalDictId < end) { + return i; + } + } + return -1; + } + + /** Returns the entity type that owns the dict ID, or null when out of bounds. */ + @Nullable + public String getEntityTypeForDictId(int globalDictId) { + int idx = findSubDictIdx(globalDictId); + return idx >= 0 ? _entityTypes[idx] : null; + } + + // ---------------- ordering / metadata ---------------- + + @Override + public boolean isSorted() { + // Multi-entity: sub-dicts are independently sorted but the buffer is not lex-sorted as + // URN strings across entities, so callers that scan in dict-ID order can't rely on + // monotone URN ordering. Single-entity: the one sub-dict is sorted, and the polymorphic + // wrapper just prepends a constant prefix on read, so the order is preserved. + return _subDicts.length == 1; + } + + @Override + public DataType getValueType() { + // Always-URN-strings on the outside, regardless of any sub-dict's underlying type. + return DataType.STRING; + } + + @Override + public int length() { + return _totalLength; + } + + @Override + public int compare(int dictId1, int dictId2) { + return Integer.compare(dictId1, dictId2); + } + + @Override + public Comparable getMinVal() { + return _totalLength == 0 ? null : getStringValue(0); + } + + @Override + public Comparable getMaxVal() { + return _totalLength == 0 ? null : getStringValue(_totalLength - 1); + } + + @Override + public Object getSortedValues() { + throw new UnsupportedOperationException(); + } + + // ---------------- per-dictId accessors ---------------- + + @Override + public Object get(int dictId) { + return getStringValue(dictId); + } + + @Override + public String getStringValue(int dictId) { + int idx = findSubDictIdx(dictId); + if (idx < 0) { + return null; + } + int localId = dictId - _offsets[idx]; + return _entityPrefixes[idx] + _subDicts[idx].getStringValue(localId); + } + + @Override + public void readStringValues(int[] dictIds, int length, String[] outValues) { + for (int i = 0; i < length; i++) { + outValues[i] = getStringValue(dictIds[i]); + } + } + + @Override + public long getLongValue(int dictId) { + int idx = findSubDictIdx(dictId); + return idx < 0 ? 0L : _subDicts[idx].getLongValue(dictId - _offsets[idx]); + } + + @Override + public int getIntValue(int dictId) { + int idx = findSubDictIdx(dictId); + return idx < 0 ? 0 : _subDicts[idx].getIntValue(dictId - _offsets[idx]); + } + + @Override + public float getFloatValue(int dictId) { + int idx = findSubDictIdx(dictId); + return idx < 0 ? 0f : _subDicts[idx].getFloatValue(dictId - _offsets[idx]); + } + + @Override + public double getDoubleValue(int dictId) { + int idx = findSubDictIdx(dictId); + return idx < 0 ? 0d : _subDicts[idx].getDoubleValue(dictId - _offsets[idx]); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + int idx = findSubDictIdx(dictId); + return idx < 0 ? null : _subDicts[idx].getBigDecimalValue(dictId - _offsets[idx]); + } + + // ---------------- lookup ---------------- + + @Override + public int indexOf(String urnString) { + int entityIdx = findEntityIdxByPrefix(urnString); + if (entityIdx < 0) { + return NULL_VALUE_INDEX; + } + String suffix = urnString.substring(_entityPrefixes[entityIdx].length()); + int localId = safeIndexOf(_subDicts[entityIdx], suffix); + return localId < 0 ? NULL_VALUE_INDEX : _offsets[entityIdx] + localId; + } + + @Override + public int insertionIndexOf(String urnString) { + int entityIdx = findEntityIdxByPrefix(urnString); + if (entityIdx < 0) { + return ~_totalLength; + } + String suffix = urnString.substring(_entityPrefixes[entityIdx].length()); + int subInsertion = safeInsertionIndexOf(_subDicts[entityIdx], suffix, _subDicts[entityIdx].length()); + if (subInsertion >= 0) { + return _offsets[entityIdx] + subInsertion; + } + int insertionPoint = ~subInsertion; + return ~(_offsets[entityIdx] + insertionPoint); + } + + @Override + public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { + int lowerEntity = findEntityIdxByPrefix(lower); + int upperEntity = findEntityIdxByPrefix(upper); + // Cross-entity ranges are semantically undefined; return empty rather than throwing so we + // behave gracefully under query generators. + if (lowerEntity < 0 || upperEntity < 0 || lowerEntity != upperEntity) { + return new IntOpenHashSet(0); + } + int entity = lowerEntity; + String lowerSuffix = lower.substring(_entityPrefixes[entity].length()); + String upperSuffix = upper.substring(_entityPrefixes[entity].length()); + // Sub-dictionaries (LongDictionary, StringDictionary, ...) are sorted but inherit the base + // class's "this method should not be called" implementation of getDictIdsInRange. Compute + // the range on the sub-dict directly via binary search using insertionIndexOf — this gives + // us a [from, to) half-open relative dict-ID range, which we shift by the entity's offset. + Dictionary sub = _subDicts[entity]; + int subLen = sub.length(); + int lowerInsert = safeInsertionIndexOf(sub, lowerSuffix, subLen); + int from; + if (lowerInsert >= 0) { + from = includeLower ? lowerInsert : lowerInsert + 1; + } else { + from = ~lowerInsert; + } + int upperInsert = safeInsertionIndexOf(sub, upperSuffix, subLen); + int to; + if (upperInsert >= 0) { + to = includeUpper ? upperInsert + 1 : upperInsert; + } else { + to = ~upperInsert; + } + if (from < 0) { + from = 0; + } + if (to > subLen) { + to = subLen; + } + if (from >= to) { + return new IntOpenHashSet(0); + } + IntSet result = new IntOpenHashSet(to - from); + int offset = _offsets[entity]; + for (int i = from; i < to; i++) { + result.add(offset + i); + } + return result; + } + + /** + * Wraps {@link Dictionary#indexOf(String)} to swallow parse errors from numeric sub-dicts — + * e.g. {@link org.apache.pinot.segment.local.segment.index.readers.LongDictionary#indexOf} + * calls {@code Long.parseLong} on the key, which throws {@link NumberFormatException} for + * malformed suffixes. The polymorphic wrapper treats that as "no such value" so a malformed + * query literal returns no matches instead of crashing the query. + */ + private static int safeIndexOf(Dictionary sub, String suffix) { + try { + return sub.indexOf(suffix); + } catch (NumberFormatException e) { + return NULL_VALUE_INDEX; + } + } + + /** + * Wraps {@link Dictionary#insertionIndexOf(String)} the same way as {@link #safeIndexOf}. + * Malformed input returns {@code ~length} — past-end, no matches — matching the contract that + * {@code getDictIdsInRange} relies on. + */ + private static int safeInsertionIndexOf(Dictionary sub, String suffix, int subLen) { + try { + return sub.insertionIndexOf(suffix); + } catch (NumberFormatException e) { + return ~subLen; + } + } + + /** Finds the entity whose prefix is a prefix of {@code urnString}, or -1. */ + private int findEntityIdxByPrefix(@Nullable String urnString) { + if (urnString == null) { + return -1; + } + for (int i = 0; i < _entityPrefixes.length; i++) { + if (urnString.startsWith(_entityPrefixes[i])) { + return i; + } + } + return -1; + } + + // ---------------- lifecycle ---------------- + + @Override + public void close() + throws IOException { + for (Dictionary subDict : _subDicts) { + subDict.close(); + } + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformerTest.java new file mode 100644 index 0000000000..756905f907 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformerTest.java @@ -0,0 +1,159 @@ +/** + * 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.segment.local.recordtransformer; + +import java.util.Arrays; +import java.util.List; +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.GenericRow; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +public class UrnKeyEncodingTransformerTest { + + private static Schema singleEntityLongSchema() { + DimensionFieldSpec dim = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + return new Schema.SchemaBuilder().addField(dim).build(); + } + + private static Schema multiEntityLongSchema() { + DimensionFieldSpec dim = new DimensionFieldSpec("resourceId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG))); + return new Schema.SchemaBuilder().addField(dim).build(); + } + + private static Schema untypedSchema() { + return new Schema.SchemaBuilder() + .addSingleValueDimension("rawUrn", FieldSpec.DataType.URN) + .build(); + } + + private static GenericRow rowWithValue(String col, Object val) { + GenericRow row = new GenericRow(); + row.putValue(col, val); + return row; + } + + // --- isNoOp --- + + @Test + public void isNoOpWhenNoUrnColumns() { + Schema schema = new Schema.SchemaBuilder() + .addSingleValueDimension("name", FieldSpec.DataType.STRING) + .build(); + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(schema); + assertTrue(t.isNoOp()); + } + + @Test + public void isNoOpForUntypedUrnColumn() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(untypedSchema()); + assertTrue(t.isNoOp(), "Untyped URN (no urnTypes) should be a no-op"); + } + + @Test + public void isNotNoOpForTypedLongUrnColumn() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(singleEntityLongSchema()); + assertTrue(!t.isNoOp(), "Typed LONG URN should not be a no-op"); + } + + // --- single-entity typed LONG column --- + + @Test + public void singleEntityLongUrnIsEncodedToLong() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(singleEntityLongSchema()); + GenericRow row = rowWithValue("memberId", "urn:li:memberId:12345"); + t.transform(row); + assertEquals(row.getValue("memberId"), 12345L); + } + + @Test + public void singleEntityNullUrnLeftAsNull() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(singleEntityLongSchema()); + GenericRow row = rowWithValue("memberId", null); + t.transform(row); + assertNull(row.getValue("memberId")); + } + + @Test + public void singleEntityMalformedUrnNulledOut() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(singleEntityLongSchema()); + GenericRow row = rowWithValue("memberId", "not-a-urn"); + t.transform(row); + assertNull(row.getValue("memberId"), "Malformed URN should result in null"); + } + + @Test + public void singleEntityNonNumericKeyNulledOut() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(singleEntityLongSchema()); + GenericRow row = rowWithValue("memberId", "urn:li:memberId:abc"); + t.transform(row); + assertNull(row.getValue("memberId"), "Non-numeric key should be nulled out"); + } + + // --- multi-entity typed LONG column (companion column written) --- + + @Test + public void multiEntityFirstEntityStashesEntityIdx0() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", "urn:li:memberId:100"); + t.transform(row); + // Multi-entity LONG storage (UrnDictionary layout): main column holds just the numeric key; + // the entity index is stashed in a transient row field consumed by the segment-build pipeline + // and is never persisted as a column. + assertEquals(row.getValue("resourceId"), 100L); + assertEquals(row.getValue("__urnEntityIdxBuild_resourceId"), 0); + } + + @Test + public void multiEntitySecondEntityStashesEntityIdx1() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", "urn:li:groupId:42"); + t.transform(row); + assertEquals(row.getValue("resourceId"), 42L); + assertEquals(row.getValue("__urnEntityIdxBuild_resourceId"), 1); + } + + @Test + public void multiEntityUnknownEntityNullsBothColumns() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", "urn:li:orgId:99"); + t.transform(row); + assertNull(row.getValue("resourceId"), "Unknown entity should be nulled"); + assertNull(row.getValue("__urnEntityIdxBuild_resourceId"), "Transient entity idx should be null too"); + } + + @Test + public void multiEntityNullUrnLeavesTransientNull() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", null); + t.transform(row); + assertNull(row.getValue("resourceId")); + assertNull(row.getValue("__urnEntityIdxBuild_resourceId")); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformerTest.java new file mode 100644 index 0000000000..8d78f18a5e --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformerTest.java @@ -0,0 +1,226 @@ +/** + * 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.segment.local.recordtransformer; + +import java.util.Arrays; +import java.util.List; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; +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.GenericRow; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +public class UrnValidationTransformerTest { + + private static TableConfig defaultTableConfig() { + return new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable").build(); + } + + private static TableConfig continueOnErrorTableConfig() { + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setContinueOnError(true); + return new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable") + .setIngestionConfig(ingestionConfig).build(); + } + + private static Schema singleEntitySchema() { + DimensionFieldSpec dim = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + return new Schema.SchemaBuilder().addField(dim).build(); + } + + private static Schema multiEntitySchema() { + DimensionFieldSpec dim = new DimensionFieldSpec("resourceId", FieldSpec.DataType.URN, true); + dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG) + )); + return new Schema.SchemaBuilder().addField(dim).build(); + } + + private static Schema untypedSchema() { + return new Schema.SchemaBuilder() + .addSingleValueDimension("rawUrn", FieldSpec.DataType.URN) + .build(); + } + + // --- isNoOp --- + + @Test + public void isNoOpWhenNoUrnColumns() { + Schema schema = new Schema.SchemaBuilder() + .addSingleValueDimension("name", FieldSpec.DataType.STRING) + .build(); + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), schema); + assertTrue(t.isNoOp()); + } + + @Test + public void isNotNoOpWhenUntypedUrnColumnPresent() { + // Untyped URN columns still need shape validation so SegmentColumnarIndexCreator's + // multi-prefix auto-detect can rely on every value being parseable. + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), untypedSchema()); + assertFalse(t.isNoOp()); + } + + @Test + public void isNotNoOpWhenUrnTypesPresent() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), singleEntitySchema()); + assertFalse(t.isNoOp()); + } + + // --- valid values pass through --- + + @Test + public void validUrnPassesThrough() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", "urn:li:memberId:123456"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertEquals(result.getValue("memberId"), "urn:li:memberId:123456"); + assertFalse(result.getFieldToValueMap().containsKey(GenericRow.INCOMPLETE_RECORD_KEY)); + } + + @Test + public void nullValuePassesThrough() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", null); + GenericRow result = t.transform(record); + assertNotNull(result); + assertNull(result.getValue("memberId")); + } + + @Test + public void multiEntityBothValid() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), multiEntitySchema()); + + GenericRow r1 = new GenericRow(); + r1.putValue("resourceId", "urn:li:memberId:1"); + assertNotNull(t.transform(r1)); + + GenericRow r2 = new GenericRow(); + r2.putValue("resourceId", "urn:li:groupId:2"); + assertNotNull(t.transform(r2)); + } + + // --- invalid values: continueOnError=false --- + + @Test(expectedExceptions = IllegalStateException.class) + public void wrongEntityTypeThrowsWhenNotContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", "urn:li:groupId:99"); + t.transform(record); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void malformedUrnThrowsWhenNotContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", "not-a-urn"); + t.transform(record); + } + + // --- invalid values: continueOnError=true --- + + @Test + public void wrongEntityTypeNullifiesAndMarksIncompleteWhenContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(continueOnErrorTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", "urn:li:groupId:99"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertNull(result.getValue("memberId")); + assertTrue((Boolean) result.getValue(GenericRow.INCOMPLETE_RECORD_KEY)); + } + + @Test + public void malformedUrnNullifiesWhenContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(continueOnErrorTableConfig(), singleEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("memberId", "bad-data"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertNull(result.getValue("memberId")); + assertTrue((Boolean) result.getValue(GenericRow.INCOMPLETE_RECORD_KEY)); + } + + @Test + public void multiEntityWrongTypeNullifiesWhenContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(continueOnErrorTableConfig(), multiEntitySchema()); + GenericRow record = new GenericRow(); + record.putValue("resourceId", "urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertNull(result.getValue("resourceId")); + } + + // --- untyped URN columns --- + + @Test + public void untypedSimpleUrnPassesThrough() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), untypedSchema()); + GenericRow record = new GenericRow(); + record.putValue("rawUrn", "urn:li:memberId:42"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertEquals(result.getValue("rawUrn"), "urn:li:memberId:42"); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void untypedNonUrnThrowsWhenNotContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), untypedSchema()); + GenericRow record = new GenericRow(); + record.putValue("rawUrn", "not-a-urn"); + t.transform(record); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void untypedNonSimpleUrnThrowsWhenNotContinueOnError() { + // Compound URN — auto-detect can't promote this to a sorted-LONG sub-dictionary, so reject. + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), untypedSchema()); + GenericRow record = new GenericRow(); + record.putValue("rawUrn", "urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)"); + t.transform(record); + } + + @Test + public void untypedNonUrnNullifiesWhenContinueOnError() { + UrnValidationTransformer t = new UrnValidationTransformer(continueOnErrorTableConfig(), untypedSchema()); + GenericRow record = new GenericRow(); + record.putValue("rawUrn", "garbage"); + GenericRow result = t.transform(record); + assertNotNull(result); + assertNull(result.getValue("rawUrn")); + assertTrue((Boolean) result.getValue(GenericRow.INCOMPLETE_RECORD_KEY)); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreatorTest.java new file mode 100644 index 0000000000..c9739af822 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreatorTest.java @@ -0,0 +1,189 @@ +/** + * 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.segment.local.segment.creator.impl.urn; + +import com.google.common.io.Files; +import java.io.File; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.DimensionFieldSpec.EntityDictRange; +import org.apache.pinot.spi.data.DimensionFieldSpec.UrnType; +import org.apache.pinot.spi.data.FieldSpec; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +public class UrnPolymorphicDictionaryCreatorTest { + + private File _indexDir; + + @BeforeMethod + public void setUp() { + _indexDir = Files.createTempDir(); + } + + @AfterMethod + public void tearDown() { + FileUtils.deleteQuietly(_indexDir); + } + + private static DimensionFieldSpec dim(String name) { + return new DimensionFieldSpec(name, FieldSpec.DataType.URN, true); + } + + private long dictionaryFileBytes(String columnName) { + Path dict = _indexDir.toPath().resolve(columnName + DictionaryIndexType.getFileExtension()); + return dict.toFile().length(); + } + + @Test + public void allLongMixedEntitiesWriteAndIndex() + throws Exception { + // Two entities, both LONG (the original sorted-LONG layout). This case must produce a + // wire-format identical to what UrnSortedDictionaryCreator emitted: per-entity LONG keys + // back-to-back, ranges encoded as the legacy 2-field metadata form. + List urnTypes = Arrays.asList( + UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG)); + UrnPolymorphicDictionaryCreator creator = + new UrnPolymorphicDictionaryCreator(dim("actorId"), _indexDir, urnTypes); + + EntityValuePair[] input = new EntityValuePair[] { + new EntityValuePair(0, 100L), + new EntityValuePair(0, 200L), + new EntityValuePair(1, 50L), + new EntityValuePair(1, 60L), + new EntityValuePair(1, 70L) + }; + Arrays.sort(input); + creator.build(input); + + List ranges = creator.getRanges(); + assertEquals(ranges.size(), 2); + assertEquals(ranges.get(0).getEntityType(), "urn:li:memberId"); + assertEquals(ranges.get(0).getEndDictIdExclusive(), 2); + assertEquals(ranges.get(0).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(ranges.get(0).getNumBytesPerValue(), Long.BYTES); + assertEquals(ranges.get(1).getEntityType(), "urn:li:groupId"); + assertEquals(ranges.get(1).getEndDictIdExclusive(), 5); + assertEquals(ranges.get(1).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(ranges.get(1).getNumBytesPerValue(), Long.BYTES); + + // 5 LONG values = 40 bytes — identical to the legacy sorted-LONG layout. + assertEquals(dictionaryFileBytes("actorId"), 40L); + assertEquals(creator.indexOfSV(new EntityValuePair(0, 100L)), 0); + assertEquals(creator.indexOfSV(new EntityValuePair(0, 200L)), 1); + assertEquals(creator.indexOfSV(new EntityValuePair(1, 50L)), 2); + assertEquals(creator.indexOfSV(new EntityValuePair(1, 70L)), 4); + + // The legacy EntityKeyPair input is accepted too — wire-format compatibility for callers + // that haven't migrated to EntityValuePair yet. + assertEquals(creator.indexOfSV(new EntityKeyPair(0, 100L)), 0); + assertEquals(creator.indexOfSV(new EntityKeyPair(1, 60L)), 3); + } + + @Test + public void mixedLongAndStringEntitiesPaddedToMaxBytes() + throws Exception { + // LONG memberId + STRING corpuser. STRING sub-dict padding width = the longest UTF-8 value + // in the entity's value set, which lets the standard StringDictionary read each entry by + // fixed offset without any per-entry length prefix. + List urnTypes = Arrays.asList( + UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + UrnType.of("urn:li:corpuser", FieldSpec.DataType.STRING)); + UrnPolymorphicDictionaryCreator creator = + new UrnPolymorphicDictionaryCreator(dim("subject"), _indexDir, urnTypes); + + EntityValuePair[] input = new EntityValuePair[] { + new EntityValuePair(0, 1L), + new EntityValuePair(0, 2L), + new EntityValuePair(1, "alice"), + new EntityValuePair(1, "robert-longest") // 14 bytes — pad width + }; + Arrays.sort(input); + creator.build(input); + + List ranges = creator.getRanges(); + assertEquals(ranges.size(), 2); + assertEquals(ranges.get(0).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(ranges.get(0).getNumBytesPerValue(), Long.BYTES); + assertEquals(ranges.get(1).getDictClass(), FieldSpec.DataType.STRING); + assertEquals(ranges.get(1).getNumBytesPerValue(), 14); + assertEquals(ranges.get(0).getEndDictIdExclusive(), 2); + assertEquals(ranges.get(1).getEndDictIdExclusive(), 4); + + // File size = 2*8 (LONG) + 2*14 (STRING) = 44 bytes + assertEquals(dictionaryFileBytes("subject"), 44L); + + assertEquals(creator.indexOfSV(new EntityValuePair(0, 1L)), 0); + assertEquals(creator.indexOfSV(new EntityValuePair(0, 2L)), 1); + assertEquals(creator.indexOfSV(new EntityValuePair(1, "alice")), 2); + assertEquals(creator.indexOfSV(new EntityValuePair(1, "robert-longest")), 3); + } + + @Test + public void writesValuesReadableByStandardDictionaryClasses() + throws Exception { + // End-to-end: write a mixed-type dictionary and read it back via the standard scalar + // dictionary classes. This is what the loader does at query time. + List urnTypes = Arrays.asList( + UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + UrnType.of("urn:li:corpuser", FieldSpec.DataType.STRING)); + UrnPolymorphicDictionaryCreator creator = + new UrnPolymorphicDictionaryCreator(dim("subject"), _indexDir, urnTypes); + + EntityValuePair[] input = new EntityValuePair[] { + new EntityValuePair(0, 100L), + new EntityValuePair(0, 200L), + new EntityValuePair(1, "alice"), + new EntityValuePair(1, "bob") + }; + Arrays.sort(input); + creator.build(input); + + File dictFile = new File(_indexDir, "subject" + DictionaryIndexType.getFileExtension()); + try (PinotDataBuffer buf = PinotDataBuffer.mapReadOnlyBigEndianFile(dictFile)) { + // LONG sub-dict at offset 0, length 2, 8 bytes each. + PinotDataBuffer longView = buf.view(0, 16); + org.apache.pinot.segment.local.segment.index.readers.LongDictionary longDict = + new org.apache.pinot.segment.local.segment.index.readers.LongDictionary(longView, 2); + assertEquals(longDict.getLongValue(0), 100L); + assertEquals(longDict.getLongValue(1), 200L); + + // STRING sub-dict at offset 16, length 2, 5 bytes each (alice is 5; bob padded to 5). + List ranges = creator.getRanges(); + int stringBytes = ranges.get(1).getNumBytesPerValue(); + assertTrue(stringBytes >= 5); + PinotDataBuffer strView = buf.view(16, 16 + 2L * stringBytes); + org.apache.pinot.segment.local.segment.index.readers.StringDictionary strDict = + new org.apache.pinot.segment.local.segment.index.readers.StringDictionary(strView, 2, stringBytes); + assertEquals(strDict.getStringValue(0), "alice"); + assertEquals(strDict.getStringValue(1), "bob"); + } + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/UrnColumnSegmentTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/UrnColumnSegmentTest.java new file mode 100644 index 0000000000..fb23e446fb --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/UrnColumnSegmentTest.java @@ -0,0 +1,562 @@ +/** + * 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.segment.local.segment.index; + +import com.google.common.io.Files; +import java.io.File; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentUtil; +import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +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.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/** + * Integration test for URN column: segment build → metadata persistence → reload. + */ +public class UrnColumnSegmentTest { + private static final String SEGMENT_NAME = "urnColumnSegmentTest"; + private static final String URN_COL = "memberId"; + private static final String TIME_COL = "t"; + + private File _segmentOutputDir; + private File _segmentDir; + + @BeforeMethod + public void setUp() { + _segmentOutputDir = Files.createTempDir(); + } + + @AfterMethod + public void tearDown() { + FileUtils.deleteQuietly(_segmentOutputDir); + } + + private static TableConfig tableConfig() { + return new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); + } + + private static Schema singleEntityTypedSchema() { + DimensionFieldSpec dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + dim.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + return new Schema.SchemaBuilder() + .addField(dim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + private static Schema multiEntityTypedSchema() { + DimensionFieldSpec dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG))); + return new Schema.SchemaBuilder() + .addField(dim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + private static Schema untypedUrnSchema() { + return new Schema.SchemaBuilder() + .addSingleValueDimension(URN_COL, FieldSpec.DataType.URN) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + private List buildRows(String... urnValues) { + List rows = new ArrayList<>(); + for (String urn : urnValues) { + GenericRow row = new GenericRow(); + row.putValue(URN_COL, urn); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + } + return rows; + } + + // --- urnTypes metadata round-trip --- + + @Test + public void singleEntityUrnTypesPersistedInMetadata() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = buildRows("urn:li:memberId:1", "urn:li:memberId:2", "urn:li:memberId:3"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + ColumnMetadata colMeta = meta.getColumnMetadataMap().get(URN_COL); + assertNotNull(colMeta); + assertEquals(colMeta.getDataType(), FieldSpec.DataType.URN); + + FieldSpec fieldSpec = colMeta.getFieldSpec(); + assertTrue(fieldSpec instanceof DimensionFieldSpec); + DimensionFieldSpec dimSpec = (DimensionFieldSpec) fieldSpec; + assertTrue(dimSpec.hasUrnTypes(), "dimSpec.hasUrnTypes() was false; urnTypes: " + dimSpec.getUrnTypes()); + assertEquals(dimSpec.getUrnTypes().size(), 1); + assertEquals(dimSpec.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(dimSpec.getUrnTypes().get(0).getValueType(), FieldSpec.DataType.LONG); + } + + @Test + public void multiEntityUrnTypesPersistedInMetadata() + throws Exception { + Schema schema = multiEntityTypedSchema(); + List rows = buildRows("urn:li:memberId:1", "urn:li:groupId:2"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + ColumnMetadata colMeta = meta.getColumnMetadataMap().get(URN_COL); + DimensionFieldSpec dimSpec = (DimensionFieldSpec) colMeta.getFieldSpec(); + assertTrue(dimSpec.hasUrnTypes()); + assertEquals(dimSpec.getUrnTypes().size(), 2); + assertEquals(dimSpec.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(dimSpec.getUrnTypes().get(1).getEntityType(), "urn:li:groupId"); + } + + @Test + public void untypedUrnAutoDetectedSegmentIsQueryable() + throws Exception { + Schema schema = untypedUrnSchema(); + List rows = buildRows("urn:li:memberId:100", "urn:li:memberId:200", + "urn:li:groupId:300", "urn:li:groupId:400"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + org.apache.pinot.segment.spi.ImmutableSegment seg = + ImmutableSegmentLoader.load(_segmentDir, ReadMode.mmap); + try { + Dictionary dict = seg.getDataSource(URN_COL).getDictionary(); + // Each of the four ingested URNs maps to a unique dict ID and round-trips on read. + assertTrue(dict.indexOf("urn:li:memberId:100") >= 0); + assertTrue(dict.indexOf("urn:li:memberId:200") >= 0); + assertTrue(dict.indexOf("urn:li:groupId:300") >= 0); + assertTrue(dict.indexOf("urn:li:groupId:400") >= 0); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:100")), "urn:li:memberId:100"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:groupId:400")), "urn:li:groupId:400"); + // Cross-entity key collisions can't happen here, but a lookup against a non-existent value + // returns the null sentinel rather than crashing. + assertEquals(dict.indexOf("urn:li:memberId:99999"), Dictionary.NULL_VALUE_INDEX); + assertEquals(dict.indexOf("urn:li:notAnEntity:1"), Dictionary.NULL_VALUE_INDEX); + } finally { + seg.destroy(); + } + } + + @Test + public void untypedUrnAutoDetectsMultiEntityFromIngestedData() + throws Exception { + Schema schema = untypedUrnSchema(); + List rows = buildRows("urn:li:memberId:1", "urn:li:groupId:2"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + ColumnMetadata colMeta = meta.getColumnMetadataMap().get(URN_COL); + DimensionFieldSpec dimSpec = (DimensionFieldSpec) colMeta.getFieldSpec(); + // Multi-prefix all-numeric data: auto-detect promotes the column to multi-entity URN with + // synthesized urnTypes and per-entity dict ID ranges. The segment behaves as if the user + // had declared the column with these urnTypes. + assertTrue(dimSpec.hasUrnTypes()); + assertEquals(dimSpec.getUrnTypes().size(), 2); + // Discovered ordering depends on the lex-sorted order of unique values from the stats + // collector; assert the *set* of entity types rather than a specific order. + java.util.Set discovered = new java.util.HashSet<>(); + discovered.add(dimSpec.getUrnTypes().get(0).getEntityType()); + discovered.add(dimSpec.getUrnTypes().get(1).getEntityType()); + assertEquals(discovered, java.util.Set.of("urn:li:memberId", "urn:li:groupId")); + assertNotNull(dimSpec.getEntityDictRanges()); + } + + // --- validation: valid URNs pass through --- + + @Test + public void validUrnValuesStoredAndRetrievable() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = + buildRows("urn:li:memberId:100", "urn:li:memberId:200", "urn:li:memberId:300"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + assertEquals(meta.getTotalDocs(), 3); + } + + @Test + public void nullUrnValueIsAllowed() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = new ArrayList<>(); + GenericRow row = new GenericRow(); + row.putValue(URN_COL, null); + row.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(row); + + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + assertEquals(meta.getTotalDocs(), 1); + } + + // --- Physical storage type --- + + @Test + public void typedUrnColumnEffectiveStoredTypeIsLong() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = buildRows("urn:li:memberId:42"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + ColumnMetadata colMeta = meta.getColumnMetadataMap().get(URN_COL); + assertEquals(colMeta.getDataType(), FieldSpec.DataType.URN); + DimensionFieldSpec dimSpec = (DimensionFieldSpec) colMeta.getFieldSpec(); + assertEquals(dimSpec.getEffectiveStoredType(), FieldSpec.DataType.LONG); + } + + // --- Segment loading and dictionary reconstruction --- + + @Test + public void singleEntityTypedUrnDictionaryReconstructsUrn() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = buildRows("urn:li:memberId:100", "urn:li:memberId:200", "urn:li:memberId:300"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + ImmutableSegment segment = ImmutableSegmentLoader.load(_segmentDir, ReadMode.mmap); + try { + Dictionary dict = segment.getDictionary(URN_COL); + assertNotNull(dict, "Dictionary should exist for typed URN column"); + assertEquals(dict.length(), 3); + // Values are stored as LONG keys; dictionary should reconstruct full URN strings. + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:100")), "urn:li:memberId:100"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:200")), "urn:li:memberId:200"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:300")), "urn:li:memberId:300"); + } finally { + segment.destroy(); + } + } + + @Test + public void singleEntityTypedUrnDictionaryIndexOfReturnsNullForWrongEntity() + throws Exception { + Schema schema = singleEntityTypedSchema(); + List rows = buildRows("urn:li:memberId:1"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + ImmutableSegment segment = ImmutableSegmentLoader.load(_segmentDir, ReadMode.mmap); + try { + Dictionary dict = segment.getDictionary(URN_COL); + assertNotNull(dict); + assertEquals(dict.indexOf("urn:li:groupId:1"), Dictionary.NULL_VALUE_INDEX); + assertEquals(dict.indexOf("urn:li:memberId:999"), Dictionary.NULL_VALUE_INDEX); + } finally { + segment.destroy(); + } + } + + // --- backwards compatibility: adding URN entity types after the fact --- + + /** + * Builds two multi-entity URN segments with different {@code urnTypes} lists (the second + * extends the first) and verifies that both can be loaded and queried. This is the contract + * for adding entity types to a URN column over time: existing segments are not rewritten and + * keep using the {@code urnTypes} ordering they were built with — newly ingested segments + * pick up the extended ordering. + */ + @Test + public void multiEntitySegmentsWithDifferentUrnTypesBothQueryable() + throws Exception { + // Original schema: two entity types. + Schema schemaV1 = multiEntityTypedSchema(); + List v1Rows = buildRows("urn:li:memberId:100", "urn:li:groupId:200"); + File v1Dir = PinotSegmentUtil.createSegment(tableConfig(), schemaV1, "v1seg", _segmentOutputDir.toString(), + new GenericRowRecordReader(v1Rows)); + + // Extended schema (entity type appended at the end). + DimensionFieldSpec v2Dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + v2Dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:applicationName", FieldSpec.DataType.LONG))); + Schema schemaV2 = new Schema.SchemaBuilder() + .addField(v2Dim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + List v2Rows = buildRows("urn:li:memberId:300", "urn:li:applicationName:400"); + File v2Dir = PinotSegmentUtil.createSegment(tableConfig(), schemaV2, "v2seg", _segmentOutputDir.toString(), + new GenericRowRecordReader(v2Rows)); + + ImmutableSegment v1 = ImmutableSegmentLoader.load(v1Dir, ReadMode.mmap); + ImmutableSegment v2 = ImmutableSegmentLoader.load(v2Dir, ReadMode.mmap); + try { + DimensionFieldSpec v1Spec = (DimensionFieldSpec) v1.getDataSource(URN_COL).getDataSourceMetadata().getFieldSpec(); + DimensionFieldSpec v2Spec = (DimensionFieldSpec) v2.getDataSource(URN_COL).getDataSourceMetadata().getFieldSpec(); + assertEquals(v1Spec.getUrnTypes().size(), 2, + "v1 segment metadata keeps its original urnTypes list — not retroactively updated"); + assertEquals(v2Spec.getUrnTypes().size(), 3, + "v2 segment metadata reflects the extended urnTypes list it was built with"); + + // Dict-level: v1's dictionary doesn't know about applicationName; lookups for that entity + // return NULL_VALUE_INDEX (zero rows in this segment have that entity type, so this is the + // semantically correct result). + Dictionary v1Dict = v1.getDataSource(URN_COL).getDictionary(); + assertEquals(v1Dict.indexOf("urn:li:applicationName:1"), Dictionary.NULL_VALUE_INDEX, + "v1 dict has no applicationName range; lookups for an unknown-entity URN must miss"); + // v1 still answers correctly for the entities it does know about. + assertTrue(v1Dict.indexOf("urn:li:memberId:100") >= 0); + assertTrue(v1Dict.indexOf("urn:li:groupId:200") >= 0); + + // v2's dictionary knows all three entity types. + Dictionary v2Dict = v2.getDataSource(URN_COL).getDictionary(); + assertTrue(v2Dict.indexOf("urn:li:memberId:300") >= 0); + assertTrue(v2Dict.indexOf("urn:li:applicationName:400") >= 0); + assertEquals(v2Dict.indexOf("urn:li:groupId:999"), Dictionary.NULL_VALUE_INDEX, + "v2 dict has a groupId range but key 999 isn't in it"); + } finally { + v1.destroy(); + v2.destroy(); + } + } + + /** + * Verifies that v1 and v2 segments built with different {@code urnTypes} orderings (entity + * types reordered) still reconstruct URN strings correctly. Each segment uses the ordering it + * was built with. + */ + @Test + public void multiEntitySegmentsWithReorderedUrnTypesReconstructCorrectly() + throws Exception { + // Original ordering: memberId first. + DimensionFieldSpec v1Dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + v1Dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG))); + Schema schemaV1 = new Schema.SchemaBuilder() + .addField(v1Dim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + // Swapped ordering: groupId first. + DimensionFieldSpec v2Dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + v2Dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + Schema schemaV2 = new Schema.SchemaBuilder() + .addField(v2Dim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + + List rows = buildRows("urn:li:memberId:1", "urn:li:groupId:2"); + File v1Dir = PinotSegmentUtil.createSegment(tableConfig(), schemaV1, "v1seg", _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + File v2Dir = PinotSegmentUtil.createSegment(tableConfig(), schemaV2, "v2seg", _segmentOutputDir.toString(), + new GenericRowRecordReader(rows)); + + ImmutableSegment v1 = ImmutableSegmentLoader.load(v1Dir, ReadMode.mmap); + ImmutableSegment v2 = ImmutableSegmentLoader.load(v2Dir, ReadMode.mmap); + try { + Dictionary v1Dict = v1.getDataSource(URN_COL).getDictionary(); + Dictionary v2Dict = v2.getDataSource(URN_COL).getDictionary(); + // Both dictionaries reconstruct the same URN strings — the per-entity ranges in segment + // metadata carry the entity type, not just the ordinal. + int v1MemberDictId = v1Dict.indexOf("urn:li:memberId:1"); + int v2MemberDictId = v2Dict.indexOf("urn:li:memberId:1"); + assertTrue(v1MemberDictId >= 0); + assertTrue(v2MemberDictId >= 0); + assertEquals(v1Dict.getStringValue(v1MemberDictId), "urn:li:memberId:1"); + assertEquals(v2Dict.getStringValue(v2MemberDictId), "urn:li:memberId:1"); + // Absolute dict IDs may differ between segments (because the ranges are in different + // orders), but the reconstructed URN strings agree. + } finally { + v1.destroy(); + v2.destroy(); + } + } + + // --- untyped polymorphic discovery (covers task 28+29 end to end) --- + + @Test + public void untypedMixedKeyTypesDiscoverPerEntityValueType() + throws Exception { + // Mixed key types across entities: memberId keys are all numeric, corpuser keys are strings. + // Auto-detect should produce a polymorphic dictionary with LONG memberId + STRING corpuser + // sub-dicts; no schema-declared urnTypes needed. + Schema schema = untypedUrnSchema(); + List rows = buildRows( + "urn:li:memberId:100", "urn:li:memberId:200", + "urn:li:corpuser:alice", "urn:li:corpuser:bob"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + DimensionFieldSpec dim = (DimensionFieldSpec) meta.getColumnMetadataMap().get(URN_COL).getFieldSpec(); + assertTrue(dim.hasUrnTypes()); + assertEquals(dim.getUrnTypes().size(), 2); + + java.util.Map byEntity = new java.util.HashMap<>(); + for (DimensionFieldSpec.UrnType ut : dim.getUrnTypes()) { + byEntity.put(ut.getEntityType(), ut.getValueType()); + } + assertEquals(byEntity.get("urn:li:memberId"), FieldSpec.DataType.LONG); + assertEquals(byEntity.get("urn:li:corpuser"), FieldSpec.DataType.STRING); + + List ranges = dim.getEntityDictRanges(); + assertNotNull(ranges); + assertEquals(ranges.size(), 2); + + ImmutableSegment seg = ImmutableSegmentLoader.load(_segmentDir, ReadMode.mmap); + try { + Dictionary dict = seg.getDictionary(URN_COL); + assertNotNull(dict); + // Round-trip through the polymorphic dictionary: each value comes back as the full URN + // string, regardless of which entity's typed sub-dict actually holds the key. + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:100")), "urn:li:memberId:100"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:memberId:200")), "urn:li:memberId:200"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:corpuser:alice")), "urn:li:corpuser:alice"); + assertEquals(dict.getStringValue(dict.indexOf("urn:li:corpuser:bob")), "urn:li:corpuser:bob"); + } finally { + seg.destroy(); + } + } + + @Test + public void untypedSinglePrefixDiscoveredAsPolymorphic() + throws Exception { + // Single-prefix untyped column: every row shares the same entity prefix and has numeric + // keys. Auto-detect should treat it as a 1-entity polymorphic column with a LONG sub-dict. + Schema schema = untypedUrnSchema(); + List rows = buildRows("urn:li:memberId:100", "urn:li:memberId:200", "urn:li:memberId:300"); + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + DimensionFieldSpec dim = (DimensionFieldSpec) meta.getColumnMetadataMap().get(URN_COL).getFieldSpec(); + assertEquals(dim.getUrnTypes().size(), 1); + assertEquals(dim.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(dim.getUrnTypes().get(0).getValueType(), FieldSpec.DataType.LONG); + + List ranges = dim.getEntityDictRanges(); + assertNotNull(ranges); + assertEquals(ranges.size(), 1); + assertEquals(ranges.get(0).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(ranges.get(0).getNumBytesPerValue(), Long.BYTES); + } + + @Test + public void declaredUrnTypesWithoutValueTypeDefaultToString() { + // setUrnTypes normalizes UrnType entries with no declared valueType to STRING. + DimensionFieldSpec dim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + dim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:corpuser"))); + assertEquals(dim.getUrnTypes().get(0).getValueType(), FieldSpec.DataType.LONG); + assertEquals(dim.getUrnTypes().get(1).getValueType(), FieldSpec.DataType.STRING); + } + + // --- regression: null URN row in a mixed-type polymorphic column --- + + @Test + public void mixedTypePolymorphicColumnAcceptsNullUrnRow() + throws Exception { + // Schema: declared polymorphic [LONG memberId, STRING corpuser]. One row has a null URN. + // The null sentinel must be the LONG default (entity 0's type), not the URN sentinel + // String, otherwise UrnPolymorphicDictionaryCreator.build casts String to Number and + // ClassCastExceptions. + DimensionFieldSpec mixedDim = new DimensionFieldSpec(URN_COL, FieldSpec.DataType.URN, true); + mixedDim.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:corpuser", FieldSpec.DataType.STRING))); + Schema schema = new Schema.SchemaBuilder() + .addField(mixedDim) + .addDateTime(TIME_COL, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + List rows = new ArrayList<>(); + GenericRow nullRow = new GenericRow(); + nullRow.putValue(URN_COL, null); + nullRow.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(nullRow); + GenericRow realRow = new GenericRow(); + realRow.putValue(URN_COL, "urn:li:memberId:42"); + realRow.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(realRow); + + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + assertEquals(meta.getTotalDocs(), 2); + } + + // --- regression: URN null sentinel must not become a discovered entity --- + + @Test + public void untypedColumnNullRowsDoNotProduceSpuriousEntity() + throws Exception { + // Untyped URN column ingesting real URN data plus one null row. The auto-detect must skip + // the URN null sentinel that NullValueTransformer fills in, so the discovered urnTypes only + // contain real entities — not a phantom "urn:pinot:null". + Schema schema = untypedUrnSchema(); + List rows = new ArrayList<>(); + rows.addAll(buildRows("urn:li:memberId:100", "urn:li:memberId:200")); + GenericRow nullRow = new GenericRow(); + nullRow.putValue(URN_COL, null); + nullRow.putValue(TIME_COL, System.currentTimeMillis()); + rows.add(nullRow); + + _segmentDir = PinotSegmentUtil.createSegment(tableConfig(), schema, SEGMENT_NAME, + _segmentOutputDir.toString(), new GenericRowRecordReader(rows)); + + SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); + DimensionFieldSpec dim = (DimensionFieldSpec) meta.getColumnMetadataMap().get(URN_COL).getFieldSpec(); + assertEquals(dim.getUrnTypes().size(), 1); + assertEquals(dim.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + for (DimensionFieldSpec.UrnType ut : dim.getUrnTypes()) { + assertTrue(!ut.getEntityType().contains("urn:pinot:null"), + "Auto-detect leaked the URN null sentinel into urnTypes: " + ut.getEntityType()); + } + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionaryTest.java new file mode 100644 index 0000000000..0b96213f5b --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionaryTest.java @@ -0,0 +1,457 @@ +/** + * 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.segment.local.segment.index.readers; + +import it.unimi.dsi.fastutil.ints.IntSet; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.Arrays; +import java.util.List; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Unit tests for {@link UrnPolymorphicDictionary} using in-memory fake delegate dictionaries. + * Exercises the per-entity sub-dictionary routing: dict-ID partitioning, prefix-based lookup + * dispatch, cross-entity range handling, and pass-through accessors. + * + *

The composition uses mixed delegate types — one LONG sub-dict, one STRING sub-dict — to + * demonstrate the polymorphic design even though the all-numeric case is what's exercised in + * production today. + */ +public class UrnPolymorphicDictionaryTest { + + /** In-memory LONG dictionary fake: an array of LONGs sorted ascending. */ + private static final class FakeLongDict implements Dictionary { + private final long[] _values; + + FakeLongDict(long[] values) { + _values = values; + } + + @Override + public boolean isSorted() { + return true; + } + + @Override + public DataType getValueType() { + return DataType.LONG; + } + + @Override + public int length() { + return _values.length; + } + + @Override + public int indexOf(String stringValue) { + return indexOf(Long.parseLong(stringValue)); + } + + @Override + public int indexOf(long longValue) { + int idx = Arrays.binarySearch(_values, longValue); + return idx >= 0 ? idx : NULL_VALUE_INDEX; + } + + @Override + public int insertionIndexOf(String stringValue) { + return Arrays.binarySearch(_values, Long.parseLong(stringValue)); + } + + @Override + public IntSet getDictIdsInRange(String l, String u, boolean il, boolean iu) { + long low = Long.parseLong(l); + long high = Long.parseLong(u); + it.unimi.dsi.fastutil.ints.IntOpenHashSet result = new it.unimi.dsi.fastutil.ints.IntOpenHashSet(); + for (int i = 0; i < _values.length; i++) { + long v = _values[i]; + boolean lowOk = il ? v >= low : v > low; + boolean highOk = iu ? v <= high : v < high; + if (lowOk && highOk) { + result.add(i); + } + } + return result; + } + + @Override + public int compare(int a, int b) { + return Long.compare(_values[a], _values[b]); + } + + @Override + public Comparable getMinVal() { + return _values[0]; + } + + @Override + public Comparable getMaxVal() { + return _values[_values.length - 1]; + } + + @Override + public Object getSortedValues() { + return _values; + } + + @Override + public Object get(int dictId) { + return _values[dictId]; + } + + @Override + public int getIntValue(int dictId) { + return (int) _values[dictId]; + } + + @Override + public long getLongValue(int dictId) { + return _values[dictId]; + } + + @Override + public float getFloatValue(int dictId) { + return _values[dictId]; + } + + @Override + public double getDoubleValue(int dictId) { + return _values[dictId]; + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return BigDecimal.valueOf(_values[dictId]); + } + + @Override + public String getStringValue(int dictId) { + return Long.toString(_values[dictId]); + } + + @Override + public void close() + throws IOException { + } + } + + /** In-memory STRING dictionary fake: an array of strings sorted lexicographically. */ + private static final class FakeStringDict implements Dictionary { + private final String[] _values; + + FakeStringDict(String[] values) { + _values = values; + } + + @Override + public boolean isSorted() { + return true; + } + + @Override + public DataType getValueType() { + return DataType.STRING; + } + + @Override + public int length() { + return _values.length; + } + + @Override + public int indexOf(String stringValue) { + int idx = Arrays.binarySearch(_values, stringValue); + return idx >= 0 ? idx : NULL_VALUE_INDEX; + } + + @Override + public int insertionIndexOf(String stringValue) { + return Arrays.binarySearch(_values, stringValue); + } + + @Override + public IntSet getDictIdsInRange(String l, String u, boolean il, boolean iu) { + it.unimi.dsi.fastutil.ints.IntOpenHashSet result = new it.unimi.dsi.fastutil.ints.IntOpenHashSet(); + for (int i = 0; i < _values.length; i++) { + int cmpLow = _values[i].compareTo(l); + int cmpHigh = _values[i].compareTo(u); + boolean lowOk = il ? cmpLow >= 0 : cmpLow > 0; + boolean highOk = iu ? cmpHigh <= 0 : cmpHigh < 0; + if (lowOk && highOk) { + result.add(i); + } + } + return result; + } + + @Override + public int compare(int a, int b) { + return _values[a].compareTo(_values[b]); + } + + @Override + public Comparable getMinVal() { + return _values[0]; + } + + @Override + public Comparable getMaxVal() { + return _values[_values.length - 1]; + } + + @Override + public Object getSortedValues() { + return _values; + } + + @Override + public Object get(int dictId) { + return _values[dictId]; + } + + @Override + public int getIntValue(int dictId) { + return Integer.parseInt(_values[dictId]); + } + + @Override + public long getLongValue(int dictId) { + return Long.parseLong(_values[dictId]); + } + + @Override + public float getFloatValue(int dictId) { + return Float.parseFloat(_values[dictId]); + } + + @Override + public double getDoubleValue(int dictId) { + return Double.parseDouble(_values[dictId]); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return new BigDecimal(_values[dictId]); + } + + @Override + public String getStringValue(int dictId) { + return _values[dictId]; + } + + @Override + public void close() + throws IOException { + } + } + + /** + * Reference layout used by the tests: + * 

+   *   urn:li:memberId  → LONG  sub-dict {100, 200, 300}     globals [0..3)
+   *   urn:li:groupId   → LONG  sub-dict {100, 400}          globals [3..5) — key 100 collides!
+   *   urn:li:corpUser  → STRING sub-dict {"alice", "bob"}    globals [5..7)
+   *
+ */ + private static UrnPolymorphicDictionary referenceDict() { + return new UrnPolymorphicDictionary(Arrays.asList( + new UrnPolymorphicDictionary.SubDict("urn:li:memberId", new FakeLongDict(new long[]{100, 200, 300})), + new UrnPolymorphicDictionary.SubDict("urn:li:groupId", new FakeLongDict(new long[]{100, 400})), + new UrnPolymorphicDictionary.SubDict("urn:li:corpUser", new FakeStringDict(new String[]{"alice", "bob"})))); + } + + // --- indexOf routes by entity prefix to the appropriate sub-dict --- + + @Test + public void indexOfRoutesLongSubDict() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:memberId:100"), 0); + assertEquals(d.indexOf("urn:li:memberId:200"), 1); + assertEquals(d.indexOf("urn:li:memberId:300"), 2); + } + + @Test + public void indexOfDisambiguatesKeyCollisionAcrossEntities() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:memberId:100"), 0); + // groupId has its own sub-dict; same key 100 lives at a different global ID. + assertEquals(d.indexOf("urn:li:groupId:100"), 3); + } + + @Test + public void indexOfRoutesStringSubDict() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:corpUser:alice"), 5); + assertEquals(d.indexOf("urn:li:corpUser:bob"), 6); + } + + @Test + public void indexOfReturnsNullForMissingKey() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:memberId:999"), Dictionary.NULL_VALUE_INDEX); + assertEquals(d.indexOf("urn:li:corpUser:charlie"), Dictionary.NULL_VALUE_INDEX); + } + + @Test + public void indexOfReturnsNullForUnknownEntity() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:notARealEntity:1"), Dictionary.NULL_VALUE_INDEX); + } + + // --- getStringValue reconstructs URN strings from each sub-dict --- + + @Test + public void getStringValueReconstructsAcrossAllRanges() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.getStringValue(0), "urn:li:memberId:100"); + assertEquals(d.getStringValue(1), "urn:li:memberId:200"); + assertEquals(d.getStringValue(2), "urn:li:memberId:300"); + assertEquals(d.getStringValue(3), "urn:li:groupId:100"); + assertEquals(d.getStringValue(4), "urn:li:groupId:400"); + assertEquals(d.getStringValue(5), "urn:li:corpUser:alice"); + assertEquals(d.getStringValue(6), "urn:li:corpUser:bob"); + } + + @Test + public void getStringValueOutOfBoundsReturnsNull() { + UrnPolymorphicDictionary d = referenceDict(); + assertNull(d.getStringValue(-1)); + assertNull(d.getStringValue(7)); + } + + @Test + public void readStringValuesBatchRespectsOrdering() { + UrnPolymorphicDictionary d = referenceDict(); + String[] out = new String[3]; + d.readStringValues(new int[]{6, 3, 0}, 3, out); + assertEquals(out[0], "urn:li:corpUser:bob"); + assertEquals(out[1], "urn:li:groupId:100"); + assertEquals(out[2], "urn:li:memberId:100"); + } + + // --- numeric accessors pass through to whichever sub-dict can answer --- + + @Test + public void getLongValuePassesThrough() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.getLongValue(0), 100L); + assertEquals(d.getLongValue(3), 100L); + assertEquals(d.getLongValue(4), 400L); + } + + @Test + public void getIntValuePassesThroughForSmallNumericSubDicts() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.getIntValue(0), 100); + assertEquals(d.getIntValue(4), 400); + } + + // --- range queries within a single entity --- + + @Test + public void getDictIdsInRangeSameEntity() { + UrnPolymorphicDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:memberId:100", "urn:li:memberId:200", true, true); + assertEquals(result.size(), 2); + assertTrue(result.contains(0)); + assertTrue(result.contains(1)); + } + + @Test + public void getDictIdsInRangeCrossEntityReturnsEmpty() { + UrnPolymorphicDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:memberId:100", "urn:li:groupId:200", true, true); + assertEquals(result.size(), 0); + } + + // --- metadata --- + + @Test + public void getValueTypeIsAlwaysString() { + assertEquals(referenceDict().getValueType(), DataType.STRING); + } + + @Test + public void isSortedReturnsFalse() { + assertTrue(!referenceDict().isSorted()); + } + + @Test + public void lengthIsSumOfSubDicts() { + assertEquals(referenceDict().length(), 7); + } + + @Test + public void getMinValIsFirstReconstructedUrn() { + assertEquals(referenceDict().getMinVal(), "urn:li:memberId:100"); + } + + @Test + public void getMaxValIsLastReconstructedUrn() { + assertEquals(referenceDict().getMaxVal(), "urn:li:corpUser:bob"); + } + + @Test + public void getEntityTypeForDictId() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.getEntityTypeForDictId(0), "urn:li:memberId"); + assertEquals(d.getEntityTypeForDictId(3), "urn:li:groupId"); + assertEquals(d.getEntityTypeForDictId(6), "urn:li:corpUser"); + assertNull(d.getEntityTypeForDictId(-1)); + assertNull(d.getEntityTypeForDictId(7)); + } + + // --- single sub-dict (degenerate case) --- + + @Test + public void singleSubDictBehavesAsExpected() { + UrnPolymorphicDictionary d = new UrnPolymorphicDictionary(List.of( + new UrnPolymorphicDictionary.SubDict("urn:li:memberId", new FakeLongDict(new long[]{1, 2, 3})))); + assertEquals(d.length(), 3); + assertEquals(d.indexOf("urn:li:memberId:2"), 1); + assertEquals(d.getStringValue(0), "urn:li:memberId:1"); + } + + // --- regression: non-numeric suffix on LONG sub-dict must not propagate NumberFormatException --- + + @Test + public void indexOfNonNumericSuffixOnLongSubDictReturnsNullIndex() { + UrnPolymorphicDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:memberId:not-a-number"), + org.apache.pinot.segment.spi.index.reader.Dictionary.NULL_VALUE_INDEX); + } + + @Test + public void getDictIdsInRangeNonNumericSuffixOnLongSubDictReturnsEmpty() { + UrnPolymorphicDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:memberId:abc", "urn:li:memberId:zzz", true, true); + assertEquals(result.size(), 0); + } +} diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java index 35b319dd60..531dcf8d14 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java @@ -129,6 +129,13 @@ public static class Column { public static final String COLUMN_PROPS_KEY_PREFIX = "column."; public static final String SCHEMA_MAX_LENGTH = "schemaMaxLength"; public static final String SCHEMA_MAX_LENGTH_EXCEED_STRATEGY = "schemaMaxLengthExceedStrategy"; + // URN column metadata: the entity types list is serialized as JSON to a single property. + public static final String URN_TYPES = "urnTypes"; + // Multi-entity URN sorted-LONG dictionary layout: per-entity dictionary-ID ranges, + // serialized as "entityType,endExclusive;entityType,endExclusive;...". Identifies how the + // main column's dictionary is partitioned by entity type at the (entityIdx, longKey)-sorted + // order produced by the segment builder. + public static final String URN_ENTITY_DICT_RANGES = "urn.entityDictRanges"; public static String getKeyFor(String column, String key) { return COLUMN_PROPS_KEY_PREFIX + column + "." + key; diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/StatsCollectorConfig.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/StatsCollectorConfig.java index fdd8abfcd7..1e0b9817ec 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/StatsCollectorConfig.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/StatsCollectorConfig.java @@ -39,6 +39,8 @@ public class StatsCollectorConfig { private final Schema _schema; private final SegmentPartitionConfig _segmentPartitionConfig; private final Map _columnFieldConfigMap; + /** Extra field specs for synthetic columns not present in the schema (e.g., companion entity-index columns). */ + private final Map _extraFieldSpecs = new HashMap<>(); /** * Constructor for the class. @@ -60,9 +62,14 @@ public StatsCollectorConfig(TableConfig tableConfig, Schema schema, } } + public void registerExtraFieldSpec(FieldSpec fieldSpec) { + _extraFieldSpecs.put(fieldSpec.getName(), fieldSpec); + } + @Nullable public FieldSpec getFieldSpecForColumn(String column) { - return _schema.getFieldSpecFor(column); + FieldSpec spec = _schema.getFieldSpecFor(column); + return spec != null ? spec : _extraFieldSpecs.get(column); } @Nullable diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java index dea80c7462..35601a83f9 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/ForwardIndexCreator.java @@ -60,6 +60,7 @@ default void add(@Nonnull Object cellValue, int dictId) { putBigDecimal((BigDecimal) cellValue); break; case STRING: + case URN: putString((String) cellValue); break; case BYTES: diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java index 91eed38295..a8f0fae0d2 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java @@ -49,6 +49,7 @@ import org.apache.pinot.spi.data.TimeGranularitySpec; import org.apache.pinot.spi.env.CommonsConfigurationUtils; import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.JsonUtils; public class ColumnMetadataImpl implements ColumnMetadata { @@ -223,7 +224,7 @@ public static ColumnMetadataImpl fromPropertiesConfiguration(String column, Prop FieldSpec fieldSpec = generateFieldSpec(column, config); builder.setFieldSpec(fieldSpec); - DataType storedType = fieldSpec.getDataType().getStoredType(); + DataType storedType = fieldSpec.getEffectiveStoredType(); if (fieldSpec instanceof ComplexFieldSpec) { // Complex field does not have min/max value @@ -303,8 +304,51 @@ public static FieldSpec generateFieldSpec(String column, PropertiesConfiguration switch (fieldType) { case DIMENSION: boolean isSingleValue = config.getBoolean(Column.getKeyFor(column, Column.IS_SINGLE_VALUED)); - fieldSpec = new DimensionFieldSpec(fieldName, dataType, isSingleValue, maxLength, - defaultNullValueString, maxLengthExceedStrategy); + // For URN columns, defer defaultNullValue construction until after urnTypes are set. + // Typed URN (LONG/INT) columns have a numeric default null that would fail URN string validation. + DimensionFieldSpec dimSpec = + new DimensionFieldSpec(fieldName, dataType, isSingleValue, maxLength, + dataType == DataType.URN ? null : defaultNullValueString, maxLengthExceedStrategy); + if (dataType == DataType.URN) { + // The urnTypes JSON may contain commas, which PropertiesConfiguration splits into multiple values. + // Re-joining the array restores the original JSON. + String[] urnTypesParts = config.getStringArray(Column.getKeyFor(column, Column.URN_TYPES)); + if (urnTypesParts.length > 0) { + try { + String urnTypesJson = String.join(",", urnTypesParts); + List urnTypes = + JsonUtils.stringToObject(urnTypesJson, new com.fasterxml.jackson.core.type.TypeReference< + List>() { + }); + dimSpec.setUrnTypes(urnTypes); + } catch (Exception e) { + // urnTypes are advisory; the column is still readable without them + } + } + // For untyped URN (STRING storage), re-apply the persisted default null value if valid + if (!dimSpec.hasTypedUrnStorage() && defaultNullValueString != null) { + try { + dimSpec.setDefaultNullValue(defaultNullValueString); + } catch (Exception ignored) { + // If the persisted value is not a valid URN, fall back to the system default + } + } + // Multi-entity URN sorted-LONG dictionary layout: load per-entity dict ID ranges if + // the segment was built with the UrnDictionary format. PropertiesConfiguration splits + // the value on commas so we rejoin to restore the original "entity,end;entity,end" + // form, mirroring the URN_TYPES handling above. + String[] entityDictRangesParts = + config.getStringArray(Column.getKeyFor(column, Column.URN_ENTITY_DICT_RANGES)); + if (entityDictRangesParts.length > 0) { + try { + dimSpec.setEntityDictRanges( + DimensionFieldSpec.parseEntityDictRanges(String.join(",", entityDictRangesParts))); + } catch (Exception ignored) { + // Malformed ranges are non-fatal; the column simply won't take the sorted-dict path. + } + } + } + fieldSpec = dimSpec; break; case METRIC: fieldSpec = diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/DimensionFieldSpec.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/DimensionFieldSpec.java index c98fcd5cf3..ba120dd4b5 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/DimensionFieldSpec.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/DimensionFieldSpec.java @@ -20,12 +20,76 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.Collections; +import java.util.List; +import java.util.Objects; import javax.annotation.Nullable; +import org.apache.pinot.spi.utils.JsonUtils; @JsonIgnoreProperties(ignoreUnknown = true) public final class DimensionFieldSpec extends FieldSpec { + /** + * Prefix for the transient row-level field that the URN encoding transformer stashes the + * 0-based entity index into during multi-entity URN segment build. For a column + * {@code "actorId"} the transient field is {@code "__urnEntityIdxBuild_actorId"}. The field + * is consumed by {@code UrnMultiEntityPreIndexStatsCollector} and + * {@code SegmentColumnarIndexCreator.indexRow} to pair {@code (entityIdx, longKey)} pairs for + * the sorted-LONG dictionary, then dropped — it is not part of the segment schema and is + * never persisted as a column. + */ + public static final String URN_TRANSIENT_ENTITY_IDX_PREFIX = "__urnEntityIdxBuild_"; + + /** + * For columns with {@code dataType = URN}: the set of (entityType, valueType) pairs that this + * column may contain. Each entry declares one valid entity type prefix and, optionally, the + * stored type for its key component. + * + *
    + *
  • A single-entry list means the column is uniform: every URN shares the same entity type + * prefix, enabling the most compact storage (prefix stripped; only the key stored).
    • + *
  • A multi-entry list means the column is heterogeneous. A small per-row entity-type index + * is stored alongside the key, keeping storage compact while supporting multiple types.
    • + *
  • {@code null} or empty list means fully untyped: URNs are stored as plain strings with + * standard dictionary encoding. Correct but least efficient.
    • + *
+ * + * Example schema snippet: + * 
+   * // Single entity type, numeric key
+   * "urnTypes": [{ "entityType": "urn:li:memberId", "valueType": "LONG" }]
+   *
+   * // Two entity types, both numeric — shared entity-type dictionary
+   * "urnTypes": [
+   *   { "entityType": "urn:li:memberId", "valueType": "LONG" },
+   *   { "entityType": "urn:li:groupId",  "valueType": "LONG" }
+   * ]
+   *
+   * // Entity type declared but key type left to auto-detection at segment build
+   * "urnTypes": [{ "entityType": "urn:li:corpUser" }]
+   *
+ */ + @Nullable + private List _urnTypes; + + /** + * For multi-entity URN columns whose effective stored type is numeric: per-entity dictionary + * ID ranges populated at segment load time from the {@code urn.entityDictRanges} metadata + * property. The main column dictionary stores LONG keys sorted by {@code (entityIdx, longKey)}; + * each entity's entries occupy a contiguous range of dict IDs. + * + *

Not part of the user-declared schema — this is a segment-build-time output. Always + * {@code null} for schemas read from disk/JSON. Always populated from + * {@link org.apache.pinot.segment.spi.ColumnMetadata} at segment load when the column uses the + * sorted-range UrnDictionary layout. + */ + @JsonIgnore + @Nullable + private List _entityDictRanges; + // Default constructor required by JSON de-serializer. DO NOT REMOVE. public DimensionFieldSpec() { super(); @@ -62,16 +126,530 @@ public DimensionFieldSpec(String name, DataType dataType, boolean isSingleValueF _virtualColumnProvider = virtualColumnProviderClass.getName(); } + @Nullable + public List getUrnTypes() { + return _urnTypes; + } + + // Required by JSON de-serializer. DO NOT REMOVE. + public void setUrnTypes(@Nullable List urnTypes) { + if (urnTypes != null && !urnTypes.isEmpty()) { + List normalized = new java.util.ArrayList<>(urnTypes.size()); + for (UrnType t : urnTypes) { + Objects.requireNonNull(t, "urnTypes entry must not be null"); + t.validate(); + // Default unspecified valueType to STRING -- every URN entity needs a concrete + // sub-dictionary type for the polymorphic layout, and STRING is the safe default that + // accepts any per-entity key shape. + if (t.getValueType() == null) { + normalized.add(UrnType.of(t.getEntityType(), DataType.STRING)); + } else { + normalized.add(t); + } + } + _urnTypes = Collections.unmodifiableList(normalized); + } else { + _urnTypes = null; + } + } + + /** + * Returns {@code true} when the column has a non-empty {@code urnTypes} list, meaning the + * segment builder can use the compact storage path rather than full-string encoding. + */ + @JsonIgnore + public boolean hasUrnTypes() { + return _urnTypes != null && !_urnTypes.isEmpty(); + } + + /** + * Returns {@code true} when the column has exactly one declared URN entity type — i.e. all + * values share the same prefix. This is the case where the entity prefix can be stripped + * entirely from per-row storage. + */ + @JsonIgnore + public boolean isSingleUrnEntityType() { + return _urnTypes != null && _urnTypes.size() == 1; + } + + /** + * Per-entity dictionary-ID ranges for multi-entity URN columns using the sorted LONG dictionary + * layout. The ranges describe how the main column's LONG dictionary is partitioned: entity + * {@code i} owns dict IDs {@code [prevEnd, ranges.get(i).endDictIdExclusive)}. + */ + @JsonIgnore + @Nullable + public List getEntityDictRanges() { + return _entityDictRanges; + } + + /** Populated at segment load time; not part of the user-declared schema. */ + @JsonIgnore + public void setEntityDictRanges(@Nullable List ranges) { + if (ranges == null || ranges.isEmpty()) { + _entityDictRanges = null; + } else { + for (EntityDictRange r : ranges) { + Objects.requireNonNull(r, "EntityDictRange must not be null"); + Objects.requireNonNull(r.getEntityType(), "EntityDictRange.entityType must not be null"); + } + _entityDictRanges = Collections.unmodifiableList(ranges); + } + } + + /** + * Parses the encoded {@code urn.entityDictRanges} metadata property value of the form + * {@code "entityType,end,dictClass,numBytesPerValue;..."} — each entity declares its + * own sub-dictionary type and per-entry byte width. Returns {@code null} for null or blank + * input. + * + * @throws IllegalArgumentException if the format is malformed. + */ + @Nullable + public static List parseEntityDictRanges(@Nullable String encoded) { + if (encoded == null) { + return null; + } + String trimmed = encoded.trim(); + if (trimmed.isEmpty()) { + return null; + } + String[] parts = trimmed.split(";", -1); + List ranges = new java.util.ArrayList<>(parts.length); + for (String part : parts) { + if (part.isEmpty()) { + continue; + } + String[] fields = part.split(",", -1); + if (fields.length != 4) { + throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); + } + String entityType = fields[0]; + if (entityType.isEmpty()) { + throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); + } + int endDictIdExclusive; + try { + endDictIdExclusive = Integer.parseInt(fields[1]); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Malformed entityDictRanges end: '" + part + "'", e); + } + DataType dictClass; + try { + dictClass = DataType.valueOf(fields[2]); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Malformed entityDictRanges dictClass: '" + part + "'", e); + } + int numBytesPerValue; + try { + numBytesPerValue = Integer.parseInt(fields[3]); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Malformed entityDictRanges numBytesPerValue: '" + part + "'", e); + } + ranges.add(new EntityDictRange(entityType, endDictIdExclusive, dictClass, numBytesPerValue)); + } + return ranges.isEmpty() ? null : ranges; + } + + /** Serializes {@link #getEntityDictRanges()} into the metadata-property form. */ + @Nullable + public static String serializeEntityDictRanges(@Nullable List ranges) { + if (ranges == null || ranges.isEmpty()) { + return null; + } + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < ranges.size(); i++) { + if (i > 0) { + sb.append(';'); + } + EntityDictRange r = ranges.get(i); + sb.append(r.getEntityType()).append(',').append(r.getEndDictIdExclusive()) + .append(',').append(r.getDictClass().name()).append(',').append(r.getNumBytesPerValue()); + } + return sb.toString(); + } + + @Override + public Object getDefaultNullValue() { + if (_dataType == DataType.URN) { + // The null sentinel ends up in entity 0's sub-dictionary at build time; pick a default + // that matches entity 0's declared sub-dict type so the value is type-compatible. For + // mixed-type columns where effectiveStoredType is STRING but entity 0 is LONG/INT, this + // avoids feeding a String sentinel into a numeric sub-dictionary. + if (_urnTypes != null && !_urnTypes.isEmpty()) { + DataType vt = _urnTypes.get(0).getValueType(); + if (vt == DataType.LONG) { + return DEFAULT_DIMENSION_NULL_VALUE_OF_LONG; + } + if (vt == DataType.INT) { + return DEFAULT_DIMENSION_NULL_VALUE_OF_INT; + } + // STRING entity 0: fall through to super (URN sentinel string). + } else { + DataType est = getEffectiveStoredType(); + if (est == DataType.LONG) { + return DEFAULT_DIMENSION_NULL_VALUE_OF_LONG; + } + if (est == DataType.INT) { + return DEFAULT_DIMENSION_NULL_VALUE_OF_INT; + } + } + } + return super.getDefaultNullValue(); + } + + /** + * Returns the physical stored type for this URN column. + * + *

    + *
  • Single-entity with a {@code LONG} or {@code INT} {@code valueType}: returns that + * numeric type — the key is stored compactly without the URN prefix.
    • + *
  • Multi-entity: always returns {@code STRING}, even when all entity types share a + * numeric {@code valueType}. A bare numeric key cannot distinguish + * {@code urn:li:memberId:100} from {@code urn:li:groupId:100}, so full URN strings must + * be stored.
    • + *
  • Otherwise {@code STRING} is returned (untyped or non-numeric key type).
    • + *
+ */ + @JsonIgnore + @Override + public DataType getEffectiveStoredType() { + if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.isEmpty()) { + return _dataType.getStoredType(); + } + // Multi-entity columns always use STRING storage: different entity types may have + // overlapping numeric key spaces (e.g. both memberId:100 and groupId:100 are valid), + // so a plain LONG key cannot uniquely identify a URN value across entity types. + if (_urnTypes.size() > 1) { + // Multi-entity URN: when every declared entity type has a numeric value type, store as + // LONG. The sorted-LONG UrnDictionary layout uses companion ranges to disambiguate + // overlapping keys across entity types (memberId:100 vs groupId:100). Falls back to + // STRING storage when any declared entity has a non-numeric value type. + for (UrnType ut : _urnTypes) { + DataType vt = ut.getValueType(); + if (vt != DataType.LONG && vt != DataType.INT) { + return DataType.STRING; + } + } + return DataType.LONG; + } + DataType common = _urnTypes.get(0).getValueType(); + if (common == null || (common != DataType.LONG && common != DataType.INT)) { + return DataType.STRING; + } + return common; + } + + /** + * Returns {@code true} when this column's keys will be stored as a typed numeric (LONG or INT) + * rather than as a full URN string. + */ + @JsonIgnore + public boolean hasTypedUrnStorage() { + DataType est = getEffectiveStoredType(); + return est == DataType.LONG || est == DataType.INT; + } + + /** + * Returns {@code true} when the column uses the multi-entity, all-numeric sorted-LONG dictionary + * layout. Two restrictions distinguish this from the broader {@link + * #usesPolymorphicUrnDictionary()}: + *
    + *
  • at least two declared entities — single-entity columns don't have the cross-entity + * overlap that motivates the sorted-LONG layout; and
    • + *
  • every entity's {@code valueType} is {@code LONG} or {@code INT}.
    • + *
+ * Callers that need only "is this column polymorphic" (entity-idx-stashing, polymorphic creator + * routing) should prefer the broader {@link #usesPolymorphicUrnDictionary()}. Callers that + * specifically want to skip the {@code urnKeyLong} pushdown rewrite for multi-entity numeric + * columns — because the underlying LONG can map to multiple entities — use this predicate. + */ + @JsonIgnore + public boolean usesSortedUrnDictionary() { + if (!usesPolymorphicUrnDictionary() || _urnTypes.size() <= 1) { + return false; + } + for (DimensionFieldSpec.UrnType ut : _urnTypes) { + DataType vt = ut.getValueType(); + if (vt != DataType.LONG && vt != DataType.INT) { + return false; + } + } + return true; + } + + /** + * Returns {@code true} when the column uses the polymorphic per-entity sub-dictionary layout + * (see {@code UrnPolymorphicDictionaryCreator}). True for any URN column — single-entity or + * multi-entity — whose declared entities all have a {@code valueType} ({@code LONG}, {@code + * INT}, or {@code STRING}). The dictionary file is then a concatenation of per-entity typed + * sub-dictionaries — {@code LongDictionary} for LONG/INT entities, {@code StringDictionary} + * for STRING entities — and the segment builder needs to (a) stash the per-row entity index + * in a transient field so the dictionary creator can lay out per-entity dict-ID ranges and + * (b) skip the cross-entity min/max metadata write since comparison across entity types isn't + * meaningful. + * + *

Single-entity columns are valid polymorphic columns with one entry — entity index 0 on + * every row, one sub-dictionary covering the whole dict-ID space, the same per-entity + * sub-dictionary reader as the multi-entity case. This removes the historical special-case + * wrappers ({@code UrnReconstructingLongDictionary}, {@code UrnReconstructingStringDictionary}) + * that prepended the entity prefix on read. + */ + @JsonIgnore + public boolean usesPolymorphicUrnDictionary() { + if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.isEmpty()) { + return false; + } + for (DimensionFieldSpec.UrnType ut : _urnTypes) { + DataType vt = ut.getValueType(); + if (vt != DataType.LONG && vt != DataType.INT && vt != DataType.STRING) { + return false; + } + } + return true; + } + + /** + * Returns the transient row-attribute name the URN encoding transformer uses to stash the + * 0-based entity index for this column during ingestion. Not a real column — never appears in + * segment metadata. + */ + @JsonIgnore + public String transientUrnEntityIdxFieldName() { + return URN_TRANSIENT_ENTITY_IDX_PREFIX + _name; + } + @JsonIgnore @Override public FieldType getFieldType() { return FieldType.DIMENSION; } + @Override + public ObjectNode toJsonObject() { + ObjectNode json = super.toJsonObject(); + if (_urnTypes != null && !_urnTypes.isEmpty()) { + ArrayNode arr = json.putArray("urnTypes"); + for (UrnType t : _urnTypes) { + arr.add(t.toJsonObject()); + } + } + return json; + } + + @Override + public boolean equals(Object o) { + if (!super.equals(o)) { + return false; + } + if (!(o instanceof DimensionFieldSpec)) { + return false; + } + DimensionFieldSpec other = (DimensionFieldSpec) o; + return Objects.equals(_urnTypes, other._urnTypes); + } + + @Override + public int hashCode() { + return 31 * super.hashCode() + Objects.hashCode(_urnTypes); + } + @Override public String toString() { - return "< field type: DIMENSION, field name: " + _name + ", data type: " + _dataType + ", is single-value field: " - + _singleValueField + ", default null value: " + _defaultNullValue + ", max length exceed strategy: " - + _maxLengthExceedStrategy + " >"; + StringBuilder sb = new StringBuilder("< field type: DIMENSION, field name: ").append(_name) + .append(", data type: ").append(_dataType) + .append(", is single-value field: ").append(_singleValueField) + .append(", default null value: ").append(_defaultNullValue) + .append(", max length exceed strategy: ").append(_maxLengthExceedStrategy); + if (_urnTypes != null && !_urnTypes.isEmpty()) { + sb.append(", urnTypes: ").append(_urnTypes); + } + sb.append(" >"); + return sb.toString(); + } + + /** + * Declares one valid entity type for a URN column, plus the optional stored type of its key + * component. + * + *

The {@code entityType} must be of the form {@code "urn::"}, matching the + * canonical entity prefix used in LinkedIn-style URNs. + * + *

The {@code valueType} is the physical storage type for the key part of the URN. When + * absent, the segment builder will auto-detect the type (preferring {@code LONG} when all keys + * in the segment are parseable as long integers, falling back to {@code STRING} otherwise). + * Explicitly setting {@code valueType} ensures consistent encoding across all segments and is + * recommended for production tables once the key type is stable. + */ + @JsonIgnoreProperties(ignoreUnknown = true) + public static final class UrnType { + private String _entityType; + @Nullable + private DataType _valueType; + + // Default constructor required by JSON de-serializer. DO NOT REMOVE. + public UrnType() { + } + + public UrnType(String entityType, @Nullable DataType valueType) { + _entityType = entityType; + _valueType = valueType; + validate(); + } + + public static UrnType of(String entityType) { + return new UrnType(entityType, null); + } + + public static UrnType of(String entityType, DataType valueType) { + return new UrnType(entityType, valueType); + } + + public String getEntityType() { + return _entityType; + } + + // Required by JSON de-serializer. DO NOT REMOVE. + public void setEntityType(String entityType) { + _entityType = entityType; + } + + @Nullable + public DataType getValueType() { + return _valueType; + } + + // Required by JSON de-serializer. DO NOT REMOVE. + public void setValueType(@Nullable DataType valueType) { + _valueType = valueType; + } + + void validate() { + Objects.requireNonNull(_entityType, "entityType is required in UrnType"); + String[] parts = _entityType.split(":", -1); + if (parts.length != 3 || !"urn".equals(parts[0]) || parts[1].isEmpty() || parts[2].isEmpty()) { + throw new IllegalArgumentException( + "entityType must be of the form 'urn::', got: " + _entityType); + } + if (_valueType != null && _valueType != DataType.INT && _valueType != DataType.LONG + && _valueType != DataType.STRING) { + throw new IllegalArgumentException("valueType must be INT, LONG, or STRING, got: " + _valueType); + } + } + + /** Returns the entity prefix with trailing colon, e.g. {@code "urn:li:memberId:"}. */ + @JsonIgnore + public String entityPrefix() { + return _entityType + ":"; + } + + public ObjectNode toJsonObject() { + ObjectNode node = JsonUtils.newObjectNode(); + node.put("entityType", _entityType); + if (_valueType != null) { + node.put("valueType", _valueType.name()); + } + return node; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof UrnType)) { + return false; + } + UrnType other = (UrnType) o; + return Objects.equals(_entityType, other._entityType) && Objects.equals(_valueType, other._valueType); + } + + @Override + public int hashCode() { + return Objects.hash(_entityType, _valueType); + } + + @Override + public String toString() { + return _valueType != null ? _entityType + "(" + _valueType + ")" : _entityType; + } + } + + /** + * One entry in a multi-entity URN column's per-entity dictionary-ID range table. Records: + *

    + *
  • {@code entityType} — the entity-type prefix (e.g. {@code "urn:li:memberId"});
    • + *
  • {@code endDictIdExclusive} — exclusive upper bound (in dict IDs) of this entity's + * region; the entity owns dict IDs {@code [prevEnd, endDictIdExclusive)};
    • + *
  • {@code dictClass} — the {@link DataType} of the sub-dictionary for this entity + * ({@link DataType#LONG}, {@link DataType#INT}, {@link DataType#STRING}, ...). Allows + * mixed-type multi-entity URN columns to store each entity in a typed sub-dictionary + * (e.g. {@code LongDictionary} for memberId, {@code StringDictionary} for corpuser);
    • + *
  • {@code numBytesPerValue} — bytes per entry in this entity's sub-dictionary buffer + * region. For fixed-width types ({@code LONG}, {@code INT}, ...) this is the type's + * primitive width; for {@code STRING} this is the column-max-length-style padding width + * chosen at build time.
    • + *
+ * + *

The per-entity byte ranges into the dictionary file are derived cumulatively at load time: + * entity {@code i} occupies bytes {@code [prevByteEnd, prevByteEnd + (endDictId - prevEndDictId) + * * numBytesPerValue)}. + * + *

This is a load-time data class — never user-declared. Populated from the + * {@code urn.entityDictRanges} segment-metadata property by {@code ColumnMetadataImpl}. + */ + public static final class EntityDictRange { + private final String _entityType; + private final int _endDictIdExclusive; + private final DataType _dictClass; + private final int _numBytesPerValue; + + public EntityDictRange(String entityType, int endDictIdExclusive, DataType dictClass, int numBytesPerValue) { + _entityType = entityType; + _endDictIdExclusive = endDictIdExclusive; + _dictClass = dictClass; + _numBytesPerValue = numBytesPerValue; + } + + public String getEntityType() { + return _entityType; + } + + public int getEndDictIdExclusive() { + return _endDictIdExclusive; + } + + public DataType getDictClass() { + return _dictClass; + } + + public int getNumBytesPerValue() { + return _numBytesPerValue; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof EntityDictRange)) { + return false; + } + EntityDictRange other = (EntityDictRange) o; + return _endDictIdExclusive == other._endDictIdExclusive + && _numBytesPerValue == other._numBytesPerValue + && Objects.equals(_entityType, other._entityType) + && _dictClass == other._dictClass; + } + + @Override + public int hashCode() { + return Objects.hash(_entityType, _endDictIdExclusive, _dictClass, _numBytesPerValue); + } + + @Override + public String toString() { + return _entityType + "[..." + _endDictIdExclusive + ", " + _dictClass + "/" + _numBytesPerValue + "B)"; + } } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java index 5d29412f5e..c118f253d0 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java @@ -80,6 +80,9 @@ public abstract class FieldSpec implements Comparable, Serializable { public static final Long DEFAULT_DIMENSION_NULL_VALUE_OF_TIMESTAMP = 0L; public static final String DEFAULT_DIMENSION_NULL_VALUE_OF_STRING = "null"; public static final String DEFAULT_DIMENSION_NULL_VALUE_OF_JSON = "null"; + // A well-formed but obviously-sentinel URN. Chosen so reading code that always calls + // Urn.parse() on the column does not have to special-case the default null value. + public static final String DEFAULT_DIMENSION_NULL_VALUE_OF_URN = "urn:pinot:null:0"; public static final byte[] DEFAULT_DIMENSION_NULL_VALUE_OF_BYTES = new byte[0]; public static final BigDecimal DEFAULT_DIMENSION_NULL_VALUE_OF_BIG_DECIMAL = BigDecimal.ZERO; @@ -192,6 +195,15 @@ public DataType getDataType() { return _dataType; } + /** + * Returns the physical stored type for this column. For most types this is identical to + * {@link DataType#getStoredType()}. Subclasses may override when the physical storage differs + * from the logical type (e.g. URN columns with a declared {@code valueType}). + */ + public DataType getEffectiveStoredType() { + return _dataType.getStoredType(); + } + // Required by JSON de-serializer. DO NOT REMOVE. public void setDataType(DataType dataType) { _dataType = dataType; @@ -340,6 +352,8 @@ public static Object getDefaultNullValue(FieldType fieldType, DataType dataType, return DEFAULT_DIMENSION_NULL_VALUE_OF_STRING; case JSON: return DEFAULT_DIMENSION_NULL_VALUE_OF_JSON; + case URN: + return DEFAULT_DIMENSION_NULL_VALUE_OF_URN; case BYTES: return DEFAULT_DIMENSION_NULL_VALUE_OF_BYTES; case BIG_DECIMAL: @@ -461,6 +475,7 @@ protected void appendDefaultNullValue(ObjectNode jsonNode) { break; case STRING: case JSON: + case URN: jsonNode.put(key, (String) _defaultNullValue); break; case BYTES: @@ -540,6 +555,7 @@ public enum DataType { TIMESTAMP(LONG, false, true), STRING(false, true), JSON(STRING, false, false), + URN(STRING, false, true), BYTES(false, false), STRUCT(false, false), MAP(false, false), @@ -638,6 +654,10 @@ public Object convert(String value) { case STRING: case JSON: return value; + case URN: + // Validate that the input parses as a URN so bad data fails at ingest, not at query. + org.apache.pinot.spi.utils.Urn.parse(value); + return value; case BYTES: return BytesUtils.toBytes(value); case MAP: @@ -677,6 +697,7 @@ public int compare(Object value1, Object value2) { return Long.compare((long) value1, (long) value2); case STRING: case JSON: + case URN: return ((String) value1).compareTo((String) value2); case BYTES: return ByteArray.compare((byte[]) value1, (byte[]) value2); @@ -731,6 +752,9 @@ public Comparable convertInternal(String value) { case STRING: case JSON: return value; + case URN: + org.apache.pinot.spi.utils.Urn.parse(value); + return value; case BYTES: return BytesUtils.toByteArray(value); case MAP: diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java index 8956ff9756..3d74e9266e 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java @@ -124,6 +124,7 @@ public static void validate(FieldType fieldType, DataType dataType) { case TIMESTAMP: case STRING: case JSON: + case URN: case BYTES: break; default: diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/JsonUtils.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/JsonUtils.java index 82bbb52eae..6e81369ca8 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/JsonUtils.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/JsonUtils.java @@ -321,6 +321,7 @@ private static Object extractSingleValue(JsonNode jsonValue, DataType dataType) case TIMESTAMP: case STRING: case JSON: + case URN: return jsonValue.asText(); case BYTES: try { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/Urn.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/Urn.java new file mode 100644 index 0000000000..978e6c8c16 --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/Urn.java @@ -0,0 +1,440 @@ +/** + * 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.spi.utils; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nullable; + + +/** + * A parsed Uniform Resource Name. Pinot recognizes the LinkedIn-style URN grammar that covers most + * production usages of the form: + * + * 

+ *   urn      ::= "urn" ":" namespace ":" entityType ":" key
+ *   key      ::= primitive | tuple
+ *   tuple    ::= "(" key ("," key)* ")"
+ *   primitive ::= any sequence that is not "," "(" ")"
+ *
+ * + * Examples: + *
    + *
  • {@code urn:li:memberId:123456} — simple URN with a primitive key
    • + *
  • {@code urn:li:dataset:(urn:li:dataPlatform:pinot,memberEducation,EI)} — composite URN with + * nested URN parts and primitive parts
    • + *
+ * + * The class is immutable and provides: + *
    + *
  • Round-trip canonical serialization via {@link #toString()}
    • + *
  • Structural equality and stable hashing for use as map keys
    • + *
  • {@link Comparable} ordering by canonical string
    • + *
+ * + *

The composite key parts are exposed as {@link Part} so callers can distinguish primitive + * strings from nested URNs without re-parsing. + */ +public final class Urn implements Comparable, Serializable { + public static final String SCHEME = "urn"; + + private final String _namespace; + private final String _entityType; + private final List _parts; + // Cached canonical serialization. Computed lazily. + private transient String _serialized; + private transient int _hashCode; + + private Urn(String namespace, String entityType, List parts) { + _namespace = namespace; + _entityType = entityType; + _parts = parts; + } + + /** Constructs a simple URN with a single primitive key. */ + public static Urn simple(String namespace, String entityType, String key) { + Objects.requireNonNull(namespace, "namespace"); + Objects.requireNonNull(entityType, "entityType"); + Objects.requireNonNull(key, "key"); + return new Urn(namespace, entityType, Collections.singletonList(Part.primitive(key))); + } + + /** Constructs a composite URN with the given parts. Must have at least one part. */ + public static Urn composite(String namespace, String entityType, List parts) { + Objects.requireNonNull(namespace, "namespace"); + Objects.requireNonNull(entityType, "entityType"); + Objects.requireNonNull(parts, "parts"); + if (parts.isEmpty()) { + throw new IllegalArgumentException("URN must have at least one key part"); + } + return new Urn(namespace, entityType, Collections.unmodifiableList(new ArrayList<>(parts))); + } + + /** + * Parses a URN from its canonical string form. Throws {@link IllegalArgumentException} if the + * input does not start with {@code urn:} or is otherwise malformed. + */ + public static Urn parse(String value) { + Objects.requireNonNull(value, "value"); + Parser p = new Parser(value); + Urn urn = p.parseUrn(); + if (!p.atEnd()) { + throw new IllegalArgumentException("Trailing characters in URN: '" + value + "'"); + } + return urn; + } + + /** + * Returns the parsed URN if {@code value} is well-formed, or {@code null} otherwise. Useful in + * ingestion paths where mixed quality input is expected. + */ + @Nullable + public static Urn tryParse(String value) { + if (value == null || value.length() < 5) { + return null; + } + try { + return parse(value); + } catch (IllegalArgumentException e) { + return null; + } + } + + /** Returns true if the given string parses as a URN. */ + public static boolean isUrn(@Nullable String value) { + return tryParse(value) != null; + } + + public String namespace() { + return _namespace; + } + + public String entityType() { + return _entityType; + } + + /** + * The fully-qualified entity prefix used as a dictionary key for efficient storage. For + * {@code urn:li:memberId:123456} this returns {@code urn:li:memberId}. + */ + public String entityPrefix() { + return SCHEME + ":" + _namespace + ":" + _entityType; + } + + /** True when the URN's key is a single primitive (no nested parts or sub-URNs). */ + public boolean isSimple() { + return _parts.size() == 1 && _parts.get(0).isPrimitive(); + } + + /** The key parts. Always non-empty. */ + public List parts() { + return _parts; + } + + /** + * Returns the simple URN's primitive key, or throws if this is composite. Use {@link #parts()} + * for composite access. + */ + public String simpleKey() { + if (!isSimple()) { + throw new IllegalStateException("Not a simple URN: " + this); + } + return _parts.get(0).primitive(); + } + + @Override + public String toString() { + String s = _serialized; + if (s == null) { + StringBuilder sb = new StringBuilder(64); + sb.append(SCHEME).append(':').append(_namespace).append(':').append(_entityType).append(':'); + appendKey(sb, _parts); + s = sb.toString(); + _serialized = s; + } + return s; + } + + private static void appendKey(StringBuilder sb, List parts) { + if (parts.size() == 1 && parts.get(0).isPrimitive()) { + sb.append(parts.get(0).primitive()); + return; + } + sb.append('('); + for (int i = 0; i < parts.size(); i++) { + if (i > 0) { + sb.append(','); + } + Part p = parts.get(i); + if (p.isPrimitive()) { + sb.append(p.primitive()); + } else { + sb.append(p.urn().toString()); + } + } + sb.append(')'); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Urn)) { + return false; + } + Urn other = (Urn) o; + return _namespace.equals(other._namespace) && _entityType.equals(other._entityType) && _parts.equals(other._parts); + } + + @Override + public int hashCode() { + int h = _hashCode; + if (h == 0) { + h = _namespace.hashCode(); + h = 31 * h + _entityType.hashCode(); + h = 31 * h + _parts.hashCode(); + _hashCode = h; + } + return h; + } + + @Override + public int compareTo(Urn other) { + return toString().compareTo(other.toString()); + } + + /** A single piece of the key — either a primitive scalar or a nested URN. */ + public static final class Part implements Serializable { + @Nullable + private final String _primitive; + @Nullable + private final Urn _urn; + + private Part(@Nullable String primitive, @Nullable Urn urn) { + _primitive = primitive; + _urn = urn; + } + + public static Part primitive(String value) { + Objects.requireNonNull(value, "value"); + return new Part(value, null); + } + + public static Part urn(Urn urn) { + Objects.requireNonNull(urn, "urn"); + return new Part(null, urn); + } + + public boolean isPrimitive() { + return _primitive != null; + } + + public boolean isUrn() { + return _urn != null; + } + + public String primitive() { + if (_primitive == null) { + throw new IllegalStateException("Not a primitive part"); + } + return _primitive; + } + + public Urn urn() { + if (_urn == null) { + throw new IllegalStateException("Not a URN part"); + } + return _urn; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Part)) { + return false; + } + Part other = (Part) o; + return Objects.equals(_primitive, other._primitive) && Objects.equals(_urn, other._urn); + } + + @Override + public int hashCode() { + return Objects.hash(_primitive, _urn); + } + + @Override + public String toString() { + return _primitive != null ? _primitive : Objects.requireNonNull(_urn).toString(); + } + } + + /** + * Recursive descent parser. URNs are small (typically tens of bytes), so a single-pass character + * scanner with a tiny depth limit is plenty. + */ + private static final class Parser { + private static final int MAX_DEPTH = 32; + private final String _input; + private int _pos; + + Parser(String input) { + _input = input; + _pos = 0; + } + + Urn parseUrn() { + return parseUrn(0); + } + + private Urn parseUrn(int depth) { + if (depth > MAX_DEPTH) { + throw new IllegalArgumentException("URN nesting too deep"); + } + expectLiteral(SCHEME); + expect(':'); + String namespace = readSegment(); + expect(':'); + String entityType = readSegment(); + expect(':'); + List parts = parseKey(depth); + return new Urn(namespace, entityType, parts); + } + + private List parseKey(int depth) { + if (peek() == '(') { + _pos++; + List parts = new ArrayList<>(); + if (peek() == ')') { + throw new IllegalArgumentException("Empty composite key at position " + _pos); + } + while (true) { + parts.add(parsePart(depth)); + char c = peek(); + if (c == ',') { + _pos++; + continue; + } + if (c == ')') { + _pos++; + return parts; + } + throw new IllegalArgumentException("Expected ',' or ')' at position " + _pos); + } + } + // simple primitive key — consume the rest of the current group / string + String prim = readPrimitive(); + if (prim.isEmpty()) { + throw new IllegalArgumentException("Empty primitive key at position " + _pos); + } + return Collections.singletonList(Part.primitive(prim)); + } + + private Part parsePart(int depth) { + if (peek() == '(') { + // Anonymous nested tuple — not standard LinkedIn URN, but we accept it for forward-compat + List nested = parseKey(depth + 1); + // Wrap as a synthetic primitive containing its serialization; this preserves round-trip + // even though we lack a namespace/entityType for it. + StringBuilder sb = new StringBuilder(); + appendKey(sb, nested); + return Part.primitive(sb.toString()); + } + // Could be either a nested URN ("urn:...") or a primitive scalar. + if (startsWithUrn()) { + Urn nested = parseUrn(depth + 1); + return Part.urn(nested); + } + String prim = readPrimitive(); + if (prim.isEmpty()) { + throw new IllegalArgumentException("Empty primitive part at position " + _pos); + } + return Part.primitive(prim); + } + + private boolean startsWithUrn() { + // matches "urn:" prefix only when followed by another ':' segment (i.e., looks like a URN) + if (_pos + 4 > _input.length()) { + return false; + } + if (!_input.regionMatches(_pos, SCHEME, 0, SCHEME.length())) { + return false; + } + return _input.charAt(_pos + SCHEME.length()) == ':'; + } + + private String readSegment() { + int start = _pos; + while (_pos < _input.length()) { + char c = _input.charAt(_pos); + if (c == ':' || c == ',' || c == '(' || c == ')') { + break; + } + _pos++; + } + if (_pos == start) { + throw new IllegalArgumentException("Empty URN segment at position " + _pos); + } + return _input.substring(start, _pos); + } + + private String readPrimitive() { + int start = _pos; + while (_pos < _input.length()) { + char c = _input.charAt(_pos); + if (c == ',' || c == '(' || c == ')') { + break; + } + _pos++; + } + return _input.substring(start, _pos); + } + + private void expectLiteral(String literal) { + if (_pos + literal.length() > _input.length() + || !_input.regionMatches(_pos, literal, 0, literal.length())) { + throw new IllegalArgumentException("Expected '" + literal + "' at position " + _pos); + } + _pos += literal.length(); + } + + private void expect(char c) { + if (_pos >= _input.length() || _input.charAt(_pos) != c) { + throw new IllegalArgumentException("Expected '" + c + "' at position " + _pos); + } + _pos++; + } + + private char peek() { + if (_pos >= _input.length()) { + return '\0'; + } + return _input.charAt(_pos); + } + + boolean atEnd() { + return _pos >= _input.length(); + } + } +} diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.java new file mode 100644 index 0000000000..bd96d4dcf9 --- /dev/null +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.java @@ -0,0 +1,255 @@ +/** + * 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.spi.data; + +import java.util.Arrays; +import java.util.List; +import org.apache.pinot.spi.utils.JsonUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +public class DimensionFieldSpecUrnTest { + + @Test + public void noUrnTypesIsUntyped() { + DimensionFieldSpec spec = new DimensionFieldSpec("col", FieldSpec.DataType.URN, true); + assertNull(spec.getUrnTypes()); + assertFalse(spec.hasUrnTypes()); + assertFalse(spec.isSingleUrnEntityType()); + } + + @Test + public void singleUrnType() { + DimensionFieldSpec spec = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + spec.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + assertTrue(spec.hasUrnTypes()); + assertTrue(spec.isSingleUrnEntityType()); + assertEquals(spec.getUrnTypes().size(), 1); + assertEquals(spec.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(spec.getUrnTypes().get(0).getValueType(), FieldSpec.DataType.LONG); + assertEquals(spec.getUrnTypes().get(0).entityPrefix(), "urn:li:memberId:"); + } + + @Test + public void multipleUrnTypes() { + DimensionFieldSpec spec = new DimensionFieldSpec("resourceId", FieldSpec.DataType.URN, true); + spec.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG) + )); + + assertTrue(spec.hasUrnTypes()); + assertFalse(spec.isSingleUrnEntityType()); + assertEquals(spec.getUrnTypes().size(), 2); + } + + @Test + public void urnTypeWithoutValueType() { + DimensionFieldSpec.UrnType t = DimensionFieldSpec.UrnType.of("urn:li:corpUser"); + assertEquals(t.getEntityType(), "urn:li:corpUser"); + assertNull(t.getValueType()); + } + + @Test + public void rejectsMalformedEntityType() { + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("li:memberId")); + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("urn:memberId")); + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("urn::memberId")); + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("urn:li:")); + } + + @Test + public void rejectsUnsupportedValueType() { + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.BYTES)); + assertThrows(IllegalArgumentException.class, + () -> DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.DOUBLE)); + } + + @Test + public void jsonRoundTrip() throws Exception { + DimensionFieldSpec spec = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + spec.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + String json = JsonUtils.objectToString(spec.toJsonObject()); + DimensionFieldSpec decoded = JsonUtils.stringToObject(json, DimensionFieldSpec.class); + + assertNotNull(decoded.getUrnTypes()); + assertEquals(decoded.getUrnTypes().size(), 1); + assertEquals(decoded.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(decoded.getUrnTypes().get(0).getValueType(), FieldSpec.DataType.LONG); + } + + @Test + public void jsonRoundTripMultipleTypes() throws Exception { + DimensionFieldSpec spec = new DimensionFieldSpec("resourceId", FieldSpec.DataType.URN, true); + spec.setUrnTypes(Arrays.asList( + DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG), + DimensionFieldSpec.UrnType.of("urn:li:groupId", FieldSpec.DataType.LONG) + )); + + String json = JsonUtils.objectToString(spec.toJsonObject()); + DimensionFieldSpec decoded = JsonUtils.stringToObject(json, DimensionFieldSpec.class); + + assertEquals(decoded.getUrnTypes().size(), 2); + assertEquals(decoded.getUrnTypes().get(0).getEntityType(), "urn:li:memberId"); + assertEquals(decoded.getUrnTypes().get(1).getEntityType(), "urn:li:groupId"); + } + + @Test + public void jsonRoundTripNoUrnTypes() throws Exception { + DimensionFieldSpec spec = new DimensionFieldSpec("col", FieldSpec.DataType.URN, true); + String json = JsonUtils.objectToString(spec.toJsonObject()); + DimensionFieldSpec decoded = JsonUtils.stringToObject(json, DimensionFieldSpec.class); + assertNull(decoded.getUrnTypes()); + } + + @Test + public void equalsAndHashCode() { + DimensionFieldSpec a = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + a.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + DimensionFieldSpec b = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + b.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + + DimensionFieldSpec c = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + c.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.STRING))); + + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + assertFalse(a.equals(c)); + } + + @Test + public void toStringIncludesUrnTypes() { + DimensionFieldSpec spec = new DimensionFieldSpec("memberId", FieldSpec.DataType.URN, true); + spec.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of("urn:li:memberId", FieldSpec.DataType.LONG))); + assertTrue(spec.toString().contains("urnTypes")); + assertTrue(spec.toString().contains("urn:li:memberId")); + } + + // --- entityDictRanges round-trip --- + + @Test + public void entityDictRangesRoundTripThreeEntities() { + String encoded = + "urn:li:corpUser,33334,LONG,8;urn:li:corpGroup,66667,LONG,8;urn:li:applicationName,100000,STRING,16"; + java.util.List parsed = + DimensionFieldSpec.parseEntityDictRanges(encoded); + assertEquals(parsed.size(), 3); + assertEquals(parsed.get(0).getEntityType(), "urn:li:corpUser"); + assertEquals(parsed.get(0).getEndDictIdExclusive(), 33334); + assertEquals(parsed.get(1).getEntityType(), "urn:li:corpGroup"); + assertEquals(parsed.get(1).getEndDictIdExclusive(), 66667); + assertEquals(parsed.get(2).getEntityType(), "urn:li:applicationName"); + assertEquals(parsed.get(2).getEndDictIdExclusive(), 100000); + assertEquals(parsed.get(2).getDictClass(), FieldSpec.DataType.STRING); + assertEquals(parsed.get(2).getNumBytesPerValue(), 16); + assertEquals(DimensionFieldSpec.serializeEntityDictRanges(parsed), encoded); + } + + @Test + public void entityDictRangesNullAndBlank() { + assertNull(DimensionFieldSpec.parseEntityDictRanges(null)); + assertNull(DimensionFieldSpec.parseEntityDictRanges("")); + assertNull(DimensionFieldSpec.parseEntityDictRanges(" ")); + assertNull(DimensionFieldSpec.serializeEntityDictRanges(null)); + assertNull(DimensionFieldSpec.serializeEntityDictRanges(java.util.Collections.emptyList())); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsMalformedNoComma() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId100"); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsMalformedNonNumericEnd() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,abc,LONG,8"); + } + + @Test + public void setEntityDictRangesPersists() { + DimensionFieldSpec spec = new DimensionFieldSpec("actorId", FieldSpec.DataType.URN, true); + java.util.List ranges = java.util.Arrays.asList( + new DimensionFieldSpec.EntityDictRange("urn:li:corpUser", 100, FieldSpec.DataType.LONG, Long.BYTES), + new DimensionFieldSpec.EntityDictRange("urn:li:corpGroup", 200, FieldSpec.DataType.LONG, Long.BYTES)); + spec.setEntityDictRanges(ranges); + assertEquals(spec.getEntityDictRanges().size(), 2); + assertEquals(spec.getEntityDictRanges().get(0).getEntityType(), "urn:li:corpUser"); + spec.setEntityDictRanges(null); + assertNull(spec.getEntityDictRanges()); + } + + @Test + public void entityDictRangeEqualsAndHashCode() { + DimensionFieldSpec.EntityDictRange a = + new DimensionFieldSpec.EntityDictRange("urn:li:x", 5, FieldSpec.DataType.LONG, Long.BYTES); + DimensionFieldSpec.EntityDictRange b = + new DimensionFieldSpec.EntityDictRange("urn:li:x", 5, FieldSpec.DataType.LONG, Long.BYTES); + DimensionFieldSpec.EntityDictRange c = + new DimensionFieldSpec.EntityDictRange("urn:li:y", 5, FieldSpec.DataType.LONG, Long.BYTES); + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + assertTrue(!a.equals(c)); + } + + @Test + public void entityDictRangesPolymorphicRoundTrip() { + String encoded = "urn:li:memberId,100,LONG,8;urn:li:corpuser,150,STRING,16"; + java.util.List parsed = + DimensionFieldSpec.parseEntityDictRanges(encoded); + assertEquals(parsed.size(), 2); + assertEquals(parsed.get(0).getEntityType(), "urn:li:memberId"); + assertEquals(parsed.get(0).getEndDictIdExclusive(), 100); + assertEquals(parsed.get(0).getDictClass(), FieldSpec.DataType.LONG); + assertEquals(parsed.get(0).getNumBytesPerValue(), Long.BYTES); + assertEquals(parsed.get(1).getEntityType(), "urn:li:corpuser"); + assertEquals(parsed.get(1).getEndDictIdExclusive(), 150); + assertEquals(parsed.get(1).getDictClass(), FieldSpec.DataType.STRING); + assertEquals(parsed.get(1).getNumBytesPerValue(), 16); + assertEquals(DimensionFieldSpec.serializeEntityDictRanges(parsed), encoded); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsTwoFieldEntry() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100"); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsThreeFieldEntry() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100,LONG"); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsUnknownDictClass() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100,WIDGET,8"); + } +} diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/UrnTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/UrnTest.java new file mode 100644 index 0000000000..e0dd3add6f --- /dev/null +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/UrnTest.java @@ -0,0 +1,144 @@ +/** + * 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.spi.utils; + +import java.util.HashMap; +import java.util.Map; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +public class UrnTest { + + @Test + public void parsesSimpleUrn() { + Urn urn = Urn.parse("urn:li:memberId:123456"); + assertEquals(urn.namespace(), "li"); + assertEquals(urn.entityType(), "memberId"); + assertTrue(urn.isSimple()); + assertEquals(urn.simpleKey(), "123456"); + assertEquals(urn.entityPrefix(), "urn:li:memberId"); + assertEquals(urn.toString(), "urn:li:memberId:123456"); + } + + @Test + public void parsesCompositeUrnWithNestedUrnAndPrimitives() { + Urn urn = Urn.parse("urn:li:dataset:(urn:li:dataPlatform:pinot,memberEducation,EI)"); + assertEquals(urn.namespace(), "li"); + assertEquals(urn.entityType(), "dataset"); + assertFalse(urn.isSimple()); + assertEquals(urn.parts().size(), 3); + assertTrue(urn.parts().get(0).isUrn()); + assertEquals(urn.parts().get(0).urn().toString(), "urn:li:dataPlatform:pinot"); + assertTrue(urn.parts().get(1).isPrimitive()); + assertEquals(urn.parts().get(1).primitive(), "memberEducation"); + assertEquals(urn.parts().get(2).primitive(), "EI"); + assertEquals(urn.toString(), "urn:li:dataset:(urn:li:dataPlatform:pinot,memberEducation,EI)"); + } + + @Test + public void roundTripsAllExamples() { + String[] examples = { + "urn:li:memberId:0", + "urn:li:memberId:123456", + "urn:li:corpuser:alice", + "urn:li:url:https://example.com/path?query=1", + "urn:li:dataset:(urn:li:dataPlatform:pinot,memberEducation,EI)", + "urn:li:assertion:(urn:li:dataset:(urn:li:dataPlatform:hive,foo,PROD),check-1)" + }; + for (String s : examples) { + assertEquals(Urn.parse(s).toString(), s, "round-trip failed for " + s); + } + } + + @Test + public void rejectsMalformedUrns() { + String[] bad = { + "not-a-urn", + "urn:", + "urn:li:", + "urn:li:memberId:", + "urn:li:memberId", + "urn:li:dataset:()", + "urn:li:dataset:(a,", + "urn:li:dataset:(a,b))", + "urn:li:dataset:(a,b)trailing", + "" + }; + for (String s : bad) { + assertThrows("expected parse failure for '" + s + "'", IllegalArgumentException.class, () -> Urn.parse(s)); + } + } + + @Test + public void tryParseReturnsNullOnFailure() { + assertNull(Urn.tryParse(null)); + assertNull(Urn.tryParse("")); + assertNull(Urn.tryParse("urn:bad")); + assertEquals(Urn.tryParse("urn:li:foo:bar").toString(), "urn:li:foo:bar"); + } + + @Test + public void simpleConstructorAndCompositeConstructor() { + Urn simple = Urn.simple("li", "memberId", "42"); + assertEquals(simple.toString(), "urn:li:memberId:42"); + + Urn nested = Urn.simple("li", "dataPlatform", "pinot"); + Urn composite = Urn.composite("li", "dataset", + java.util.Arrays.asList(Urn.Part.urn(nested), Urn.Part.primitive("foo"), Urn.Part.primitive("PROD"))); + assertEquals(composite.toString(), "urn:li:dataset:(urn:li:dataPlatform:pinot,foo,PROD)"); + } + + @Test + public void equalsHashCodeAndMapUsage() { + Urn a = Urn.parse("urn:li:memberId:1"); + Urn b = Urn.parse("urn:li:memberId:1"); + Urn c = Urn.parse("urn:li:memberId:2"); + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + assertNotEquals(a, c); + + Map m = new HashMap<>(); + m.put(a, 1); + assertEquals(m.get(b), Integer.valueOf(1)); + } + + @Test + public void comparesByCanonicalString() { + Urn a = Urn.parse("urn:li:memberId:1"); + Urn b = Urn.parse("urn:li:memberId:2"); + assertTrue(a.compareTo(b) < 0); + assertTrue(b.compareTo(a) > 0); + assertEquals(a.compareTo(Urn.parse("urn:li:memberId:1")), 0); + } + + @Test + public void isUrnUtility() { + assertTrue(Urn.isUrn("urn:li:memberId:1")); + assertTrue(Urn.isUrn("urn:li:dataset:(urn:li:dataPlatform:pinot,foo,bar)")); + assertFalse(Urn.isUrn("foo:bar:baz:qux")); + assertFalse(Urn.isUrn(null)); + } +}