From 0f0a79768efb63cd0b0c52822566daf931182c3c Mon Sep 17 00:00:00 2001 From: dinoocch Date: Tue, 26 May 2026 21:55:37 +0000 Subject: [PATCH 01/24] URN data type: typed/untyped columns, predicates, transforms, benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the URN logical DataType with two storage paths: - Single-entity typed (e.g., urn:li:memberId) — stores compact LONG keys with UrnReconstructingLongDictionary for transparent SELECT. - Multi-entity / untyped — stores full URN strings; multi-entity also writes a companion __urnEntityIdx INT column for fast urnEntity() lookup at query time. Includes ingestion transformers (UrnKeyEncoding, UrnValidation), scalar functions (urnEntity, urnKeyLong + underscore aliases), predicate evaluator support (EQ/IN/RANGE/REGEXP_LIKE), and a JMH benchmark showing 2.1x–7.4x speedups vs untyped STRING storage across four operations. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../common/datablock/DataBlockEquals.java | 1 + .../pinot/common/function/FunctionUtils.java | 1 + .../common/function/scalar/UrnFunctions.java | 190 ++++++++ .../encoder/ArrowResponseEncoder.java | 3 + .../response/encoder/JsonResponseEncoder.java | 1 + .../apache/pinot/common/utils/DataSchema.java | 12 + .../pinot/common/utils/PinotDataType.java | 7 + .../function/scalar/UrnFunctionsTest.java | 202 ++++++++ .../EqualsPredicateEvaluatorFactory.java | 1 + .../InPredicateEvaluatorFactory.java | 3 +- .../NotEqualsPredicateEvaluatorFactory.java | 1 + .../NotInPredicateEvaluatorFactory.java | 3 +- .../filter/predicate/PredicateUtils.java | 1 + .../RangePredicateEvaluatorFactory.java | 1 + .../RegexpLikePredicateEvaluatorFactory.java | 5 +- .../function/CaseTransformFunction.java | 1 + .../function/TransformFunctionFactory.java | 5 + .../function/UrnEntityTransformFunction.java | 123 +++++ .../function/UrnKeyLongTransformFunction.java | 116 +++++ .../function/AggregationFunctionUtils.java | 1 + .../ParentExprMinMaxAggregationFunction.java | 6 + .../utils/exprminmax/ExprMinMaxObject.java | 2 + .../ExprMinMaxProjectionValSetWrapper.java | 1 + .../exprminmax/ExprMinMaxWrapperValSet.java | 1 + .../query/reduce/GroupByDataTableReducer.java | 1 + .../apache/pinot/core/util/GapfillUtils.java | 1 + .../function/UrnTransformFunctionTest.java | 332 +++++++++++++ .../pinot/perf/BenchmarkUrnTransform.java | 212 +++++++++ .../planner/logical/RexExpressionUtils.java | 3 +- .../serde/RexExpressionToProtoExpression.java | 2 + .../apache/pinot/query/type/TypeFactory.java | 1 + .../CompositeTransformer.java | 2 + .../SchemaConformingTransformer.java | 1 + .../UrnKeyEncodingTransformer.java | 165 +++++++ .../UrnValidationTransformer.java | 134 ++++++ .../impl/SegmentColumnarIndexCreator.java | 91 +++- .../impl/SegmentDictionaryCreator.java | 2 +- .../impl/SegmentIndexCreationDriverImpl.java | 20 +- .../SegmentPreIndexStatsCollectorImpl.java | 17 +- .../index/dictionary/DictionaryIndexType.java | 23 +- .../forward/ForwardIndexCreatorFactory.java | 2 +- .../forward/ForwardIndexReaderFactory.java | 3 +- .../index/forward/ForwardIndexType.java | 2 +- .../UrnReconstructingLongDictionary.java | 182 ++++++++ .../UrnKeyEncodingTransformerTest.java | 158 +++++++ .../UrnValidationTransformerTest.java | 184 ++++++++ .../segment/index/UrnColumnSegmentTest.java | 261 +++++++++++ .../apache/pinot/segment/spi/V1Constants.java | 2 + .../spi/creator/StatsCollectorConfig.java | 9 +- .../index/creator/ForwardIndexCreator.java | 1 + .../index/metadata/ColumnMetadataImpl.java | 36 +- .../pinot/spi/data/DimensionFieldSpec.java | 327 ++++++++++++- .../org/apache/pinot/spi/data/FieldSpec.java | 24 + .../org/apache/pinot/spi/data/Schema.java | 1 + .../org/apache/pinot/spi/utils/JsonUtils.java | 1 + .../java/org/apache/pinot/spi/utils/Urn.java | 440 ++++++++++++++++++ .../spi/data/DimensionFieldSpecUrnTest.java | 158 +++++++ .../org/apache/pinot/spi/utils/UrnTest.java | 144 ++++++ 58 files changed, 3603 insertions(+), 27 deletions(-) create mode 100644 pinot-common/src/main/java/org/apache/pinot/common/function/scalar/UrnFunctions.java create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/function/scalar/UrnFunctionsTest.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnEntityTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnKeyLongTransformFunction.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformer.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformerTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformerTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/UrnColumnSegmentTest.java create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/utils/Urn.java create mode 100644 pinot-spi/src/test/java/org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.java create mode 100644 pinot-spi/src/test/java/org/apache/pinot/spi/utils/UrnTest.java 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..c5be2c9d4e --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnEntityTransformFunction.java @@ -0,0 +1,123 @@ +/** + * 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.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 urnEntity(col)} transform function — returns the entity-type prefix of a URN column. + * + *

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

    + *
  1. Constant — single-entity typed URN column with LONG storage: every row in the + * block gets the same literal string, so the block is filled with one {@link Arrays#fill} + * call and no per-row URN parsing.
    2. + *
  2. Parse-string — all other cases (untyped URN or multi-entity URN, both stored as + * STRING): reads each URN string value and calls {@link Urn#tryParse} per row. Multi-entity + * columns use STRING storage because a bare LONG key cannot distinguish entity types with + * overlapping key spaces (e.g., {@code urn:li:memberId:100} vs {@code urn:li:groupId:100}). + * A faster companion-column strategy (reading entity index from {@code col__urnEntityIdx}) + * is left for a future optimization once projection-plan wiring is in place.
    3. + *
+ */ +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, PARSE_STRING + } + + private Strategy _strategy; + private String _constantEntityType; + + @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"); + + 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() && !dim.needsCompanionEntityIdxColumn()) { + _strategy = Strategy.CONSTANT; + _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); + return; + } + } + } + } + } + _strategy = Strategy.PARSE_STRING; + } + + @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 PARSE_STRING: + default: + String[] strValues = _arguments.get(0).transformToStringValuesSV(valueBlock); + for (int i = 0; i < length; i++) { + String urnStr = strValues[i]; + if (urnStr == null) { + _stringValuesSV[i] = null; + } else { + Urn urn = Urn.tryParse(urnStr); + _stringValuesSV[i] = urn != null ? urn.entityPrefix() : 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/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/UrnTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java new file mode 100644 index 0000000000..dc28d2cc82 --- /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 STRING column (PARSE_STRING fallback) --- + + @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 STRING column (PARSE_STRING fallback) --- + + @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-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..69db445dee --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java @@ -0,0 +1,212 @@ +/** + * 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.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.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.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +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.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; + + +@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 ENTITY = "urn:li:memberId"; + private static final String TYPED_COL = "typedId"; + private static final String UNTYPED_COL = "untypedId"; + private static final String TIME_COL = "t"; + private static final String SEGMENT_NAME = "benchmarkUrnSegment"; + + private File _segmentDir; + private ImmutableSegment _segment; + private Map _dataSourceMap; + private ProjectionBlock _projectionBlock; + private Dictionary _typedDict; + private Dictionary _untypedDict; + private String[] _lookupValues; + + public static void main(String[] args) throws Exception { + new Runner(new OptionsBuilder() + .include(BenchmarkUrnTransform.class.getSimpleName()) + .build()).run(); + } + + @Setup + public void setUp() throws Exception { + _segmentDir = Files.createTempDir(); + + DimensionFieldSpec typed = new DimensionFieldSpec(TYPED_COL, FieldSpec.DataType.URN, true); + typed.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of(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 = 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, SEGMENT_NAME, + _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] = ENTITY + ":" + i; + } + } + + @TearDown + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + FileUtils.deleteQuietly(_segmentDir); + } + + @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)); + } + } +} 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, MapTwo encoding modes are selected at construction time per column: + *
    + *
  • Key-encoding ({@link DimensionFieldSpec#hasTypedUrnStorage()}): single-entity LONG + * or INT column. The URN string is parsed and its numeric key replaces the string value on + * the row. No companion column is written.
    • + *
  • Companion-only ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}): + * multi-entity URN column where all declared types have numeric value types. The main column + * retains the full URN string (STRING storage); the 0-based entity-type index is written to + * the companion column {@code __urnEntityIdx} for fast per-row entity-type lookup + * without string parsing at query time.
    • + *
+ * + *

If the URN value is {@code null} or unparseable, both the main and companion columns are set + * to {@code null}; downstream {@link NullValueTransformer} fills in default null values. + * + *

Must run after {@link UrnValidationTransformer} so that the entity type has already + * been verified to be in the declared set. + */ +public class UrnKeyEncodingTransformer implements RecordTransformer { + + /** Per-column encoding plan. */ + private final List _encodings; + + public UrnKeyEncodingTransformer(Schema schema) { + List encodings = new ArrayList<>(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() != FieldSpec.DataType.URN) { + continue; + } + if (!(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.hasTypedUrnStorage()) { + // Single-entity typed LONG/INT: encode the key to numeric; no companion column. + encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), null)); + } else if (dim.needsCompanionEntityIdxColumn()) { + // Multi-entity with all-numeric value types: leave main column as STRING URN; + // write 0-based entity-type index to companion column for fast urnEntity() lookup. + encodings.add(new ColumnEncoding(dim.getName(), null, dim.getUrnTypes(), dim.companionEntityIdxColumnName())); + } + } + _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._companionColumnName != null) { + record.putValue(enc._companionColumnName, null); + } + continue; + } + + String urnStr = raw.toString(); + Urn urn = Urn.tryParse(urnStr); + if (urn == null || !urn.isSimple()) { + // Malformed or composite URN: null out; NullValueTransformer will fill the default + record.putValue(enc._columnName, null); + if (enc._companionColumnName != null) { + record.putValue(enc._companionColumnName, null); + } + continue; + } + + // For multi-entity: find and write entity-type index + if (enc._companionColumnName != null) { + String entityType = urn.entityPrefix(); // e.g. "urn:li:memberId" + int entityIdx = -1; + for (int i = 0; i < enc._urnTypes.size(); i++) { + if (enc._urnTypes.get(i).getEntityType().equals(entityType)) { + entityIdx = i; + break; + } + } + if (entityIdx < 0) { + // Entity not in declared set (validation should have caught this); null out both + record.putValue(enc._columnName, null); + record.putValue(enc._companionColumnName, null); + continue; + } + record.putValue(enc._companionColumnName, entityIdx); + } + + // Convert the key to the declared numeric type (skipped in companion-only mode) + if (enc._storedType != null) { + String keyStr = urn.simpleKey(); + try { + if (enc._storedType == FieldSpec.DataType.LONG) { + record.putValue(enc._columnName, Long.parseLong(keyStr)); + } else { + record.putValue(enc._columnName, Integer.parseInt(keyStr)); + } + } catch (NumberFormatException e) { + record.putValue(enc._columnName, null); + if (enc._companionColumnName != null) { + record.putValue(enc._companionColumnName, null); + } + } + } + } + return record; + } + + private static final class ColumnEncoding { + final String _columnName; + @Nullable + final FieldSpec.DataType _storedType; + final List _urnTypes; + @Nullable + final String _companionColumnName; + + ColumnEncoding(String columnName, @Nullable FieldSpec.DataType storedType, + List urnTypes, @Nullable String companionColumnName) { + _columnName = columnName; + _storedType = storedType; + _urnTypes = urnTypes; + _companionColumnName = companionColumnName; + } + } +} 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..6329f31ee0 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformer.java @@ -0,0 +1,134 @@ +/** + * 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 against the {@link DimensionFieldSpec#getUrnTypes()} list + * declared in the schema. + * + *

For each URN column that has a non-empty {@code urnTypes} list: + *

    + *
  • The value must be parseable as a URN.
    • + *
  • The URN's entity type prefix must appear in the declared {@code urnTypes}.
    • + *
+ * + *

If a value fails validation: + *

    + *
  • With {@code continueOnError = false} (default): an {@link IllegalStateException} is thrown + * and ingestion of this record fails.
    • + *
  • With {@code continueOnError = true}: the value is replaced with {@code null} (the + * {@link NullValueTransformer} that follows will fill in the column's default-null-value), + * and the row is marked incomplete via {@link GenericRow#INCOMPLETE_RECORD_KEY}.
    • + *
+ * + *

Columns with {@code dataType = URN} but no {@code urnTypes} declaration are not validated + * here (the {@link DataTypeTransformer} already verifies basic URN parseability via + * {@link Urn#parse(String)}). + * + *

This transformer is a no-op when the schema contains no URN columns with declared types. + * + *

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 (e.g. "urn:li:memberId"). */ + private final Map> _allowedEntityPrefixes; + private final boolean _continueOnError; + + public UrnValidationTransformer(TableConfig tableConfig, Schema schema) { + Map> columnEntityPrefixes = new HashMap<>(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() != FieldSpec.DataType.URN) { + continue; + } + if (!(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + List urnTypes = dim.getUrnTypes(); + if (urnTypes == null || urnTypes.isEmpty()) { + continue; + } + Set prefixes = new LinkedHashSet<>(); + for (DimensionFieldSpec.UrnType t : urnTypes) { + prefixes.add(t.getEntityType()); + } + columnEntityPrefixes.put(fieldSpec.getName(), Collections.unmodifiableSet(prefixes)); + } + _allowedEntityPrefixes = Collections.unmodifiableMap(columnEntityPrefixes); + + IngestionConfig ingestionConfig = tableConfig.getIngestionConfig(); + _continueOnError = ingestionConfig != null && ingestionConfig.isContinueOnError(); + } + + @Override + public boolean isNoOp() { + return _allowedEntityPrefixes.isEmpty(); + } + + @Nullable + @Override + public GenericRow transform(GenericRow record) { + for (Map.Entry> entry : _allowedEntityPrefixes.entrySet()) { + String column = entry.getKey(); + Set allowed = entry.getValue(); + + Object rawValue = record.getValue(column); + if (rawValue == null) { + continue; + } + String valueStr = rawValue.toString(); + Urn urn = Urn.tryParse(valueStr); + if (urn == null || !allowed.contains(urn.entityPrefix())) { + String errorMessage = String.format( + "URN value '%s' in column '%s' does not match any declared urnTypes %s", valueStr, column, allowed); + if (_continueOnError) { + LOGGER.debug(errorMessage); + record.putValue(column, null); + record.putValue(GenericRow.INCOMPLETE_RECORD_KEY, true); + } else { + throw new IllegalStateException(errorMessage); + } + } + } + return record; + } +} 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..4baf217e08 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 @@ -64,12 +64,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 +106,7 @@ 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_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 +143,8 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { */ private Map, IndexCreator>> _creatorsByColAndIndex = new HashMap<>(); private final Map _nullValueVectorCreatorMap = new HashMap<>(); + /** Synthetic field specs for companion entity-index columns (not in schema). */ + private final Map _companionFieldSpecs = new HashMap<>(); private String _segmentName; private Schema _schema; private File _indexDir; @@ -277,12 +282,64 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio LOGGER.info("Column: {} is not nullable", columnName); } } + + // Initialize companion entity-index columns for typed multi-entity URN dimensions + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (!(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (!dim.needsCompanionEntityIdxColumn()) { + continue; + } + String compCol = dim.companionEntityIdxColumnName(); + ColumnIndexCreationInfo compInfo = indexCreationInfoMap.get(compCol); + if (compInfo == null) { + continue; + } + DimensionFieldSpec compSpec = dim.createCompanionEntityIdxFieldSpec(); + compSpec.setNullable(fieldSpec.isNullable()); + _companionFieldSpecs.put(compCol, compSpec); + + IndexCreationContext.Common compContext = IndexCreationContext.builder() + .withIndexDir(_indexDir) + .withDictionary(true) + .withFieldSpec(compSpec) + .withTotalDocs(segmentIndexCreationInfo.getTotalDocs()) + .withColumnIndexCreationInfo(compInfo) + .onHeap(segmentCreationSpec.isOnHeap()) + .build(); + + DictionaryIndexConfig dictConfig = DictionaryIndexConfig.DEFAULT; + SegmentDictionaryCreator dictCreator = + new DictionaryIndexPlugin().getIndexType().createIndexCreator(compContext, dictConfig); + try { + dictCreator.build(compContext.getSortedUniqueElementsArray()); + } catch (Exception e) { + LOGGER.error("Error building dictionary for companion column: {}", compCol); + throw e; + } + _dictionaryCreatorMap.put(compCol, dictCreator); + + Map, IndexCreator> compCreators = new HashMap<>(); + tryCreateIndexCreator(compCreators, StandardIndexes.forward(), compContext, FieldIndexConfigs.EMPTY); + _creatorsByColAndIndex.put(compCol, compCreators); + + if (isNullable(compSpec)) { + _nullValueVectorCreatorMap.put(compCol, new NullValueVectorCreator(_indexDir, compCol)); + } + } } private boolean isNullable(FieldSpec fieldSpec) { return _schema.isEnableColumnBasedNullHandling() ? fieldSpec.isNullable() : _config.isDefaultNullHandlingEnabled(); } + private FieldSpec resolveFieldSpec(String columnName) { + FieldSpec fs = _schema.getFieldSpecFor(columnName); + return fs != null ? fs : _companionFieldSpecs.get(columnName); + } + private FieldIndexConfigs adaptConfig(String columnName, FieldIndexConfigs config, ColumnIndexCreationInfo columnIndexCreationInfo, SegmentGeneratorConfig segmentCreationSpec) { FieldIndexConfigs.Builder builder = new FieldIndexConfigs.Builder(config); @@ -369,7 +426,7 @@ public void indexRow(GenericRow row) Map, IndexCreator> creatorsByIndex = byColEntry.getValue(); - FieldSpec fieldSpec = _schema.getFieldSpecFor(columnName); + FieldSpec fieldSpec = resolveFieldSpec(columnName); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); try { if (fieldSpec.isSingleValueField()) { @@ -404,7 +461,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 +661,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 +723,18 @@ 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()); + } + } + } + // complex field if (fieldSpec.getFieldType() == FieldType.COMPLEX) { ComplexFieldSpec complexFieldSpec = (ComplexFieldSpec) fieldSpec; @@ -681,12 +750,12 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str 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 +803,18 @@ 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()); + } + } + } + // 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..f54cd4595b 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 @@ -70,6 +70,7 @@ import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.spi.config.table.StarTreeIndexConfig; 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.IngestionSchemaValidator; @@ -528,7 +529,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(); @@ -542,6 +543,23 @@ void collectStatsAndIndexCreationInfo() new ColumnIndexCreationInfo(columnProfile, createDictionary, useVarLengthDictionary, false/*isAutoGenerated*/, defaultNullValue)); } + // Add companion entity-index columns for typed multi-entity URN columns + for (FieldSpec fieldSpec : _dataSchema.getAllFieldSpecs()) { + if (!(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (!dim.needsCompanionEntityIdxColumn()) { + continue; + } + String compCol = dim.companionEntityIdxColumnName(); + ColumnStatistics compProfile = _segmentStats.getColumnProfileFor(compCol); + if (compProfile == null) { + continue; + } + _indexCreationInfoMap.put(compCol, + new ColumnIndexCreationInfo(compProfile, true, false, true /*isAutoGenerated*/, Integer.MIN_VALUE)); + } _segmentIndexCreationInfo.setTotalDocs(_totalDocs); _totalStatsCollectorTimeNs = System.nanoTime() - statsCollectorStartTime; } 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..aecdd73631 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,8 @@ public void init() { Schema dataSchema = _statsCollectorConfig.getSchema(); for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) { String column = fieldSpec.getName(); - switch (fieldSpec.getDataType().getStoredType()) { + // 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; @@ -78,6 +80,19 @@ public void init() { throw new IllegalStateException("Unsupported data type: " + fieldSpec.getDataType()); } } + // Register companion entity-index columns for typed multi-entity URN columns + for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) { + if (!(fieldSpec instanceof DimensionFieldSpec)) { + continue; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.needsCompanionEntityIdxColumn()) { + DimensionFieldSpec compSpec = dim.createCompanionEntityIdxFieldSpec(); + _statsCollectorConfig.registerExtraFieldSpec(compSpec); + _columnStatsCollectorMap.put(compSpec.getName(), + new IntColumnPreIndexStatsCollector(compSpec.getName(), _statsCollectorConfig)); + } + } } @Override 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..e63209fea4 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.UrnReconstructingLongDictionary; 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,14 +321,27 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat } } + // 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); - case LONG: - return loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) + case LONG: { + Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) : new LongDictionary(dataBuffer, length); + // For single-entity typed URN LONG columns, wrap to reconstruct full URN strings on read. + FieldSpec fieldSpec = metadata.getFieldSpec(); + if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.hasTypedUrnStorage() && !dim.needsCompanionEntityIdxColumn()) { + String entityType = dim.getUrnTypes().get(0).getEntityType(); + return new UrnReconstructingLongDictionary(dict, entityType); + } + } + return dict; + } case FLOAT: return loadOnHeap ? new OnHeapFloatDictionary(dataBuffer, length) : new FloatDictionary(dataBuffer, length); 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/UrnReconstructingLongDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java new file mode 100644 index 0000000000..1b471b3753 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java @@ -0,0 +1,182 @@ +/** + * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/** + * Wraps a LONG dictionary for a single-entity typed URN column and reconstructs full URN strings + * (e.g., "urn:li:memberId:100") on read. Predicates using full URN string literals are also + * transparently resolved by stripping the entity prefix before delegating to the inner dictionary. + * + *

This allows {@code SELECT col} on a typed URN LONG column to return the original URN strings + * while the segment physically stores compact LONG keys. + */ +public class UrnReconstructingLongDictionary implements Dictionary { + private final Dictionary _delegate; + private final String _entityTypePrefix; + + /** + * @param delegate underlying LONG dictionary (off-heap or on-heap) + * @param entityType the entity type without trailing colon, e.g. {@code "urn:li:memberId"} + */ + public UrnReconstructingLongDictionary(Dictionary delegate, String entityType) { + _delegate = delegate; + _entityTypePrefix = entityType + ":"; + } + + @Override + public boolean isSorted() { + return true; + } + + @Override + public DataType getValueType() { + return DataType.LONG; + } + + @Override + public int length() { + return _delegate.length(); + } + + /** Resolves a full URN string to its dict ID by stripping the entity prefix. */ + @Override + public int indexOf(String stringValue) { + if (stringValue.startsWith(_entityTypePrefix)) { + try { + long key = Long.parseLong(stringValue.substring(_entityTypePrefix.length())); + return _delegate.indexOf(key); + } catch (NumberFormatException e) { + return NULL_VALUE_INDEX; + } + } + return NULL_VALUE_INDEX; + } + + @Override + public int indexOf(long longValue) { + return _delegate.indexOf(longValue); + } + + /** + * Returns the binary-search insertion index for range predicate evaluation. + * Strips the entity prefix before delegating. + */ + @Override + public int insertionIndexOf(String stringValue) { + if (stringValue.startsWith(_entityTypePrefix)) { + try { + long key = Long.parseLong(stringValue.substring(_entityTypePrefix.length())); + return _delegate.insertionIndexOf(Long.toString(key)); + } catch (NumberFormatException e) { + return ~0; + } + } + return ~0; + } + + @Override + public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { + String lowerKey = toKeyString(lower); + String upperKey = toKeyString(upper); + return _delegate.getDictIdsInRange( + lowerKey != null ? lowerKey : lower, + upperKey != null ? upperKey : upper, + includeLower, includeUpper); + } + + private String toKeyString(String urnString) { + if (urnString != null && urnString.startsWith(_entityTypePrefix)) { + try { + return Long.toString(Long.parseLong(urnString.substring(_entityTypePrefix.length()))); + } catch (NumberFormatException e) { + return null; + } + } + return null; + } + + @Override + public int compare(int dictId1, int dictId2) { + return Integer.compare(dictId1, dictId2); + } + + @Override + public Comparable getMinVal() { + return _delegate.getMinVal(); + } + + @Override + public Comparable getMaxVal() { + return _delegate.getMaxVal(); + } + + @Override + public Object getSortedValues() { + return _delegate.getSortedValues(); + } + + @Override + public Object get(int dictId) { + return _delegate.get(dictId); + } + + @Override + public int getIntValue(int dictId) { + return _delegate.getIntValue(dictId); + } + + @Override + public long getLongValue(int dictId) { + return _delegate.getLongValue(dictId); + } + + @Override + public float getFloatValue(int dictId) { + return _delegate.getFloatValue(dictId); + } + + @Override + public double getDoubleValue(int dictId) { + return _delegate.getDoubleValue(dictId); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return _delegate.getBigDecimalValue(dictId); + } + + /** Returns the full URN string, e.g., {@code "urn:li:memberId:100"}. */ + @Override + public String getStringValue(int dictId) { + return _entityTypePrefix + _delegate.getLongValue(dictId); + } + + @Override + public void close() + throws IOException { + _delegate.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..07b6a64947 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformerTest.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.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 multiEntityFirstEntityWritesCompanionIdx0() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", "urn:li:memberId:100"); + t.transform(row); + // Multi-entity uses STRING storage; main column retains the full URN string + assertEquals(row.getValue("resourceId"), "urn:li:memberId:100"); + assertEquals(row.getValue("resourceId__urnEntityIdx"), 0); + } + + @Test + public void multiEntitySecondEntityWritesCompanionIdx1() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", "urn:li:groupId:42"); + t.transform(row); + // Multi-entity uses STRING storage; main column retains the full URN string + assertEquals(row.getValue("resourceId"), "urn:li:groupId:42"); + assertEquals(row.getValue("resourceId__urnEntityIdx"), 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("resourceId__urnEntityIdx"), "Companion column should be null too"); + } + + @Test + public void multiEntityNullUrnLeavesCompanionNull() { + UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); + GenericRow row = rowWithValue("resourceId", null); + t.transform(row); + assertNull(row.getValue("resourceId")); + assertNull(row.getValue("resourceId__urnEntityIdx")); + } +} 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..1a4836e207 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/UrnValidationTransformerTest.java @@ -0,0 +1,184 @@ +/** + * 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 isNoOpWhenUrnColumnHasNoUrnTypes() { + UrnValidationTransformer t = new UrnValidationTransformer(defaultTableConfig(), untypedSchema()); + assertTrue(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")); + } +} 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..dbacee18ed --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/UrnColumnSegmentTest.java @@ -0,0 +1,261 @@ +/** + * 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.assertFalse; +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 untypedUrnHasNoUrnTypesInMetadata() + 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(); + assertFalse(dimSpec.hasUrnTypes()); + } + + // --- 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(); + } + } +} 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..173afed11a 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,8 @@ 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"; 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..cd557c8d9f 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,37 @@ 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 + } + } + } + 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..cb3cfd480e 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,58 @@ 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 { + /** + * Suffix appended to a typed multi-entity URN column name to form the companion entity-index + * column. For a column {@code "actorId"} with multiple URN types, the companion column is + * {@code "actorId__urnEntityIdx"} and stores a 0-based index into the column's {@code urnTypes} + * list, allowing efficient per-row entity-type lookup without string parsing. + */ + public static final String COMPANION_ENTITY_IDX_SUFFIX = "__urnEntityIdx"; + + /** + * 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; + // Default constructor required by JSON de-serializer. DO NOT REMOVE. public DimensionFieldSpec() { super(); @@ -62,16 +108,291 @@ 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()) { + for (UrnType t : urnTypes) { + Objects.requireNonNull(t, "urnTypes entry must not be null"); + t.validate(); + } + _urnTypes = Collections.unmodifiableList(urnTypes); + } 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; + } + + @Override + public Object getDefaultNullValue() { + if (_dataType == DataType.URN) { + 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) { + return DataType.STRING; + } + 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 should have a companion entity-index INT column written + * alongside the main column. This is true for multi-entity URN columns whose entity types all + * declare a numeric ({@code LONG} or {@code INT}) {@code valueType}; the companion column stores + * the 0-based entity-type index per row, enabling fast {@code urnEntity()} resolution without + * parsing the full URN string. + */ + @JsonIgnore + public boolean needsCompanionEntityIdxColumn() { + if (_dataType != DataType.URN || _urnTypes == null || _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 the name of the companion entity-index column for this column, i.e. + * {@code name + COMPANION_ENTITY_IDX_SUFFIX}. + */ + @JsonIgnore + public String companionEntityIdxColumnName() { + return _name + COMPANION_ENTITY_IDX_SUFFIX; + } + + /** + * Creates and returns the {@link DimensionFieldSpec} for the companion entity-index column. + * The companion column has type {@code INT} and is a single-value dimension. + */ + @JsonIgnore + public DimensionFieldSpec createCompanionEntityIdxFieldSpec() { + return new DimensionFieldSpec(companionEntityIdxColumnName(), DataType.INT, true); + } + @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; + } } } 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..1cca03f620 --- /dev/null +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.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.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")); + } +} 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)); + } +} From aa59293505fa3e1a80bb51e319e9f759a94d0c9c Mon Sep 17 00:00:00 2001 From: dinoocch Date: Wed, 27 May 2026 20:37:10 +0000 Subject: [PATCH 02/24] URN: multi-entity urnEntity() fast path + urnKeyLong filter pushdown Closes the two query-time optimizations that the previous URN commit deferred. Multi-entity companion-column fast path for urnEntity(): - UrnEntityTransformFunction gains a COMPANION_LOOKUP strategy that reads the pre-computed entity index from the companion __urnEntityIdx INT column and indexes into a String[] of entity types. No URN parsing per row. - ProjectPlanNode auto-projects the companion column whenever urnEntity(col) is referenced on a multi-entity URN column. - SegmentColumnarIndexCreator now includes synthetic companion columns in the DIMENSIONS list written to metadata.properties so they are loaded as physical columns when the segment is opened. Filter pushdown for urnKeyLong(typed col) OP literal: - FilterPlanNode rewrites predicates of the form urnKeyLong(typed_long_urn_col) OP literal into col OP entityPrefix:literal, sending them through the existing identifier-based dictionary fast path with prefix stripping inside UrnReconstructingLongDictionary. Supports EQ/NEQ/IN/NOT_IN/RANGE. Multi-entity URN columns are skipped because the literal alone cannot identify which entity prefix to apply. Adds 13 new tests (UrnEntityMultiEntityTest x3, FilterPlanNodeUrnRewriteTest x10). All 119 URN tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../function/UrnEntityTransformFunction.java | 42 +++- .../pinot/core/plan/FilterPlanNode.java | 94 ++++++++- .../pinot/core/plan/ProjectPlanNode.java | 47 +++++ .../function/UrnEntityMultiEntityTest.java | 191 ++++++++++++++++++ .../plan/FilterPlanNodeUrnRewriteTest.java | 185 +++++++++++++++++ .../impl/SegmentColumnarIndexCreator.java | 10 +- 6 files changed, 559 insertions(+), 10 deletions(-) create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityMultiEntityTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeUrnRewriteTest.java 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 index c5be2c9d4e..4162ec42c1 100644 --- 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 @@ -22,6 +22,7 @@ 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; @@ -36,15 +37,16 @@ * *

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

    - *
  1. Constant — single-entity typed URN column with LONG storage: every row in the + *
  2. Constant — single-entity typed URN column with LONG/INT storage: every row in the * block gets the same literal string, so the block is filled with one {@link Arrays#fill} * call and no per-row URN parsing.
    3. - *
  3. Parse-string — all other cases (untyped URN or multi-entity URN, both stored as - * STRING): reads each URN string value and calls {@link Urn#tryParse} per row. Multi-entity - * columns use STRING storage because a bare LONG key cannot distinguish entity types with - * overlapping key spaces (e.g., {@code urn:li:memberId:100} vs {@code urn:li:groupId:100}). - * A faster companion-column strategy (reading entity index from {@code col__urnEntityIdx}) - * is left for a future optimization once projection-plan wiring is in place.
    4. + *
  4. Companion-lookup — multi-entity typed URN column with the companion + * {@code __urnEntityIdx} INT column projected: reads the per-row entity index and + * indexes into a pre-resolved {@code String[]} of entity types. No URN string parsing per + * row.
    5. + *
  5. Parse-string — fallback: untyped URN columns, or multi-entity columns whose + * companion column was not included in the projection. Reads each URN string value and + * calls {@link Urn#tryParse} per row.
    6. *
*/ public class UrnEntityTransformFunction extends BaseTransformFunction { @@ -53,11 +55,13 @@ public class UrnEntityTransformFunction extends BaseTransformFunction { public static final String FUNCTION_NAME_ALT = "urn_entity"; private enum Strategy { - CONSTANT, PARSE_STRING + CONSTANT, COMPANION_LOOKUP, PARSE_STRING } private Strategy _strategy; private String _constantEntityType; + private String _companionColumnName; + private String[] _entityTypesByIdx; @Override public String getName() { @@ -84,6 +88,19 @@ public void init(List arguments, Map c _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); return; } + if (dim.needsCompanionEntityIdxColumn()) { + String companionCol = dim.companionEntityIdxColumnName(); + if (columnContextMap.containsKey(companionCol)) { + List urnTypes = dim.getUrnTypes(); + _entityTypesByIdx = new String[urnTypes.size()]; + for (int i = 0; i < urnTypes.size(); i++) { + _entityTypesByIdx[i] = urnTypes.get(i).getEntityType(); + } + _companionColumnName = companionCol; + _strategy = Strategy.COMPANION_LOOKUP; + return; + } + } } } } @@ -104,6 +121,15 @@ public String[] transformToStringValuesSV(ValueBlock valueBlock) { case CONSTANT: Arrays.fill(_stringValuesSV, 0, length, _constantEntityType); break; + case COMPANION_LOOKUP: + BlockValSet companionBlock = valueBlock.getBlockValueSet(_companionColumnName); + int[] entityIdxes = companionBlock.getIntValuesSV(); + int numTypes = _entityTypesByIdx.length; + for (int i = 0; i < length; i++) { + int idx = entityIdxes[i]; + _stringValuesSV[i] = (idx >= 0 && idx < numTypes) ? _entityTypesByIdx[idx] : null; + } + break; case PARSE_STRING: default: String[] strValues = _arguments.get(0).transformToStringValuesSV(valueBlock); 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..2a3fdd5994 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.needsCompanionEntityIdxColumn()) { + 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/plan/ProjectPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java index d9b4cb8aa8..3ac58236dd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java @@ -25,6 +25,7 @@ import java.util.Set; import javax.annotation.Nullable; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.FunctionContext; import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.DocIdSetOperator; @@ -32,10 +33,14 @@ import org.apache.pinot.core.operator.ProjectionOperatorUtils; import org.apache.pinot.core.operator.filter.BaseFilterOperator; import org.apache.pinot.core.operator.transform.TransformOperator; +import org.apache.pinot.core.operator.transform.function.UrnEntityTransformFunction; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentContext; 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.data.Schema; /** @@ -74,6 +79,7 @@ public BaseProjectOperator run() { hasNonIdentifierExpression = true; } } + addUrnCompanionColumns(_expressions, _queryContext.getSchema(), projectionColumns); Map dataSourceMap = new HashMap<>(HashUtil.getHashMapCapacity(projectionColumns.size())); projectionColumns.forEach( column -> dataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema()))); @@ -88,4 +94,45 @@ public BaseProjectOperator run() { return hasNonIdentifierExpression ? new TransformOperator(_queryContext, projectionOperator, _expressions) : projectionOperator; } + + /** + * Walks expressions for {@code urnEntity(col)} calls on multi-entity typed URN columns and adds + * the companion {@code __urnEntityIdx} INT column to the projection so the transform + * function can take its fast COMPANION_LOOKUP path instead of per-row URN string parsing. + */ + private static void addUrnCompanionColumns(Collection expressions, @Nullable Schema schema, + Set projectionColumns) { + if (schema == null) { + return; + } + for (ExpressionContext expression : expressions) { + collectUrnCompanions(expression, schema, projectionColumns); + } + } + + private static void collectUrnCompanions(ExpressionContext expression, Schema schema, Set projectionColumns) { + if (expression.getType() != ExpressionContext.Type.FUNCTION) { + return; + } + FunctionContext function = expression.getFunction(); + String name = function.getFunctionName(); + if ((UrnEntityTransformFunction.FUNCTION_NAME.equalsIgnoreCase(name) + || UrnEntityTransformFunction.FUNCTION_NAME_ALT.equalsIgnoreCase(name)) + && function.getArguments().size() == 1) { + ExpressionContext arg = function.getArguments().get(0); + if (arg.getType() == ExpressionContext.Type.IDENTIFIER) { + String colName = arg.getIdentifier(); + FieldSpec fieldSpec = schema.getFieldSpecFor(colName); + if (fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.needsCompanionEntityIdxColumn()) { + projectionColumns.add(dim.companionEntityIdxColumnName()); + } + } + } + } + for (ExpressionContext child : function.getArguments()) { + collectUrnCompanions(child, schema, projectionColumns); + } + } } 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..394bcbbf9c --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityMultiEntityTest.java @@ -0,0 +1,191 @@ +/** + * 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 + * pre-computed entity indices from the companion {@code col__urnEntityIdx} INT column. + */ +public class UrnEntityMultiEntityTest { + + private static final String SEGMENT_NAME = "urnEntityMultiSegment"; + private static final String URN_COL = "actorId"; + private static final String COMPANION_COL = URN_COL + "__urnEntityIdx"; + 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 ProjectionBlock _projectionBlockNoCompanion; + 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(COMPANION_COL, _segment.getDataSource(COMPANION_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(); + + Map noCompanionMap = new HashMap<>(); + noCompanionMap.put(URN_COL, _segment.getDataSource(URN_COL)); + noCompanionMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); + _projectionBlockNoCompanion = new ProjectionOperator(noCompanionMap, + 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 companionColumnIsPresent() { + DataSource ds = _segment.getDataSource(COMPANION_COL); + assertNotNull(ds, "Companion column should be auto-created during segment build"); + assertEquals(ds.getDataSourceMetadata().getDataType(), FieldSpec.DataType.INT); + } + + @Test + public void urnEntityUsesCompanionFastPath() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); + assertNotNull(fn); + + String[] result = fn.transformToStringValuesSV(_projectionBlock); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + + // Result order matches doc order in the segment. Multi-entity STRING storage preserves + // ingestion order for the main column; the companion column is written in lockstep. + 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); + } + + @Test + public void urnEntityFallsBackToParseStringWhenCompanionNotProjected() { + TransformFunction fn = TransformFunctionFactory.get( + RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), + new HashMap() {{ + put(URN_COL, _segment.getDataSource(URN_COL)); + put(TIME_COL, _segment.getDataSource(TIME_COL)); + }}); + assertNotNull(fn); + + String[] result = fn.transformToStringValuesSV(_projectionBlockNoCompanion); + assertNotNull(result); + assertEquals(result.length, NUM_ROWS); + + 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/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-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 4baf217e08..369cf098cb 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 @@ -581,7 +581,15 @@ private void writeMetadata() properties.setProperty(SEGMENT_PADDING_CHARACTER, String.valueOf(V1Constants.Str.DEFAULT_STRING_PAD_CHAR)); properties.setProperty(SEGMENT_NAME, _segmentName); properties.setProperty(TABLE_NAME, _config.getTableName()); - properties.setProperty(DIMENSIONS, _config.getDimensions()); + // Augment user-declared dimensions with synthetic companion entity-index columns so they + // are loaded as physical dimensions when the segment is opened. + List dimensions = new ArrayList<>(_config.getDimensions()); + for (String compCol : _companionFieldSpecs.keySet()) { + if (!dimensions.contains(compCol)) { + dimensions.add(compCol); + } + } + properties.setProperty(DIMENSIONS, dimensions); properties.setProperty(METRICS, _config.getMetrics()); properties.setProperty(DATETIME_COLUMNS, _config.getDateTimeColumnNames()); properties.setProperty(COMPLEX_COLUMNS, _config.getComplexColumnNames()); From 9f77015424470190cdf88e333ed1872ca91aa87b Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 00:33:50 +0000 Subject: [PATCH 03/24] URN: extend benchmark and wire companion projection into ExpressionFilterOperator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends BenchmarkUrnTransform with end-to-end filter benchmarks that compare typed vs untyped URN columns for "col = literal" and "urnEntity(col) = literal" queries on both single-entity and multi-entity columns, and prints on-disk segment sizes at setup. The first benchmark run showed that "WHERE urnEntity(actorId) = ..." queries on multi-entity typed columns were no faster than untyped — the companion column was only being projected by ProjectPlanNode (used by SELECT), not by ExpressionFilterOperator (used by WHERE on function expressions). Hoisted the companion-projection helper into UrnProjectionUtils so both call sites apply it. Multi-entity urnEntity filter is now ~11x faster on typed vs untyped. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../filter/ExpressionFilterOperator.java | 2 + .../function/UrnProjectionUtils.java | 83 ++++++ .../pinot/core/plan/ProjectPlanNode.java | 49 +--- .../pinot/perf/BenchmarkUrnTransform.java | 271 +++++++++++++++++- 4 files changed, 351 insertions(+), 54 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java index 27889f8ec3..3673cb6bbe 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java @@ -39,6 +39,7 @@ import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluatorProvider; import org.apache.pinot.core.operator.transform.function.TransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +import org.apache.pinot.core.operator.transform.function.UrnProjectionUtils; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.datasource.DataSource; @@ -60,6 +61,7 @@ public ExpressionFilterOperator(IndexSegment segment, QueryContext queryContext, Set columns = new HashSet<>(); ExpressionContext lhs = predicate.getLhs(); lhs.getColumns(columns); + UrnProjectionUtils.addUrnCompanionColumns(Collections.singletonList(lhs), queryContext.getSchema(), columns); int mapCapacity = HashUtil.getHashMapCapacity(columns.size()); _dataSourceMap = new HashMap<>(mapCapacity); Map columnContextMap = new HashMap<>(mapCapacity); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java new file mode 100644 index 0000000000..1ba20eb721 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java @@ -0,0 +1,83 @@ +/** + * 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 java.util.Collection; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.FunctionContext; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; + + +/** + * Utilities for ensuring URN companion entity-index columns are included in a projection so the + * {@link UrnEntityTransformFunction} can take its fast COMPANION_LOOKUP path on multi-entity typed + * URN columns instead of falling back to per-row URN string parsing. + */ +public final class UrnProjectionUtils { + + private UrnProjectionUtils() { + } + + /** + * Walks {@code expressions} for {@code urnEntity(col)} calls on multi-entity typed URN columns + * (see {@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}) and adds the corresponding + * companion {@code __urnEntityIdx} INT column names to {@code projectionColumns}. + * + *

No-op when {@code schema} is null or no matching URN expressions are present. + */ + public static void addUrnCompanionColumns(Collection expressions, @Nullable Schema schema, + Set projectionColumns) { + if (schema == null) { + return; + } + for (ExpressionContext expression : expressions) { + collect(expression, schema, projectionColumns); + } + } + + private static void collect(ExpressionContext expression, Schema schema, Set projectionColumns) { + if (expression.getType() != ExpressionContext.Type.FUNCTION) { + return; + } + FunctionContext function = expression.getFunction(); + String name = function.getFunctionName(); + if ((UrnEntityTransformFunction.FUNCTION_NAME.equalsIgnoreCase(name) + || UrnEntityTransformFunction.FUNCTION_NAME_ALT.equalsIgnoreCase(name)) + && function.getArguments().size() == 1) { + ExpressionContext arg = function.getArguments().get(0); + if (arg.getType() == ExpressionContext.Type.IDENTIFIER) { + String colName = arg.getIdentifier(); + FieldSpec fieldSpec = schema.getFieldSpecFor(colName); + if (fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.needsCompanionEntityIdxColumn()) { + projectionColumns.add(dim.companionEntityIdxColumnName()); + } + } + } + } + for (ExpressionContext child : function.getArguments()) { + collect(child, schema, projectionColumns); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java index 3ac58236dd..68dc63491c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java @@ -25,7 +25,6 @@ import java.util.Set; import javax.annotation.Nullable; import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.request.context.FunctionContext; import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.DocIdSetOperator; @@ -33,14 +32,11 @@ import org.apache.pinot.core.operator.ProjectionOperatorUtils; import org.apache.pinot.core.operator.filter.BaseFilterOperator; import org.apache.pinot.core.operator.transform.TransformOperator; -import org.apache.pinot.core.operator.transform.function.UrnEntityTransformFunction; +import org.apache.pinot.core.operator.transform.function.UrnProjectionUtils; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentContext; 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.data.Schema; /** @@ -79,7 +75,7 @@ public BaseProjectOperator run() { hasNonIdentifierExpression = true; } } - addUrnCompanionColumns(_expressions, _queryContext.getSchema(), projectionColumns); + UrnProjectionUtils.addUrnCompanionColumns(_expressions, _queryContext.getSchema(), projectionColumns); Map dataSourceMap = new HashMap<>(HashUtil.getHashMapCapacity(projectionColumns.size())); projectionColumns.forEach( column -> dataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema()))); @@ -94,45 +90,4 @@ public BaseProjectOperator run() { return hasNonIdentifierExpression ? new TransformOperator(_queryContext, projectionOperator, _expressions) : projectionOperator; } - - /** - * Walks expressions for {@code urnEntity(col)} calls on multi-entity typed URN columns and adds - * the companion {@code __urnEntityIdx} INT column to the projection so the transform - * function can take its fast COMPANION_LOOKUP path instead of per-row URN string parsing. - */ - private static void addUrnCompanionColumns(Collection expressions, @Nullable Schema schema, - Set projectionColumns) { - if (schema == null) { - return; - } - for (ExpressionContext expression : expressions) { - collectUrnCompanions(expression, schema, projectionColumns); - } - } - - private static void collectUrnCompanions(ExpressionContext expression, Schema schema, Set projectionColumns) { - if (expression.getType() != ExpressionContext.Type.FUNCTION) { - return; - } - FunctionContext function = expression.getFunction(); - String name = function.getFunctionName(); - if ((UrnEntityTransformFunction.FUNCTION_NAME.equalsIgnoreCase(name) - || UrnEntityTransformFunction.FUNCTION_NAME_ALT.equalsIgnoreCase(name)) - && function.getArguments().size() == 1) { - ExpressionContext arg = function.getArguments().get(0); - if (arg.getType() == ExpressionContext.Type.IDENTIFIER) { - String colName = arg.getIdentifier(); - FieldSpec fieldSpec = schema.getFieldSpecFor(colName); - if (fieldSpec instanceof DimensionFieldSpec) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.needsCompanionEntityIdxColumn()) { - projectionColumns.add(dim.companionEntityIdxColumnName()); - } - } - } - } - for (ExpressionContext child : function.getArguments()) { - collectUrnCompanions(child, schema, projectionColumns); - } - } } 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 index 69db445dee..de760e30e0 100644 --- a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java @@ -21,23 +21,31 @@ 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; @@ -64,6 +72,16 @@ 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) @@ -74,11 +92,15 @@ public class BenchmarkUrnTransform { private static final int NUM_ROWS = 100_000; private static final int CARDINALITY = 10_000; - private static final String ENTITY = "urn:li:memberId"; + + 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 static final String SEGMENT_NAME = "benchmarkUrnSegment"; private File _segmentDir; private ImmutableSegment _segment; @@ -88,6 +110,24 @@ public class BenchmarkUrnTransform { 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()) @@ -96,10 +136,17 @@ public static void main(String[] args) throws Exception { @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(ENTITY, FieldSpec.DataType.LONG))); + typed.setUrnTypes(List.of(DimensionFieldSpec.UrnType.of(SINGLE_ENTITY, FieldSpec.DataType.LONG))); Schema schema = new Schema.SchemaBuilder() .addField(typed) @@ -112,7 +159,7 @@ public void setUp() throws Exception { List rows = new ArrayList<>(NUM_ROWS); for (int i = 0; i < NUM_ROWS; i++) { long key = i % CARDINALITY; - String urn = ENTITY + ":" + key; + String urn = SINGLE_ENTITY + ":" + key; GenericRow row = new GenericRow(); row.putValue(TYPED_COL, urn); row.putValue(UNTYPED_COL, urn); @@ -120,9 +167,8 @@ public void setUp() throws Exception { rows.add(row); } - File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, SEGMENT_NAME, + File segDir = PinotSegmentUtil.createSegment(tableConfig, schema, "benchmarkUrnSegment", _segmentDir.toString(), new GenericRowRecordReader(rows)); - _segment = ImmutableSegmentLoader.load(segDir, ReadMode.mmap); _dataSourceMap = new HashMap<>(); @@ -140,8 +186,129 @@ public void setUp() throws Exception { _lookupValues = new String[CARDINALITY]; for (int i = 0; i < CARDINALITY; i++) { - _lookupValues[i] = ENTITY + ":" + 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.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("================================================================"); + System.out.println(); + } + + 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 @@ -150,8 +317,24 @@ public void tearDown() { _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( @@ -209,4 +392,78 @@ public void dictIndexOfUntyped(Blackhole bh) { 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; + } + } } From cdc788277222383fee065bf743e2805763a39f7c Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 00:54:53 +0000 Subject: [PATCH 04/24] URN: strip prefix bytes from multi-entity dictionary entries Multi-entity typed URN columns previously stored full URN strings in the main dictionary (e.g. "urn:li:corpUser:docchial"). With ~16-20 bytes of repeated prefix per dictionary entry, this dwarfed the actual key bytes. UrnKeyEncodingTransformer now writes the main column as ":" so the dictionary only carries the entity index plus the URN key. UrnReconstructingMultiEntityDictionary wraps the loaded STRING dictionary to reconstruct full URN strings on read and to encode incoming full-URN predicate literals before delegating to the underlying STRING lookup. Benchmark (100K rows, 10K unique URNs, 3 entity types): multi-entity typed before: 472,879 bytes multi-entity typed after: 262,862 bytes (44% reduction) vs untyped: 446,692 bytes (1.70x smaller now) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../UrnKeyEncodingTransformer.java | 28 ++- .../index/dictionary/DictionaryIndexType.java | 23 ++- ...rnReconstructingMultiEntityDictionary.java | 193 ++++++++++++++++++ .../UrnKeyEncodingTransformerTest.java | 9 +- 4 files changed, 238 insertions(+), 15 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java index b42ba0175c..dea03d4e4f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java @@ -39,11 +39,14 @@ *
  • Key-encoding ({@link DimensionFieldSpec#hasTypedUrnStorage()}): single-entity LONG * or INT column. The URN string is parsed and its numeric key replaces the string value on * the row. No companion column is written.
    • - *
  • Companion-only ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}): + *
  • Prefix-stripped ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}): * multi-entity URN column where all declared types have numeric value types. The main column - * retains the full URN string (STRING storage); the 0-based entity-type index is written to - * the companion column {@code __urnEntityIdx} for fast per-row entity-type lookup - * without string parsing at query time.
    • + * is rewritten as {@code ":"} (e.g. {@code "0:12345"} for the first + * declared entity type), stripping the repeated URN prefix from every dictionary entry. + * The 0-based entity index is also written to the companion {@code __urnEntityIdx} + * column for fast per-row entity-type lookup. At read time, the column dictionary is + * wrapped so {@code getStringValue()} reconstructs the full URN string and predicate + * lookups encode their literals before delegating to the underlying dictionary. * * *

    If the URN value is {@code null} or unparseable, both the main and companion columns are set @@ -71,9 +74,11 @@ public UrnKeyEncodingTransformer(Schema schema) { // Single-entity typed LONG/INT: encode the key to numeric; no companion column. encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), null)); } else if (dim.needsCompanionEntityIdxColumn()) { - // Multi-entity with all-numeric value types: leave main column as STRING URN; - // write 0-based entity-type index to companion column for fast urnEntity() lookup. - encodings.add(new ColumnEncoding(dim.getName(), null, dim.getUrnTypes(), dim.companionEntityIdxColumnName())); + // Multi-entity with all-numeric value types: encode main column as ":" + // to strip the repeated URN prefix from every dictionary entry. The companion column + // stores the 0-based entity-type index for fast urnEntity() lookup. + encodings.add(new ColumnEncoding(dim.getName(), FieldSpec.DataType.STRING, dim.getUrnTypes(), + dim.companionEntityIdxColumnName())); } } _encodings = Collections.unmodifiableList(encodings); @@ -126,8 +131,13 @@ public GenericRow transform(GenericRow record) { record.putValue(enc._companionColumnName, entityIdx); } - // Convert the key to the declared numeric type (skipped in companion-only mode) - if (enc._storedType != null) { + // Encode the main column. + if (enc._storedType == FieldSpec.DataType.STRING) { + // Multi-entity prefix-stripped encoding: ":". + record.putValue(enc._columnName, + ((Integer) record.getValue(enc._companionColumnName)).intValue() + ":" + urn.simpleKey()); + } else if (enc._storedType != null) { + // Single-entity typed numeric: keep just the LONG/INT key. String keyStr = urn.simpleKey(); try { if (enc._storedType == FieldSpec.DataType.LONG) { 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 e63209fea4..f6514ecdfc 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 @@ -50,6 +50,7 @@ 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.UrnReconstructingLongDictionary; +import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingMultiEntityDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -352,10 +353,28 @@ 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) + Dictionary dict = loadOnHeap + ? new OnHeapStringDictionary(dataBuffer, length, numBytesPerValue, strInterner, byteInterner) : new StringDictionary(dataBuffer, length, numBytesPerValue); + // For multi-entity typed URN columns the dictionary entries are stored as + // ":"; wrap to reconstruct full URN strings on read and to encode + // full-URN predicate literals before delegating to the underlying STRING dictionary. + FieldSpec fieldSpec = metadata.getFieldSpec(); + if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.needsCompanionEntityIdxColumn()) { + List urnTypes = dim.getUrnTypes(); + String[] entityTypes = new String[urnTypes.size()]; + for (int i = 0; i < urnTypes.size(); i++) { + entityTypes[i] = urnTypes.get(i).getEntityType(); + } + return new UrnReconstructingMultiEntityDictionary(dict, entityTypes); + } + } + return dict; + } case BYTES: numBytesPerValue = metadata.getColumnMaxLength(); return loadOnHeap ? new OnHeapBytesDictionary(dataBuffer, length, numBytesPerValue, byteInterner) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java new file mode 100644 index 0000000000..76281ab531 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java @@ -0,0 +1,193 @@ +/** + * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/** + * Wraps a STRING dictionary for a multi-entity typed URN column whose entries are physically + * stored as {@code ":"} after the {@link + * org.apache.pinot.segment.local.recordtransformer.UrnKeyEncodingTransformer} stripped the URN + * prefix. Reconstructs full URN strings (e.g. {@code "urn:li:corpUser:docchial"}) on read, and + * encodes incoming full-URN predicate literals back into the stored form before delegating to the + * underlying STRING dictionary. + * + *

    Dictionary entries that do not match the {@code ":..."} encoding (for example the + * default null value {@code "urn:pinot:null:0"}) are passed through unchanged in both directions. + */ +public class UrnReconstructingMultiEntityDictionary implements Dictionary { + private final Dictionary _delegate; + private final String[] _entityTypes; + + /** + * @param delegate underlying STRING dictionary of encoded {@code ":"} + * values + * @param entityTypes the declared entity types in order; the index used in the encoded form + * refers into this array + */ + public UrnReconstructingMultiEntityDictionary(Dictionary delegate, String[] entityTypes) { + _delegate = delegate; + _entityTypes = entityTypes; + } + + @Override + public boolean isSorted() { + return false; + } + + @Override + public DataType getValueType() { + return DataType.STRING; + } + + @Override + public int length() { + return _delegate.length(); + } + + /** Resolves a full URN string to its dict ID by encoding it to the stored form first. */ + @Override + public int indexOf(String urnString) { + String encoded = encode(urnString); + return _delegate.indexOf(encoded != null ? encoded : urnString); + } + + /** Range lookup is best-effort on multi-entity columns; delegates after attempted encoding. */ + @Override + public int insertionIndexOf(String urnString) { + String encoded = encode(urnString); + return _delegate.insertionIndexOf(encoded != null ? encoded : urnString); + } + + @Override + public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { + String encodedLower = encode(lower); + String encodedUpper = encode(upper); + return _delegate.getDictIdsInRange( + encodedLower != null ? encodedLower : lower, + encodedUpper != null ? encodedUpper : upper, + includeLower, includeUpper); + } + + /** + * Encode a full URN string ({@code "urn:li:corpUser:foo"}) to the stored {@code "0:foo"} form + * by finding which declared entity type the URN starts with. Returns {@code null} when no + * declared entity matches (e.g. a predicate literal for an entity the segment does not store). + */ + private String encode(String urnString) { + if (urnString == null) { + return null; + } + for (int i = 0; i < _entityTypes.length; i++) { + String prefix = _entityTypes[i] + ":"; + if (urnString.startsWith(prefix)) { + return i + ":" + urnString.substring(prefix.length()); + } + } + return null; + } + + @Override + public int compare(int dictId1, int dictId2) { + return _delegate.compare(dictId1, dictId2); + } + + @Override + public Comparable getMinVal() { + return _delegate.getMinVal(); + } + + @Override + public Comparable getMaxVal() { + return _delegate.getMaxVal(); + } + + @Override + public Object getSortedValues() { + return _delegate.getSortedValues(); + } + + @Override + public Object get(int dictId) { + return getStringValue(dictId); + } + + @Override + public int getIntValue(int dictId) { + return _delegate.getIntValue(dictId); + } + + @Override + public long getLongValue(int dictId) { + return _delegate.getLongValue(dictId); + } + + @Override + public float getFloatValue(int dictId) { + return _delegate.getFloatValue(dictId); + } + + @Override + public double getDoubleValue(int dictId) { + return _delegate.getDoubleValue(dictId); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return _delegate.getBigDecimalValue(dictId); + } + + /** Returns the full URN string reconstructed from the stored {@code ":"}. */ + @Override + public String getStringValue(int dictId) { + String raw = _delegate.getStringValue(dictId); + return decode(raw); + } + + private String decode(String raw) { + if (raw == null) { + return null; + } + int colon = raw.indexOf(':'); + if (colon <= 0) { + return raw; + } + int idx; + try { + idx = Integer.parseInt(raw, 0, colon, 10); + } catch (NumberFormatException e) { + return raw; + } + if (idx < 0 || idx >= _entityTypes.length) { + return raw; + } + return _entityTypes[idx] + ":" + raw.substring(colon + 1); + } + + @Override + public void close() + throws IOException { + _delegate.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 index 07b6a64947..f69e312e96 100644 --- 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 @@ -123,8 +123,9 @@ public void multiEntityFirstEntityWritesCompanionIdx0() { UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); GenericRow row = rowWithValue("resourceId", "urn:li:memberId:100"); t.transform(row); - // Multi-entity uses STRING storage; main column retains the full URN string - assertEquals(row.getValue("resourceId"), "urn:li:memberId:100"); + // Multi-entity strips the URN prefix and encodes the main column as ":". + // UrnReconstructingMultiEntityDictionary reverses this at read time. + assertEquals(row.getValue("resourceId"), "0:100"); assertEquals(row.getValue("resourceId__urnEntityIdx"), 0); } @@ -133,8 +134,8 @@ public void multiEntitySecondEntityWritesCompanionIdx1() { UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); GenericRow row = rowWithValue("resourceId", "urn:li:groupId:42"); t.transform(row); - // Multi-entity uses STRING storage; main column retains the full URN string - assertEquals(row.getValue("resourceId"), "urn:li:groupId:42"); + // Multi-entity strips the URN prefix and encodes the main column as ":". + assertEquals(row.getValue("resourceId"), "1:42"); assertEquals(row.getValue("resourceId__urnEntityIdx"), 1); } From d0edee5622d703e283d3283bce90fb6157e6efe0 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 01:24:35 +0000 Subject: [PATCH 05/24] URN: auto-detect single prefix on untyped columns; extend CONSTANT path Untyped URN columns now get the same prefix-stripping treatment as typed single-entity columns, at segment build time, whenever every unique value happens to share a common urn::: prefix. SegmentColumnarIndexCreator scans the sorted unique values when building a URN column's dictionary; if a single common prefix is detected, the dictionary stores only the suffixes and the discovered entity type is recorded in column metadata as urnTypes. Row-time indexing strips the same prefix so forward-index dictId lookups still match. UrnReconstructingStringDictionary wraps the loaded STRING dict to reconstruct full URN strings on read and to strip the prefix from incoming full-URN predicate literals. UrnEntityTransformFunction now takes the CONSTANT fast path for *any* single-entity URN column (declared or auto-detected), not just those with LONG/INT storage. This makes urnEntity() effectively free on both declared single-entity columns and on untyped columns where a single prefix was discovered. Benchmark (100K rows, 10K unique URNs): single-entity untyped storage: 376,736 -> 216,812 bytes (-42%) urnEntity() filter on untyped: 14,994 -> 1,543 us/op (~10x faster) urnEntity() transform untyped: 1,050 -> 156 us/op (~7x faster) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../function/UrnEntityTransformFunction.java | 4 +- .../impl/SegmentColumnarIndexCreator.java | 119 ++++++++++++- .../index/dictionary/DictionaryIndexType.java | 7 + ...rnReconstructingMultiEntityDictionary.java | 8 + .../UrnReconstructingStringDictionary.java | 158 ++++++++++++++++++ 5 files changed, 293 insertions(+), 3 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java 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 index 4162ec42c1..268ea1e02e 100644 --- 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 @@ -83,7 +83,9 @@ public void init(List arguments, Map c FieldSpec fieldSpec = dataSource.getDataSourceMetadata().getFieldSpec(); if (fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.hasTypedUrnStorage() && !dim.needsCompanionEntityIdxColumn()) { + // Single-entity URN columns (declared single-entity or auto-detected single-prefix) + // always have the same entity prefix on every row — fill the block with one literal. + if (dim.isSingleUrnEntityType()) { _strategy = Strategy.CONSTANT; _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); return; 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 369cf098cb..b8cb6b2837 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; @@ -145,6 +146,13 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { private final Map _nullValueVectorCreatorMap = new HashMap<>(); /** Synthetic field specs for companion entity-index columns (not in schema). */ private final Map _companionFieldSpecs = new HashMap<>(); + /** + * For untyped URN columns in which every observed value shares the same {@code urn:::} + * prefix, records the discovered entity type (without trailing colon). The prefix is stripped + * from the dictionary entries and persisted in column metadata so the load-time path treats the + * column as single-entity typed. + */ + private final Map _discoveredUrnEntities = new HashMap<>(); private String _segmentName; private Schema _schema; private File _indexDir; @@ -238,8 +246,9 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio SegmentDictionaryCreator creator = new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); + Object dictValues = maybeStripUrnPrefix(columnName, fieldSpec, context.getSortedUniqueElementsArray()); 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()); @@ -337,7 +346,110 @@ private boolean isNullable(FieldSpec fieldSpec) { private FieldSpec resolveFieldSpec(String columnName) { FieldSpec fs = _schema.getFieldSpecFor(columnName); - return fs != null ? fs : _companionFieldSpecs.get(columnName); + if (fs == null) { + return _companionFieldSpecs.get(columnName); + } + // For untyped URN columns where a single common prefix was discovered during dictionary + // build, surface a synthetic urnTypes so the column behaves like single-entity typed when the + // segment is loaded (prefix reconstruction, predicate prefix-strip). + String discoveredEntity = _discoveredUrnEntities.get(columnName); + if (discoveredEntity != null && fs instanceof DimensionFieldSpec + && fs.getDataType() == DataType.URN && !((DimensionFieldSpec) fs).hasUrnTypes()) { + DimensionFieldSpec original = (DimensionFieldSpec) fs; + DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, + original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), + original.getMaxLengthExceedStrategy()); + synthetic.setUrnTypes(Collections.singletonList(DimensionFieldSpec.UrnType.of(discoveredEntity))); + return synthetic; + } + return fs; + } + + /** + * For untyped URN columns whose unique values all share the same {@code urn:::} + * prefix, strips the prefix from every value before it goes into the dictionary, records the + * discovered entity type so {@link #resolveFieldSpec} can surface it at metadata-write time, + * and returns the stripped array. For all other cases (typed URN, non-URN, or untyped URN with + * mixed prefixes) returns {@code values} unchanged. + */ + private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Object values) { + if (!(fieldSpec instanceof DimensionFieldSpec) || fieldSpec.getDataType() != DataType.URN) { + return values; + } + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + if (dim.hasUrnTypes() || !(values instanceof Object[])) { + return values; + } + Object[] arr = (Object[]) values; + if (arr.length == 0) { + return values; + } + String entityType = discoverCommonUrnEntity(arr); + if (entityType == null) { + return values; + } + String prefix = entityType + ":"; + int prefixLen = prefix.length(); + String[] stripped = new String[arr.length]; + for (int i = 0; i < arr.length; i++) { + String s = arr[i].toString(); + stripped[i] = s.substring(prefixLen); + } + _discoveredUrnEntities.put(columnName, entityType); + return stripped; + } + + /** + * Strips the discovered URN entity prefix from a row's value at index time so the dictionary + * lookup hits the suffix-only entry. + */ + private Object maybeStripUrnPrefixForIndex(String columnName, Object value) { + String entityType = _discoveredUrnEntities.get(columnName); + if (entityType == null || !(value instanceof String)) { + return value; + } + String s = (String) value; + String prefix = entityType + ":"; + return s.startsWith(prefix) ? s.substring(prefix.length()) : value; + } + + /** + * Returns the entity type {@code urn::} when every value in {@code values} starts with + * the same {@code urn:::} prefix, or {@code null} otherwise. + */ + @Nullable + private static String discoverCommonUrnEntity(Object[] values) { + String prefix = extractEntityPrefix(values[0].toString()); + if (prefix == null) { + return null; + } + for (int i = 1; i < values.length; i++) { + if (!values[i].toString().startsWith(prefix)) { + return null; + } + } + return prefix.substring(0, prefix.length() - 1); + } + + /** + * Extracts the entity prefix {@code urn:::} (with trailing colon) from a URN string. + * Returns {@code null} when the input is not in canonical URN form. + */ + @Nullable + private static String extractEntityPrefix(String urnString) { + if (!urnString.startsWith("urn:")) { + return null; + } + int colons = 0; + for (int i = 0; i < urnString.length(); i++) { + if (urnString.charAt(i) == ':') { + colons++; + if (colons == 3) { + return urnString.substring(0, i + 1); + } + } + } + return null; } private FieldIndexConfigs adaptConfig(String columnName, FieldIndexConfigs config, @@ -423,6 +535,9 @@ public void indexRow(GenericRow row) if (columnValueToIndex == null) { throw new RuntimeException("Null value for column:" + columnName); } + // If we discovered a common URN entity prefix for this column at dictionary-build time, + // strip it from each row's value too so the dictionary lookup hits the suffix-only entries. + columnValueToIndex = maybeStripUrnPrefixForIndex(columnName, columnValueToIndex); Map, IndexCreator> creatorsByIndex = byColEntry.getValue(); 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 f6514ecdfc..4c7d742819 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 @@ -51,6 +51,7 @@ import org.apache.pinot.segment.local.segment.index.readers.StringDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingLongDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingMultiEntityDictionary; +import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingStringDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -372,6 +373,12 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat } return new UrnReconstructingMultiEntityDictionary(dict, entityTypes); } + if (dim.isSingleUrnEntityType()) { + // Single-entity URN with STRING storage: the dictionary stores suffixes only. + // Covers both schema-declared single-entity URN columns and untyped URN columns where + // a single common prefix was discovered at segment build time. + return new UrnReconstructingStringDictionary(dict, dim.getUrnTypes().get(0).getEntityType()); + } } return dict; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java index 76281ab531..c08e61639a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java @@ -165,6 +165,14 @@ public String getStringValue(int dictId) { return decode(raw); } + @Override + public void readStringValues(int[] dictIds, int length, String[] outValues) { + _delegate.readStringValues(dictIds, length, outValues); + for (int i = 0; i < length; i++) { + outValues[i] = decode(outValues[i]); + } + } + private String decode(String raw) { if (raw == null) { return null; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java new file mode 100644 index 0000000000..09e9324d56 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.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.segment.index.readers; + +import it.unimi.dsi.fastutil.ints.IntSet; +import java.io.IOException; +import java.math.BigDecimal; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/** + * Wraps a STRING dictionary for a single-entity URN column whose entries were stored with the + * {@code urn:::} prefix stripped (because every value in the segment shared that + * prefix). Reconstructs the full URN string on read and strips the entity prefix from incoming + * full-URN predicate literals before delegating to the underlying STRING dictionary. + */ +public class UrnReconstructingStringDictionary implements Dictionary { + private final Dictionary _delegate; + private final String _entityTypePrefix; + + /** + * @param delegate underlying STRING dictionary of URN suffix-only values + * @param entityType the entity type without trailing colon, e.g. {@code "urn:li:corpUser"} + */ + public UrnReconstructingStringDictionary(Dictionary delegate, String entityType) { + _delegate = delegate; + _entityTypePrefix = entityType + ":"; + } + + @Override + public boolean isSorted() { + return _delegate.isSorted(); + } + + @Override + public DataType getValueType() { + return DataType.STRING; + } + + @Override + public int length() { + return _delegate.length(); + } + + @Override + public int indexOf(String stringValue) { + String suffix = stripPrefix(stringValue); + return suffix != null ? _delegate.indexOf(suffix) : _delegate.indexOf(stringValue); + } + + @Override + public int insertionIndexOf(String stringValue) { + String suffix = stripPrefix(stringValue); + return suffix != null ? _delegate.insertionIndexOf(suffix) : _delegate.insertionIndexOf(stringValue); + } + + @Override + public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { + String lowerKey = stripPrefix(lower); + String upperKey = stripPrefix(upper); + return _delegate.getDictIdsInRange( + lowerKey != null ? lowerKey : lower, + upperKey != null ? upperKey : upper, + includeLower, includeUpper); + } + + private String stripPrefix(String stringValue) { + return stringValue != null && stringValue.startsWith(_entityTypePrefix) + ? stringValue.substring(_entityTypePrefix.length()) : null; + } + + @Override + public int compare(int dictId1, int dictId2) { + return _delegate.compare(dictId1, dictId2); + } + + @Override + public Comparable getMinVal() { + return _delegate.getMinVal(); + } + + @Override + public Comparable getMaxVal() { + return _delegate.getMaxVal(); + } + + @Override + public Object getSortedValues() { + return _delegate.getSortedValues(); + } + + @Override + public Object get(int dictId) { + return getStringValue(dictId); + } + + @Override + public int getIntValue(int dictId) { + return _delegate.getIntValue(dictId); + } + + @Override + public long getLongValue(int dictId) { + return _delegate.getLongValue(dictId); + } + + @Override + public float getFloatValue(int dictId) { + return _delegate.getFloatValue(dictId); + } + + @Override + public double getDoubleValue(int dictId) { + return _delegate.getDoubleValue(dictId); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return _delegate.getBigDecimalValue(dictId); + } + + /** Returns the full URN string by prepending the entity prefix. */ + @Override + public String getStringValue(int dictId) { + return _entityTypePrefix + _delegate.getStringValue(dictId); + } + + @Override + public void readStringValues(int[] dictIds, int length, String[] outValues) { + _delegate.readStringValues(dictIds, length, outValues); + for (int i = 0; i < length; i++) { + outValues[i] = _entityTypePrefix + outValues[i]; + } + } + + @Override + public void close() + throws IOException { + _delegate.close(); + } +} From 8b5445a20a819a8a63bb759d51872dfd6a6ad921 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 04:56:32 +0000 Subject: [PATCH 06/24] UrnDictionary commit 1/5: data model + metadata round-trip Defines the storage-level data class and segment-metadata property for the sorted-LONG UrnDictionary design described in ~/urn_sorted_dict_design.md. No behavior change yet -- segment build still emits the existing encoded-STRING multi-entity layout. Commit 2 wires up the new dict creator; commits 3-5 add the read path, drop the companion column, and update tests/benchmark. Adds: * DimensionFieldSpec.EntityDictRange data class: (entityType, endDictIdExclusive) * DimensionFieldSpec.parseEntityDictRanges / serializeEntityDictRanges for the "entityType,end;entityType,end;..." header form * V1Constants.URN_ENTITY_DICT_RANGES metadata key * SegmentColumnarIndexCreator writes urn.entityDictRanges when set * ColumnMetadataImpl reads urn.entityDictRanges at segment load * 6 unit tests in DimensionFieldSpecUrnTest Co-Authored-By: Claude Opus 4.7 (1M context) --- .../impl/SegmentColumnarIndexCreator.java | 17 ++ .../apache/pinot/segment/spi/V1Constants.java | 5 + .../index/metadata/ColumnMetadataImpl.java | 11 ++ .../pinot/spi/data/DimensionFieldSpec.java | 145 ++++++++++++++++++ .../spi/data/DimensionFieldSpecUrnTest.java | 69 +++++++++ 5 files changed, 247 insertions(+) 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 b8cb6b2837..9459aec965 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 @@ -107,6 +107,7 @@ 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; @@ -856,6 +857,14 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str 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 @@ -936,6 +945,14 @@ public static void addFieldSpec(PropertiesConfiguration properties, String colum 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 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 173afed11a..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 @@ -131,6 +131,11 @@ public static class Column { 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/index/metadata/ColumnMetadataImpl.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java index cd557c8d9f..b01e716e43 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 @@ -333,6 +333,17 @@ public static FieldSpec generateFieldSpec(String column, PropertiesConfiguration // 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. + String entityDictRangesStr = + config.getString(Column.getKeyFor(column, Column.URN_ENTITY_DICT_RANGES), null); + if (entityDictRangesStr != null) { + try { + dimSpec.setEntityDictRanges(DimensionFieldSpec.parseEntityDictRanges(entityDictRangesStr)); + } catch (Exception ignored) { + // Malformed ranges are non-fatal; the column simply won't take the sorted-dict path. + } + } } fieldSpec = dimSpec; break; 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 cb3cfd480e..3c2246095a 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 @@ -72,6 +72,21 @@ public final class DimensionFieldSpec extends FieldSpec { @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(); @@ -145,6 +160,86 @@ 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 "entityType1,end1;entityType2,end2;..."} where each {@code end} is the exclusive upper + * bound (in dict IDs) of that entity's range. 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; + } + int comma = part.lastIndexOf(','); + if (comma <= 0 || comma == part.length() - 1) { + throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); + } + String entityType = part.substring(0, comma); + int endDictIdExclusive; + try { + endDictIdExclusive = Integer.parseInt(part.substring(comma + 1)); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Malformed entityDictRanges end: '" + part + "'", e); + } + ranges.add(new EntityDictRange(entityType, endDictIdExclusive)); + } + 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()); + } + return sb.toString(); + } + @Override public Object getDefaultNullValue() { if (_dataType == DataType.URN) { @@ -395,4 +490,54 @@ public String toString() { return _valueType != null ? _entityType + "(" + _valueType + ")" : _entityType; } } + + /** + * One entry in a multi-entity URN column's per-entity dictionary-ID range table. Identifies + * which entity type owns the dictionary IDs ending at (exclusive) {@code endDictIdExclusive}. + * + *

    The {@code i}-th range covers dict IDs {@code [prevEnd, endDictIdExclusive)}, where + * {@code prevEnd} is the previous range's exclusive end (or zero for the first range). + * + *

    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; + + public EntityDictRange(String entityType, int endDictIdExclusive) { + _entityType = entityType; + _endDictIdExclusive = endDictIdExclusive; + } + + public String getEntityType() { + return _entityType; + } + + public int getEndDictIdExclusive() { + return _endDictIdExclusive; + } + + @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 && Objects.equals(_entityType, other._entityType); + } + + @Override + public int hashCode() { + return Objects.hash(_entityType, _endDictIdExclusive); + } + + @Override + public String toString() { + return _entityType + "[..." + _endDictIdExclusive + ")"; + } + } } 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 index 1cca03f620..4e94572de0 100644 --- 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 @@ -155,4 +155,73 @@ public void toStringIncludesUrnTypes() { 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;urn:li:corpGroup,66667;urn:li:applicationName,100000"; + 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(DimensionFieldSpec.serializeEntityDictRanges(parsed), encoded); + } + + @Test + public void entityDictRangesUserDocExample() { + String encoded = "urn:li:memberId,100;urn:li:companyPage,120"; + java.util.List parsed = + DimensionFieldSpec.parseEntityDictRanges(encoded); + assertEquals(parsed.size(), 2); + assertEquals(parsed.get(0).getEndDictIdExclusive(), 100); + assertEquals(parsed.get(1).getEndDictIdExclusive(), 120); + } + + @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"); + } + + @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), + new DimensionFieldSpec.EntityDictRange("urn:li:corpGroup", 200)); + 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); + DimensionFieldSpec.EntityDictRange b = new DimensionFieldSpec.EntityDictRange("urn:li:x", 5); + DimensionFieldSpec.EntityDictRange c = new DimensionFieldSpec.EntityDictRange("urn:li:y", 5); + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + assertTrue(!a.equals(c)); + } } From 54b67bdc81604961a238502e10d84b6bccd39165 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 17:17:23 +0000 Subject: [PATCH 07/24] UrnDictionary commit 2/5: UrnSortedDictionary reader Adds the dictionary reader class for the sorted-LONG UrnDictionary layout described in ~/urn_sorted_dict_design.md. Wraps a LONG buffer delegate and uses the per-entity dictId ranges loaded in commit 1. Lookups: * indexOf("urn:li:corpUser:100") parses the entity prefix to find the range, then range-restricted binary searches the LONG key. * getStringValue(dictId) finds which entity range owns the id and reconstructs the URN string from delegate.getLongValue. * getLongValue / getIntValue pass straight through to the delegate (urnKeyLong-style consumers get the raw key with zero overhead). * getDictIdsInRange supports same-entity range queries; cross-entity range queries return empty (semantically undefined). No integration yet -- DictionaryIndexType still returns the existing LongDictionary / UrnReconstructingMultiEntityDictionary depending on layout. Commit 3 wires up the build + load paths. Unit tests cover: range-restricted indexOf, key-collision disambiguation across entities, reconstruction across all ranges, same-entity range queries (inclusive/exclusive), cross-entity / unknown-entity rejection, batch readStringValues, min/max, and pass-through accessors. 20 new tests in UrnSortedDictionaryTest (145 URN tests pass total). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../index/readers/UrnSortedDictionary.java | 313 +++++++++++++++++ .../readers/UrnSortedDictionaryTest.java | 323 ++++++++++++++++++ 2 files changed, 636 insertions(+) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java new file mode 100644 index 0000000000..9923c717f9 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java @@ -0,0 +1,313 @@ +/** + * 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.DimensionFieldSpec.EntityDictRange; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/** + * Dictionary wrapper for multi-entity URN columns using the sorted-LONG layout. + * + *

    The underlying delegate is a buffer of LONG keys, partitioned by entity type. Each entity + * type owns a contiguous, key-sorted range of dictionary IDs identified by {@link + * EntityDictRange#getEndDictIdExclusive()}. Within a range, entries are sorted ascending by LONG + * value; across ranges, the same LONG value may appear multiple times (once per entity that has + * that key). + * + *

    This wrapper exposes the column as URN strings: + *

      + *
    • {@link #indexOf(String) indexOf("urn:li:corpUser:100")} parses the entity prefix, finds + * the corresponding entity range, then runs a range-restricted binary search for the LONG + * key. Returns the absolute dict ID, or {@link #NULL_VALUE_INDEX} if the entity prefix is + * unknown or the key isn't present.
      • + *
    • {@link #getStringValue(int) getStringValue(dictId)} looks up which entity range owns + * {@code dictId}, reads the LONG key from the delegate, and returns + * {@code ":"}.
      • + *
    • {@link #getLongValue(int) getLongValue(dictId)} passes through to the delegate — the raw + * LONG key is exposed for {@code urnKeyLong()} and LONG-typed predicates.
      • + *
    + * + *

    Cross-entity range filters are explicitly unsupported: {@code WHERE col BETWEEN + * 'urn:li:a:100' AND 'urn:li:b:200'} returns the empty set rather than throwing, to be robust to + * generated queries. Same-entity range filters are supported via {@link #getDictIdsInRange}. + */ +public class UrnSortedDictionary implements Dictionary { + + private final Dictionary _delegate; + private final EntityDictRange[] _ranges; + private final String[] _entityTypes; + private final String[] _entityPrefixes; + private final int[] _rangeStarts; + + public UrnSortedDictionary(Dictionary delegate, List ranges) { + if (delegate.getValueType() != DataType.LONG && delegate.getValueType() != DataType.INT) { + throw new IllegalArgumentException( + "UrnSortedDictionary delegate must be a LONG or INT dictionary, got: " + delegate.getValueType()); + } + _delegate = delegate; + _ranges = ranges.toArray(new EntityDictRange[0]); + _entityTypes = new String[_ranges.length]; + _entityPrefixes = new String[_ranges.length]; + _rangeStarts = new int[_ranges.length]; + int prev = 0; + for (int i = 0; i < _ranges.length; i++) { + _entityTypes[i] = _ranges[i].getEntityType(); + _entityPrefixes[i] = _entityTypes[i] + ":"; + _rangeStarts[i] = prev; + prev = _ranges[i].getEndDictIdExclusive(); + } + } + + @Override + public boolean isSorted() { + // The buffer is sorted only within each entity range, not globally. + return false; + } + + @Override + public DataType getValueType() { + return DataType.STRING; + } + + @Override + public int length() { + return _delegate.length(); + } + + // ---------------- lookup ---------------- + + @Override + public int indexOf(String urnString) { + int entityIdx = findEntityIdx(urnString); + if (entityIdx < 0) { + return NULL_VALUE_INDEX; + } + long key; + try { + key = Long.parseLong(urnString, _entityPrefixes[entityIdx].length(), urnString.length(), 10); + } catch (NumberFormatException e) { + return NULL_VALUE_INDEX; + } + return rangeRestrictedIndexOf(entityIdx, key); + } + + @Override + public int insertionIndexOf(String urnString) { + int entityIdx = findEntityIdx(urnString); + if (entityIdx < 0) { + return ~_delegate.length(); + } + long key; + try { + key = Long.parseLong(urnString, _entityPrefixes[entityIdx].length(), urnString.length(), 10); + } catch (NumberFormatException e) { + return ~_delegate.length(); + } + int found = rangeBinarySearch(entityIdx, key); + return found >= 0 ? found : found; + } + + @Override + public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { + int lowerEntity = findEntityIdx(lower); + int upperEntity = findEntityIdx(upper); + // Cross-entity ranges are undefined; return empty per the design doc. + if (lowerEntity < 0 || upperEntity < 0 || lowerEntity != upperEntity) { + return new IntOpenHashSet(0); + } + long lowerKey; + long upperKey; + try { + lowerKey = Long.parseLong(lower, _entityPrefixes[lowerEntity].length(), lower.length(), 10); + upperKey = Long.parseLong(upper, _entityPrefixes[upperEntity].length(), upper.length(), 10); + } catch (NumberFormatException e) { + return new IntOpenHashSet(0); + } + int rangeStart = _rangeStarts[lowerEntity]; + int rangeEnd = _ranges[lowerEntity].getEndDictIdExclusive(); + int firstDictId = -1; + int lastDictIdInclusive = -1; + // Linear pass; range is typically much smaller than the full dict. + for (int dictId = rangeStart; dictId < rangeEnd; dictId++) { + long v = _delegate.getLongValue(dictId); + boolean ge = includeLower ? v >= lowerKey : v > lowerKey; + boolean le = includeUpper ? v <= upperKey : v < upperKey; + if (ge && le) { + if (firstDictId < 0) { + firstDictId = dictId; + } + lastDictIdInclusive = dictId; + } + } + if (firstDictId < 0) { + return new IntOpenHashSet(0); + } + IntSet result = new IntOpenHashSet(lastDictIdInclusive - firstDictId + 1); + for (int i = firstDictId; i <= lastDictIdInclusive; i++) { + result.add(i); + } + return result; + } + + // ---------------- per-dictId accessors ---------------- + + @Override + public Object get(int dictId) { + return getStringValue(dictId); + } + + @Override + public int getIntValue(int dictId) { + return _delegate.getIntValue(dictId); + } + + @Override + public long getLongValue(int dictId) { + return _delegate.getLongValue(dictId); + } + + @Override + public float getFloatValue(int dictId) { + return _delegate.getFloatValue(dictId); + } + + @Override + public double getDoubleValue(int dictId) { + return _delegate.getDoubleValue(dictId); + } + + @Override + public BigDecimal getBigDecimalValue(int dictId) { + return _delegate.getBigDecimalValue(dictId); + } + + @Override + public String getStringValue(int dictId) { + int entityIdx = entityIdxForDictId(dictId); + if (entityIdx < 0) { + return null; + } + return _entityPrefixes[entityIdx] + _delegate.getLongValue(dictId); + } + + @Override + public void readStringValues(int[] dictIds, int length, String[] outValues) { + for (int i = 0; i < length; i++) { + outValues[i] = getStringValue(dictIds[i]); + } + } + + // ---------------- ordering / metadata ---------------- + + @Override + public int compare(int dictId1, int dictId2) { + return _delegate.compare(dictId1, dictId2); + } + + @Override + public Comparable getMinVal() { + // Composite ordering (entityType, then key) — return as string for cross-type stability. + return _delegate.length() == 0 ? null : getStringValue(0); + } + + @Override + public Comparable getMaxVal() { + int len = _delegate.length(); + return len == 0 ? null : getStringValue(len - 1); + } + + @Override + public Object getSortedValues() { + throw new UnsupportedOperationException(); + } + + @Override + public void close() + throws IOException { + _delegate.close(); + } + + // ---------------- internal helpers ---------------- + + /** Finds the entity index whose prefix is a prefix of {@code urnString}, or -1. */ + private int findEntityIdx(@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; + } + + /** Returns the entity index that owns {@code dictId}, or -1 if out of bounds. */ + int entityIdxForDictId(int dictId) { + if (dictId < 0 || dictId >= _delegate.length()) { + return -1; + } + // Linear scan — entity counts are small (typically <= 16). Switch to binary search if needed. + for (int i = 0; i < _ranges.length; i++) { + if (dictId < _ranges[i].getEndDictIdExclusive()) { + return i; + } + } + return -1; + } + + /** + * Range-restricted lookup for {@code longKey} within entity {@code entityIdx}'s dict-ID range. + * Returns the absolute dict ID or {@link #NULL_VALUE_INDEX}. + */ + private int rangeRestrictedIndexOf(int entityIdx, long longKey) { + int result = rangeBinarySearch(entityIdx, longKey); + return result >= 0 ? result : NULL_VALUE_INDEX; + } + + /** + * Returns the absolute dict ID of {@code longKey} within entity {@code entityIdx}'s range, + * or {@code -(insertionPoint+1)} when not found (standard {@code Arrays.binarySearch} contract, + * with the insertion point referring to an absolute dict ID). + */ + private int rangeBinarySearch(int entityIdx, long longKey) { + int low = _rangeStarts[entityIdx]; + int high = _ranges[entityIdx].getEndDictIdExclusive() - 1; + while (low <= high) { + int mid = (low + high) >>> 1; + long midValue = _delegate.getLongValue(mid); + if (midValue < longKey) { + low = mid + 1; + } else if (midValue > longKey) { + high = mid - 1; + } else { + return mid; + } + } + return -(low + 1); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java new file mode 100644 index 0000000000..74907afc59 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java @@ -0,0 +1,323 @@ +/** + * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.DimensionFieldSpec.EntityDictRange; +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 UrnSortedDictionary} using an in-memory fake LONG delegate. Verifies + * range-restricted indexOf, range-bound getStringValue reconstruction, cross-entity range + * handling, and pass-through accessors. + */ +public class UrnSortedDictionaryTest { + + /** + * In-memory stand-in for a LONG dictionary. Exposes a {@code long[]} buffer directly so tests + * can construct exotic layouts without spinning up the full segment-build machinery. + */ + private static final class FakeLongDictionary implements Dictionary { + private final long[] _values; + + FakeLongDictionary(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) { + throw new UnsupportedOperationException(); + } + + @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 { + } + } + + private static UrnSortedDictionary build(long[] keys, EntityDictRange... ranges) { + return new UrnSortedDictionary(new FakeLongDictionary(keys), Arrays.asList(ranges)); + } + + // Reference layout used by most tests: + // corpUser: dict IDs [0, 3), keys 100, 200, 300 + // groupId: dict IDs [3, 5), keys 100, 400 (note 100 collides with corpUser) + // appName: dict IDs [5, 6), key 50 + private static UrnSortedDictionary referenceDict() { + return build(new long[]{100, 200, 300, 100, 400, 50}, + new EntityDictRange("urn:li:corpUser", 3), + new EntityDictRange("urn:li:groupId", 5), + new EntityDictRange("urn:li:applicationName", 6)); + } + + // --- indexOf --- + + @Test + public void indexOfReturnsDictIdInCorpUserRange() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:corpUser:100"), 0); + assertEquals(d.indexOf("urn:li:corpUser:200"), 1); + assertEquals(d.indexOf("urn:li:corpUser:300"), 2); + } + + @Test + public void indexOfDisambiguatesCollidingKeyAcrossEntities() { + UrnSortedDictionary d = referenceDict(); + // Same key 100, different entities -> different dict IDs. + assertEquals(d.indexOf("urn:li:corpUser:100"), 0); + assertEquals(d.indexOf("urn:li:groupId:100"), 3); + } + + @Test + public void indexOfReturnsNullForKeyMissingFromEntity() { + UrnSortedDictionary d = referenceDict(); + // 999 isn't in any entity's range. + assertEquals(d.indexOf("urn:li:corpUser:999"), Dictionary.NULL_VALUE_INDEX); + // 50 exists in appName but not in groupId. + assertEquals(d.indexOf("urn:li:groupId:50"), Dictionary.NULL_VALUE_INDEX); + } + + @Test + public void indexOfReturnsNullForUnknownEntity() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:notAnEntity:100"), Dictionary.NULL_VALUE_INDEX); + } + + @Test + public void indexOfReturnsNullForMalformedKey() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.indexOf("urn:li:corpUser:notANumber"), Dictionary.NULL_VALUE_INDEX); + } + + // --- getStringValue --- + + @Test + public void getStringValueReconstructsAcrossAllRanges() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.getStringValue(0), "urn:li:corpUser:100"); + assertEquals(d.getStringValue(1), "urn:li:corpUser:200"); + assertEquals(d.getStringValue(2), "urn:li:corpUser:300"); + assertEquals(d.getStringValue(3), "urn:li:groupId:100"); + assertEquals(d.getStringValue(4), "urn:li:groupId:400"); + assertEquals(d.getStringValue(5), "urn:li:applicationName:50"); + } + + @Test + public void getStringValueReturnsNullForOutOfBounds() { + UrnSortedDictionary d = referenceDict(); + assertNull(d.getStringValue(-1)); + assertNull(d.getStringValue(6)); + } + + @Test + public void readStringValuesBatch() { + UrnSortedDictionary d = referenceDict(); + String[] out = new String[3]; + d.readStringValues(new int[]{5, 3, 0}, 3, out); + assertEquals(out[0], "urn:li:applicationName:50"); + assertEquals(out[1], "urn:li:groupId:100"); + assertEquals(out[2], "urn:li:corpUser:100"); + } + + // --- pass-through accessors --- + + @Test + public void getLongValuePassesThroughToDelegate() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.getLongValue(0), 100L); + assertEquals(d.getLongValue(3), 100L); + assertEquals(d.getLongValue(5), 50L); + } + + @Test + public void getIntValuePassesThroughForSmallKeys() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.getIntValue(0), 100); + assertEquals(d.getIntValue(5), 50); + } + + @Test + public void getReturnsReconstructedUrnString() { + UrnSortedDictionary d = referenceDict(); + assertEquals(d.get(3), "urn:li:groupId:100"); + } + + // --- metadata --- + + @Test + public void getValueTypeIsString() { + assertEquals(referenceDict().getValueType(), DataType.STRING); + } + + @Test + public void isSortedReturnsFalse() { + // The buffer is sorted only within each entity range; not globally lex-sorted as URNs. + assertTrue(!referenceDict().isSorted()); + } + + @Test + public void lengthMatchesDelegate() { + assertEquals(referenceDict().length(), 6); + } + + // --- getDictIdsInRange --- + + @Test + public void getDictIdsInRangeSameEntityInclusive() { + UrnSortedDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:corpUser:200", true, true); + assertEquals(result.size(), 2); + assertTrue(result.contains(0)); + assertTrue(result.contains(1)); + } + + @Test + public void getDictIdsInRangeSameEntityExclusiveBounds() { + UrnSortedDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:corpUser:300", false, false); + assertEquals(result.size(), 1); + assertTrue(result.contains(1)); + } + + @Test + public void getDictIdsInRangeCrossEntityReturnsEmpty() { + UrnSortedDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:groupId:200", true, true); + assertEquals(result.size(), 0); + } + + @Test + public void getDictIdsInRangeUnknownEntityReturnsEmpty() { + UrnSortedDictionary d = referenceDict(); + IntSet result = + d.getDictIdsInRange("urn:li:notAnEntity:100", "urn:li:notAnEntity:200", true, true); + assertEquals(result.size(), 0); + } + + // --- min/max --- + + @Test + public void getMinValIsFirstReconstructedUrn() { + assertEquals(referenceDict().getMinVal(), "urn:li:corpUser:100"); + } + + @Test + public void getMaxValIsLastReconstructedUrn() { + assertEquals(referenceDict().getMaxVal(), "urn:li:applicationName:50"); + } +} From bce50fff50f6d274be4ee5e0bf84f8e5b57af04d Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 21:31:14 +0000 Subject: [PATCH 08/24] UrnDictionary commit 3/5: switch multi-entity URN main column to sorted LONG End-to-end wiring of the sorted-LONG UrnDictionary layout for multi-entity URN columns whose declared entity types all have numeric value types. The encoded ":" STRING dictionary is replaced with a LONG dictionary partitioned by entity range; UrnSortedDictionary handles all read-time reconstruction and lookup. Storage / build path: * DimensionFieldSpec.getEffectiveStoredType -> LONG for multi-entity URN (was STRING). * UrnKeyEncodingTransformer multi-entity branch now encodes only the numeric key into the main column. The companion entity-idx INT column is still written (parallel-write; commit 4 will drop it). * UrnMultiEntityPreIndexStatsCollector tracks unique (entityIdx, longKey) pairs by reading the companion column from the row. Reports isSorted=false so the sorted-forward-index path isn't triggered (dict IDs are not monotonic across entity boundaries even when the LONG values are). * SegmentPreIndexStatsCollectorImpl dispatches the row to the new collector via collect(value, row) for multi-entity URN columns. * UrnSortedDictionaryCreator (extends SegmentDictionaryCreator): consumes EntityKeyPair[] from the stats collector, writes a LONG buffer sorted by (entityIdx, longKey), and exposes the per-entity dict-ID ranges. indexOfSV takes an EntityKeyPair from the row indexer. * SegmentColumnarIndexCreator branches to UrnSortedDictionaryCreator for multi-entity URN columns, plumbs EntityKeyPair through indexRow, and surfaces the computed ranges on the FieldSpec so the metadata writer persists them. Read path: * DictionaryIndexType.read returns UrnSortedDictionary (commit 2) when the loaded FieldSpec has entityDictRanges. * ColumnMetadataImpl.parseEntityDictRanges (commit 1) handles the PropertiesConfiguration comma-split via getStringArray + rejoin. Tests / benchmarks: * UrnReconstructingMultiEntityDictionary still exists but is now unused (commit 4 will delete it). * UrnKeyEncodingTransformerTest multi-entity assertions updated to expect numeric LONG keys on the main column. All 145 URN tests pass. Companion column is still auto-created in this commit -- the query layer still uses COMPANION_LOOKUP for urnEntity(). Commit 4 will drop the companion column entirely and switch urnEntity() to RANGE_LOOKUP over the entity ranges. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../UrnKeyEncodingTransformer.java | 43 ++---- .../impl/SegmentColumnarIndexCreator.java | 31 +++- .../SegmentPreIndexStatsCollectorImpl.java | 31 +++- .../UrnMultiEntityPreIndexStatsCollector.java | 118 ++++++++++++++ .../creator/impl/urn/EntityKeyPair.java | 71 +++++++++ .../impl/urn/UrnSortedDictionaryCreator.java | 144 ++++++++++++++++++ .../index/dictionary/DictionaryIndexType.java | 8 +- .../UrnKeyEncodingTransformerTest.java | 10 +- .../index/metadata/ColumnMetadataImpl.java | 13 +- .../pinot/spi/data/DimensionFieldSpec.java | 12 +- 10 files changed, 431 insertions(+), 50 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/UrnMultiEntityPreIndexStatsCollector.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityKeyPair.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java index dea03d4e4f..f0f91012e0 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java @@ -34,20 +34,13 @@ * Converts URN string values in typed URN columns to their compact physical representation and * writes per-row entity-type indices for multi-entity URN columns. * - *

    Two encoding modes are selected at construction time per column: - *

      - *
    • Key-encoding ({@link DimensionFieldSpec#hasTypedUrnStorage()}): single-entity LONG - * or INT column. The URN string is parsed and its numeric key replaces the string value on - * the row. No companion column is written.
      • - *
    • Prefix-stripped ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}): - * multi-entity URN column where all declared types have numeric value types. The main column - * is rewritten as {@code ":"} (e.g. {@code "0:12345"} for the first - * declared entity type), stripping the repeated URN prefix from every dictionary entry. - * The 0-based entity index is also written to the companion {@code __urnEntityIdx} - * column for fast per-row entity-type lookup. At read time, the column dictionary is - * wrapped so {@code getStringValue()} reconstructs the full URN string and predicate - * lookups encode their literals before delegating to the underlying dictionary.
      • - *
    + *

    For URN columns whose effective stored type is LONG or INT + * ({@link DimensionFieldSpec#hasTypedUrnStorage()}), the URN string is parsed and its numeric key + * replaces the string value on the row. For multi-entity URN columns + * ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}), the 0-based entity index is also + * written to the companion {@code __urnEntityIdx} column. The stats collector and + * {@code UrnSortedDictionaryCreator} pair the main column LONG with the companion column entity + * index to lay out a per-entity sorted-LONG dictionary at segment build time. * *

    If the URN value is {@code null} or unparseable, both the main and companion columns are set * to {@code null}; downstream {@link NullValueTransformer} fills in default null values. @@ -71,14 +64,11 @@ public UrnKeyEncodingTransformer(Schema schema) { } DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; if (dim.hasTypedUrnStorage()) { - // Single-entity typed LONG/INT: encode the key to numeric; no companion column. - encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), null)); - } else if (dim.needsCompanionEntityIdxColumn()) { - // Multi-entity with all-numeric value types: encode main column as ":" - // to strip the repeated URN prefix from every dictionary entry. The companion column - // stores the 0-based entity-type index for fast urnEntity() lookup. - encodings.add(new ColumnEncoding(dim.getName(), FieldSpec.DataType.STRING, dim.getUrnTypes(), - dim.companionEntityIdxColumnName())); + // Single-entity typed LONG/INT and multi-entity LONG/INT (UrnDictionary layout): encode + // the numeric key into the main column. Multi-entity also writes the entity index to a + // companion INT column so the segment builder can pair (entityIdx, longKey). + String companion = dim.needsCompanionEntityIdxColumn() ? dim.companionEntityIdxColumnName() : null; + encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), companion)); } } _encodings = Collections.unmodifiableList(encodings); @@ -131,13 +121,8 @@ public GenericRow transform(GenericRow record) { record.putValue(enc._companionColumnName, entityIdx); } - // Encode the main column. - if (enc._storedType == FieldSpec.DataType.STRING) { - // Multi-entity prefix-stripped encoding: ":". - record.putValue(enc._columnName, - ((Integer) record.getValue(enc._companionColumnName)).intValue() + ":" + urn.simpleKey()); - } else if (enc._storedType != null) { - // Single-entity typed numeric: keep just the LONG/INT key. + // Encode the main column as the typed numeric key. + if (enc._storedType != null) { String keyStr = urn.simpleKey(); try { if (enc._storedType == FieldSpec.DataType.LONG) { 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 9459aec965..774b03128c 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 @@ -38,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.UrnSortedDictionaryCreator; 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; @@ -244,8 +245,16 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio dictConfig = new DictionaryIndexConfig(dictConfig, columnIndexCreationInfo.isUseVarLengthDictionary()); } - SegmentDictionaryCreator creator = - new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); + // Multi-entity URN columns get a custom dict creator that lays out a per-entity sorted + // LONG buffer plus dictId range metadata. All other columns go through the standard path. + SegmentDictionaryCreator creator; + if (fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).needsCompanionEntityIdxColumn()) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + creator = new UrnSortedDictionaryCreator(fieldSpec, _indexDir, dim.getUrnTypes()); + } else { + creator = new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); + } Object dictValues = maybeStripUrnPrefix(columnName, fieldSpec, context.getSortedUniqueElementsArray()); try { @@ -256,6 +265,12 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio 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 (read in commit 1). + if (creator instanceof UrnSortedDictionaryCreator && fieldSpec instanceof DimensionFieldSpec) { + ((DimensionFieldSpec) fieldSpec).setEntityDictRanges(((UrnSortedDictionaryCreator) creator).getRanges()); + } + _dictionaryCreatorMap.put(columnName, creator); } @@ -544,6 +559,18 @@ public void indexRow(GenericRow row) FieldSpec fieldSpec = resolveFieldSpec(columnName); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); + // For multi-entity URN columns, pair the row's main column LONG with the companion column's + // entity index so the UrnSortedDictionaryCreator can map (entityIdx, longKey) → dict ID. + if (dictionaryCreator instanceof UrnSortedDictionaryCreator) { + Object companionVal = row.getValue( + ((DimensionFieldSpec) fieldSpec).companionEntityIdxColumnName()); + if (companionVal != null) { + int entityIdx = ((Number) companionVal).intValue(); + long longKey = ((Number) columnValueToIndex).longValue(); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( + entityIdx, longKey); + } + } try { if (fieldSpec.isSingleValueField()) { indexSingleValueRow(dictionaryCreator, columnValueToIndex, creatorsByIndex); 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 aecdd73631..583873c700 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 @@ -55,7 +55,16 @@ public void init() { _columnStatsCollectorMap.put(column, new IntColumnPreIndexStatsCollector(column, _statsCollectorConfig)); break; case LONG: - _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig)); + // Multi-entity URN columns get a stats collector that tracks (entityIdx, longKey) + // pairs so the UrnSortedDictionaryCreator can lay out per-entity dict ID ranges. + if (fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).needsCompanionEntityIdxColumn()) { + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + _columnStatsCollectorMap.put(column, new UrnMultiEntityPreIndexStatsCollector(column, + dim.companionEntityIdxColumnName(), _statsCollectorConfig)); + } else { + _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig)); + } break; case FLOAT: _columnStatsCollectorMap.put(column, new FloatColumnPreIndexStatsCollector(column, _statsCollectorConfig)); @@ -113,13 +122,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..7f56be0fef --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/UrnMultiEntityPreIndexStatsCollector.java @@ -0,0 +1,118 @@ +/** + * 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.EntityKeyPair; +import org.apache.pinot.segment.spi.creator.StatsCollectorConfig; +import org.apache.pinot.spi.data.readers.GenericRow; + + +/** + * Stats collector for multi-entity URN columns whose effective stored type is LONG. + * + *

    Unlike the standard {@link LongColumnPreIndexStatsCollector}, which dedupes by LONG value + * alone, this collector tracks unique {@code (entityIdx, longKey)} pairs. Two URNs that share a + * key but belong to different entity types remain distinct dictionary entries, which is what the + * {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnSortedDictionaryCreator} + * needs to lay out the per-entity dict ID ranges. + * + *

    Per-column dispatch in {@code SegmentPreIndexStatsCollectorImpl.collectRow} routes to + * {@link #collect(Object, GenericRow)} so the collector can pick the entity index up from the + * row's companion column. The standard {@link #collect(Object)} method falls back to ignoring + * the entity dimension (effectively LONG-only) — only used by paths that don't have row context. + */ +public class UrnMultiEntityPreIndexStatsCollector extends LongColumnPreIndexStatsCollector { + private final String _companionColumnName; + private final Set _pairs = new HashSet<>(); + private EntityKeyPair[] _sortedPairs; + private boolean _sealed = false; + + public UrnMultiEntityPreIndexStatsCollector(String column, String companionColumnName, + StatsCollectorConfig statsCollectorConfig) { + super(column, statsCollectorConfig); + _companionColumnName = companionColumnName; + } + + public String getCompanionColumnName() { + return _companionColumnName; + } + + /** + * Collects a row's main-column LONG value together with the companion column entity index. + * Used by {@code SegmentPreIndexStatsCollectorImpl.collectRow} for multi-entity URN columns. + */ + public void collect(Object value, GenericRow row) { + assert !_sealed; + if (value == null) { + return; + } + Object companion = row.getValue(_companionColumnName); + if (companion == null) { + // Row was nulled by the encoder (e.g. unknown entity); ignore for pair tracking. + return; + } + int entityIdx = ((Number) companion).intValue(); + long key = ((Number) value).longValue(); + _pairs.add(new EntityKeyPair(entityIdx, key)); + super.collect(value); // keep LongColumn min/max/cardinality stats in sync + } + + /** + * Returns the sorted unique {@code (entityIdx, longKey)} pairs after {@link #seal()}. + * Consumed by the {@code UrnSortedDictionaryCreator}. + */ + @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 void seal() { + if (_sealed) { + return; + } + super.seal(); + _sortedPairs = _pairs.toArray(new EntityKeyPair[0]); + Arrays.sort(_sortedPairs); + _pairs.clear(); + _sealed = true; + } + + /** + * Multi-entity URN columns are never considered sorted for forward-index purposes: the LONG + * key buffer is partitioned by entity, so consecutive rows may jump across entity ranges in + * dict-ID space even when the LONG values themselves come in monotone order. Letting the + * sorted forward-index path kick in would corrupt the per-row dict-ID readback. + */ + @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/UrnSortedDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java new file mode 100644 index 0000000000..82d1c0c87d --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.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.segment.local.segment.creator.impl.urn; + +import com.google.common.base.Preconditions; +import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap; +import java.io.File; +import java.io.IOException; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.io.util.FixedByteValueReaderWriter; +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.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Custom dictionary creator for multi-entity URN columns. Writes a LONG buffer of keys sorted by + * {@code (entityIdx, longKey)} — allowing duplicate LONG values across entity boundaries — plus + * per-entity dictionary-ID ranges. The companion {@link + * org.apache.pinot.segment.local.segment.index.readers.UrnSortedDictionary} reader uses the + * ranges to do range-restricted lookups at query time. + */ +public class UrnSortedDictionaryCreator extends SegmentDictionaryCreator { + + private static final Logger LOGGER = LoggerFactory.getLogger(UrnSortedDictionaryCreator.class); + + private final String _columnName; + private final File _dictionaryFile; + private final List _urnTypes; + private List _ranges; + // (entityIdx, longKey) → absolute dictId. Two-level: array per entity, then long→dictId map. + private Long2IntOpenHashMap[] _pairToDictId; + private int _numEntries; + + public UrnSortedDictionaryCreator(FieldSpec fieldSpec, File indexDir, + List urnTypes) { + super(fieldSpec, indexDir, false); + _columnName = fieldSpec.getName(); + _dictionaryFile = new File(indexDir, _columnName + DictionaryIndexType.getFileExtension()); + _urnTypes = urnTypes; + } + + /** + * Builds the dictionary from a pre-sorted array of {@link EntityKeyPair}s. Input must be sorted + * by {@link EntityKeyPair#compareTo} (i.e. ascending entity index, then ascending long key) and + * contain no duplicates. + */ + @Override + public void build(Object sortedValues) + throws IOException { + Preconditions.checkArgument(sortedValues instanceof EntityKeyPair[], + "UrnSortedDictionaryCreator expects an EntityKeyPair[] from the stats collector"); + EntityKeyPair[] pairs = (EntityKeyPair[]) sortedValues; + Preconditions.checkState(pairs.length > 0, "UrnSortedDictionaryCreator received empty input"); + + FileUtils.touch(_dictionaryFile); + _numEntries = pairs.length; + _pairToDictId = new Long2IntOpenHashMap[_urnTypes.size()]; + for (int i = 0; i < _pairToDictId.length; i++) { + _pairToDictId[i] = new Long2IntOpenHashMap(); + _pairToDictId[i].defaultReturnValue(-1); + } + List ranges = new ArrayList<>(); + int currentEntity = -1; + + try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapFile(_dictionaryFile, false, 0, + (long) _numEntries * Long.BYTES, ByteOrder.BIG_ENDIAN, getClass().getSimpleName()); + FixedByteValueReaderWriter writer = new FixedByteValueReaderWriter(dataBuffer)) { + for (int i = 0; i < _numEntries; i++) { + EntityKeyPair pair = pairs[i]; + int entityIdx = pair.getEntityIdx(); + Preconditions.checkState(entityIdx >= currentEntity, + "UrnSortedDictionaryCreator input must be sorted by entityIdx"); + // Cross an entity boundary: close out the previous entity's range. + while (currentEntity < entityIdx) { + if (currentEntity >= 0) { + ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), i)); + } + currentEntity++; + } + writer.writeLong(i, pair.getLongKey()); + _pairToDictId[entityIdx].put(pair.getLongKey(), i); + } + } + // Close out the trailing range(s). Entities with no rows still get an entry with the same + // endDictId as the previous one (empty range), so readers can use a uniform lookup table. + while (currentEntity < _urnTypes.size() - 1) { + ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), _numEntries)); + currentEntity++; + } + ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), _numEntries)); + _ranges = ranges; + + LOGGER.info("Built UrnSortedDictionary for column: {} with cardinality: {} across {} entities", _columnName, + _numEntries, _urnTypes.size()); + } + + /** {@inheritDoc} For multi-entity URN, {@code value} must be an {@link EntityKeyPair}. */ + @Override + public int indexOfSV(Object value) { + Preconditions.checkState(_pairToDictId != null, "build() must be called before indexOfSV()"); + EntityKeyPair pair = (EntityKeyPair) value; + int dictId = _pairToDictId[pair.getEntityIdx()].get(pair.getLongKey()); + Preconditions.checkState(dictId >= 0, + "Row references (entity=%s, key=%s) that wasn't seen during stats collection", + _urnTypes.get(pair.getEntityIdx()).getEntityType(), pair.getLongKey()); + return dictId; + } + + @Override + public int getNumBytesPerEntry() { + return Long.BYTES; + } + + /** 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 4c7d742819..c0d0fffc1b 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 @@ -52,6 +52,7 @@ import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingLongDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingMultiEntityDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingStringDictionary; +import org.apache.pinot.segment.local.segment.index.readers.UrnSortedDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -333,10 +334,15 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat case LONG: { Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) : new LongDictionary(dataBuffer, length); - // For single-entity typed URN LONG columns, wrap to reconstruct full URN strings on read. FieldSpec fieldSpec = metadata.getFieldSpec(); if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + // Multi-entity URN with per-entity dict ID ranges (UrnDictionary layout): wrap with a + // range-aware reader that reconstructs URN strings using the appropriate entity prefix. + if (dim.getEntityDictRanges() != null) { + return new UrnSortedDictionary(dict, dim.getEntityDictRanges()); + } + // Single-entity typed URN LONG: wrap to prepend the entity prefix. if (dim.hasTypedUrnStorage() && !dim.needsCompanionEntityIdxColumn()) { String entityType = dim.getUrnTypes().get(0).getEntityType(); return new UrnReconstructingLongDictionary(dict, entityType); 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 index f69e312e96..b11e457fef 100644 --- 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 @@ -123,9 +123,10 @@ public void multiEntityFirstEntityWritesCompanionIdx0() { UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); GenericRow row = rowWithValue("resourceId", "urn:li:memberId:100"); t.transform(row); - // Multi-entity strips the URN prefix and encodes the main column as ":". - // UrnReconstructingMultiEntityDictionary reverses this at read time. - assertEquals(row.getValue("resourceId"), "0:100"); + // Multi-entity LONG storage (UrnDictionary layout): main column holds just the numeric key; + // the companion column carries the entity index for the segment-build pipeline to pair + // them into (entityIdx, longKey) dictionary entries. + assertEquals(row.getValue("resourceId"), 100L); assertEquals(row.getValue("resourceId__urnEntityIdx"), 0); } @@ -134,8 +135,7 @@ public void multiEntitySecondEntityWritesCompanionIdx1() { UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); GenericRow row = rowWithValue("resourceId", "urn:li:groupId:42"); t.transform(row); - // Multi-entity strips the URN prefix and encodes the main column as ":". - assertEquals(row.getValue("resourceId"), "1:42"); + assertEquals(row.getValue("resourceId"), 42L); assertEquals(row.getValue("resourceId__urnEntityIdx"), 1); } 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 b01e716e43..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 @@ -334,12 +334,15 @@ public static FieldSpec generateFieldSpec(String column, PropertiesConfiguration } } // Multi-entity URN sorted-LONG dictionary layout: load per-entity dict ID ranges if - // the segment was built with the UrnDictionary format. - String entityDictRangesStr = - config.getString(Column.getKeyFor(column, Column.URN_ENTITY_DICT_RANGES), null); - if (entityDictRangesStr != null) { + // 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(entityDictRangesStr)); + 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. } 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 3c2246095a..801e070ed1 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 @@ -277,7 +277,17 @@ public DataType getEffectiveStoredType() { // 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) { - return DataType.STRING; + // 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)) { From 486fc88fd1f1a537b865d0e9f7d860edf2342da1 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 22:11:47 +0000 Subject: [PATCH 09/24] UrnDictionary commit 4/5: drop companion column, switch to RANGE_LOOKUP Eliminates the per-row __urnEntityIdx INT companion column for multi-entity URN columns. Entity index is now stashed on the row in a transient field at ingest time (UrnKeyEncodingTransformer), consumed by the stats collector and forward-index code, and dropped before metadata write. No new physical column lives in the segment. SPI: * DimensionFieldSpec.needsCompanionEntityIdxColumn -> usesSortedUrnDictionary * companionEntityIdxColumnName / createCompanionEntityIdxFieldSpec / COMPANION_ENTITY_IDX_SUFFIX -> deleted, replaced by transientUrnEntityIdxFieldName / URN_TRANSIENT_ENTITY_IDX_PREFIX Build path: * UrnKeyEncodingTransformer writes the entity index to __urnEntityIdxBuild_ (transient) instead of a column. * SegmentColumnarIndexCreator / SegmentIndexCreationDriverImpl / SegmentPreIndexStatsCollectorImpl no longer auto-create, register, or write the companion column. _companionFieldSpecs map removed, DIMENSIONS list no longer gets the synthetic suffix appended. * UrnMultiEntityPreIndexStatsCollector reads the transient row field instead of a sibling column. Query path: * UrnEntityTransformFunction: COMPANION_LOOKUP strategy replaced with RANGE_LOOKUP. It reads the main column's dict IDs and finds the owning entity range with a linear scan over the metadata range table (typically <=16 entries). No per-row companion read, no URN string parsing. * UrnProjectionUtils deleted -- nothing to inject into projections. Callers in ProjectPlanNode and ExpressionFilterOperator updated. * DictionaryIndexType: branch returning UrnReconstructingMultiEntityDictionary removed (obsoleted by UrnSortedDictionary returned via the LONG path). * UrnReconstructingMultiEntityDictionary file deleted. Tests: * UrnEntityMultiEntityTest: companionColumnIsPresent -> entityDictRangesPersistedInMetadata; urnEntityUsesCompanionFastPath -> urnEntityUsesRangeLookupFastPath; the urnEntityFallsBackToParseStringWhenCompanionNotProjected case is obsolete (no companion column to omit) and was dropped. * UrnKeyEncodingTransformerTest multi-entity asserts updated to read the transient field instead of a companion column. * All 144 URN tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../filter/ExpressionFilterOperator.java | 2 - .../function/UrnEntityTransformFunction.java | 76 ++++--- .../function/UrnProjectionUtils.java | 83 -------- .../pinot/core/plan/FilterPlanNode.java | 2 +- .../pinot/core/plan/ProjectPlanNode.java | 2 - .../function/UrnEntityMultiEntityTest.java | 58 ++--- .../UrnKeyEncodingTransformer.java | 37 ++-- .../impl/SegmentColumnarIndexCreator.java | 76 +------ .../impl/SegmentIndexCreationDriverImpl.java | 18 -- .../SegmentPreIndexStatsCollectorImpl.java | 20 +- .../UrnMultiEntityPreIndexStatsCollector.java | 21 +- .../index/dictionary/DictionaryIndexType.java | 11 +- ...rnReconstructingMultiEntityDictionary.java | 201 ------------------ .../UrnKeyEncodingTransformerTest.java | 18 +- .../pinot/spi/data/DimensionFieldSpec.java | 40 ++-- 15 files changed, 129 insertions(+), 536 deletions(-) delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java delete mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java index 3673cb6bbe..27889f8ec3 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExpressionFilterOperator.java @@ -39,7 +39,6 @@ import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluatorProvider; import org.apache.pinot.core.operator.transform.function.TransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; -import org.apache.pinot.core.operator.transform.function.UrnProjectionUtils; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.datasource.DataSource; @@ -61,7 +60,6 @@ public ExpressionFilterOperator(IndexSegment segment, QueryContext queryContext, Set columns = new HashSet<>(); ExpressionContext lhs = predicate.getLhs(); lhs.getColumns(columns); - UrnProjectionUtils.addUrnCompanionColumns(Collections.singletonList(lhs), queryContext.getSchema(), columns); int mapCapacity = HashUtil.getHashMapCapacity(columns.size()); _dataSourceMap = new HashMap<>(mapCapacity); Map columnContextMap = new HashMap<>(mapCapacity); 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 index 268ea1e02e..a20398b6aa 100644 --- 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 @@ -28,6 +28,7 @@ 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.DimensionFieldSpec.EntityDictRange; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.utils.Urn; @@ -37,16 +38,16 @@ * *

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

      - *
    1. Constant — single-entity typed URN column with LONG/INT storage: every row in the - * block gets the same literal string, so the block is filled with one {@link Arrays#fill} - * call and no per-row URN parsing.
      2. - *
    2. Companion-lookup — multi-entity typed URN column with the companion - * {@code __urnEntityIdx} INT column projected: reads the per-row entity index and - * indexes into a pre-resolved {@code String[]} of entity types. No URN string parsing per - * row.
      3. - *
    3. Parse-string — fallback: untyped URN columns, or multi-entity columns whose - * companion column was not included in the projection. Reads each URN string value and - * calls {@link Urn#tryParse} per row.
      4. + *
    4. Constant — single-entity typed URN column (declared or auto-detected): every row + * in the block gets the same literal string, so the block is filled with one + * {@link Arrays#fill} call and no per-row URN parsing.
      5. + *
    5. Range-lookup — multi-entity typed URN column using the sorted-LONG UrnDictionary + * layout (see {@link DimensionFieldSpec#getEntityDictRanges()}): reads main-column dict IDs + * and finds the owning entity range with a linear scan over the per-entity range table + * (typically <=16 entries). No URN string parsing per row.
      6. + *
    6. Parse-string — fallback: untyped URN columns, or multi-entity columns whose entity + * dict ranges aren't available. Reads each URN string value and calls {@link Urn#tryParse} + * per row.
      7. *
    */ public class UrnEntityTransformFunction extends BaseTransformFunction { @@ -55,13 +56,14 @@ public class UrnEntityTransformFunction extends BaseTransformFunction { public static final String FUNCTION_NAME_ALT = "urn_entity"; private enum Strategy { - CONSTANT, COMPANION_LOOKUP, PARSE_STRING + CONSTANT, RANGE_LOOKUP, PARSE_STRING } private Strategy _strategy; private String _constantEntityType; - private String _companionColumnName; - private String[] _entityTypesByIdx; + private String _mainColumnName; + private int[] _rangeEnds; + private String[] _entityTypesByRange; @Override public String getName() { @@ -83,25 +85,26 @@ public void init(List arguments, Map c FieldSpec fieldSpec = dataSource.getDataSourceMetadata().getFieldSpec(); if (fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - // Single-entity URN columns (declared single-entity or auto-detected single-prefix) - // always have the same entity prefix on every row — fill the block with one literal. + // Single-entity URN columns always have the same entity prefix on every row — fill + // the block with one literal. if (dim.isSingleUrnEntityType()) { _strategy = Strategy.CONSTANT; _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); return; } - if (dim.needsCompanionEntityIdxColumn()) { - String companionCol = dim.companionEntityIdxColumnName(); - if (columnContextMap.containsKey(companionCol)) { - List urnTypes = dim.getUrnTypes(); - _entityTypesByIdx = new String[urnTypes.size()]; - for (int i = 0; i < urnTypes.size(); i++) { - _entityTypesByIdx[i] = urnTypes.get(i).getEntityType(); - } - _companionColumnName = companionCol; - _strategy = Strategy.COMPANION_LOOKUP; - return; + // Multi-entity URN columns using the sorted-LONG UrnDictionary layout: the per-entity + // dict ID ranges in column metadata tell us which entity owns each row's dict ID. + 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; } } } @@ -123,13 +126,22 @@ public String[] transformToStringValuesSV(ValueBlock valueBlock) { case CONSTANT: Arrays.fill(_stringValuesSV, 0, length, _constantEntityType); break; - case COMPANION_LOOKUP: - BlockValSet companionBlock = valueBlock.getBlockValueSet(_companionColumnName); - int[] entityIdxes = companionBlock.getIntValuesSV(); - int numTypes = _entityTypesByIdx.length; + case RANGE_LOOKUP: + BlockValSet mainBlock = valueBlock.getBlockValueSet(_mainColumnName); + int[] dictIds = mainBlock.getDictionaryIdsSV(); + int numRanges = _rangeEnds.length; for (int i = 0; i < length; i++) { - int idx = entityIdxes[i]; - _stringValuesSV[i] = (idx >= 0 && idx < numTypes) ? _entityTypesByIdx[idx] : null; + int dictId = dictIds[i]; + // Linear scan — range count is small (typically <=16) and the scan is branch-friendly. + // For higher cardinalities, switch to binary search on _rangeEnds. + String entityType = null; + for (int r = 0; r < numRanges; r++) { + if (dictId < _rangeEnds[r]) { + entityType = _entityTypesByRange[r]; + break; + } + } + _stringValuesSV[i] = entityType; } break; case PARSE_STRING: diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java deleted file mode 100644 index 1ba20eb721..0000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/UrnProjectionUtils.java +++ /dev/null @@ -1,83 +0,0 @@ -/** - * 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 java.util.Collection; -import java.util.Set; -import javax.annotation.Nullable; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.request.context.FunctionContext; -import org.apache.pinot.spi.data.DimensionFieldSpec; -import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.Schema; - - -/** - * Utilities for ensuring URN companion entity-index columns are included in a projection so the - * {@link UrnEntityTransformFunction} can take its fast COMPANION_LOOKUP path on multi-entity typed - * URN columns instead of falling back to per-row URN string parsing. - */ -public final class UrnProjectionUtils { - - private UrnProjectionUtils() { - } - - /** - * Walks {@code expressions} for {@code urnEntity(col)} calls on multi-entity typed URN columns - * (see {@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}) and adds the corresponding - * companion {@code __urnEntityIdx} INT column names to {@code projectionColumns}. - * - *

    No-op when {@code schema} is null or no matching URN expressions are present. - */ - public static void addUrnCompanionColumns(Collection expressions, @Nullable Schema schema, - Set projectionColumns) { - if (schema == null) { - return; - } - for (ExpressionContext expression : expressions) { - collect(expression, schema, projectionColumns); - } - } - - private static void collect(ExpressionContext expression, Schema schema, Set projectionColumns) { - if (expression.getType() != ExpressionContext.Type.FUNCTION) { - return; - } - FunctionContext function = expression.getFunction(); - String name = function.getFunctionName(); - if ((UrnEntityTransformFunction.FUNCTION_NAME.equalsIgnoreCase(name) - || UrnEntityTransformFunction.FUNCTION_NAME_ALT.equalsIgnoreCase(name)) - && function.getArguments().size() == 1) { - ExpressionContext arg = function.getArguments().get(0); - if (arg.getType() == ExpressionContext.Type.IDENTIFIER) { - String colName = arg.getIdentifier(); - FieldSpec fieldSpec = schema.getFieldSpecFor(colName); - if (fieldSpec instanceof DimensionFieldSpec) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.needsCompanionEntityIdxColumn()) { - projectionColumns.add(dim.companionEntityIdxColumnName()); - } - } - } - } - for (ExpressionContext child : function.getArguments()) { - collect(child, schema, projectionColumns); - } - } -} 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 2a3fdd5994..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 @@ -396,7 +396,7 @@ Predicate rewriteUrnKeyLongPredicate(Predicate predicate) { return predicate; } DimensionFieldSpec dim = (DimensionFieldSpec) fs; - if (!dim.hasTypedUrnStorage() || dim.needsCompanionEntityIdxColumn()) { + if (!dim.hasTypedUrnStorage() || dim.usesSortedUrnDictionary()) { return predicate; } String prefix = dim.getUrnTypes().get(0).getEntityType() + ":"; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java index 68dc63491c..d9b4cb8aa8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/ProjectPlanNode.java @@ -32,7 +32,6 @@ import org.apache.pinot.core.operator.ProjectionOperatorUtils; import org.apache.pinot.core.operator.filter.BaseFilterOperator; import org.apache.pinot.core.operator.transform.TransformOperator; -import org.apache.pinot.core.operator.transform.function.UrnProjectionUtils; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentContext; @@ -75,7 +74,6 @@ public BaseProjectOperator run() { hasNonIdentifierExpression = true; } } - UrnProjectionUtils.addUrnCompanionColumns(_expressions, _queryContext.getSchema(), projectionColumns); Map dataSourceMap = new HashMap<>(HashUtil.getHashMapCapacity(projectionColumns.size())); projectionColumns.forEach( column -> dataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema()))); 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 index 394bcbbf9c..5bda752857 100644 --- 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 @@ -55,13 +55,13 @@ /** * Tests the multi-entity {@code urnEntity()} fast path: COMPANION_LOOKUP strategy that reads - * pre-computed entity indices from the companion {@code col__urnEntityIdx} INT column. + * 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 COMPANION_COL = URN_COL + "__urnEntityIdx"; private static final String TIME_COL = "t"; private static final int NUM_ROWS = 60; @@ -69,7 +69,6 @@ public class UrnEntityMultiEntityTest { private ImmutableSegment _segment; private Map _dataSourceMap; private ProjectionBlock _projectionBlock; - private ProjectionBlock _projectionBlockNoCompanion; private final String[] _expectedEntities = new String[NUM_ROWS]; @BeforeClass @@ -106,20 +105,12 @@ public void setUp() _dataSourceMap = new HashMap<>(); _dataSourceMap.put(URN_COL, _segment.getDataSource(URN_COL)); - _dataSourceMap.put(COMPANION_COL, _segment.getDataSource(COMPANION_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(); - - Map noCompanionMap = new HashMap<>(); - noCompanionMap.put(URN_COL, _segment.getDataSource(URN_COL)); - noCompanionMap.put(TIME_COL, _segment.getDataSource(TIME_COL)); - _projectionBlockNoCompanion = new ProjectionOperator(noCompanionMap, - new DocIdSetOperator(new MatchAllFilterOperator(numDocs), - DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock(); } @AfterClass @@ -131,14 +122,18 @@ public void tearDown() { } @Test - public void companionColumnIsPresent() { - DataSource ds = _segment.getDataSource(COMPANION_COL); - assertNotNull(ds, "Companion column should be auto-created during segment build"); - assertEquals(ds.getDataSourceMetadata().getDataType(), FieldSpec.DataType.INT); + 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 urnEntityUsesCompanionFastPath() { + public void urnEntityUsesRangeLookupFastPath() { TransformFunction fn = TransformFunctionFactory.get( RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), _dataSourceMap); assertNotNull(fn); @@ -147,35 +142,8 @@ public void urnEntityUsesCompanionFastPath() { assertNotNull(result); assertEquals(result.length, NUM_ROWS); - // Result order matches doc order in the segment. Multi-entity STRING storage preserves - // ingestion order for the main column; the companion column is written in lockstep. - 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); - } - - @Test - public void urnEntityFallsBackToParseStringWhenCompanionNotProjected() { - TransformFunction fn = TransformFunctionFactory.get( - RequestContextUtils.getExpression("urnEntity(" + URN_COL + ")"), - new HashMap() {{ - put(URN_COL, _segment.getDataSource(URN_COL)); - put(TIME_COL, _segment.getDataSource(TIME_COL)); - }}); - assertNotNull(fn); - - String[] result = fn.transformToStringValuesSV(_projectionBlockNoCompanion); - 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) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java index f0f91012e0..8168ff167e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java @@ -37,7 +37,7 @@ *

    For URN columns whose effective stored type is LONG or INT * ({@link DimensionFieldSpec#hasTypedUrnStorage()}), the URN string is parsed and its numeric key * replaces the string value on the row. For multi-entity URN columns - * ({@link DimensionFieldSpec#needsCompanionEntityIdxColumn()}), the 0-based entity index is also + * ({@link DimensionFieldSpec#usesSortedUrnDictionary()}), the 0-based entity index is also * written to the companion {@code __urnEntityIdx} column. The stats collector and * {@code UrnSortedDictionaryCreator} pair the main column LONG with the companion column entity * index to lay out a per-entity sorted-LONG dictionary at segment build time. @@ -65,10 +65,13 @@ public UrnKeyEncodingTransformer(Schema schema) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; if (dim.hasTypedUrnStorage()) { // Single-entity typed LONG/INT and multi-entity LONG/INT (UrnDictionary layout): encode - // the numeric key into the main column. Multi-entity also writes the entity index to a - // companion INT column so the segment builder can pair (entityIdx, longKey). - String companion = dim.needsCompanionEntityIdxColumn() ? dim.companionEntityIdxColumnName() : null; - encodings.add(new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), companion)); + // the numeric key into the main column. For multi-entity, also stash the entity index + // into a transient row field so the segment builder can pair (entityIdx, longKey). The + // transient field is consumed at stats collection / forward-index time and is never + // persisted as a column. + String transientField = dim.usesSortedUrnDictionary() ? dim.transientUrnEntityIdxFieldName() : null; + encodings.add( + new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), transientField)); } } _encodings = Collections.unmodifiableList(encodings); @@ -85,8 +88,8 @@ public GenericRow transform(GenericRow record) { for (ColumnEncoding enc : _encodings) { Object raw = record.getValue(enc._columnName); if (raw == null) { - if (enc._companionColumnName != null) { - record.putValue(enc._companionColumnName, null); + if (enc._entityIdxFieldName != null) { + record.putValue(enc._entityIdxFieldName, null); } continue; } @@ -96,14 +99,14 @@ public GenericRow transform(GenericRow record) { if (urn == null || !urn.isSimple()) { // Malformed or composite URN: null out; NullValueTransformer will fill the default record.putValue(enc._columnName, null); - if (enc._companionColumnName != null) { - record.putValue(enc._companionColumnName, null); + if (enc._entityIdxFieldName != null) { + record.putValue(enc._entityIdxFieldName, null); } continue; } // For multi-entity: find and write entity-type index - if (enc._companionColumnName != null) { + if (enc._entityIdxFieldName != null) { String entityType = urn.entityPrefix(); // e.g. "urn:li:memberId" int entityIdx = -1; for (int i = 0; i < enc._urnTypes.size(); i++) { @@ -115,10 +118,10 @@ public GenericRow transform(GenericRow record) { if (entityIdx < 0) { // Entity not in declared set (validation should have caught this); null out both record.putValue(enc._columnName, null); - record.putValue(enc._companionColumnName, null); + record.putValue(enc._entityIdxFieldName, null); continue; } - record.putValue(enc._companionColumnName, entityIdx); + record.putValue(enc._entityIdxFieldName, entityIdx); } // Encode the main column as the typed numeric key. @@ -132,8 +135,8 @@ public GenericRow transform(GenericRow record) { } } catch (NumberFormatException e) { record.putValue(enc._columnName, null); - if (enc._companionColumnName != null) { - record.putValue(enc._companionColumnName, null); + if (enc._entityIdxFieldName != null) { + record.putValue(enc._entityIdxFieldName, null); } } } @@ -147,14 +150,14 @@ private static final class ColumnEncoding { final FieldSpec.DataType _storedType; final List _urnTypes; @Nullable - final String _companionColumnName; + final String _entityIdxFieldName; ColumnEncoding(String columnName, @Nullable FieldSpec.DataType storedType, - List urnTypes, @Nullable String companionColumnName) { + List urnTypes, @Nullable String entityIdxFieldName) { _columnName = columnName; _storedType = storedType; _urnTypes = urnTypes; - _companionColumnName = companionColumnName; + _entityIdxFieldName = entityIdxFieldName; } } } 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 774b03128c..8b6e7f7242 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 @@ -146,8 +146,6 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { */ private Map, IndexCreator>> _creatorsByColAndIndex = new HashMap<>(); private final Map _nullValueVectorCreatorMap = new HashMap<>(); - /** Synthetic field specs for companion entity-index columns (not in schema). */ - private final Map _companionFieldSpecs = new HashMap<>(); /** * For untyped URN columns in which every observed value shares the same {@code urn:::} * prefix, records the discovered entity type (without trailing colon). The prefix is stripped @@ -249,7 +247,7 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio // LONG buffer plus dictId range metadata. All other columns go through the standard path. SegmentDictionaryCreator creator; if (fieldSpec instanceof DimensionFieldSpec - && ((DimensionFieldSpec) fieldSpec).needsCompanionEntityIdxColumn()) { + && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; creator = new UrnSortedDictionaryCreator(fieldSpec, _indexDir, dim.getUrnTypes()); } else { @@ -307,53 +305,6 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio LOGGER.info("Column: {} is not nullable", columnName); } } - - // Initialize companion entity-index columns for typed multi-entity URN dimensions - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!(fieldSpec instanceof DimensionFieldSpec)) { - continue; - } - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (!dim.needsCompanionEntityIdxColumn()) { - continue; - } - String compCol = dim.companionEntityIdxColumnName(); - ColumnIndexCreationInfo compInfo = indexCreationInfoMap.get(compCol); - if (compInfo == null) { - continue; - } - DimensionFieldSpec compSpec = dim.createCompanionEntityIdxFieldSpec(); - compSpec.setNullable(fieldSpec.isNullable()); - _companionFieldSpecs.put(compCol, compSpec); - - IndexCreationContext.Common compContext = IndexCreationContext.builder() - .withIndexDir(_indexDir) - .withDictionary(true) - .withFieldSpec(compSpec) - .withTotalDocs(segmentIndexCreationInfo.getTotalDocs()) - .withColumnIndexCreationInfo(compInfo) - .onHeap(segmentCreationSpec.isOnHeap()) - .build(); - - DictionaryIndexConfig dictConfig = DictionaryIndexConfig.DEFAULT; - SegmentDictionaryCreator dictCreator = - new DictionaryIndexPlugin().getIndexType().createIndexCreator(compContext, dictConfig); - try { - dictCreator.build(compContext.getSortedUniqueElementsArray()); - } catch (Exception e) { - LOGGER.error("Error building dictionary for companion column: {}", compCol); - throw e; - } - _dictionaryCreatorMap.put(compCol, dictCreator); - - Map, IndexCreator> compCreators = new HashMap<>(); - tryCreateIndexCreator(compCreators, StandardIndexes.forward(), compContext, FieldIndexConfigs.EMPTY); - _creatorsByColAndIndex.put(compCol, compCreators); - - if (isNullable(compSpec)) { - _nullValueVectorCreatorMap.put(compCol, new NullValueVectorCreator(_indexDir, compCol)); - } - } } private boolean isNullable(FieldSpec fieldSpec) { @@ -363,7 +314,7 @@ private boolean isNullable(FieldSpec fieldSpec) { private FieldSpec resolveFieldSpec(String columnName) { FieldSpec fs = _schema.getFieldSpecFor(columnName); if (fs == null) { - return _companionFieldSpecs.get(columnName); + return null; } // For untyped URN columns where a single common prefix was discovered during dictionary // build, surface a synthetic urnTypes so the column behaves like single-entity typed when the @@ -559,13 +510,14 @@ public void indexRow(GenericRow row) FieldSpec fieldSpec = resolveFieldSpec(columnName); SegmentDictionaryCreator dictionaryCreator = _dictionaryCreatorMap.get(columnName); - // For multi-entity URN columns, pair the row's main column LONG with the companion column's - // entity index so the UrnSortedDictionaryCreator can map (entityIdx, longKey) → dict ID. + // 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 UrnSortedDictionaryCreator + // can map (entityIdx, longKey) → dict ID. The transient field is dropped after build. if (dictionaryCreator instanceof UrnSortedDictionaryCreator) { - Object companionVal = row.getValue( - ((DimensionFieldSpec) fieldSpec).companionEntityIdxColumnName()); - if (companionVal != null) { - int entityIdx = ((Number) companionVal).intValue(); + Object entityIdxVal = row.getValue( + ((DimensionFieldSpec) fieldSpec).transientUrnEntityIdxFieldName()); + if (entityIdxVal != null) { + int entityIdx = ((Number) entityIdxVal).intValue(); long longKey = ((Number) columnValueToIndex).longValue(); columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( entityIdx, longKey); @@ -724,15 +676,7 @@ private void writeMetadata() properties.setProperty(SEGMENT_PADDING_CHARACTER, String.valueOf(V1Constants.Str.DEFAULT_STRING_PAD_CHAR)); properties.setProperty(SEGMENT_NAME, _segmentName); properties.setProperty(TABLE_NAME, _config.getTableName()); - // Augment user-declared dimensions with synthetic companion entity-index columns so they - // are loaded as physical dimensions when the segment is opened. - List dimensions = new ArrayList<>(_config.getDimensions()); - for (String compCol : _companionFieldSpecs.keySet()) { - if (!dimensions.contains(compCol)) { - dimensions.add(compCol); - } - } - properties.setProperty(DIMENSIONS, dimensions); + properties.setProperty(DIMENSIONS, _config.getDimensions()); properties.setProperty(METRICS, _config.getMetrics()); properties.setProperty(DATETIME_COLUMNS, _config.getDateTimeColumnNames()); properties.setProperty(COMPLEX_COLUMNS, _config.getComplexColumnNames()); 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 f54cd4595b..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 @@ -70,7 +70,6 @@ import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.spi.config.table.StarTreeIndexConfig; 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.IngestionSchemaValidator; @@ -543,23 +542,6 @@ void collectStatsAndIndexCreationInfo() new ColumnIndexCreationInfo(columnProfile, createDictionary, useVarLengthDictionary, false/*isAutoGenerated*/, defaultNullValue)); } - // Add companion entity-index columns for typed multi-entity URN columns - for (FieldSpec fieldSpec : _dataSchema.getAllFieldSpecs()) { - if (!(fieldSpec instanceof DimensionFieldSpec)) { - continue; - } - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (!dim.needsCompanionEntityIdxColumn()) { - continue; - } - String compCol = dim.companionEntityIdxColumnName(); - ColumnStatistics compProfile = _segmentStats.getColumnProfileFor(compCol); - if (compProfile == null) { - continue; - } - _indexCreationInfoMap.put(compCol, - new ColumnIndexCreationInfo(compProfile, true, false, true /*isAutoGenerated*/, Integer.MIN_VALUE)); - } _segmentIndexCreationInfo.setTotalDocs(_totalDocs); _totalStatsCollectorTimeNs = System.nanoTime() - statsCollectorStartTime; } 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 583873c700..cabde7e598 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 @@ -56,12 +56,13 @@ public void init() { break; case LONG: // Multi-entity URN columns get a stats collector that tracks (entityIdx, longKey) - // pairs so the UrnSortedDictionaryCreator can lay out per-entity dict ID ranges. + // pairs so the UrnSortedDictionaryCreator can lay out per-entity dict ID ranges. The + // entity index comes from the transient field stashed by UrnKeyEncodingTransformer. if (fieldSpec instanceof DimensionFieldSpec - && ((DimensionFieldSpec) fieldSpec).needsCompanionEntityIdxColumn()) { + && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; _columnStatsCollectorMap.put(column, new UrnMultiEntityPreIndexStatsCollector(column, - dim.companionEntityIdxColumnName(), _statsCollectorConfig)); + dim.transientUrnEntityIdxFieldName(), _statsCollectorConfig)); } else { _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig)); } @@ -89,19 +90,6 @@ public void init() { throw new IllegalStateException("Unsupported data type: " + fieldSpec.getDataType()); } } - // Register companion entity-index columns for typed multi-entity URN columns - for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) { - if (!(fieldSpec instanceof DimensionFieldSpec)) { - continue; - } - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.needsCompanionEntityIdxColumn()) { - DimensionFieldSpec compSpec = dim.createCompanionEntityIdxFieldSpec(); - _statsCollectorConfig.registerExtraFieldSpec(compSpec); - _columnStatsCollectorMap.put(compSpec.getName(), - new IntColumnPreIndexStatsCollector(compSpec.getName(), _statsCollectorConfig)); - } - } } @Override 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 index 7f56be0fef..82cffd0bf5 100644 --- 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 @@ -41,36 +41,33 @@ * the entity dimension (effectively LONG-only) — only used by paths that don't have row context. */ public class UrnMultiEntityPreIndexStatsCollector extends LongColumnPreIndexStatsCollector { - private final String _companionColumnName; + private final String _transientEntityIdxFieldName; private final Set _pairs = new HashSet<>(); private EntityKeyPair[] _sortedPairs; private boolean _sealed = false; - public UrnMultiEntityPreIndexStatsCollector(String column, String companionColumnName, + public UrnMultiEntityPreIndexStatsCollector(String column, String transientEntityIdxFieldName, StatsCollectorConfig statsCollectorConfig) { super(column, statsCollectorConfig); - _companionColumnName = companionColumnName; - } - - public String getCompanionColumnName() { - return _companionColumnName; + _transientEntityIdxFieldName = transientEntityIdxFieldName; } /** - * Collects a row's main-column LONG value together with the companion column entity index. - * Used by {@code SegmentPreIndexStatsCollectorImpl.collectRow} for multi-entity URN columns. + * Collects a row's main-column LONG value paired with the entity index stashed by + * {@code UrnKeyEncodingTransformer} in the transient row field. Used by + * {@code SegmentPreIndexStatsCollectorImpl.collectRow} for multi-entity URN columns. */ public void collect(Object value, GenericRow row) { assert !_sealed; if (value == null) { return; } - Object companion = row.getValue(_companionColumnName); - if (companion == null) { + Object entityIdxVal = row.getValue(_transientEntityIdxFieldName); + if (entityIdxVal == null) { // Row was nulled by the encoder (e.g. unknown entity); ignore for pair tracking. return; } - int entityIdx = ((Number) companion).intValue(); + int entityIdx = ((Number) entityIdxVal).intValue(); long key = ((Number) value).longValue(); _pairs.add(new EntityKeyPair(entityIdx, key)); super.collect(value); // keep LongColumn min/max/cardinality stats in sync 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 c0d0fffc1b..07029dbd66 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 @@ -50,7 +50,6 @@ 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.UrnReconstructingLongDictionary; -import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingMultiEntityDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingStringDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnSortedDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -343,7 +342,7 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat return new UrnSortedDictionary(dict, dim.getEntityDictRanges()); } // Single-entity typed URN LONG: wrap to prepend the entity prefix. - if (dim.hasTypedUrnStorage() && !dim.needsCompanionEntityIdxColumn()) { + if (dim.hasTypedUrnStorage() && !dim.usesSortedUrnDictionary()) { String entityType = dim.getUrnTypes().get(0).getEntityType(); return new UrnReconstructingLongDictionary(dict, entityType); } @@ -371,14 +370,6 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat FieldSpec fieldSpec = metadata.getFieldSpec(); if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.needsCompanionEntityIdxColumn()) { - List urnTypes = dim.getUrnTypes(); - String[] entityTypes = new String[urnTypes.size()]; - for (int i = 0; i < urnTypes.size(); i++) { - entityTypes[i] = urnTypes.get(i).getEntityType(); - } - return new UrnReconstructingMultiEntityDictionary(dict, entityTypes); - } if (dim.isSingleUrnEntityType()) { // Single-entity URN with STRING storage: the dictionary stores suffixes only. // Covers both schema-declared single-entity URN columns and untyped URN columns where diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java deleted file mode 100644 index c08e61639a..0000000000 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingMultiEntityDictionary.java +++ /dev/null @@ -1,201 +0,0 @@ -/** - * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; - - -/** - * Wraps a STRING dictionary for a multi-entity typed URN column whose entries are physically - * stored as {@code ":"} after the {@link - * org.apache.pinot.segment.local.recordtransformer.UrnKeyEncodingTransformer} stripped the URN - * prefix. Reconstructs full URN strings (e.g. {@code "urn:li:corpUser:docchial"}) on read, and - * encodes incoming full-URN predicate literals back into the stored form before delegating to the - * underlying STRING dictionary. - * - *

    Dictionary entries that do not match the {@code ":..."} encoding (for example the - * default null value {@code "urn:pinot:null:0"}) are passed through unchanged in both directions. - */ -public class UrnReconstructingMultiEntityDictionary implements Dictionary { - private final Dictionary _delegate; - private final String[] _entityTypes; - - /** - * @param delegate underlying STRING dictionary of encoded {@code ":"} - * values - * @param entityTypes the declared entity types in order; the index used in the encoded form - * refers into this array - */ - public UrnReconstructingMultiEntityDictionary(Dictionary delegate, String[] entityTypes) { - _delegate = delegate; - _entityTypes = entityTypes; - } - - @Override - public boolean isSorted() { - return false; - } - - @Override - public DataType getValueType() { - return DataType.STRING; - } - - @Override - public int length() { - return _delegate.length(); - } - - /** Resolves a full URN string to its dict ID by encoding it to the stored form first. */ - @Override - public int indexOf(String urnString) { - String encoded = encode(urnString); - return _delegate.indexOf(encoded != null ? encoded : urnString); - } - - /** Range lookup is best-effort on multi-entity columns; delegates after attempted encoding. */ - @Override - public int insertionIndexOf(String urnString) { - String encoded = encode(urnString); - return _delegate.insertionIndexOf(encoded != null ? encoded : urnString); - } - - @Override - public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { - String encodedLower = encode(lower); - String encodedUpper = encode(upper); - return _delegate.getDictIdsInRange( - encodedLower != null ? encodedLower : lower, - encodedUpper != null ? encodedUpper : upper, - includeLower, includeUpper); - } - - /** - * Encode a full URN string ({@code "urn:li:corpUser:foo"}) to the stored {@code "0:foo"} form - * by finding which declared entity type the URN starts with. Returns {@code null} when no - * declared entity matches (e.g. a predicate literal for an entity the segment does not store). - */ - private String encode(String urnString) { - if (urnString == null) { - return null; - } - for (int i = 0; i < _entityTypes.length; i++) { - String prefix = _entityTypes[i] + ":"; - if (urnString.startsWith(prefix)) { - return i + ":" + urnString.substring(prefix.length()); - } - } - return null; - } - - @Override - public int compare(int dictId1, int dictId2) { - return _delegate.compare(dictId1, dictId2); - } - - @Override - public Comparable getMinVal() { - return _delegate.getMinVal(); - } - - @Override - public Comparable getMaxVal() { - return _delegate.getMaxVal(); - } - - @Override - public Object getSortedValues() { - return _delegate.getSortedValues(); - } - - @Override - public Object get(int dictId) { - return getStringValue(dictId); - } - - @Override - public int getIntValue(int dictId) { - return _delegate.getIntValue(dictId); - } - - @Override - public long getLongValue(int dictId) { - return _delegate.getLongValue(dictId); - } - - @Override - public float getFloatValue(int dictId) { - return _delegate.getFloatValue(dictId); - } - - @Override - public double getDoubleValue(int dictId) { - return _delegate.getDoubleValue(dictId); - } - - @Override - public BigDecimal getBigDecimalValue(int dictId) { - return _delegate.getBigDecimalValue(dictId); - } - - /** Returns the full URN string reconstructed from the stored {@code ":"}. */ - @Override - public String getStringValue(int dictId) { - String raw = _delegate.getStringValue(dictId); - return decode(raw); - } - - @Override - public void readStringValues(int[] dictIds, int length, String[] outValues) { - _delegate.readStringValues(dictIds, length, outValues); - for (int i = 0; i < length; i++) { - outValues[i] = decode(outValues[i]); - } - } - - private String decode(String raw) { - if (raw == null) { - return null; - } - int colon = raw.indexOf(':'); - if (colon <= 0) { - return raw; - } - int idx; - try { - idx = Integer.parseInt(raw, 0, colon, 10); - } catch (NumberFormatException e) { - return raw; - } - if (idx < 0 || idx >= _entityTypes.length) { - return raw; - } - return _entityTypes[idx] + ":" + raw.substring(colon + 1); - } - - @Override - public void close() - throws IOException { - _delegate.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 index b11e457fef..756905f907 100644 --- 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 @@ -119,24 +119,24 @@ public void singleEntityNonNumericKeyNulledOut() { // --- multi-entity typed LONG column (companion column written) --- @Test - public void multiEntityFirstEntityWritesCompanionIdx0() { + 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 companion column carries the entity index for the segment-build pipeline to pair - // them into (entityIdx, longKey) dictionary entries. + // 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("resourceId__urnEntityIdx"), 0); + assertEquals(row.getValue("__urnEntityIdxBuild_resourceId"), 0); } @Test - public void multiEntitySecondEntityWritesCompanionIdx1() { + 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("resourceId__urnEntityIdx"), 1); + assertEquals(row.getValue("__urnEntityIdxBuild_resourceId"), 1); } @Test @@ -145,15 +145,15 @@ public void multiEntityUnknownEntityNullsBothColumns() { GenericRow row = rowWithValue("resourceId", "urn:li:orgId:99"); t.transform(row); assertNull(row.getValue("resourceId"), "Unknown entity should be nulled"); - assertNull(row.getValue("resourceId__urnEntityIdx"), "Companion column should be null too"); + assertNull(row.getValue("__urnEntityIdxBuild_resourceId"), "Transient entity idx should be null too"); } @Test - public void multiEntityNullUrnLeavesCompanionNull() { + public void multiEntityNullUrnLeavesTransientNull() { UrnKeyEncodingTransformer t = new UrnKeyEncodingTransformer(multiEntityLongSchema()); GenericRow row = rowWithValue("resourceId", null); t.transform(row); assertNull(row.getValue("resourceId")); - assertNull(row.getValue("resourceId__urnEntityIdx")); + assertNull(row.getValue("__urnEntityIdxBuild_resourceId")); } } 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 801e070ed1..859390f613 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 @@ -33,12 +33,15 @@ public final class DimensionFieldSpec extends FieldSpec { /** - * Suffix appended to a typed multi-entity URN column name to form the companion entity-index - * column. For a column {@code "actorId"} with multiple URN types, the companion column is - * {@code "actorId__urnEntityIdx"} and stores a 0-based index into the column's {@code urnTypes} - * list, allowing efficient per-row entity-type lookup without string parsing. + * 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 COMPANION_ENTITY_IDX_SUFFIX = "__urnEntityIdx"; + public static final String URN_TRANSIENT_ENTITY_IDX_PREFIX = "__urnEntityIdxBuild_"; /** * For columns with {@code dataType = URN}: the set of (entityType, valueType) pairs that this @@ -309,12 +312,13 @@ public boolean hasTypedUrnStorage() { /** * Returns {@code true} when the column should have a companion entity-index INT column written * alongside the main column. This is true for multi-entity URN columns whose entity types all - * declare a numeric ({@code LONG} or {@code INT}) {@code valueType}; the companion column stores - * the 0-based entity-type index per row, enabling fast {@code urnEntity()} resolution without - * parsing the full URN string. + * declare a numeric ({@code LONG} or {@code INT}) {@code valueType}. The segment builder uses + * a sorted LONG dictionary partitioned by entity-type ranges (see + * {@link #URN_TRANSIENT_ENTITY_IDX_PREFIX}). The companion field is built only as a transient + * row attribute during ingestion and is never persisted as a column. */ @JsonIgnore - public boolean needsCompanionEntityIdxColumn() { + public boolean usesSortedUrnDictionary() { if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.size() <= 1) { return false; } @@ -328,21 +332,13 @@ public boolean needsCompanionEntityIdxColumn() { } /** - * Returns the name of the companion entity-index column for this column, i.e. - * {@code name + COMPANION_ENTITY_IDX_SUFFIX}. + * 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 companionEntityIdxColumnName() { - return _name + COMPANION_ENTITY_IDX_SUFFIX; - } - - /** - * Creates and returns the {@link DimensionFieldSpec} for the companion entity-index column. - * The companion column has type {@code INT} and is a single-value dimension. - */ - @JsonIgnore - public DimensionFieldSpec createCompanionEntityIdxFieldSpec() { - return new DimensionFieldSpec(companionEntityIdxColumnName(), DataType.INT, true); + public String transientUrnEntityIdxFieldName() { + return URN_TRANSIENT_ENTITY_IDX_PREFIX + _name; } @JsonIgnore From 44725019b6c086062e2a82081e4d91b5429de3d5 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 28 May 2026 23:59:42 +0000 Subject: [PATCH 10/24] UrnDictionary commit 5/5: tests for adding entity types over time Backwards-compatibility tests confirming that segments built with different urnTypes lists (extended set, reordered set) both load and query correctly. Each segment uses the urnTypes ordering it was built with -- existing segments do not need to be rewritten when a column's urnTypes list is extended in the table schema. Tests: * multiEntitySegmentsWithDifferentUrnTypesBothQueryable: a v1 segment with [memberId, groupId] and a v2 segment with [memberId, groupId, applicationName] both load; v1 metadata keeps its 2-entry list and queries for the new applicationName entity miss (zero rows have it, semantically correct), v2 queries for all three. * multiEntitySegmentsWithReorderedUrnTypesReconstructCorrectly: a v1 segment with [memberId, groupId] and a v2 segment with [groupId, memberId] both reconstruct the same URN strings on read -- per- entity ranges in segment metadata carry the entity type, not just its ordinal. 146 URN tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../segment/index/UrnColumnSegmentTest.java | 119 ++++++++++++++++++ 1 file changed, 119 insertions(+) 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 index dbacee18ed..8d1840d45d 100644 --- 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 @@ -258,4 +258,123 @@ public void singleEntityTypedUrnDictionaryIndexOfReturnsNullForWrongEntity() 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(); + } + } } From 83e70a4df5e15a1fb23ec58e617a258c0c937bc7 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 00:12:45 +0000 Subject: [PATCH 11/24] BenchmarkUrnTransform: report dictionary file size separately Adds a v3-layout index_map parser so the benchmark prints the dictionary file size for each segment alongside the total. Total segment size is dominated by the forward index (~70% of bytes), so the dict-only number is what actually shows the URN feature's storage win. Numbers from the run (100K rows, 10K unique keys): multi-entity untyped (STRING dict): 269,981 bytes multi-entity typed (sorted LONG): 80,000 bytes (3.37x smaller) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../pinot/perf/BenchmarkUrnTransform.java | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) 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 index de760e30e0..a3251e0fed 100644 --- a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUrnTransform.java @@ -226,14 +226,53 @@ private void printStorageSizes() { 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; } From 89f9acb7e5c41186c745f95e0e0507189ae72b78 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 00:51:59 +0000 Subject: [PATCH 12/24] URN: refactor sorted-LONG dictionary into UrnPolymorphicDictionary Replaces UrnSortedDictionary with UrnPolymorphicDictionary: a list of per-entity sub-dictionaries with cumulative dict-ID offsets. Each sub-dict is a standard Pinot Dictionary (LongDictionary, StringDictionary, IntDictionary, ...). The wrapper routes every operation -- indexOf, getStringValue, getLongValue, getDictIdsInRange -- to the appropriate sub-dictionary and re-attaches the entity prefix on read. This is the architecturally-cleaner restatement of the same design: the current all-numeric-keys case still produces one composite LONG buffer sliced into per-entity LongDictionary views, with the existing urn.entityDictRanges metadata recording the slice boundaries. Wire format is unchanged; this commit is a reader-side refactor. The next commits extend this to genuinely polymorphic sub-dictionaries (mixed LONG / STRING / INT per entity), wire up the build path to choose the natural delegate type per entity, add multi-prefix auto-detect for untyped URN columns, and introduce a fail/skip/pause ingest mode for non-URN-shaped data so PARSE_STRING goes away entirely. * New: UrnPolymorphicDictionary + 19 unit tests (mixed LONG + STRING sub-dicts exercise the polymorphic routing). * Removed: UrnSortedDictionary and its 20 tests (superseded). * DictionaryIndexType.read constructs per-entity LongDictionary views by slicing the main LONG buffer at the entity-range offsets, then composes them into a UrnPolymorphicDictionary. 129 URN tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../index/dictionary/DictionaryIndexType.java | 36 +- .../readers/UrnPolymorphicDictionary.java | 296 ++++++++++++ .../index/readers/UrnSortedDictionary.java | 313 ------------- .../readers/UrnPolymorphicDictionaryTest.java | 440 ++++++++++++++++++ .../readers/UrnSortedDictionaryTest.java | 323 ------------- 5 files changed, 764 insertions(+), 644 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionary.java delete mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionaryTest.java delete mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java 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 07029dbd66..69c1ca4c79 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,9 +49,9 @@ 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.local.segment.index.readers.UrnReconstructingLongDictionary; import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingStringDictionary; -import org.apache.pinot.segment.local.segment.index.readers.UrnSortedDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -331,23 +331,43 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat return loadOnHeap ? new OnHeapIntDictionary(dataBuffer, length) : new IntDictionary(dataBuffer, length); case LONG: { - Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) - : new LongDictionary(dataBuffer, length); FieldSpec fieldSpec = metadata.getFieldSpec(); if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - // Multi-entity URN with per-entity dict ID ranges (UrnDictionary layout): wrap with a - // range-aware reader that reconstructs URN strings using the appropriate entity prefix. - if (dim.getEntityDictRanges() != null) { - return new UrnSortedDictionary(dict, dim.getEntityDictRanges()); + // Multi-entity URN: slice the LONG buffer per entity range and wrap each slice in a + // standard LongDictionary, then compose them into a UrnPolymorphicDictionary. The + // wrapper routes by entity prefix so each sub-dict gets to use its native lookup. + List ranges = dim.getEntityDictRanges(); + if (ranges != null) { + List subDicts = new ArrayList<>(ranges.size()); + int prevEnd = 0; + for (DimensionFieldSpec.EntityDictRange r : ranges) { + int subLen = r.getEndDictIdExclusive() - prevEnd; + if (subLen > 0) { + PinotDataBuffer subBuffer = dataBuffer.view( + (long) prevEnd * Long.BYTES, (long) r.getEndDictIdExclusive() * Long.BYTES); + Dictionary subDict = loadOnHeap + ? new OnHeapLongDictionary(subBuffer, subLen) + : new LongDictionary(subBuffer, subLen); + subDicts.add(new UrnPolymorphicDictionary.SubDict(r.getEntityType(), subDict)); + } + prevEnd = r.getEndDictIdExclusive(); + } + if (!subDicts.isEmpty()) { + return new UrnPolymorphicDictionary(subDicts); + } } + Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) + : new LongDictionary(dataBuffer, length); // Single-entity typed URN LONG: wrap to prepend the entity prefix. if (dim.hasTypedUrnStorage() && !dim.usesSortedUrnDictionary()) { String entityType = dim.getUrnTypes().get(0).getEntityType(); return new UrnReconstructingLongDictionary(dict, entityType); } + return dict; } - return dict; + return loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) + : new LongDictionary(dataBuffer, length); } case FLOAT: return loadOnHeap ? new OnHeapFloatDictionary(dataBuffer, length) 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..6593222bec --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionary.java @@ -0,0 +1,296 @@ +/** + * 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() { + // Sub-dicts are independently sorted; globally the buffer is NOT lex-sorted as URN strings. + return false; + } + + @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 = _subDicts[entityIdx].indexOf(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 = _subDicts[entityIdx].insertionIndexOf(suffix); + 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); + } + String lowerSuffix = lower.substring(_entityPrefixes[lowerEntity].length()); + String upperSuffix = upper.substring(_entityPrefixes[upperEntity].length()); + IntSet subResult = + _subDicts[lowerEntity].getDictIdsInRange(lowerSuffix, upperSuffix, includeLower, includeUpper); + IntSet result = new IntOpenHashSet(subResult.size()); + int offset = _offsets[lowerEntity]; + subResult.forEach((java.util.function.IntConsumer) id -> result.add(offset + id)); + return result; + } + + /** 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/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java deleted file mode 100644 index 9923c717f9..0000000000 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionary.java +++ /dev/null @@ -1,313 +0,0 @@ -/** - * 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.DimensionFieldSpec.EntityDictRange; -import org.apache.pinot.spi.data.FieldSpec.DataType; - - -/** - * Dictionary wrapper for multi-entity URN columns using the sorted-LONG layout. - * - *

    The underlying delegate is a buffer of LONG keys, partitioned by entity type. Each entity - * type owns a contiguous, key-sorted range of dictionary IDs identified by {@link - * EntityDictRange#getEndDictIdExclusive()}. Within a range, entries are sorted ascending by LONG - * value; across ranges, the same LONG value may appear multiple times (once per entity that has - * that key). - * - *

    This wrapper exposes the column as URN strings: - *

      - *
    • {@link #indexOf(String) indexOf("urn:li:corpUser:100")} parses the entity prefix, finds - * the corresponding entity range, then runs a range-restricted binary search for the LONG - * key. Returns the absolute dict ID, or {@link #NULL_VALUE_INDEX} if the entity prefix is - * unknown or the key isn't present.
      • - *
    • {@link #getStringValue(int) getStringValue(dictId)} looks up which entity range owns - * {@code dictId}, reads the LONG key from the delegate, and returns - * {@code ":"}.
      • - *
    • {@link #getLongValue(int) getLongValue(dictId)} passes through to the delegate — the raw - * LONG key is exposed for {@code urnKeyLong()} and LONG-typed predicates.
      • - *
    - * - *

    Cross-entity range filters are explicitly unsupported: {@code WHERE col BETWEEN - * 'urn:li:a:100' AND 'urn:li:b:200'} returns the empty set rather than throwing, to be robust to - * generated queries. Same-entity range filters are supported via {@link #getDictIdsInRange}. - */ -public class UrnSortedDictionary implements Dictionary { - - private final Dictionary _delegate; - private final EntityDictRange[] _ranges; - private final String[] _entityTypes; - private final String[] _entityPrefixes; - private final int[] _rangeStarts; - - public UrnSortedDictionary(Dictionary delegate, List ranges) { - if (delegate.getValueType() != DataType.LONG && delegate.getValueType() != DataType.INT) { - throw new IllegalArgumentException( - "UrnSortedDictionary delegate must be a LONG or INT dictionary, got: " + delegate.getValueType()); - } - _delegate = delegate; - _ranges = ranges.toArray(new EntityDictRange[0]); - _entityTypes = new String[_ranges.length]; - _entityPrefixes = new String[_ranges.length]; - _rangeStarts = new int[_ranges.length]; - int prev = 0; - for (int i = 0; i < _ranges.length; i++) { - _entityTypes[i] = _ranges[i].getEntityType(); - _entityPrefixes[i] = _entityTypes[i] + ":"; - _rangeStarts[i] = prev; - prev = _ranges[i].getEndDictIdExclusive(); - } - } - - @Override - public boolean isSorted() { - // The buffer is sorted only within each entity range, not globally. - return false; - } - - @Override - public DataType getValueType() { - return DataType.STRING; - } - - @Override - public int length() { - return _delegate.length(); - } - - // ---------------- lookup ---------------- - - @Override - public int indexOf(String urnString) { - int entityIdx = findEntityIdx(urnString); - if (entityIdx < 0) { - return NULL_VALUE_INDEX; - } - long key; - try { - key = Long.parseLong(urnString, _entityPrefixes[entityIdx].length(), urnString.length(), 10); - } catch (NumberFormatException e) { - return NULL_VALUE_INDEX; - } - return rangeRestrictedIndexOf(entityIdx, key); - } - - @Override - public int insertionIndexOf(String urnString) { - int entityIdx = findEntityIdx(urnString); - if (entityIdx < 0) { - return ~_delegate.length(); - } - long key; - try { - key = Long.parseLong(urnString, _entityPrefixes[entityIdx].length(), urnString.length(), 10); - } catch (NumberFormatException e) { - return ~_delegate.length(); - } - int found = rangeBinarySearch(entityIdx, key); - return found >= 0 ? found : found; - } - - @Override - public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { - int lowerEntity = findEntityIdx(lower); - int upperEntity = findEntityIdx(upper); - // Cross-entity ranges are undefined; return empty per the design doc. - if (lowerEntity < 0 || upperEntity < 0 || lowerEntity != upperEntity) { - return new IntOpenHashSet(0); - } - long lowerKey; - long upperKey; - try { - lowerKey = Long.parseLong(lower, _entityPrefixes[lowerEntity].length(), lower.length(), 10); - upperKey = Long.parseLong(upper, _entityPrefixes[upperEntity].length(), upper.length(), 10); - } catch (NumberFormatException e) { - return new IntOpenHashSet(0); - } - int rangeStart = _rangeStarts[lowerEntity]; - int rangeEnd = _ranges[lowerEntity].getEndDictIdExclusive(); - int firstDictId = -1; - int lastDictIdInclusive = -1; - // Linear pass; range is typically much smaller than the full dict. - for (int dictId = rangeStart; dictId < rangeEnd; dictId++) { - long v = _delegate.getLongValue(dictId); - boolean ge = includeLower ? v >= lowerKey : v > lowerKey; - boolean le = includeUpper ? v <= upperKey : v < upperKey; - if (ge && le) { - if (firstDictId < 0) { - firstDictId = dictId; - } - lastDictIdInclusive = dictId; - } - } - if (firstDictId < 0) { - return new IntOpenHashSet(0); - } - IntSet result = new IntOpenHashSet(lastDictIdInclusive - firstDictId + 1); - for (int i = firstDictId; i <= lastDictIdInclusive; i++) { - result.add(i); - } - return result; - } - - // ---------------- per-dictId accessors ---------------- - - @Override - public Object get(int dictId) { - return getStringValue(dictId); - } - - @Override - public int getIntValue(int dictId) { - return _delegate.getIntValue(dictId); - } - - @Override - public long getLongValue(int dictId) { - return _delegate.getLongValue(dictId); - } - - @Override - public float getFloatValue(int dictId) { - return _delegate.getFloatValue(dictId); - } - - @Override - public double getDoubleValue(int dictId) { - return _delegate.getDoubleValue(dictId); - } - - @Override - public BigDecimal getBigDecimalValue(int dictId) { - return _delegate.getBigDecimalValue(dictId); - } - - @Override - public String getStringValue(int dictId) { - int entityIdx = entityIdxForDictId(dictId); - if (entityIdx < 0) { - return null; - } - return _entityPrefixes[entityIdx] + _delegate.getLongValue(dictId); - } - - @Override - public void readStringValues(int[] dictIds, int length, String[] outValues) { - for (int i = 0; i < length; i++) { - outValues[i] = getStringValue(dictIds[i]); - } - } - - // ---------------- ordering / metadata ---------------- - - @Override - public int compare(int dictId1, int dictId2) { - return _delegate.compare(dictId1, dictId2); - } - - @Override - public Comparable getMinVal() { - // Composite ordering (entityType, then key) — return as string for cross-type stability. - return _delegate.length() == 0 ? null : getStringValue(0); - } - - @Override - public Comparable getMaxVal() { - int len = _delegate.length(); - return len == 0 ? null : getStringValue(len - 1); - } - - @Override - public Object getSortedValues() { - throw new UnsupportedOperationException(); - } - - @Override - public void close() - throws IOException { - _delegate.close(); - } - - // ---------------- internal helpers ---------------- - - /** Finds the entity index whose prefix is a prefix of {@code urnString}, or -1. */ - private int findEntityIdx(@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; - } - - /** Returns the entity index that owns {@code dictId}, or -1 if out of bounds. */ - int entityIdxForDictId(int dictId) { - if (dictId < 0 || dictId >= _delegate.length()) { - return -1; - } - // Linear scan — entity counts are small (typically <= 16). Switch to binary search if needed. - for (int i = 0; i < _ranges.length; i++) { - if (dictId < _ranges[i].getEndDictIdExclusive()) { - return i; - } - } - return -1; - } - - /** - * Range-restricted lookup for {@code longKey} within entity {@code entityIdx}'s dict-ID range. - * Returns the absolute dict ID or {@link #NULL_VALUE_INDEX}. - */ - private int rangeRestrictedIndexOf(int entityIdx, long longKey) { - int result = rangeBinarySearch(entityIdx, longKey); - return result >= 0 ? result : NULL_VALUE_INDEX; - } - - /** - * Returns the absolute dict ID of {@code longKey} within entity {@code entityIdx}'s range, - * or {@code -(insertionPoint+1)} when not found (standard {@code Arrays.binarySearch} contract, - * with the insertion point referring to an absolute dict ID). - */ - private int rangeBinarySearch(int entityIdx, long longKey) { - int low = _rangeStarts[entityIdx]; - int high = _ranges[entityIdx].getEndDictIdExclusive() - 1; - while (low <= high) { - int mid = (low + high) >>> 1; - long midValue = _delegate.getLongValue(mid); - if (midValue < longKey) { - low = mid + 1; - } else if (midValue > longKey) { - high = mid - 1; - } else { - return mid; - } - } - return -(low + 1); - } -} 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..4dda731ff0 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnPolymorphicDictionaryTest.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.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"); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java deleted file mode 100644 index 74907afc59..0000000000 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/UrnSortedDictionaryTest.java +++ /dev/null @@ -1,323 +0,0 @@ -/** - * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.DimensionFieldSpec.EntityDictRange; -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 UrnSortedDictionary} using an in-memory fake LONG delegate. Verifies - * range-restricted indexOf, range-bound getStringValue reconstruction, cross-entity range - * handling, and pass-through accessors. - */ -public class UrnSortedDictionaryTest { - - /** - * In-memory stand-in for a LONG dictionary. Exposes a {@code long[]} buffer directly so tests - * can construct exotic layouts without spinning up the full segment-build machinery. - */ - private static final class FakeLongDictionary implements Dictionary { - private final long[] _values; - - FakeLongDictionary(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) { - throw new UnsupportedOperationException(); - } - - @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 { - } - } - - private static UrnSortedDictionary build(long[] keys, EntityDictRange... ranges) { - return new UrnSortedDictionary(new FakeLongDictionary(keys), Arrays.asList(ranges)); - } - - // Reference layout used by most tests: - // corpUser: dict IDs [0, 3), keys 100, 200, 300 - // groupId: dict IDs [3, 5), keys 100, 400 (note 100 collides with corpUser) - // appName: dict IDs [5, 6), key 50 - private static UrnSortedDictionary referenceDict() { - return build(new long[]{100, 200, 300, 100, 400, 50}, - new EntityDictRange("urn:li:corpUser", 3), - new EntityDictRange("urn:li:groupId", 5), - new EntityDictRange("urn:li:applicationName", 6)); - } - - // --- indexOf --- - - @Test - public void indexOfReturnsDictIdInCorpUserRange() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.indexOf("urn:li:corpUser:100"), 0); - assertEquals(d.indexOf("urn:li:corpUser:200"), 1); - assertEquals(d.indexOf("urn:li:corpUser:300"), 2); - } - - @Test - public void indexOfDisambiguatesCollidingKeyAcrossEntities() { - UrnSortedDictionary d = referenceDict(); - // Same key 100, different entities -> different dict IDs. - assertEquals(d.indexOf("urn:li:corpUser:100"), 0); - assertEquals(d.indexOf("urn:li:groupId:100"), 3); - } - - @Test - public void indexOfReturnsNullForKeyMissingFromEntity() { - UrnSortedDictionary d = referenceDict(); - // 999 isn't in any entity's range. - assertEquals(d.indexOf("urn:li:corpUser:999"), Dictionary.NULL_VALUE_INDEX); - // 50 exists in appName but not in groupId. - assertEquals(d.indexOf("urn:li:groupId:50"), Dictionary.NULL_VALUE_INDEX); - } - - @Test - public void indexOfReturnsNullForUnknownEntity() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.indexOf("urn:li:notAnEntity:100"), Dictionary.NULL_VALUE_INDEX); - } - - @Test - public void indexOfReturnsNullForMalformedKey() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.indexOf("urn:li:corpUser:notANumber"), Dictionary.NULL_VALUE_INDEX); - } - - // --- getStringValue --- - - @Test - public void getStringValueReconstructsAcrossAllRanges() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.getStringValue(0), "urn:li:corpUser:100"); - assertEquals(d.getStringValue(1), "urn:li:corpUser:200"); - assertEquals(d.getStringValue(2), "urn:li:corpUser:300"); - assertEquals(d.getStringValue(3), "urn:li:groupId:100"); - assertEquals(d.getStringValue(4), "urn:li:groupId:400"); - assertEquals(d.getStringValue(5), "urn:li:applicationName:50"); - } - - @Test - public void getStringValueReturnsNullForOutOfBounds() { - UrnSortedDictionary d = referenceDict(); - assertNull(d.getStringValue(-1)); - assertNull(d.getStringValue(6)); - } - - @Test - public void readStringValuesBatch() { - UrnSortedDictionary d = referenceDict(); - String[] out = new String[3]; - d.readStringValues(new int[]{5, 3, 0}, 3, out); - assertEquals(out[0], "urn:li:applicationName:50"); - assertEquals(out[1], "urn:li:groupId:100"); - assertEquals(out[2], "urn:li:corpUser:100"); - } - - // --- pass-through accessors --- - - @Test - public void getLongValuePassesThroughToDelegate() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.getLongValue(0), 100L); - assertEquals(d.getLongValue(3), 100L); - assertEquals(d.getLongValue(5), 50L); - } - - @Test - public void getIntValuePassesThroughForSmallKeys() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.getIntValue(0), 100); - assertEquals(d.getIntValue(5), 50); - } - - @Test - public void getReturnsReconstructedUrnString() { - UrnSortedDictionary d = referenceDict(); - assertEquals(d.get(3), "urn:li:groupId:100"); - } - - // --- metadata --- - - @Test - public void getValueTypeIsString() { - assertEquals(referenceDict().getValueType(), DataType.STRING); - } - - @Test - public void isSortedReturnsFalse() { - // The buffer is sorted only within each entity range; not globally lex-sorted as URNs. - assertTrue(!referenceDict().isSorted()); - } - - @Test - public void lengthMatchesDelegate() { - assertEquals(referenceDict().length(), 6); - } - - // --- getDictIdsInRange --- - - @Test - public void getDictIdsInRangeSameEntityInclusive() { - UrnSortedDictionary d = referenceDict(); - IntSet result = - d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:corpUser:200", true, true); - assertEquals(result.size(), 2); - assertTrue(result.contains(0)); - assertTrue(result.contains(1)); - } - - @Test - public void getDictIdsInRangeSameEntityExclusiveBounds() { - UrnSortedDictionary d = referenceDict(); - IntSet result = - d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:corpUser:300", false, false); - assertEquals(result.size(), 1); - assertTrue(result.contains(1)); - } - - @Test - public void getDictIdsInRangeCrossEntityReturnsEmpty() { - UrnSortedDictionary d = referenceDict(); - IntSet result = - d.getDictIdsInRange("urn:li:corpUser:100", "urn:li:groupId:200", true, true); - assertEquals(result.size(), 0); - } - - @Test - public void getDictIdsInRangeUnknownEntityReturnsEmpty() { - UrnSortedDictionary d = referenceDict(); - IntSet result = - d.getDictIdsInRange("urn:li:notAnEntity:100", "urn:li:notAnEntity:200", true, true); - assertEquals(result.size(), 0); - } - - // --- min/max --- - - @Test - public void getMinValIsFirstReconstructedUrn() { - assertEquals(referenceDict().getMinVal(), "urn:li:corpUser:100"); - } - - @Test - public void getMaxValIsLastReconstructedUrn() { - assertEquals(referenceDict().getMaxVal(), "urn:li:applicationName:50"); - } -} From bb3dd558c1ed440ec55f6230203f7232f30fa600 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 02:56:56 +0000 Subject: [PATCH 13/24] URN: multi-prefix auto-detect for untyped columns Extends segment-build URN auto-detection from single-prefix-only to multi-prefix. If every value in an untyped URN column parses as a simple URN with a numeric key and at least two distinct entity prefixes appear in the data, the column is promoted to multi-entity URN for the rest of the build: urnTypes are synthesized in observed order, the sorted-LONG UrnDictionary layout is used, and the segment loads as if the user had declared the column with these urnTypes. Build-path changes: * maybeStripUrnPrefix tries single-prefix first; on failure, runs tryMultiPrefixUrnDetection. When successful, returns EntityKeyPair[] so UrnSortedDictionaryCreator handles the rest. * SegmentColumnarIndexCreator picks the dict creator after preprocessing so EntityKeyPair[] from auto-detect routes through UrnSortedDictionaryCreator just like declared multi-entity. * indexRow handles auto-detected multi-entity columns by parsing the URN string inline (no transformer ran on the untyped column). * resolveFieldSpec caches synthetic FieldSpecs so per-build mutations (setEntityDictRanges) persist across resolve calls. * Skip min/max metadata for multi-entity URN columns: cross-entity comparisons aren't semantically meaningful and the values won't round-trip through LONG parsing at load. Test updates: * UrnColumnSegmentTest.untypedUrnHasNoUrnTypesInMetadata renamed to untypedUrnAutoDetectsMultiEntityFromIngestedData; asserts that multi-prefix all-numeric data triggers auto-detection. * New untypedUrnAutoDetectedSegmentIsQueryable verifies the segment round-trips: all four ingested URNs are indexable, getStringValue reconstructs correctly, and unknown lookups return NULL_VALUE_INDEX. 148 URN tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../impl/SegmentColumnarIndexCreator.java | 215 ++++++++++++++---- .../segment/index/UrnColumnSegmentTest.java | 45 +++- 2 files changed, 217 insertions(+), 43 deletions(-) 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 8b6e7f7242..0ed22bb22a 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 @@ -153,6 +153,16 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { * column as single-entity typed. */ private final Map _discoveredUrnEntities = new HashMap<>(); + /** + * For untyped URN columns whose unique values span multiple {@code urn:::} prefixes + * with numeric keys, records the discovered entity types in observed order. The column is + * promoted to multi-entity typed for the rest of the build (sorted-LONG UrnDictionary layout) + * and load-time {@code urnTypes} surface the discovered list. + */ + 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; @@ -243,18 +253,24 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio dictConfig = new DictionaryIndexConfig(dictConfig, columnIndexCreationInfo.isUseVarLengthDictionary()); } - // Multi-entity URN columns get a custom dict creator that lays out a per-entity sorted - // LONG buffer plus dictId range metadata. All other columns go through the standard path. + // 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; - if (fieldSpec instanceof DimensionFieldSpec - && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - creator = new UrnSortedDictionaryCreator(fieldSpec, _indexDir, dim.getUrnTypes()); + boolean usesSortedUrnDict = dictValues instanceof org.apache.pinot.segment.local.segment.creator.impl.urn + .EntityKeyPair[] + || (fieldSpec instanceof DimensionFieldSpec + && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()); + if (usesSortedUrnDict) { + List urnTypesForCreator = + _discoveredMultiEntityUrnTypes.containsKey(columnName) + ? _discoveredMultiEntityUrnTypes.get(columnName) + : ((DimensionFieldSpec) fieldSpec).getUrnTypes(); + creator = new UrnSortedDictionaryCreator(fieldSpec, _indexDir, urnTypesForCreator); } else { creator = new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); } - Object dictValues = maybeStripUrnPrefix(columnName, fieldSpec, context.getSortedUniqueElementsArray()); try { creator.build(dictValues); } catch (Exception e) { @@ -264,9 +280,10 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio } // For multi-entity URN columns, surface the computed per-entity dict ID ranges so the - // metadata write can persist them under urn.entityDictRanges (read in commit 1). - if (creator instanceof UrnSortedDictionaryCreator && fieldSpec instanceof DimensionFieldSpec) { - ((DimensionFieldSpec) fieldSpec).setEntityDictRanges(((UrnSortedDictionaryCreator) creator).getRanges()); + // metadata write can persist them under urn.entityDictRanges. + if (creator instanceof UrnSortedDictionaryCreator) { + DimensionFieldSpec dim = (DimensionFieldSpec) resolveFieldSpec(columnName); + dim.setEntityDictRanges(((UrnSortedDictionaryCreator) creator).getRanges()); } _dictionaryCreatorMap.put(columnName, creator); @@ -316,34 +333,66 @@ private FieldSpec resolveFieldSpec(String columnName) { if (fs == null) { return null; } - // For untyped URN columns where a single common prefix was discovered during dictionary - // build, surface a synthetic urnTypes so the column behaves like single-entity typed when the - // segment is loaded (prefix reconstruction, predicate prefix-strip). - String discoveredEntity = _discoveredUrnEntities.get(columnName); - if (discoveredEntity != null && fs instanceof DimensionFieldSpec - && fs.getDataType() == DataType.URN && !((DimensionFieldSpec) fs).hasUrnTypes()) { + 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; - DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, - original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), - original.getMaxLengthExceedStrategy()); - synthetic.setUrnTypes(Collections.singletonList(DimensionFieldSpec.UrnType.of(discoveredEntity))); - return synthetic; + // 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 discoveredMulti = _discoveredMultiEntityUrnTypes.get(columnName); + if (discoveredMulti != null) { + DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, + original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), + original.getMaxLengthExceedStrategy()); + synthetic.setUrnTypes(discoveredMulti); + _syntheticFieldSpecCache.put(columnName, synthetic); + return synthetic; + } + // Single-prefix discovery: synthesize single-entry urnTypes (no valueType so the column + // stays STRING-stored and gets the suffix-stripped reconstruction wrapper at load time). + String discoveredEntity = _discoveredUrnEntities.get(columnName); + if (discoveredEntity != null) { + DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, + original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), + original.getMaxLengthExceedStrategy()); + synthetic.setUrnTypes(Collections.singletonList(DimensionFieldSpec.UrnType.of(discoveredEntity))); + _syntheticFieldSpecCache.put(columnName, synthetic); + return synthetic; + } } return fs; } /** - * For untyped URN columns whose unique values all share the same {@code urn:::} - * prefix, strips the prefix from every value before it goes into the dictionary, records the - * discovered entity type so {@link #resolveFieldSpec} can surface it at metadata-write time, - * and returns the stripped array. For all other cases (typed URN, non-URN, or untyped URN with - * mixed prefixes) returns {@code values} unchanged. + * Preprocesses URN-column dictionary values. Returns either: + *
      + *
    • {@code EntityKeyPair[]} when the column is declared multi-entity or auto-detected as + * multi-prefix with all-numeric keys — feeds {@link + * org.apache.pinot.segment.local.segment.creator.impl.urn.UrnSortedDictionaryCreator}.
      • + *
    • {@code String[]} of stripped suffixes when an untyped column was auto-detected as + * single-entity — feeds the standard string dict creator.
      • + *
    • {@code values} unchanged for non-URN columns and untyped URN columns whose data isn't + * URN-shaped enough to optimize.
      • + *
    + * Side effects: populates {@link #_discoveredUrnEntities} or {@link + * #_discoveredMultiEntityUrnTypes} when a pattern is detected so {@link #resolveFieldSpec} + * can surface the synthetic 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 multi-entity URN: the stats collector already produced EntityKeyPair[] (see + // UrnMultiEntityPreIndexStatsCollector). Pass straight through to the + // UrnSortedDictionaryCreator. + if (dim.usesSortedUrnDictionary()) { + return values; + } if (dim.hasUrnTypes() || !(values instanceof Object[])) { return values; } @@ -351,19 +400,70 @@ private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Objec if (arr.length == 0) { return values; } + // Try single-prefix detection first; it's the most compact form. String entityType = discoverCommonUrnEntity(arr); - if (entityType == null) { - return values; + if (entityType != null) { + String prefix = entityType + ":"; + int prefixLen = prefix.length(); + String[] stripped = new String[arr.length]; + for (int i = 0; i < arr.length; i++) { + stripped[i] = arr[i].toString().substring(prefixLen); + } + _discoveredUrnEntities.put(columnName, entityType); + return stripped; } - String prefix = entityType + ":"; - int prefixLen = prefix.length(); - String[] stripped = new String[arr.length]; - for (int i = 0; i < arr.length; i++) { - String s = arr[i].toString(); - stripped[i] = s.substring(prefixLen); - } - _discoveredUrnEntities.put(columnName, entityType); - return stripped; + // Try multi-prefix auto-detect (all-numeric keys). When successful, the column is treated as + // multi-entity URN for the rest of the build and load paths. + org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] pairs = + tryMultiPrefixUrnDetection(columnName, arr); + if (pairs != null) { + return pairs; + } + return values; + } + + /** + * If every value in {@code values} parses as a simple URN with a numeric key and the column + * spans at least two distinct entity prefixes, returns the sorted {@code (entityIdx, key)} + * pair array and records the discovered {@code urnTypes}. Returns {@code null} otherwise. + */ + @Nullable + private org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] tryMultiPrefixUrnDetection( + String columnName, Object[] values) { + java.util.LinkedHashMap entityToIdx = new java.util.LinkedHashMap<>(); + List pairs = new ArrayList<>(values.length); + for (Object obj : values) { + String s = obj.toString(); + org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(s); + if (urn == null || !urn.isSimple()) { + return null; + } + long key; + try { + key = Long.parseLong(urn.simpleKey()); + } catch (NumberFormatException e) { + return null; + } + String entityType = urn.entityPrefix(); + Integer existing = entityToIdx.get(entityType); + int idx = existing != null ? existing : entityToIdx.size(); + if (existing == null) { + entityToIdx.put(entityType, idx); + } + pairs.add(new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair(idx, key)); + } + if (entityToIdx.size() < 2) { + return null; + } + org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] sorted = + pairs.toArray(new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[0]); + java.util.Arrays.sort(sorted); + List urnTypes = new ArrayList<>(entityToIdx.size()); + for (String entityType : entityToIdx.keySet()) { + urnTypes.add(DimensionFieldSpec.UrnType.of(entityType, DataType.LONG)); + } + _discoveredMultiEntityUrnTypes.put(columnName, urnTypes); + return sorted; } /** @@ -513,14 +613,44 @@ public void indexRow(GenericRow row) // 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 UrnSortedDictionaryCreator // 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 UrnSortedDictionaryCreator) { - Object entityIdxVal = row.getValue( - ((DimensionFieldSpec) fieldSpec).transientUrnEntityIdxFieldName()); + DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; + Object entityIdxVal = row.getValue(dim.transientUrnEntityIdxFieldName()); if (entityIdxVal != null) { + // Declared multi-entity path: transient field carries the entity idx. int entityIdx = ((Number) entityIdxVal).intValue(); long longKey = ((Number) columnValueToIndex).longValue(); columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( entityIdx, longKey); + } else if (columnValueToIndex instanceof String) { + // Auto-detected multi-entity path: parse the URN string against the discovered types. + org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse((String) columnValueToIndex); + 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) { + try { + long longKey = Long.parseLong(urn.simpleKey()); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( + entityIdx, longKey); + } catch (NumberFormatException ignored) { + // Auto-detection only synthesizes urnTypes when all keys parsed as LONG, so this + // branch should be unreachable; if it happens, fall through and let indexOfSV + // raise its own error. + } + } + } } } try { @@ -849,7 +979,12 @@ 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).usesSortedUrnDictionary(); + if ((fieldSpec.getFieldType() != FieldType.COMPLEX) && (totalDocs > 0) && !isMultiEntityUrn) { Object min = columnIndexCreationInfo.getMin(); Object max = columnIndexCreationInfo.getMax(); if (min != null && max != null) { 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 index 8d1840d45d..c76d04329d 100644 --- 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 @@ -44,7 +44,6 @@ 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.assertTrue; @@ -154,7 +153,36 @@ public void multiEntityUrnTypesPersistedInMetadata() } @Test - public void untypedUrnHasNoUrnTypesInMetadata() + 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"); @@ -164,7 +192,18 @@ public void untypedUrnHasNoUrnTypesInMetadata() SegmentMetadataImpl meta = new SegmentMetadataImpl(_segmentDir); ColumnMetadata colMeta = meta.getColumnMetadataMap().get(URN_COL); DimensionFieldSpec dimSpec = (DimensionFieldSpec) colMeta.getFieldSpec(); - assertFalse(dimSpec.hasUrnTypes()); + // 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 --- From 1d54f8ffd721b0d97325ecf97b7361d998e2de74 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 04:02:45 +0000 Subject: [PATCH 14/24] URN: validate ingest shape, replace runtime parsing with dict lookup UrnValidationTransformer now rejects any value that is not a simple URN in untyped URN columns, mirroring the existing typed-column check. With shape guaranteed at ingest time, every URN column reaches query time with structured metadata, so urnEntity() never needs to parse strings per row. UrnEntityTransformFunction's PARSE_STRING strategy is replaced with DICT_LOOKUP for multi-entity STRING-stored columns (mixed-type entities): init walks the dictionary once to build a dictId -> entityPrefix array, then per row reads the dict ID and indexes into the array. The CONSTANT and RANGE_LOOKUP strategies are unchanged. Rejection behaves like JSON ingest: throw when continueOnError=false, or null the value and mark the row INCOMPLETE when continueOnError=true. --- .../function/UrnEntityTransformFunction.java | 132 +++++++++------- .../function/UrnEntityDictLookupTest.java | 147 ++++++++++++++++++ .../function/UrnTransformFunctionTest.java | 4 +- .../UrnValidationTransformer.java | 120 ++++++++------ .../UrnValidationTransformerTest.java | 46 +++++- 5 files changed, 340 insertions(+), 109 deletions(-) create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityDictLookupTest.java 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 index a20398b6aa..9e5d65571e 100644 --- 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 @@ -27,6 +27,7 @@ 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; @@ -36,18 +37,20 @@ /** * {@code urnEntity(col)} transform function — returns the entity-type prefix of a URN column. * - *

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

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

      - *
    1. Constant — single-entity typed URN column (declared or auto-detected): every row - * in the block gets the same literal string, so the block is filled with one - * {@link Arrays#fill} call and no per-row URN parsing.
      2. - *
    2. Range-lookup — multi-entity typed URN column using the sorted-LONG UrnDictionary - * layout (see {@link DimensionFieldSpec#getEntityDictRanges()}): reads main-column dict IDs - * and finds the owning entity range with a linear scan over the per-entity range table - * (typically <=16 entries). No URN string parsing per row.
      3. - *
    3. Parse-string — fallback: untyped URN columns, or multi-entity columns whose entity - * dict ranges aren't available. Reads each URN string value and calls {@link Urn#tryParse} - * per row.
      4. + *
    4. 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.
      5. + *
    5. 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).
      6. + *
    6. 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.
      7. *
    */ public class UrnEntityTransformFunction extends BaseTransformFunction { @@ -56,7 +59,7 @@ public class UrnEntityTransformFunction extends BaseTransformFunction { public static final String FUNCTION_NAME_ALT = "urn_entity"; private enum Strategy { - CONSTANT, RANGE_LOOKUP, PARSE_STRING + CONSTANT, RANGE_LOOKUP, DICT_LOOKUP } private Strategy _strategy; @@ -64,6 +67,7 @@ private enum Strategy { private String _mainColumnName; private int[] _rangeEnds; private String[] _entityTypesByRange; + private String[] _entityTypeByDictId; @Override public String getName() { @@ -75,42 +79,62 @@ public void init(List arguments, Map c super.init(arguments, columnContextMap); Preconditions.checkArgument(arguments.size() == 1, "urnEntity 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; - // Single-entity URN columns always have the same entity prefix on every row — fill - // the block with one literal. - if (dim.isSingleUrnEntityType()) { - _strategy = Strategy.CONSTANT; - _constantEntityType = dim.getUrnTypes().get(0).getEntityType(); - return; - } - // Multi-entity URN columns using the sorted-LONG UrnDictionary layout: the per-entity - // dict ID ranges in column metadata tell us which entity owns each row's dict ID. - 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; - } - } - } + 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(); } - _strategy = Strategy.PARSE_STRING; + _mainColumnName = colName; + _strategy = Strategy.DICT_LOOKUP; } @Override @@ -132,8 +156,6 @@ public String[] transformToStringValuesSV(ValueBlock valueBlock) { int numRanges = _rangeEnds.length; for (int i = 0; i < length; i++) { int dictId = dictIds[i]; - // Linear scan — range count is small (typically <=16) and the scan is branch-friendly. - // For higher cardinalities, switch to binary search on _rangeEnds. String entityType = null; for (int r = 0; r < numRanges; r++) { if (dictId < _rangeEnds[r]) { @@ -144,17 +166,13 @@ public String[] transformToStringValuesSV(ValueBlock valueBlock) { _stringValuesSV[i] = entityType; } break; - case PARSE_STRING: + case DICT_LOOKUP: default: - String[] strValues = _arguments.get(0).transformToStringValuesSV(valueBlock); + BlockValSet dictBlock = valueBlock.getBlockValueSet(_mainColumnName); + int[] dictLookupIds = dictBlock.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { - String urnStr = strValues[i]; - if (urnStr == null) { - _stringValuesSV[i] = null; - } else { - Urn urn = Urn.tryParse(urnStr); - _stringValuesSV[i] = urn != null ? urn.entityPrefix() : null; - } + int id = dictLookupIds[i]; + _stringValuesSV[i] = id >= 0 ? _entityTypeByDictId[id] : null; } break; } 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..9a2c57c23a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnEntityDictLookupTest.java @@ -0,0 +1,147 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.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; + + +/** + * Covers the {@code DICT_LOOKUP} strategy of {@link UrnEntityTransformFunction} — the path used + * for multi-entity URN columns that fall back to STRING storage because at least one declared + * entity type has a non-numeric value type. The function precomputes entity-per-dict-id at init + * time and then does an O(1) array lookup per row, with no per-row URN parsing. + */ +public class UrnEntityDictLookupTest { + + private static final String SEGMENT_NAME = "urnEntityDictLookupSegment"; + 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: one LONG entity (memberId), one STRING entity (corpuser). + // getEffectiveStoredType() returns STRING here, so this column is NOT a sorted-LONG layout, + // entityDictRanges is null, and the function falls through to the DICT_LOOKUP strategy. + 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 mixedTypeUrnColumnRoutesThroughDictLookup() { + DimensionFieldSpec dim = (DimensionFieldSpec) _segment.getDataSource(URN_COL) + .getDataSourceMetadata().getFieldSpec(); + // Pre-condition for DICT_LOOKUP: STRING storage and no per-entity dict ranges. + assertEquals(dim.getEffectiveStoredType(), FieldSpec.DataType.STRING); + assertEquals(dim.getEntityDictRanges(), null); + + 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/UrnTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/UrnTransformFunctionTest.java index dc28d2cc82..5a00496d8a 100644 --- 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 @@ -163,7 +163,7 @@ public void testUrnEntityUnderscoreAliasOnTypedLongColumn() { } } - // --- urnEntity on untyped STRING column (PARSE_STRING fallback) --- + // --- urnEntity on untyped URN column (auto-detected to CONSTANT here: single-prefix data) --- @Test public void testUrnEntityOnUntypedUrnColumn() { @@ -213,7 +213,7 @@ public void testUrnKeyLongResultIsLongType() { assertEquals(fn.getResultMetadata().getDataType(), FieldSpec.DataType.LONG); } - // --- urnKeyLong on untyped STRING column (PARSE_STRING fallback) --- + // --- urnKeyLong on untyped URN column --- @Test public void testUrnKeyLongOnUntypedUrnColumn() { 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 index 6329f31ee0..41f9ef090e 100644 --- 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 @@ -38,30 +38,31 @@ /** - * Validates values in URN-typed columns against the {@link DimensionFieldSpec#getUrnTypes()} list - * declared in the schema. + * Validates values in URN-typed columns and rejects non-URN-shaped input. * - *

    For each URN column that has a non-empty {@code urnTypes} list: + *

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

      - *
    • The value must be parseable as a URN.
      • - *
    • The URN's entity type prefix must appear in the declared {@code urnTypes}.
      • + *
    • 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.
      • *
    * - *

    If a value fails validation: + *

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

      - *
    • With {@code continueOnError = false} (default): an {@link IllegalStateException} is thrown - * and ingestion of this record fails.
      • - *
    • With {@code continueOnError = true}: the value is replaced with {@code null} (the - * {@link NullValueTransformer} that follows will fill in the column's default-null-value), - * and the row is marked incomplete via {@link GenericRow#INCOMPLETE_RECORD_KEY}.
      • + *
    • {@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.
      • *
    * - *

    Columns with {@code dataType = URN} but no {@code urnTypes} declaration are not validated - * here (the {@link DataTypeTransformer} already verifies basic URN parseability via - * {@link Urn#parse(String)}). - * - *

    This transformer is a no-op when the schema contains no URN columns with declared types. - * *

    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. @@ -69,31 +70,34 @@ public class UrnValidationTransformer implements RecordTransformer { private static final Logger LOGGER = LoggerFactory.getLogger(UrnValidationTransformer.class); - /** Maps column name → set of allowed entity prefixes (e.g. "urn:li:memberId"). */ - private final Map> _allowedEntityPrefixes; + /** 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> columnEntityPrefixes = new HashMap<>(); + Map> typed = new HashMap<>(); + Set untyped = new LinkedHashSet<>(); for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (fieldSpec.getDataType() != FieldSpec.DataType.URN) { - continue; - } - if (!(fieldSpec instanceof DimensionFieldSpec)) { + 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()); } - columnEntityPrefixes.put(fieldSpec.getName(), Collections.unmodifiableSet(prefixes)); + typed.put(fieldSpec.getName(), Collections.unmodifiableSet(prefixes)); } - _allowedEntityPrefixes = Collections.unmodifiableMap(columnEntityPrefixes); + _columnAllowedPrefixes = Collections.unmodifiableMap(typed); + _untypedUrnColumns = Collections.unmodifiableSet(untyped); IngestionConfig ingestionConfig = tableConfig.getIngestionConfig(); _continueOnError = ingestionConfig != null && ingestionConfig.isContinueOnError(); @@ -101,34 +105,54 @@ public UrnValidationTransformer(TableConfig tableConfig, Schema schema) { @Override public boolean isNoOp() { - return _allowedEntityPrefixes.isEmpty(); + return _columnAllowedPrefixes.isEmpty() && _untypedUrnColumns.isEmpty(); } @Nullable @Override public GenericRow transform(GenericRow record) { - for (Map.Entry> entry : _allowedEntityPrefixes.entrySet()) { - String column = entry.getKey(); - Set allowed = entry.getValue(); - - Object rawValue = record.getValue(column); - if (rawValue == null) { - continue; - } - String valueStr = rawValue.toString(); - Urn urn = Urn.tryParse(valueStr); - if (urn == null || !allowed.contains(urn.entityPrefix())) { - String errorMessage = String.format( - "URN value '%s' in column '%s' does not match any declared urnTypes %s", valueStr, column, allowed); - if (_continueOnError) { - LOGGER.debug(errorMessage); - record.putValue(column, null); - record.putValue(GenericRow.INCOMPLETE_RECORD_KEY, true); - } else { - throw new IllegalStateException(errorMessage); - } - } + 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/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 index 1a4836e207..8d78f18a5e 100644 --- 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 @@ -83,9 +83,11 @@ public void isNoOpWhenNoUrnColumns() { } @Test - public void isNoOpWhenUrnColumnHasNoUrnTypes() { + 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()); - assertTrue(t.isNoOp()); + assertFalse(t.isNoOp()); } @Test @@ -181,4 +183,44 @@ public void multiEntityWrongTypeNullifiesWhenContinueOnError() { 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)); + } } From 8159f73a33aab444b3a5623ea8f98bb0a7e91f63 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 18:00:13 +0000 Subject: [PATCH 15/24] URN: polymorphic wire format -- per-entity typed sub-dicts (data + reader + creator) Extends multi-entity URN dictionary storage so each entity can use its own typed sub-dictionary instead of a single uniform LONG buffer. Wire format adds two fields per EntityDictRange: - dictClass: INT/LONG/STRING/...; selects the sub-dictionary reader - numBytesPerValue: fixed-width entry size in the entity's sub-buffer Metadata serialization stays backwards compatible. Ranges that are all default LONG/8-byte still emit the legacy two-field "entityType,end" form; ranges with mixed sub-dict types emit the four-field "entityType,end,dictClass,numBytesPerValue" form. The parser accepts both. DictionaryIndexType.read() hoists multi-entity URN columns above the effective-stored-type switch and slices the file into per-entity byte regions, wrapping each region in its declared sub-dict reader (LongDictionary, StringDictionary, ...). Per-region byte ranges are derived cumulatively from numBytesPerValue. Existing all-LONG segments still load through the LongDictionary path with identical byte layouts. UrnPolymorphicDictionaryCreator writes the polymorphic file: per-entity typed sub-dictionaries concatenated back-to-back. When every entity is LONG/8-byte the file is byte-identical to the prior sorted-LONG layout, so swapping it in for UrnSortedDictionaryCreator is wire-compatible. EntityValuePair generalizes EntityKeyPair to any Comparable value so mixed-type rows can flow through the same sort/index pipeline. Not yet wired into SegmentColumnarIndexCreator -- that swap (and the stats-collector generalization needed to feed mixed-type values into the new creator) is the next step. --- .../creator/impl/urn/EntityValuePair.java | 79 +++++ .../urn/UrnPolymorphicDictionaryCreator.java | 301 ++++++++++++++++++ .../index/dictionary/DictionaryIndexType.java | 87 +++-- .../UrnPolymorphicDictionaryCreatorTest.java | 189 +++++++++++ .../pinot/spi/data/DimensionFieldSpec.java | 113 ++++++- .../spi/data/DimensionFieldSpecUrnTest.java | 55 ++++ 6 files changed, 784 insertions(+), 40 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/EntityValuePair.java create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreator.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreatorTest.java 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..66661b0450 --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnPolymorphicDictionaryCreator.java @@ -0,0 +1,301 @@ +/** + * 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; + } + Preconditions.checkState(pairs.length > 0, "UrnPolymorphicDictionaryCreator received empty input"); + + 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); + } + + /** {@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 69c1ca4c79..4609cd0ef8 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 @@ -323,6 +323,20 @@ 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(); @@ -334,29 +348,6 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat FieldSpec fieldSpec = metadata.getFieldSpec(); if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - // Multi-entity URN: slice the LONG buffer per entity range and wrap each slice in a - // standard LongDictionary, then compose them into a UrnPolymorphicDictionary. The - // wrapper routes by entity prefix so each sub-dict gets to use its native lookup. - List ranges = dim.getEntityDictRanges(); - if (ranges != null) { - List subDicts = new ArrayList<>(ranges.size()); - int prevEnd = 0; - for (DimensionFieldSpec.EntityDictRange r : ranges) { - int subLen = r.getEndDictIdExclusive() - prevEnd; - if (subLen > 0) { - PinotDataBuffer subBuffer = dataBuffer.view( - (long) prevEnd * Long.BYTES, (long) r.getEndDictIdExclusive() * Long.BYTES); - Dictionary subDict = loadOnHeap - ? new OnHeapLongDictionary(subBuffer, subLen) - : new LongDictionary(subBuffer, subLen); - subDicts.add(new UrnPolymorphicDictionary.SubDict(r.getEntityType(), subDict)); - } - prevEnd = r.getEndDictIdExclusive(); - } - if (!subDicts.isEmpty()) { - return new UrnPolymorphicDictionary(subDicts); - } - } Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) : new LongDictionary(dataBuffer, length); // Single-entity typed URN LONG: wrap to prepend the entity prefix. @@ -408,6 +399,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/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-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 859390f613..91051eb337 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 @@ -189,9 +189,15 @@ public void setEntityDictRanges(@Nullable List ranges) { } /** - * Parses the encoded {@code urn.entityDictRanges} metadata property value of the form - * {@code "entityType1,end1;entityType2,end2;..."} where each {@code end} is the exclusive upper - * bound (in dict IDs) of that entity's range. Returns {@code null} for null or blank input. + * Parses the encoded {@code urn.entityDictRanges} metadata property value. Two formats are + * accepted for backwards compatibility: + *

      + *
    • Legacy: {@code "entityType,end;entityType,end;..."} — all sub-dicts are LONG + * with 8-byte entries (the original sorted-LONG layout).
      • + *
    • Polymorphic: {@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. */ @@ -210,28 +216,60 @@ public static List parseEntityDictRanges(@Nullable String encod if (part.isEmpty()) { continue; } - int comma = part.lastIndexOf(','); - if (comma <= 0 || comma == part.length() - 1) { + String[] fields = part.split(",", -1); + if (fields.length != 2 && fields.length != 4) { + throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); + } + String entityType = fields[0]; + if (entityType.isEmpty()) { throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); } - String entityType = part.substring(0, comma); int endDictIdExclusive; try { - endDictIdExclusive = Integer.parseInt(part.substring(comma + 1)); + endDictIdExclusive = Integer.parseInt(fields[1]); } catch (NumberFormatException e) { throw new IllegalArgumentException("Malformed entityDictRanges end: '" + part + "'", e); } - ranges.add(new EntityDictRange(entityType, endDictIdExclusive)); + if (fields.length == 2) { + // Legacy format — defaults to LONG + 8 bytes per value. + ranges.add(new EntityDictRange(entityType, endDictIdExclusive)); + } else { + 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. */ + /** + * Serializes {@link #getEntityDictRanges()} into the metadata-property form. Writes the legacy + * 2-field form when every range is a default LONG/8-byte sub-dict (the all-numeric sorted-LONG + * layout), preserving wire compatibility with segments built before the polymorphic format. + * Otherwise writes the 4-field polymorphic form. + */ @Nullable public static String serializeEntityDictRanges(@Nullable List ranges) { if (ranges == null || ranges.isEmpty()) { return null; } + boolean polymorphic = false; + for (EntityDictRange r : ranges) { + if (r.getDictClass() != DataType.LONG || r.getNumBytesPerValue() != Long.BYTES) { + polymorphic = true; + break; + } + } StringBuilder sb = new StringBuilder(); for (int i = 0; i < ranges.size(); i++) { if (i > 0) { @@ -239,6 +277,9 @@ public static String serializeEntityDictRanges(@Nullable List r } EntityDictRange r = ranges.get(i); sb.append(r.getEntityType()).append(',').append(r.getEndDictIdExclusive()); + if (polymorphic) { + sb.append(',').append(r.getDictClass().name()).append(',').append(r.getNumBytesPerValue()); + } } return sb.toString(); } @@ -498,11 +539,24 @@ public String toString() { } /** - * One entry in a multi-entity URN column's per-entity dictionary-ID range table. Identifies - * which entity type owns the dictionary IDs ending at (exclusive) {@code endDictIdExclusive}. + * 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 {@code i}-th range covers dict IDs {@code [prevEnd, endDictIdExclusive)}, where - * {@code prevEnd} is the previous range's exclusive end (or zero for the first range). + *

    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}. @@ -510,10 +564,24 @@ public String toString() { public static final class EntityDictRange { private final String _entityType; private final int _endDictIdExclusive; - + private final DataType _dictClass; + private final int _numBytesPerValue; + + /** + * Legacy 2-arg constructor for backwards compatibility with segments built before the + * polymorphic wire format. Defaults to {@code LONG} sub-dictionary with 8-byte entries — + * which matches the all-LONG sorted-dictionary layout that was the only multi-entity URN + * storage format prior to this extension. + */ public EntityDictRange(String entityType, int endDictIdExclusive) { + this(entityType, endDictIdExclusive, DataType.LONG, Long.BYTES); + } + + public EntityDictRange(String entityType, int endDictIdExclusive, DataType dictClass, int numBytesPerValue) { _entityType = entityType; _endDictIdExclusive = endDictIdExclusive; + _dictClass = dictClass; + _numBytesPerValue = numBytesPerValue; } public String getEntityType() { @@ -524,6 +592,14 @@ public int getEndDictIdExclusive() { return _endDictIdExclusive; } + public DataType getDictClass() { + return _dictClass; + } + + public int getNumBytesPerValue() { + return _numBytesPerValue; + } + @Override public boolean equals(Object o) { if (this == o) { @@ -533,17 +609,20 @@ public boolean equals(Object o) { return false; } EntityDictRange other = (EntityDictRange) o; - return _endDictIdExclusive == other._endDictIdExclusive && Objects.equals(_entityType, other._entityType); + return _endDictIdExclusive == other._endDictIdExclusive + && _numBytesPerValue == other._numBytesPerValue + && Objects.equals(_entityType, other._entityType) + && _dictClass == other._dictClass; } @Override public int hashCode() { - return Objects.hash(_entityType, _endDictIdExclusive); + return Objects.hash(_entityType, _endDictIdExclusive, _dictClass, _numBytesPerValue); } @Override public String toString() { - return _entityType + "[..." + _endDictIdExclusive + ")"; + return _entityType + "[..." + _endDictIdExclusive + ", " + _dictClass + "/" + _numBytesPerValue + "B)"; } } } 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 index 4e94572de0..8060ed4f3e 100644 --- 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 @@ -224,4 +224,59 @@ public void entityDictRangeEqualsAndHashCode() { assertEquals(a.hashCode(), b.hashCode()); assertTrue(!a.equals(c)); } + + // --- polymorphic entityDictRanges (per-entity dictClass + numBytesPerValue) --- + + @Test + public void entityDictRangesLegacyFormatDefaultsToLong() { + // The 2-field legacy format must read back as LONG/8-byte sub-dicts so that segments built + // before the polymorphic wire format keep loading exactly as before. + java.util.List parsed = + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100;urn:li:groupId,200"); + assertEquals(parsed.size(), 2); + for (DimensionFieldSpec.EntityDictRange r : parsed) { + assertEquals(r.getDictClass(), FieldSpec.DataType.LONG); + assertEquals(r.getNumBytesPerValue(), Long.BYTES); + } + } + + @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 + public void serializeAllDefaultLongPrefersLegacyForm() { + // When every range is the default LONG/8-byte sub-dict, the serialized form stays in the + // legacy 2-field shape so a segment built today reads byte-identical metadata to one built + // before the polymorphic wire format existed. + java.util.List ranges = java.util.Arrays.asList( + new DimensionFieldSpec.EntityDictRange("urn:li:memberId", 100), + new DimensionFieldSpec.EntityDictRange("urn:li:groupId", 200, + FieldSpec.DataType.LONG, Long.BYTES)); + assertEquals(DimensionFieldSpec.serializeEntityDictRanges(ranges), + "urn:li:memberId,100;urn:li:groupId,200"); + } + + @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"); + } } From 59e4b6e4d515bf8783dc875fb17d6ddf08cd81d8 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 18:15:39 +0000 Subject: [PATCH 16/24] URN: replace UrnSortedDictionaryCreator with UrnPolymorphicDictionaryCreator Transparent swap -- the polymorphic creator writes byte-identical output to the legacy sorted-LONG creator when every entity is LONG/8-byte, so all 121 existing URN tests pass without modification. This unifies the multi-entity URN build path around one creator. Mixing in non-numeric entities (next step) just requires generalizing the stats collector and per-row encoder to feed EntityValuePair[] instead of EntityKeyPair[]; the creator side is already ready. --- .claude/scheduled_tasks.lock | 1 + .../UrnKeyEncodingTransformer.java | 2 +- .../impl/SegmentColumnarIndexCreator.java | 16 +- .../SegmentPreIndexStatsCollectorImpl.java | 2 +- .../UrnMultiEntityPreIndexStatsCollector.java | 4 +- .../impl/urn/UrnSortedDictionaryCreator.java | 144 ------------------ 6 files changed, 13 insertions(+), 156 deletions(-) create mode 100644 .claude/scheduled_tasks.lock delete mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java diff --git a/.claude/scheduled_tasks.lock b/.claude/scheduled_tasks.lock new file mode 100644 index 0000000000..e1db557749 --- /dev/null +++ b/.claude/scheduled_tasks.lock @@ -0,0 +1 @@ +{"sessionId":"39e586ce-401c-49f0-808e-1d397531da99","pid":2573156,"procStart":"49717156","acquiredAt":1780026967388} \ No newline at end of file diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java index 8168ff167e..eaa67d5580 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java @@ -39,7 +39,7 @@ * replaces the string value on the row. For multi-entity URN columns * ({@link DimensionFieldSpec#usesSortedUrnDictionary()}), the 0-based entity index is also * written to the companion {@code __urnEntityIdx} column. The stats collector and - * {@code UrnSortedDictionaryCreator} pair the main column LONG with the companion column entity + * {@code UrnPolymorphicDictionaryCreator} pair the main column LONG with the companion column entity * index to lay out a per-entity sorted-LONG dictionary at segment build time. * *

    If the URN value is {@code null} or unparseable, both the main and companion columns are set 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 0ed22bb22a..2d129b10d1 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 @@ -38,7 +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.UrnSortedDictionaryCreator; +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; @@ -266,7 +266,7 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio _discoveredMultiEntityUrnTypes.containsKey(columnName) ? _discoveredMultiEntityUrnTypes.get(columnName) : ((DimensionFieldSpec) fieldSpec).getUrnTypes(); - creator = new UrnSortedDictionaryCreator(fieldSpec, _indexDir, urnTypesForCreator); + creator = new UrnPolymorphicDictionaryCreator(fieldSpec, _indexDir, urnTypesForCreator); } else { creator = new DictionaryIndexPlugin().getIndexType().createIndexCreator(context, dictConfig); } @@ -281,9 +281,9 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio // 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 UrnSortedDictionaryCreator) { + if (creator instanceof UrnPolymorphicDictionaryCreator) { DimensionFieldSpec dim = (DimensionFieldSpec) resolveFieldSpec(columnName); - dim.setEntityDictRanges(((UrnSortedDictionaryCreator) creator).getRanges()); + dim.setEntityDictRanges(((UrnPolymorphicDictionaryCreator) creator).getRanges()); } _dictionaryCreatorMap.put(columnName, creator); @@ -372,7 +372,7 @@ private FieldSpec resolveFieldSpec(String columnName) { *

      *
    • {@code EntityKeyPair[]} when the column is declared multi-entity or auto-detected as * multi-prefix with all-numeric keys — feeds {@link - * org.apache.pinot.segment.local.segment.creator.impl.urn.UrnSortedDictionaryCreator}.
      • + * org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator}. *
    • {@code String[]} of stripped suffixes when an untyped column was auto-detected as * single-entity — feeds the standard string dict creator.
      • *
    • {@code values} unchanged for non-URN columns and untyped URN columns whose data isn't @@ -389,7 +389,7 @@ private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Objec DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; // Declared multi-entity URN: the stats collector already produced EntityKeyPair[] (see // UrnMultiEntityPreIndexStatsCollector). Pass straight through to the - // UrnSortedDictionaryCreator. + // UrnPolymorphicDictionaryCreator. if (dim.usesSortedUrnDictionary()) { return values; } @@ -611,13 +611,13 @@ public void indexRow(GenericRow row) 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 UrnSortedDictionaryCreator + // 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 UrnSortedDictionaryCreator) { + if (dictionaryCreator instanceof UrnPolymorphicDictionaryCreator) { DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; Object entityIdxVal = row.getValue(dim.transientUrnEntityIdxFieldName()); if (entityIdxVal != null) { 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 cabde7e598..99d617303c 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 @@ -56,7 +56,7 @@ public void init() { break; case LONG: // Multi-entity URN columns get a stats collector that tracks (entityIdx, longKey) - // pairs so the UrnSortedDictionaryCreator can lay out per-entity dict ID ranges. The + // pairs so the UrnPolymorphicDictionaryCreator can lay out per-entity dict ID ranges. The // entity index comes from the transient field stashed by UrnKeyEncodingTransformer. if (fieldSpec instanceof DimensionFieldSpec && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()) { 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 index 82cffd0bf5..deab734a62 100644 --- 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 @@ -32,7 +32,7 @@ *

      Unlike the standard {@link LongColumnPreIndexStatsCollector}, which dedupes by LONG value * alone, this collector tracks unique {@code (entityIdx, longKey)} pairs. Two URNs that share a * key but belong to different entity types remain distinct dictionary entries, which is what the - * {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnSortedDictionaryCreator} + * {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator} * needs to lay out the per-entity dict ID ranges. * *

      Per-column dispatch in {@code SegmentPreIndexStatsCollectorImpl.collectRow} routes to @@ -75,7 +75,7 @@ public void collect(Object value, GenericRow row) { /** * Returns the sorted unique {@code (entityIdx, longKey)} pairs after {@link #seal()}. - * Consumed by the {@code UrnSortedDictionaryCreator}. + * Consumed by the {@code UrnPolymorphicDictionaryCreator}. */ @Override public Object getUniqueValuesSet() { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java deleted file mode 100644 index 82d1c0c87d..0000000000 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/urn/UrnSortedDictionaryCreator.java +++ /dev/null @@ -1,144 +0,0 @@ -/** - * 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.longs.Long2IntOpenHashMap; -import java.io.File; -import java.io.IOException; -import java.nio.ByteOrder; -import java.util.ArrayList; -import java.util.List; -import org.apache.commons.io.FileUtils; -import org.apache.pinot.segment.local.io.util.FixedByteValueReaderWriter; -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.slf4j.Logger; -import org.slf4j.LoggerFactory; - - -/** - * Custom dictionary creator for multi-entity URN columns. Writes a LONG buffer of keys sorted by - * {@code (entityIdx, longKey)} — allowing duplicate LONG values across entity boundaries — plus - * per-entity dictionary-ID ranges. The companion {@link - * org.apache.pinot.segment.local.segment.index.readers.UrnSortedDictionary} reader uses the - * ranges to do range-restricted lookups at query time. - */ -public class UrnSortedDictionaryCreator extends SegmentDictionaryCreator { - - private static final Logger LOGGER = LoggerFactory.getLogger(UrnSortedDictionaryCreator.class); - - private final String _columnName; - private final File _dictionaryFile; - private final List _urnTypes; - private List _ranges; - // (entityIdx, longKey) → absolute dictId. Two-level: array per entity, then long→dictId map. - private Long2IntOpenHashMap[] _pairToDictId; - private int _numEntries; - - public UrnSortedDictionaryCreator(FieldSpec fieldSpec, File indexDir, - List urnTypes) { - super(fieldSpec, indexDir, false); - _columnName = fieldSpec.getName(); - _dictionaryFile = new File(indexDir, _columnName + DictionaryIndexType.getFileExtension()); - _urnTypes = urnTypes; - } - - /** - * Builds the dictionary from a pre-sorted array of {@link EntityKeyPair}s. Input must be sorted - * by {@link EntityKeyPair#compareTo} (i.e. ascending entity index, then ascending long key) and - * contain no duplicates. - */ - @Override - public void build(Object sortedValues) - throws IOException { - Preconditions.checkArgument(sortedValues instanceof EntityKeyPair[], - "UrnSortedDictionaryCreator expects an EntityKeyPair[] from the stats collector"); - EntityKeyPair[] pairs = (EntityKeyPair[]) sortedValues; - Preconditions.checkState(pairs.length > 0, "UrnSortedDictionaryCreator received empty input"); - - FileUtils.touch(_dictionaryFile); - _numEntries = pairs.length; - _pairToDictId = new Long2IntOpenHashMap[_urnTypes.size()]; - for (int i = 0; i < _pairToDictId.length; i++) { - _pairToDictId[i] = new Long2IntOpenHashMap(); - _pairToDictId[i].defaultReturnValue(-1); - } - List ranges = new ArrayList<>(); - int currentEntity = -1; - - try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapFile(_dictionaryFile, false, 0, - (long) _numEntries * Long.BYTES, ByteOrder.BIG_ENDIAN, getClass().getSimpleName()); - FixedByteValueReaderWriter writer = new FixedByteValueReaderWriter(dataBuffer)) { - for (int i = 0; i < _numEntries; i++) { - EntityKeyPair pair = pairs[i]; - int entityIdx = pair.getEntityIdx(); - Preconditions.checkState(entityIdx >= currentEntity, - "UrnSortedDictionaryCreator input must be sorted by entityIdx"); - // Cross an entity boundary: close out the previous entity's range. - while (currentEntity < entityIdx) { - if (currentEntity >= 0) { - ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), i)); - } - currentEntity++; - } - writer.writeLong(i, pair.getLongKey()); - _pairToDictId[entityIdx].put(pair.getLongKey(), i); - } - } - // Close out the trailing range(s). Entities with no rows still get an entry with the same - // endDictId as the previous one (empty range), so readers can use a uniform lookup table. - while (currentEntity < _urnTypes.size() - 1) { - ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), _numEntries)); - currentEntity++; - } - ranges.add(new EntityDictRange(_urnTypes.get(currentEntity).getEntityType(), _numEntries)); - _ranges = ranges; - - LOGGER.info("Built UrnSortedDictionary for column: {} with cardinality: {} across {} entities", _columnName, - _numEntries, _urnTypes.size()); - } - - /** {@inheritDoc} For multi-entity URN, {@code value} must be an {@link EntityKeyPair}. */ - @Override - public int indexOfSV(Object value) { - Preconditions.checkState(_pairToDictId != null, "build() must be called before indexOfSV()"); - EntityKeyPair pair = (EntityKeyPair) value; - int dictId = _pairToDictId[pair.getEntityIdx()].get(pair.getLongKey()); - Preconditions.checkState(dictId >= 0, - "Row references (entity=%s, key=%s) that wasn't seen during stats collection", - _urnTypes.get(pair.getEntityIdx()).getEntityType(), pair.getLongKey()); - return dictId; - } - - @Override - public int getNumBytesPerEntry() { - return Long.BYTES; - } - - /** 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; - } -} From f2b0c16bd86b654b56d44b6eb82ac7eeb414ccd1 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 18:34:41 +0000 Subject: [PATCH 17/24] URN: route mixed-type multi-entity columns to polymorphic creator End-to-end wire-up for the polymorphic dictionary layout. Mixed-type multi-entity URN columns (e.g. LONG memberId + STRING corpuser) now write per-entity typed sub-dictionaries instead of falling back to plain STRING storage. DimensionFieldSpec gains usesPolymorphicUrnDictionary(): true for any multi-entity URN with declared LONG/INT/STRING valueTypes (superset of the existing usesSortedUrnDictionary which stays as the narrower all-numeric predicate). Wiring uses the broader predicate everywhere except where the all-numeric fast path is actually load-bearing. UrnKeyEncodingTransformer encodes per-entity-typed values for polymorphic columns: Long for LONG entities, Integer for INT, String for STRING. The column value type is heterogeneous across rows, but the transient entity index keeps the type paired with its entity so downstream readers stay type-safe via EntityValuePair. UrnMultiEntityPreIndexStatsCollector decoupled from LongColumnPreIndexStats and rebuilt around EntityValuePair so mixed-type values flow through the same dedup/sort pipeline as the all-LONG case. SegmentColumnarIndexCreator.indexRow builds EntityValuePair (not EntityKeyPair) for polymorphic columns, with the value pulled directly from the row's typed column value. UrnEntityDictLookupTest now verifies the mixed-type column persists typed EntityDictRanges (LONG memberId + STRING corpuser) and that urnEntity() returns the right prefix per row via RANGE_LOOKUP. DICT_LOOKUP becomes unreachable for new segments since every URN column now has structured per-entity metadata. --- .../function/UrnEntityDictLookupTest.java | 34 ++-- .../UrnKeyEncodingTransformer.java | 145 +++++++++++------- .../impl/SegmentColumnarIndexCreator.java | 34 ++-- .../SegmentPreIndexStatsCollectorImpl.java | 24 +-- .../UrnMultiEntityPreIndexStatsCollector.java | 91 +++++++---- .../pinot/spi/data/DimensionFieldSpec.java | 38 ++++- 6 files changed, 234 insertions(+), 132 deletions(-) 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 index 9a2c57c23a..849912bd83 100644 --- 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 @@ -54,14 +54,15 @@ /** - * Covers the {@code DICT_LOOKUP} strategy of {@link UrnEntityTransformFunction} — the path used - * for multi-entity URN columns that fall back to STRING storage because at least one declared - * entity type has a non-numeric value type. The function precomputes entity-per-dict-id at init - * time and then does an O(1) array lookup per row, with no per-row URN parsing. + * 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 = "urnEntityDictLookupSegment"; + 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; @@ -78,9 +79,9 @@ public void setUp() _segmentOutputDir = Files.createTempDir(); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); - // Mixed-type multi-entity URN: one LONG entity (memberId), one STRING entity (corpuser). - // getEffectiveStoredType() returns STRING here, so this column is NOT a sorted-LONG layout, - // entityDictRanges is null, and the function falls through to the DICT_LOOKUP strategy. + // 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), @@ -126,17 +127,24 @@ public void tearDown() { } @Test - public void mixedTypeUrnColumnRoutesThroughDictLookup() { + public void mixedTypeColumnPersistsTypedDictRanges() { DimensionFieldSpec dim = (DimensionFieldSpec) _segment.getDataSource(URN_COL) .getDataSourceMetadata().getFieldSpec(); - // Pre-condition for DICT_LOOKUP: STRING storage and no per-entity dict ranges. - assertEquals(dim.getEffectiveStoredType(), FieldSpec.DataType.STRING); - assertEquals(dim.getEntityDictRanges(), null); + 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); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java index eaa67d5580..b7c4aacb5d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/UrnKeyEncodingTransformer.java @@ -31,47 +31,53 @@ /** - * Converts URN string values in typed URN columns to their compact physical representation and - * writes per-row entity-type indices for multi-entity URN columns. + * Parses URN-string row values into their compact per-entity-typed key and (for multi-entity + * URN columns) stashes the entity index in a transient row field. * - *

      For URN columns whose effective stored type is LONG or INT - * ({@link DimensionFieldSpec#hasTypedUrnStorage()}), the URN string is parsed and its numeric key - * replaces the string value on the row. For multi-entity URN columns - * ({@link DimensionFieldSpec#usesSortedUrnDictionary()}), the 0-based entity index is also - * written to the companion {@code __urnEntityIdx} column. The stats collector and - * {@code UrnPolymorphicDictionaryCreator} pair the main column LONG with the companion column entity - * index to lay out a per-entity sorted-LONG dictionary at segment build time. + *

      Three 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 value is {@code null} or unparseable, both the main and companion columns are set - * to {@code null}; downstream {@link NullValueTransformer} fills in default null values. - * - *

      Must run after {@link UrnValidationTransformer} so that the entity type has already - * been verified to be in the declared set. + *

      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 { - /** Per-column encoding plan. */ private final List _encodings; public UrnKeyEncodingTransformer(Schema schema) { List encodings = new ArrayList<>(); for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (fieldSpec.getDataType() != FieldSpec.DataType.URN) { - continue; - } - if (!(fieldSpec instanceof DimensionFieldSpec)) { + if (fieldSpec.getDataType() != FieldSpec.DataType.URN + || !(fieldSpec instanceof DimensionFieldSpec)) { continue; } DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.hasTypedUrnStorage()) { - // Single-entity typed LONG/INT and multi-entity LONG/INT (UrnDictionary layout): encode - // the numeric key into the main column. For multi-entity, also stash the entity index - // into a transient row field so the segment builder can pair (entityIdx, longKey). The - // transient field is consumed at stats collection / forward-index time and is never - // persisted as a column. - String transientField = dim.usesSortedUrnDictionary() ? dim.transientUrnEntityIdxFieldName() : null; - encodings.add( - new ColumnEncoding(dim.getName(), dim.getEffectiveStoredType(), dim.getUrnTypes(), transientField)); + 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); @@ -94,21 +100,16 @@ public GenericRow transform(GenericRow record) { continue; } - String urnStr = raw.toString(); - Urn urn = Urn.tryParse(urnStr); + Urn urn = Urn.tryParse(raw.toString()); if (urn == null || !urn.isSimple()) { - // Malformed or composite URN: null out; NullValueTransformer will fill the default - record.putValue(enc._columnName, null); - if (enc._entityIdxFieldName != null) { - record.putValue(enc._entityIdxFieldName, null); - } + nullOut(record, enc); continue; } - // For multi-entity: find and write entity-type index if (enc._entityIdxFieldName != null) { - String entityType = urn.entityPrefix(); // e.g. "urn:li:memberId" + // 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; @@ -116,46 +117,74 @@ public GenericRow transform(GenericRow record) { } } if (entityIdx < 0) { - // Entity not in declared set (validation should have caught this); null out both - record.putValue(enc._columnName, null); - record.putValue(enc._entityIdxFieldName, null); + 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); - } - - // Encode the main column as the typed numeric key. - if (enc._storedType != null) { - String keyStr = urn.simpleKey(); - try { - if (enc._storedType == FieldSpec.DataType.LONG) { - record.putValue(enc._columnName, Long.parseLong(keyStr)); - } else { - record.putValue(enc._columnName, Integer.parseInt(keyStr)); - } - } catch (NumberFormatException e) { + 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); - if (enc._entityIdxFieldName != null) { - record.putValue(enc._entityIdxFieldName, 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 _storedType; + 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 storedType, + ColumnEncoding(String columnName, @Nullable FieldSpec.DataType singleEntityStoredType, List urnTypes, @Nullable String entityIdxFieldName) { _columnName = columnName; - _storedType = storedType; + _singleEntityStoredType = singleEntityStoredType; _urnTypes = urnTypes; _entityIdxFieldName = entityIdxFieldName; } 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 2d129b10d1..a6f0d52606 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 @@ -257,11 +257,12 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio // right dict creator. The standard path handles everything else. Object dictValues = maybeStripUrnPrefix(columnName, fieldSpec, context.getSortedUniqueElementsArray()); SegmentDictionaryCreator creator; - boolean usesSortedUrnDict = dictValues instanceof org.apache.pinot.segment.local.segment.creator.impl.urn + 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).usesSortedUrnDictionary()); - if (usesSortedUrnDict) { + && ((DimensionFieldSpec) fieldSpec).usesPolymorphicUrnDictionary()); + if (usesPolymorphicUrnDict) { List urnTypesForCreator = _discoveredMultiEntityUrnTypes.containsKey(columnName) ? _discoveredMultiEntityUrnTypes.get(columnName) @@ -387,10 +388,10 @@ private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Objec return values; } DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - // Declared multi-entity URN: the stats collector already produced EntityKeyPair[] (see - // UrnMultiEntityPreIndexStatsCollector). Pass straight through to the - // UrnPolymorphicDictionaryCreator. - if (dim.usesSortedUrnDictionary()) { + // Declared multi-entity polymorphic URN: the stats collector already produced + // EntityKeyPair[] (all-numeric path) or EntityValuePair[] (mixed-type path). Pass straight + // through to UrnPolymorphicDictionaryCreator. + if (dim.usesPolymorphicUrnDictionary()) { return values; } if (dim.hasUrnTypes() || !(values instanceof Object[])) { @@ -621,13 +622,18 @@ public void indexRow(GenericRow row) DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; Object entityIdxVal = row.getValue(dim.transientUrnEntityIdxFieldName()); if (entityIdxVal != null) { - // Declared multi-entity path: transient field carries the entity idx. + // 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(); - long longKey = ((Number) columnValueToIndex).longValue(); - columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( - entityIdx, longKey); + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( + entityIdx, (Comparable) columnValueToIndex); } else if (columnValueToIndex instanceof String) { - // Auto-detected multi-entity path: parse the URN string against the discovered types. + // Auto-detected multi-entity path: UrnKeyEncodingTransformer didn't run (column was + // untyped at schema time) — parse the URN string here against the discovered urnTypes. + // Auto-detect only fires when every value is parseable as LONG, so the sub-dictionary + // for each entity is LONG and we emit a Long-valued EntityValuePair. org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse((String) columnValueToIndex); if (urn != null && urn.isSimple()) { List urnTypes = dim.getUrnTypes(); @@ -642,7 +648,7 @@ public void indexRow(GenericRow row) if (entityIdx >= 0) { try { long longKey = Long.parseLong(urn.simpleKey()); - columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair( + columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( entityIdx, longKey); } catch (NumberFormatException ignored) { // Auto-detection only synthesizes urnTypes when all keys parsed as LONG, so this @@ -983,7 +989,7 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str // 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).usesSortedUrnDictionary(); + && ((DimensionFieldSpec) fieldSpec).usesPolymorphicUrnDictionary(); if ((fieldSpec.getFieldType() != FieldType.COMPLEX) && (totalDocs > 0) && !isMultiEntityUrn) { Object min = columnIndexCreationInfo.getMin(); Object max = columnIndexCreationInfo.getMax(); 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 99d617303c..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 @@ -49,23 +49,25 @@ public void init() { Schema dataSchema = _statsCollectorConfig.getSchema(); for (FieldSpec fieldSpec : dataSchema.getAllFieldSpecs()) { String column = fieldSpec.getName(); + // 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; case LONG: - // Multi-entity URN columns get a stats collector that tracks (entityIdx, longKey) - // pairs so the UrnPolymorphicDictionaryCreator can lay out per-entity dict ID ranges. The - // entity index comes from the transient field stashed by UrnKeyEncodingTransformer. - if (fieldSpec instanceof DimensionFieldSpec - && ((DimensionFieldSpec) fieldSpec).usesSortedUrnDictionary()) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - _columnStatsCollectorMap.put(column, new UrnMultiEntityPreIndexStatsCollector(column, - dim.transientUrnEntityIdxFieldName(), _statsCollectorConfig)); - } else { - _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig)); - } + _columnStatsCollectorMap.put(column, new LongColumnPreIndexStatsCollector(column, _statsCollectorConfig)); break; case FLOAT: _columnStatsCollectorMap.put(column, new FloatColumnPreIndexStatsCollector(column, _statsCollectorConfig)); 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 index deab734a62..28aa4288e2 100644 --- 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 @@ -21,29 +21,33 @@ import java.util.Arrays; import java.util.HashSet; import java.util.Set; -import org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair; +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 whose effective stored type is LONG. + * Stats collector for multi-entity URN columns using the polymorphic sub-dictionary layout. * - *

      Unlike the standard {@link LongColumnPreIndexStatsCollector}, which dedupes by LONG value - * alone, this collector tracks unique {@code (entityIdx, longKey)} pairs. Two URNs that share a - * key but belong to different entity types remain distinct dictionary entries, which is what the - * {@link org.apache.pinot.segment.local.segment.creator.impl.urn.UrnPolymorphicDictionaryCreator} - * needs to lay out the per-entity dict ID ranges. + *

      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. * - *

      Per-column dispatch in {@code SegmentPreIndexStatsCollectorImpl.collectRow} routes to - * {@link #collect(Object, GenericRow)} so the collector can pick the entity index up from the - * row's companion column. The standard {@link #collect(Object)} method falls back to ignoring - * the entity dimension (effectively LONG-only) — only used by paths that don't have row context. + *

      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 LongColumnPreIndexStatsCollector { +public class UrnMultiEntityPreIndexStatsCollector extends AbstractColumnStatisticsCollector { private final String _transientEntityIdxFieldName; - private final Set _pairs = new HashSet<>(); - private EntityKeyPair[] _sortedPairs; + private final Set _pairs = new HashSet<>(); + private EntityValuePair[] _sortedPairs; + private int _totalNumberOfEntries; private boolean _sealed = false; public UrnMultiEntityPreIndexStatsCollector(String column, String transientEntityIdxFieldName, @@ -53,9 +57,9 @@ public UrnMultiEntityPreIndexStatsCollector(String column, String transientEntit } /** - * Collects a row's main-column LONG value paired with the entity index stashed by - * {@code UrnKeyEncodingTransformer} in the transient row field. Used by - * {@code SegmentPreIndexStatsCollectorImpl.collectRow} for multi-entity URN columns. + * 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; @@ -64,18 +68,28 @@ public void collect(Object value, GenericRow row) { } Object entityIdxVal = row.getValue(_transientEntityIdxFieldName); if (entityIdxVal == null) { - // Row was nulled by the encoder (e.g. unknown entity); ignore for pair tracking. return; } int entityIdx = ((Number) entityIdxVal).intValue(); - long key = ((Number) value).longValue(); - _pairs.add(new EntityKeyPair(entityIdx, key)); - super.collect(value); // keep LongColumn min/max/cardinality stats in sync + _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, longKey)} pairs after {@link #seal()}. - * Consumed by the {@code UrnPolymorphicDictionaryCreator}. + * 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() { @@ -90,23 +104,42 @@ 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; } - super.seal(); - _sortedPairs = _pairs.toArray(new EntityKeyPair[0]); + _sortedPairs = _pairs.toArray(new EntityValuePair[0]); Arrays.sort(_sortedPairs); _pairs.clear(); _sealed = true; } /** - * Multi-entity URN columns are never considered sorted for forward-index purposes: the LONG - * key buffer is partitioned by entity, so consecutive rows may jump across entity ranges in - * dict-ID space even when the LONG values themselves come in monotone order. Letting the - * sorted forward-index path kick in would corrupt the per-row dict-ID readback. + * 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() { 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 91051eb337..32e6c9a29d 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 @@ -351,16 +351,16 @@ public boolean hasTypedUrnStorage() { } /** - * Returns {@code true} when the column should have a companion entity-index INT column written - * alongside the main column. This is true for multi-entity URN columns whose entity types all - * declare a numeric ({@code LONG} or {@code INT}) {@code valueType}. The segment builder uses - * a sorted LONG dictionary partitioned by entity-type ranges (see - * {@link #URN_TRANSIENT_ENTITY_IDX_PREFIX}). The companion field is built only as a transient - * row attribute during ingestion and is never persisted as a column. + * Returns {@code true} when the column uses the all-numeric sorted-LONG dictionary layout — a + * multi-entity URN column whose declared entities are all {@code LONG} or {@code INT}. The + * segment builder writes a single LONG buffer partitioned by per-entity dict-ID ranges. This is + * the narrower case of {@link #usesPolymorphicUrnDictionary()}; callers that don't depend on + * the all-numeric property (e.g. entity-idx-stashing, polymorphic creator routing) should + * prefer the broader predicate. */ @JsonIgnore public boolean usesSortedUrnDictionary() { - if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.size() <= 1) { + if (!usesPolymorphicUrnDictionary()) { return false; } for (DimensionFieldSpec.UrnType ut : _urnTypes) { @@ -372,6 +372,30 @@ public boolean usesSortedUrnDictionary() { return true; } + /** + * Returns {@code true} when the column uses the polymorphic per-entity sub-dictionary layout + * (see {@code UrnPolymorphicDictionaryCreator}). True for any multi-entity URN column 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. + */ + @JsonIgnore + public boolean usesPolymorphicUrnDictionary() { + if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.size() <= 1) { + 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 From 469a6e322f03a420d54bd95af2f698282ea92442 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 21:10:55 +0000 Subject: [PATCH 18/24] URN: unify single-entity columns into the polymorphic path Drops the size >= 2 restriction on usesPolymorphicUrnDictionary(): a single-entity URN column becomes a polymorphic dictionary with one EntityDictRange entry. The standalone UrnReconstructingLongDictionary and UrnReconstructingStringDictionary wrappers no longer need to fire for declared single-entity columns -- the polymorphic wrapper already prepends the entity prefix on read, and the metadata overhead is one extra range row. To make this work end to end: - UrnPolymorphicDictionaryCreator handles empty input (a column whose only rows are null) by emitting zero-width per-entity ranges. - UrnMultiEntityPreIndexStatsCollector associates the column's default null sentinel with entity 0 so it gets a dict ID and the polymorphic forward index can resolve null rows. SegmentColumnarIndexCreator.indexRow mirrors that, wrapping the bare typed value as EntityValuePair(0, ...) when the transient entity index is absent. - UrnPolymorphicDictionary.isSorted() reports true for single-entity columns -- one sub-dict, prefix prepended on read, ordering preserved -- so the sorted range-predicate evaluator runs. For multi-entity it stays false (cross-entity isn't lex-sorted as URN strings). - UrnPolymorphicDictionary.getDictIdsInRange computes range bounds via the sub-dict's insertionIndexOf instead of delegating to the sub-dict's getDictIdsInRange (which inherits the base "should not be called" implementation from BaseImmutableDictionary). The result is the half-open [from, to) relative dict-ID range, shifted by the entity's offset. --- .../impl/SegmentColumnarIndexCreator.java | 7 +++ .../UrnMultiEntityPreIndexStatsCollector.java | 9 ++-- .../urn/UrnPolymorphicDictionaryCreator.java | 36 ++++++++++++- .../readers/UrnPolymorphicDictionary.java | 51 +++++++++++++++---- .../pinot/spi/data/DimensionFieldSpec.java | 23 ++++++--- 5 files changed, 104 insertions(+), 22 deletions(-) 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 a6f0d52606..29e8b908ae 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 @@ -629,6 +629,13 @@ public void indexRow(GenericRow row) 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 multi-entity path: UrnKeyEncodingTransformer didn't run (column was // untyped at schema time) — parse the URN string here against the discovered urnTypes. 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 index 28aa4288e2..9a4c1716ca 100644 --- 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 @@ -67,10 +67,11 @@ public void collect(Object value, GenericRow row) { return; } Object entityIdxVal = row.getValue(_transientEntityIdxFieldName); - if (entityIdxVal == null) { - return; - } - int entityIdx = ((Number) entityIdxVal).intValue(); + // 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++; } 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 index 66661b0450..ae82e5e150 100644 --- 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 @@ -101,7 +101,13 @@ public void build(Object sortedValues) "UrnPolymorphicDictionaryCreator expects EntityValuePair[] or EntityKeyPair[]"); pairs = (EntityValuePair[]) sortedValues; } - Preconditions.checkState(pairs.length > 0, "UrnPolymorphicDictionaryCreator received empty input"); + // 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]; @@ -249,6 +255,34 @@ ByteOrder.BIG_ENDIAN, getClass().getSimpleName())) { _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) { 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 index 6593222bec..57a5b54a27 100644 --- 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 @@ -135,8 +135,11 @@ public String getEntityTypeForDictId(int globalDictId) { @Override public boolean isSorted() { - // Sub-dicts are independently sorted; globally the buffer is NOT lex-sorted as URN strings. - return false; + // 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 @@ -261,13 +264,43 @@ public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower if (lowerEntity < 0 || upperEntity < 0 || lowerEntity != upperEntity) { return new IntOpenHashSet(0); } - String lowerSuffix = lower.substring(_entityPrefixes[lowerEntity].length()); - String upperSuffix = upper.substring(_entityPrefixes[upperEntity].length()); - IntSet subResult = - _subDicts[lowerEntity].getDictIdsInRange(lowerSuffix, upperSuffix, includeLower, includeUpper); - IntSet result = new IntOpenHashSet(subResult.size()); - int offset = _offsets[lowerEntity]; - subResult.forEach((java.util.function.IntConsumer) id -> result.add(offset + id)); + 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 = sub.insertionIndexOf(lowerSuffix); + int from; + if (lowerInsert >= 0) { + from = includeLower ? lowerInsert : lowerInsert + 1; + } else { + from = ~lowerInsert; + } + int upperInsert = sub.insertionIndexOf(upperSuffix); + 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; } 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 32e6c9a29d..cbebc8eded 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 @@ -374,17 +374,24 @@ public boolean usesSortedUrnDictionary() { /** * Returns {@code true} when the column uses the polymorphic per-entity sub-dictionary layout - * (see {@code UrnPolymorphicDictionaryCreator}). True for any multi-entity URN column 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. + * (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.size() <= 1) { + if (_dataType != DataType.URN || _urnTypes == null || _urnTypes.isEmpty()) { return false; } for (DimensionFieldSpec.UrnType ut : _urnTypes) { From 9d7568ef34965a3ea46b0204475b2628b450d30c Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 21:16:37 +0000 Subject: [PATCH 19/24] URN: untyped columns route to polymorphic; unspecified valueType defaults to STRING All URN columns with dictionaries now use the polymorphic layout. Two unifications land here: 1. DimensionFieldSpec.setUrnTypes defaults any UrnType with no declared valueType to STRING. The polymorphic layout needs a concrete sub-dictionary type per entity; STRING accepts any key shape and is the safe default. Schemas that explicitly set valueType=LONG/INT/STRING are unchanged. 2. Untyped URN columns (no schema-declared urnTypes) flow through a single auto-detect path. tryUrnAutoDetection parses every value as a URN, groups by entity prefix, and chooses a per-entity sub-dict type: LONG when every key for that entity parses as long, otherwise STRING. It then emits an EntityValuePair[] for the polymorphic creator and stashes the discovered urnTypes (with their per-entity valueType) so resolveFieldSpec surfaces a synthetic polymorphic FieldSpec at metadata-write time. The old single-prefix-strip path is gone: _discoveredUrnEntities, maybeStripUrnPrefixForIndex, discoverCommonUrnEntity, and extractEntityPrefix are all removed. Single-prefix columns now run through the same polymorphic codepath as multi-prefix. --- .../impl/SegmentColumnarIndexCreator.java | 236 +++++++----------- .../pinot/spi/data/DimensionFieldSpec.java | 11 +- 2 files changed, 95 insertions(+), 152 deletions(-) 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 29e8b908ae..265b51f490 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 @@ -147,17 +147,11 @@ public class SegmentColumnarIndexCreator implements SegmentCreator { private Map, IndexCreator>> _creatorsByColAndIndex = new HashMap<>(); private final Map _nullValueVectorCreatorMap = new HashMap<>(); /** - * For untyped URN columns in which every observed value shares the same {@code urn:::} - * prefix, records the discovered entity type (without trailing colon). The prefix is stripped - * from the dictionary entries and persisted in column metadata so the load-time path treats the - * column as single-entity typed. - */ - private final Map _discoveredUrnEntities = new HashMap<>(); - /** - * For untyped URN columns whose unique values span multiple {@code urn:::} prefixes - * with numeric keys, records the discovered entity types in observed order. The column is - * promoted to multi-entity typed for the rest of the build (sorted-LONG UrnDictionary layout) - * and load-time {@code urnTypes} surface the discovered list. + * 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 @@ -344,23 +338,12 @@ private FieldSpec resolveFieldSpec(String columnName) { // 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 discoveredMulti = _discoveredMultiEntityUrnTypes.get(columnName); - if (discoveredMulti != null) { + 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(discoveredMulti); - _syntheticFieldSpecCache.put(columnName, synthetic); - return synthetic; - } - // Single-prefix discovery: synthesize single-entry urnTypes (no valueType so the column - // stays STRING-stored and gets the suffix-stripped reconstruction wrapper at load time). - String discoveredEntity = _discoveredUrnEntities.get(columnName); - if (discoveredEntity != null) { - DimensionFieldSpec synthetic = new DimensionFieldSpec(original.getName(), DataType.URN, - original.isSingleValueField(), original.getMaxLength(), original.getDefaultNullValue(), - original.getMaxLengthExceedStrategy()); - synthetic.setUrnTypes(Collections.singletonList(DimensionFieldSpec.UrnType.of(discoveredEntity))); + synthetic.setUrnTypes(discovered); _syntheticFieldSpecCache.put(columnName, synthetic); return synthetic; } @@ -371,26 +354,26 @@ private FieldSpec resolveFieldSpec(String columnName) { /** * Preprocesses URN-column dictionary values. Returns either: *

        - *
      • {@code EntityKeyPair[]} when the column is declared multi-entity or auto-detected as - * multi-prefix with all-numeric keys — feeds {@link + *
      • {@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 String[]} of stripped suffixes when an untyped column was auto-detected as - * single-entity — feeds the standard string dict creator.
        • + *
      • {@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 enough to optimize.
        • + * URN-shaped. *
      - * Side effects: populates {@link #_discoveredUrnEntities} or {@link - * #_discoveredMultiEntityUrnTypes} when a pattern is detected so {@link #resolveFieldSpec} - * can surface the synthetic FieldSpec at metadata-write time. + * 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 multi-entity polymorphic URN: the stats collector already produced - // EntityKeyPair[] (all-numeric path) or EntityValuePair[] (mixed-type path). Pass straight - // through to UrnPolymorphicDictionaryCreator. + // 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; } @@ -401,125 +384,84 @@ private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Objec if (arr.length == 0) { return values; } - // Try single-prefix detection first; it's the most compact form. - String entityType = discoverCommonUrnEntity(arr); - if (entityType != null) { - String prefix = entityType + ":"; - int prefixLen = prefix.length(); - String[] stripped = new String[arr.length]; - for (int i = 0; i < arr.length; i++) { - stripped[i] = arr[i].toString().substring(prefixLen); - } - _discoveredUrnEntities.put(columnName, entityType); - return stripped; - } - // Try multi-prefix auto-detect (all-numeric keys). When successful, the column is treated as - // multi-entity URN for the rest of the build and load paths. - org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] pairs = - tryMultiPrefixUrnDetection(columnName, arr); - if (pairs != null) { - return pairs; - } - 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); } /** - * If every value in {@code values} parses as a simple URN with a numeric key and the column - * spans at least two distinct entity prefixes, returns the sorted {@code (entityIdx, key)} - * pair array and records the discovered {@code urnTypes}. Returns {@code null} otherwise. + * 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. */ - @Nullable - private org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] tryMultiPrefixUrnDetection( - String columnName, Object[] values) { + private Object tryUrnAutoDetection(String columnName, Object[] values) { java.util.LinkedHashMap entityToIdx = new java.util.LinkedHashMap<>(); - List pairs = new ArrayList<>(values.length); - for (Object obj : values) { - String s = obj.toString(); - org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(s); + List perEntityValues = new ArrayList<>(); // entry per entity, holding (Long or String) per row + boolean[] perEntityAllLong = new boolean[8]; // grows as needed + Object[] parsedSuffixes = new Object[values.length]; // String suffix per row + int[] perRowEntityIdx = new int[values.length]; + + for (int rowIdx = 0; rowIdx < values.length; rowIdx++) { + org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(values[rowIdx].toString()); if (urn == null || !urn.isSimple()) { - return null; - } - long key; - try { - key = Long.parseLong(urn.simpleKey()); - } catch (NumberFormatException e) { - return null; + return values; } String entityType = urn.entityPrefix(); Integer existing = entityToIdx.get(entityType); - int idx = existing != null ? existing : entityToIdx.size(); - if (existing == null) { + 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; + } } - pairs.add(new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair(idx, key)); - } - if (entityToIdx.size() < 2) { - return null; } - org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[] sorted = - pairs.toArray(new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityKeyPair[0]); - java.util.Arrays.sort(sorted); - List urnTypes = new ArrayList<>(entityToIdx.size()); + + int numEntities = entityToIdx.size(); + List urnTypes = new ArrayList<>(numEntities); + DataType[] perEntityType = new DataType[numEntities]; + int e = 0; for (String entityType : entityToIdx.keySet()) { - urnTypes.add(DimensionFieldSpec.UrnType.of(entityType, DataType.LONG)); + DataType vt = perEntityAllLong[e] ? DataType.LONG : DataType.STRING; + perEntityType[e] = vt; + urnTypes.add(DimensionFieldSpec.UrnType.of(entityType, vt)); + e++; } _discoveredMultiEntityUrnTypes.put(columnName, urnTypes); - return sorted; - } - /** - * Strips the discovered URN entity prefix from a row's value at index time so the dictionary - * lookup hits the suffix-only entry. - */ - private Object maybeStripUrnPrefixForIndex(String columnName, Object value) { - String entityType = _discoveredUrnEntities.get(columnName); - if (entityType == null || !(value instanceof String)) { - return value; - } - String s = (String) value; - String prefix = entityType + ":"; - return s.startsWith(prefix) ? s.substring(prefix.length()) : value; - } - - /** - * Returns the entity type {@code urn::} when every value in {@code values} starts with - * the same {@code urn:::} prefix, or {@code null} otherwise. - */ - @Nullable - private static String discoverCommonUrnEntity(Object[] values) { - String prefix = extractEntityPrefix(values[0].toString()); - if (prefix == null) { - return null; - } - for (int i = 1; i < values.length; i++) { - if (!values[i].toString().startsWith(prefix)) { - return null; - } + org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair[] pairs = + new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair[values.length]; + for (int rowIdx = 0; rowIdx < values.length; rowIdx++) { + int idx = perRowEntityIdx[rowIdx]; + String suffix = (String) parsedSuffixes[rowIdx]; + Comparable typedValue = perEntityType[idx] == DataType.LONG ? Long.valueOf(suffix) : suffix; + pairs[rowIdx] = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair(idx, typedValue); } - return prefix.substring(0, prefix.length() - 1); + java.util.Arrays.sort(pairs); + return pairs; } /** - * Extracts the entity prefix {@code urn:::} (with trailing colon) from a URN string. - * Returns {@code null} when the input is not in canonical URN form. + * Strips the discovered URN entity prefix from a row's value at index time so the dictionary + * lookup hits the suffix-only entry. */ - @Nullable - private static String extractEntityPrefix(String urnString) { - if (!urnString.startsWith("urn:")) { - return null; - } - int colons = 0; - for (int i = 0; i < urnString.length(); i++) { - if (urnString.charAt(i) == ':') { - colons++; - if (colons == 3) { - return urnString.substring(0, i + 1); - } - } - } - return null; - } - private FieldIndexConfigs adaptConfig(String columnName, FieldIndexConfigs config, ColumnIndexCreationInfo columnIndexCreationInfo, SegmentGeneratorConfig segmentCreationSpec) { FieldIndexConfigs.Builder builder = new FieldIndexConfigs.Builder(config); @@ -603,10 +545,6 @@ public void indexRow(GenericRow row) if (columnValueToIndex == null) { throw new RuntimeException("Null value for column:" + columnName); } - // If we discovered a common URN entity prefix for this column at dictionary-build time, - // strip it from each row's value too so the dictionary lookup hits the suffix-only entries. - columnValueToIndex = maybeStripUrnPrefixForIndex(columnName, columnValueToIndex); - Map, IndexCreator> creatorsByIndex = byColEntry.getValue(); FieldSpec fieldSpec = resolveFieldSpec(columnName); @@ -637,10 +575,10 @@ public void indexRow(GenericRow row) columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( 0, (Comparable) columnValueToIndex); } else if (columnValueToIndex instanceof String) { - // Auto-detected multi-entity path: UrnKeyEncodingTransformer didn't run (column was - // untyped at schema time) — parse the URN string here against the discovered urnTypes. - // Auto-detect only fires when every value is parseable as LONG, so the sub-dictionary - // for each entity is LONG and we emit a Long-valued EntityValuePair. + // 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). org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse((String) columnValueToIndex); if (urn != null && urn.isSimple()) { List urnTypes = dim.getUrnTypes(); @@ -653,15 +591,11 @@ public void indexRow(GenericRow row) } } if (entityIdx >= 0) { - try { - long longKey = Long.parseLong(urn.simpleKey()); - columnValueToIndex = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair( - entityIdx, longKey); - } catch (NumberFormatException ignored) { - // Auto-detection only synthesizes urnTypes when all keys parsed as LONG, so this - // branch should be unreachable; if it happens, fall through and let indexOfSV - // raise its own error. - } + 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); } } } 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 cbebc8eded..a416e01cef 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 @@ -134,11 +134,20 @@ public List getUrnTypes() { // 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(urnTypes); + _urnTypes = Collections.unmodifiableList(normalized); } else { _urnTypes = null; } From 2730f6f85486a566c6f54f7bbc934a725a3776da Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 21:19:06 +0000 Subject: [PATCH 20/24] URN: delete UrnReconstructing{Long,String}Dictionary wrappers With every URN column now flowing through the polymorphic dictionary, the legacy single-entity reconstructing wrappers are dead code. The polymorphic dictionary's getStringValue() already prepends the entity prefix on read and its indexOf()/insertionIndexOf() already accept full URN strings, so wrapping a plain LongDictionary or StringDictionary is no longer needed. DictionaryIndexType.read drops both special-case branches and the wrapper imports; the LONG and STRING cases become plain scalar loads again. Single-entity URN columns either populate entityDictRanges and take the polymorphic load path at the top of read(), or fall through to the scalar load and behave like any other column. --- .../index/dictionary/DictionaryIndexType.java | 33 +--- .../UrnReconstructingLongDictionary.java | 182 ------------------ .../UrnReconstructingStringDictionary.java | 158 --------------- 3 files changed, 2 insertions(+), 371 deletions(-) delete mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java delete mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java 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 4609cd0ef8..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 @@ -50,8 +50,6 @@ 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.local.segment.index.readers.UrnReconstructingLongDictionary; -import org.apache.pinot.segment.local.segment.index.readers.UrnReconstructingStringDictionary; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.ColumnStatistics; @@ -344,22 +342,9 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat case INT: return loadOnHeap ? new OnHeapIntDictionary(dataBuffer, length) : new IntDictionary(dataBuffer, length); - case LONG: { - FieldSpec fieldSpec = metadata.getFieldSpec(); - if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - Dictionary dict = loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) - : new LongDictionary(dataBuffer, length); - // Single-entity typed URN LONG: wrap to prepend the entity prefix. - if (dim.hasTypedUrnStorage() && !dim.usesSortedUrnDictionary()) { - String entityType = dim.getUrnTypes().get(0).getEntityType(); - return new UrnReconstructingLongDictionary(dict, entityType); - } - return dict; - } + case LONG: return loadOnHeap ? new OnHeapLongDictionary(dataBuffer, length) : new LongDictionary(dataBuffer, length); - } case FLOAT: return loadOnHeap ? new OnHeapFloatDictionary(dataBuffer, length) : new FloatDictionary(dataBuffer, length); @@ -372,23 +357,9 @@ public static Dictionary read(PinotDataBuffer dataBuffer, ColumnMetadata metadat : new BigDecimalDictionary(dataBuffer, length, numBytesPerValue); case STRING: { numBytesPerValue = metadata.getColumnMaxLength(); - Dictionary dict = loadOnHeap + return loadOnHeap ? new OnHeapStringDictionary(dataBuffer, length, numBytesPerValue, strInterner, byteInterner) : new StringDictionary(dataBuffer, length, numBytesPerValue); - // For multi-entity typed URN columns the dictionary entries are stored as - // ":"; wrap to reconstruct full URN strings on read and to encode - // full-URN predicate literals before delegating to the underlying STRING dictionary. - FieldSpec fieldSpec = metadata.getFieldSpec(); - if (dataType == DataType.URN && fieldSpec instanceof DimensionFieldSpec) { - DimensionFieldSpec dim = (DimensionFieldSpec) fieldSpec; - if (dim.isSingleUrnEntityType()) { - // Single-entity URN with STRING storage: the dictionary stores suffixes only. - // Covers both schema-declared single-entity URN columns and untyped URN columns where - // a single common prefix was discovered at segment build time. - return new UrnReconstructingStringDictionary(dict, dim.getUrnTypes().get(0).getEntityType()); - } - } - return dict; } case BYTES: numBytesPerValue = metadata.getColumnMaxLength(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java deleted file mode 100644 index 1b471b3753..0000000000 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingLongDictionary.java +++ /dev/null @@ -1,182 +0,0 @@ -/** - * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; - - -/** - * Wraps a LONG dictionary for a single-entity typed URN column and reconstructs full URN strings - * (e.g., "urn:li:memberId:100") on read. Predicates using full URN string literals are also - * transparently resolved by stripping the entity prefix before delegating to the inner dictionary. - * - *

      This allows {@code SELECT col} on a typed URN LONG column to return the original URN strings - * while the segment physically stores compact LONG keys. - */ -public class UrnReconstructingLongDictionary implements Dictionary { - private final Dictionary _delegate; - private final String _entityTypePrefix; - - /** - * @param delegate underlying LONG dictionary (off-heap or on-heap) - * @param entityType the entity type without trailing colon, e.g. {@code "urn:li:memberId"} - */ - public UrnReconstructingLongDictionary(Dictionary delegate, String entityType) { - _delegate = delegate; - _entityTypePrefix = entityType + ":"; - } - - @Override - public boolean isSorted() { - return true; - } - - @Override - public DataType getValueType() { - return DataType.LONG; - } - - @Override - public int length() { - return _delegate.length(); - } - - /** Resolves a full URN string to its dict ID by stripping the entity prefix. */ - @Override - public int indexOf(String stringValue) { - if (stringValue.startsWith(_entityTypePrefix)) { - try { - long key = Long.parseLong(stringValue.substring(_entityTypePrefix.length())); - return _delegate.indexOf(key); - } catch (NumberFormatException e) { - return NULL_VALUE_INDEX; - } - } - return NULL_VALUE_INDEX; - } - - @Override - public int indexOf(long longValue) { - return _delegate.indexOf(longValue); - } - - /** - * Returns the binary-search insertion index for range predicate evaluation. - * Strips the entity prefix before delegating. - */ - @Override - public int insertionIndexOf(String stringValue) { - if (stringValue.startsWith(_entityTypePrefix)) { - try { - long key = Long.parseLong(stringValue.substring(_entityTypePrefix.length())); - return _delegate.insertionIndexOf(Long.toString(key)); - } catch (NumberFormatException e) { - return ~0; - } - } - return ~0; - } - - @Override - public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { - String lowerKey = toKeyString(lower); - String upperKey = toKeyString(upper); - return _delegate.getDictIdsInRange( - lowerKey != null ? lowerKey : lower, - upperKey != null ? upperKey : upper, - includeLower, includeUpper); - } - - private String toKeyString(String urnString) { - if (urnString != null && urnString.startsWith(_entityTypePrefix)) { - try { - return Long.toString(Long.parseLong(urnString.substring(_entityTypePrefix.length()))); - } catch (NumberFormatException e) { - return null; - } - } - return null; - } - - @Override - public int compare(int dictId1, int dictId2) { - return Integer.compare(dictId1, dictId2); - } - - @Override - public Comparable getMinVal() { - return _delegate.getMinVal(); - } - - @Override - public Comparable getMaxVal() { - return _delegate.getMaxVal(); - } - - @Override - public Object getSortedValues() { - return _delegate.getSortedValues(); - } - - @Override - public Object get(int dictId) { - return _delegate.get(dictId); - } - - @Override - public int getIntValue(int dictId) { - return _delegate.getIntValue(dictId); - } - - @Override - public long getLongValue(int dictId) { - return _delegate.getLongValue(dictId); - } - - @Override - public float getFloatValue(int dictId) { - return _delegate.getFloatValue(dictId); - } - - @Override - public double getDoubleValue(int dictId) { - return _delegate.getDoubleValue(dictId); - } - - @Override - public BigDecimal getBigDecimalValue(int dictId) { - return _delegate.getBigDecimalValue(dictId); - } - - /** Returns the full URN string, e.g., {@code "urn:li:memberId:100"}. */ - @Override - public String getStringValue(int dictId) { - return _entityTypePrefix + _delegate.getLongValue(dictId); - } - - @Override - public void close() - throws IOException { - _delegate.close(); - } -} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java deleted file mode 100644 index 09e9324d56..0000000000 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/UrnReconstructingStringDictionary.java +++ /dev/null @@ -1,158 +0,0 @@ -/** - * 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 org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; - - -/** - * Wraps a STRING dictionary for a single-entity URN column whose entries were stored with the - * {@code urn:::} prefix stripped (because every value in the segment shared that - * prefix). Reconstructs the full URN string on read and strips the entity prefix from incoming - * full-URN predicate literals before delegating to the underlying STRING dictionary. - */ -public class UrnReconstructingStringDictionary implements Dictionary { - private final Dictionary _delegate; - private final String _entityTypePrefix; - - /** - * @param delegate underlying STRING dictionary of URN suffix-only values - * @param entityType the entity type without trailing colon, e.g. {@code "urn:li:corpUser"} - */ - public UrnReconstructingStringDictionary(Dictionary delegate, String entityType) { - _delegate = delegate; - _entityTypePrefix = entityType + ":"; - } - - @Override - public boolean isSorted() { - return _delegate.isSorted(); - } - - @Override - public DataType getValueType() { - return DataType.STRING; - } - - @Override - public int length() { - return _delegate.length(); - } - - @Override - public int indexOf(String stringValue) { - String suffix = stripPrefix(stringValue); - return suffix != null ? _delegate.indexOf(suffix) : _delegate.indexOf(stringValue); - } - - @Override - public int insertionIndexOf(String stringValue) { - String suffix = stripPrefix(stringValue); - return suffix != null ? _delegate.insertionIndexOf(suffix) : _delegate.insertionIndexOf(stringValue); - } - - @Override - public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower, boolean includeUpper) { - String lowerKey = stripPrefix(lower); - String upperKey = stripPrefix(upper); - return _delegate.getDictIdsInRange( - lowerKey != null ? lowerKey : lower, - upperKey != null ? upperKey : upper, - includeLower, includeUpper); - } - - private String stripPrefix(String stringValue) { - return stringValue != null && stringValue.startsWith(_entityTypePrefix) - ? stringValue.substring(_entityTypePrefix.length()) : null; - } - - @Override - public int compare(int dictId1, int dictId2) { - return _delegate.compare(dictId1, dictId2); - } - - @Override - public Comparable getMinVal() { - return _delegate.getMinVal(); - } - - @Override - public Comparable getMaxVal() { - return _delegate.getMaxVal(); - } - - @Override - public Object getSortedValues() { - return _delegate.getSortedValues(); - } - - @Override - public Object get(int dictId) { - return getStringValue(dictId); - } - - @Override - public int getIntValue(int dictId) { - return _delegate.getIntValue(dictId); - } - - @Override - public long getLongValue(int dictId) { - return _delegate.getLongValue(dictId); - } - - @Override - public float getFloatValue(int dictId) { - return _delegate.getFloatValue(dictId); - } - - @Override - public double getDoubleValue(int dictId) { - return _delegate.getDoubleValue(dictId); - } - - @Override - public BigDecimal getBigDecimalValue(int dictId) { - return _delegate.getBigDecimalValue(dictId); - } - - /** Returns the full URN string by prepending the entity prefix. */ - @Override - public String getStringValue(int dictId) { - return _entityTypePrefix + _delegate.getStringValue(dictId); - } - - @Override - public void readStringValues(int[] dictIds, int length, String[] outValues) { - _delegate.readStringValues(dictIds, length, outValues); - for (int i = 0; i < length; i++) { - outValues[i] = _entityTypePrefix + outValues[i]; - } - } - - @Override - public void close() - throws IOException { - _delegate.close(); - } -} From 06c22d40c6ba291711291183fee6417d681529a3 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 21:20:48 +0000 Subject: [PATCH 21/24] URN: end-to-end tests for polymorphic untyped + STRING-default paths - untypedMixedKeyTypesDiscoverPerEntityValueType: untyped column with mixed LONG (memberId) + STRING (corpuser) keys auto-detects per-entity sub-dict types and round-trips full URN strings through the polymorphic dictionary. - untypedSinglePrefixDiscoveredAsPolymorphic: single-prefix untyped column becomes a 1-entity polymorphic column with LONG sub-dict; no more special-case prefix-stripping path. - declaredUrnTypesWithoutValueTypeDefaultToString: setUrnTypes normalizes any UrnType missing a valueType to STRING. --- .../segment/index/UrnColumnSegmentTest.java | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) 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 index c76d04329d..36200b86b4 100644 --- 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 @@ -416,4 +416,84 @@ public void multiEntitySegmentsWithReorderedUrnTypesReconstructCorrectly() 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); + } } From 78ce260b2d37986ac0937a80350380cd55eeeec0 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 21:29:37 +0000 Subject: [PATCH 22/24] URN: drop legacy 2-field wire format -- always polymorphic The all-LONG byte-identical wire compat was a refactor safety net while swapping UrnSortedDictionaryCreator -> UrnPolymorphicDictionaryCreator; it's not needed in production because the URN feature itself is new in this PR, so no existing segments use the legacy form. Collapsing to one format: - serializeEntityDictRanges always emits the 4-field "entityType,end,dictClass,numBytesPerValue" form. - parseEntityDictRanges only accepts the 4-field form; 2-field input is rejected as malformed. - EntityDictRange drops the 2-arg legacy constructor; all callers now pass dictClass and numBytesPerValue explicitly. Segments are now self-describing -- the sub-dict type for each entity is stated in the metadata rather than inferred from defaults. --- .../pinot/spi/data/DimensionFieldSpec.java | 74 +++++-------------- .../spi/data/DimensionFieldSpecUrnTest.java | 58 ++++----------- 2 files changed, 35 insertions(+), 97 deletions(-) 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 a416e01cef..0f6ac2aff6 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 @@ -198,15 +198,10 @@ public void setEntityDictRanges(@Nullable List ranges) { } /** - * Parses the encoded {@code urn.entityDictRanges} metadata property value. Two formats are - * accepted for backwards compatibility: - *

        - *
      • Legacy: {@code "entityType,end;entityType,end;..."} — all sub-dicts are LONG - * with 8-byte entries (the original sorted-LONG layout).
        • - *
      • Polymorphic: {@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. + * 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. */ @@ -226,7 +221,7 @@ public static List parseEntityDictRanges(@Nullable String encod continue; } String[] fields = part.split(",", -1); - if (fields.length != 2 && fields.length != 4) { + if (fields.length != 4) { throw new IllegalArgumentException("Malformed entityDictRanges entry: '" + part + "'"); } String entityType = fields[0]; @@ -239,56 +234,37 @@ public static List parseEntityDictRanges(@Nullable String encod } catch (NumberFormatException e) { throw new IllegalArgumentException("Malformed entityDictRanges end: '" + part + "'", e); } - if (fields.length == 2) { - // Legacy format — defaults to LONG + 8 bytes per value. - ranges.add(new EntityDictRange(entityType, endDictIdExclusive)); - } else { - 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)); + 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. Writes the legacy - * 2-field form when every range is a default LONG/8-byte sub-dict (the all-numeric sorted-LONG - * layout), preserving wire compatibility with segments built before the polymorphic format. - * Otherwise writes the 4-field polymorphic form. - */ + /** Serializes {@link #getEntityDictRanges()} into the metadata-property form. */ @Nullable public static String serializeEntityDictRanges(@Nullable List ranges) { if (ranges == null || ranges.isEmpty()) { return null; } - boolean polymorphic = false; - for (EntityDictRange r : ranges) { - if (r.getDictClass() != DataType.LONG || r.getNumBytesPerValue() != Long.BYTES) { - polymorphic = true; - break; - } - } 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()); - if (polymorphic) { - sb.append(',').append(r.getDictClass().name()).append(',').append(r.getNumBytesPerValue()); - } + sb.append(r.getEntityType()).append(',').append(r.getEndDictIdExclusive()) + .append(',').append(r.getDictClass().name()).append(',').append(r.getNumBytesPerValue()); } return sb.toString(); } @@ -607,16 +583,6 @@ public static final class EntityDictRange { private final DataType _dictClass; private final int _numBytesPerValue; - /** - * Legacy 2-arg constructor for backwards compatibility with segments built before the - * polymorphic wire format. Defaults to {@code LONG} sub-dictionary with 8-byte entries — - * which matches the all-LONG sorted-dictionary layout that was the only multi-entity URN - * storage format prior to this extension. - */ - public EntityDictRange(String entityType, int endDictIdExclusive) { - this(entityType, endDictIdExclusive, DataType.LONG, Long.BYTES); - } - public EntityDictRange(String entityType, int endDictIdExclusive, DataType dictClass, int numBytesPerValue) { _entityType = entityType; _endDictIdExclusive = endDictIdExclusive; 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 index 8060ed4f3e..e3686ac569 100644 --- 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 @@ -160,7 +160,7 @@ public void toStringIncludesUrnTypes() { @Test public void entityDictRangesRoundTripThreeEntities() { - String encoded = "urn:li:corpUser,33334;urn:li:corpGroup,66667;urn:li:applicationName,100000"; + 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); @@ -170,19 +170,11 @@ public void entityDictRangesRoundTripThreeEntities() { 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 entityDictRangesUserDocExample() { - String encoded = "urn:li:memberId,100;urn:li:companyPage,120"; - java.util.List parsed = - DimensionFieldSpec.parseEntityDictRanges(encoded); - assertEquals(parsed.size(), 2); - assertEquals(parsed.get(0).getEndDictIdExclusive(), 100); - assertEquals(parsed.get(1).getEndDictIdExclusive(), 120); - } - @Test public void entityDictRangesNullAndBlank() { assertNull(DimensionFieldSpec.parseEntityDictRanges(null)); @@ -199,15 +191,15 @@ public void entityDictRangesRejectsMalformedNoComma() { @Test(expectedExceptions = IllegalArgumentException.class) public void entityDictRangesRejectsMalformedNonNumericEnd() { - DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,abc"); + 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), - new DimensionFieldSpec.EntityDictRange("urn:li:corpGroup", 200)); + 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"); @@ -217,29 +209,17 @@ public void setEntityDictRangesPersists() { @Test public void entityDictRangeEqualsAndHashCode() { - DimensionFieldSpec.EntityDictRange a = new DimensionFieldSpec.EntityDictRange("urn:li:x", 5); - DimensionFieldSpec.EntityDictRange b = new DimensionFieldSpec.EntityDictRange("urn:li:x", 5); - DimensionFieldSpec.EntityDictRange c = new DimensionFieldSpec.EntityDictRange("urn:li:y", 5); + 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)); } - // --- polymorphic entityDictRanges (per-entity dictClass + numBytesPerValue) --- - - @Test - public void entityDictRangesLegacyFormatDefaultsToLong() { - // The 2-field legacy format must read back as LONG/8-byte sub-dicts so that segments built - // before the polymorphic wire format keep loading exactly as before. - java.util.List parsed = - DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100;urn:li:groupId,200"); - assertEquals(parsed.size(), 2); - for (DimensionFieldSpec.EntityDictRange r : parsed) { - assertEquals(r.getDictClass(), FieldSpec.DataType.LONG); - assertEquals(r.getNumBytesPerValue(), Long.BYTES); - } - } - @Test public void entityDictRangesPolymorphicRoundTrip() { String encoded = "urn:li:memberId,100,LONG,8;urn:li:corpuser,150,STRING,16"; @@ -257,17 +237,9 @@ public void entityDictRangesPolymorphicRoundTrip() { assertEquals(DimensionFieldSpec.serializeEntityDictRanges(parsed), encoded); } - @Test - public void serializeAllDefaultLongPrefersLegacyForm() { - // When every range is the default LONG/8-byte sub-dict, the serialized form stays in the - // legacy 2-field shape so a segment built today reads byte-identical metadata to one built - // before the polymorphic wire format existed. - java.util.List ranges = java.util.Arrays.asList( - new DimensionFieldSpec.EntityDictRange("urn:li:memberId", 100), - new DimensionFieldSpec.EntityDictRange("urn:li:groupId", 200, - FieldSpec.DataType.LONG, Long.BYTES)); - assertEquals(DimensionFieldSpec.serializeEntityDictRanges(ranges), - "urn:li:memberId,100;urn:li:groupId,200"); + @Test(expectedExceptions = IllegalArgumentException.class) + public void entityDictRangesRejectsTwoFieldEntry() { + DimensionFieldSpec.parseEntityDictRanges("urn:li:memberId,100"); } @Test(expectedExceptions = IllegalArgumentException.class) From 8c443d36b9ae0531b5813605fb1793e7ced510aa Mon Sep 17 00:00:00 2001 From: dinoocch Date: Fri, 29 May 2026 23:52:30 +0000 Subject: [PATCH 23/24] URN: fix 4 confirmed bugs from code review + remove stale RAT-flagged file 1. usesSortedUrnDictionary() lost its size > 1 guard, causing FilterPlanNode to skip the urnKeyLong-pushdown rewrite for single-entity LONG URN columns. Restored the guard; the FilterPlanNodeUrnRewriteTest single-entity cases pass again. 2. Mixed-type polymorphic columns + null URN row caused a build-time ClassCastException because NullValueTransformer filled the column with the URN sentinel String, but the polymorphic creator's LONG/INT branch unconditionally cast the value to Number. Made DimensionFieldSpec.getDefaultNullValue() return the entity-0-typed default for polymorphic columns, so the sentinel matches entity 0's sub-dictionary type. 3. UrnPolymorphicDictionary.indexOf / insertionIndexOf / getDictIdsInRange would propagate NumberFormatException out of the query path when a range bound's suffix was non-numeric on a LONG sub-dict (e.g. WHERE col BETWEEN 'urn:li:memberId:abc' AND 'urn:li:memberId:zzz'). Wrapped sub-dict calls in safeIndexOf / safeInsertionIndexOf helpers that treat parse failure as "no such value" (NULL_VALUE_INDEX / ~length). 4. tryUrnAutoDetection promoted the URN null sentinel (DEFAULT_DIMENSION_NULL_VALUE_OF_URN = "urn:pinot:null:0") to a phantom "urn:pinot:null" entity in every untyped URN column with null rows. Filter the sentinel before entity discovery and reserve a type-appropriate null placeholder in entity 0's sub-dictionary so per-row indexer can still resolve null rows. indexRow's auto-detect branch maps sentinels to the same placeholder. Plus: removed .claude/scheduled_tasks.lock from the index and added .claude/ to .gitignore. The lockfile was an accidentally-committed Claude Code runtime artifact with no license header, causing apache-rat:check to fail with "Too many files with unapproved license". RAT now reports Unapproved: 0. Tests: 4 new regression tests, all 112 URN tests pass. --- .claude/scheduled_tasks.lock | 1 - .gitignore | 1 + .../impl/SegmentColumnarIndexCreator.java | 93 ++++++++++++++----- .../readers/UrnPolymorphicDictionary.java | 36 ++++++- .../segment/index/UrnColumnSegmentTest.java | 63 +++++++++++++ .../readers/UrnPolymorphicDictionaryTest.java | 17 ++++ .../pinot/spi/data/DimensionFieldSpec.java | 47 +++++++--- 7 files changed, 219 insertions(+), 39 deletions(-) delete mode 100644 .claude/scheduled_tasks.lock diff --git a/.claude/scheduled_tasks.lock b/.claude/scheduled_tasks.lock deleted file mode 100644 index e1db557749..0000000000 --- a/.claude/scheduled_tasks.lock +++ /dev/null @@ -1 +0,0 @@ -{"sessionId":"39e586ce-401c-49f0-808e-1d397531da99","pid":2573156,"procStart":"49717156","acquiredAt":1780026967388} \ No newline at end of file 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-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 265b51f490..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 @@ -400,13 +400,22 @@ private Object maybeStripUrnPrefix(String columnName, FieldSpec fieldSpec, Objec */ private Object tryUrnAutoDetection(String columnName, Object[] values) { java.util.LinkedHashMap entityToIdx = new java.util.LinkedHashMap<>(); - List perEntityValues = new ArrayList<>(); // entry per entity, holding (Long or String) per row boolean[] perEntityAllLong = new boolean[8]; // grows as needed - Object[] parsedSuffixes = new Object[values.length]; // String suffix per row + 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++) { - org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse(values[rowIdx].toString()); + 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; } @@ -435,6 +444,10 @@ private Object tryUrnAutoDetection(String columnName, Object[] values) { } 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; @@ -446,18 +459,44 @@ private Object tryUrnAutoDetection(String columnName, Object[] values) { } _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[values.length]; + 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[rowIdx] = new org.apache.pinot.segment.local.segment.creator.impl.urn.EntityValuePair(idx, typedValue); + 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. @@ -578,24 +617,36 @@ public void indexRow(GenericRow row) // 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). - org.apache.pinot.spi.utils.Urn urn = org.apache.pinot.spi.utils.Urn.tryParse((String) columnValueToIndex); - if (urn != null && urn.isSimple()) { + // 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(); - 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(); + 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( - entityIdx, typedKey); + 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); + } } } } 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 index 57a5b54a27..6a2df92caf 100644 --- 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 @@ -236,7 +236,7 @@ public int indexOf(String urnString) { return NULL_VALUE_INDEX; } String suffix = urnString.substring(_entityPrefixes[entityIdx].length()); - int localId = _subDicts[entityIdx].indexOf(suffix); + int localId = safeIndexOf(_subDicts[entityIdx], suffix); return localId < 0 ? NULL_VALUE_INDEX : _offsets[entityIdx] + localId; } @@ -247,7 +247,7 @@ public int insertionIndexOf(String urnString) { return ~_totalLength; } String suffix = urnString.substring(_entityPrefixes[entityIdx].length()); - int subInsertion = _subDicts[entityIdx].insertionIndexOf(suffix); + int subInsertion = safeInsertionIndexOf(_subDicts[entityIdx], suffix, _subDicts[entityIdx].length()); if (subInsertion >= 0) { return _offsets[entityIdx] + subInsertion; } @@ -273,14 +273,14 @@ public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower // 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 = sub.insertionIndexOf(lowerSuffix); + int lowerInsert = safeInsertionIndexOf(sub, lowerSuffix, subLen); int from; if (lowerInsert >= 0) { from = includeLower ? lowerInsert : lowerInsert + 1; } else { from = ~lowerInsert; } - int upperInsert = sub.insertionIndexOf(upperSuffix); + int upperInsert = safeInsertionIndexOf(sub, upperSuffix, subLen); int to; if (upperInsert >= 0) { to = includeUpper ? upperInsert + 1 : upperInsert; @@ -304,6 +304,34 @@ public IntSet getDictIdsInRange(String lower, String upper, boolean includeLower 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) { 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 index 36200b86b4..fb23e446fb 100644 --- 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 @@ -496,4 +496,67 @@ public void declaredUrnTypesWithoutValueTypeDefaultToString() { 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 index 4dda731ff0..0b96213f5b 100644 --- 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 @@ -437,4 +437,21 @@ public void singleSubDictBehavesAsExpected() { 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-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 0f6ac2aff6..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 @@ -272,12 +272,27 @@ public static String serializeEntityDictRanges(@Nullable List r @Override public Object getDefaultNullValue() { if (_dataType == DataType.URN) { - 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; + // 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(); @@ -336,16 +351,22 @@ public boolean hasTypedUrnStorage() { } /** - * Returns {@code true} when the column uses the all-numeric sorted-LONG dictionary layout — a - * multi-entity URN column whose declared entities are all {@code LONG} or {@code INT}. The - * segment builder writes a single LONG buffer partitioned by per-entity dict-ID ranges. This is - * the narrower case of {@link #usesPolymorphicUrnDictionary()}; callers that don't depend on - * the all-numeric property (e.g. entity-idx-stashing, polymorphic creator routing) should - * prefer the broader predicate. + * 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()) { + if (!usesPolymorphicUrnDictionary() || _urnTypes.size() <= 1) { return false; } for (DimensionFieldSpec.UrnType ut : _urnTypes) { From ef7a0e7b4f7a0dd97f6d27de05d967eae1197c3d Mon Sep 17 00:00:00 2001 From: dinoocch Date: Mon, 1 Jun 2026 23:31:51 +0000 Subject: [PATCH 24/24] URN: fix checkstyle line-length violation in DimensionFieldSpecUrnTest --- .../org/apache/pinot/spi/data/DimensionFieldSpecUrnTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index e3686ac569..bd96d4dcf9 100644 --- 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 @@ -160,7 +160,8 @@ public void toStringIncludesUrnTypes() { @Test public void entityDictRangesRoundTripThreeEntities() { - String encoded = "urn:li:corpUser,33334,LONG,8;urn:li:corpGroup,66667,LONG,8;urn:li:applicationName,100000,STRING,16"; + 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);