xiangfu0 commented on code in PR #19595:
URL: https://github.com/apache/pinot/pull/19595#discussion_r4051857619
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/array/BaseArrayAggStringFunction.java:
##########
@@ -40,6 +43,22 @@ public BaseArrayAggStringFunction(ExpressionContext
expression, boolean nullHand
public void aggregateGroupBySV(int length, int[] groupKeyArray,
GroupByResultHolder groupByResultHolder,
Map<ExpressionContext, BlockValSet> blockValSetMap) {
BlockValSet blockValSet = blockValSetMap.get(_expression);
+ // Star-tree pre-aggregated column: each single-value BYTES entry is a
serialized distinct set; add each of its
+ // elements to the group's accumulator.
+ if (blockValSet.getValueType() == DataType.BYTES &&
blockValSet.isSingleValue()) {
Review Comment:
Fixed. Commit 41bb55a replaces the BYTES type sniffing with explicit
provenance: the star-tree overload of
`AggregationFunctionUtils.getBlockValSetMap` wraps the pre-aggregated column in
a new `StarTreePreAggregatedBlockValSet`, and every arrayAgg star-tree branch
(aggregate + group-by, all element types) keys on `instanceof` that marker
instead of the block type. Ordinary BYTES source columns are no longer
intercepted, and this also made `ARRAY_AGG` over BYTES columns star-tree
eligible. Regression coverage added in d12afc0:
`testArrayAggBytesAsStringMatchesRawScan` runs `arrayAgg(bytesCol, 'STRING')`
(distinct and non-distinct) with a star-tree present, asserting hex-conversion
results identical to a scan with the star-tree disabled.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/array/ArrayAggDistinctLongFunction.java:
##########
@@ -36,12 +37,29 @@ public ArrayAggDistinctLongFunction(ExpressionContext
expression, DataType dataT
super(expression, dataType, nullHandlingEnabled);
}
+ @Override
+ public boolean canUseStarTree(Map<String, Object> functionParameters) {
+ return true;
Review Comment:
Fixed in d12afc0, using the rejection approach:
`StarTreeUtils.createStarTreeBasedProjectOperator` now rejects an arrayAgg
whose declared element stored type differs from the source column's stored
type, so `arrayAgg(intCol, 'LONG', true)` falls back to the raw scan and its
read-time conversion (same for BYTES read as 'STRING'). Stored-type comparison
keeps logical types working (TIMESTAMP/LONG, BOOLEAN/INT, JSON/STRING,
UUID/BYTES). Tests: `testDistinctArrayAggTypeMismatchFallsBackToRawScan`
(star-tree not used + results identical to raw scan) and
`testDistinctArrayAggIntUsesStarTree` (matching type still selected). Read-time
conversion inside the star-tree path could be layered on later if a use case
demands the speedup for mismatched types.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ArrayAggDistinctValueAggregator.java:
##########
@@ -0,0 +1,335 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.segment.local.aggregator;
+
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
+import it.unimi.dsi.fastutil.objects.ObjectSet;
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.BigDecimalUtils;
+import org.apache.pinot.spi.utils.ByteArray;
+
+
+/// Value aggregator for distinct `ARRAYAGG` (i.e. `arrayAgg(column, 'type',
true)`) used by the star-tree index.
+///
+/// The aggregated value is the set of distinct (non-null) raw values seen
under a star-tree node. Distinct semantics
+/// make the aggregation associative and idempotent (merge = set-union), which
is what the star-tree requires; the
+/// non-distinct variant would need an unbounded per-node multiset and is
intentionally not supported here.
+///
+/// The set is stored as a variable-length `BYTES` cell in the star-tree
forward index. The serialized layout is
+/// intentionally byte-for-byte identical to `ObjectSerDeUtils.*_SET_SER_DE`
(in pinot-core) so the query-time
+/// `arrayAgg` function can deserialize the stored cell directly. The wire
layout cannot be shared as code because
+/// pinot-segment-local must not depend on pinot-core; any change here must be
mirrored there.
+///
+/// NOTE: Only single-value source columns are supported (enforced at
star-tree build config validation time), so the
+/// raw value handed in by the builder is always a single boxed scalar, never
an array.
+public class ArrayAggDistinctValueAggregator implements
ValueAggregator<Object, ObjectSet<Object>> {
+ public static final DataType AGGREGATED_VALUE_TYPE = DataType.BYTES;
+
+ // Runtime element kind, inferred lazily from the first raw value (mirrors
DistinctCountThetaSketchValueAggregator's
+ // runtime dispatch). All values in a given aggregator instance come from
the same column, so the kind is stable.
+ // Each value carries an explicit stable on-disk `tag` that is persisted as
a leading byte in the serialized cell,
+ // making deserialization self-describing (independent of aggregator
instance state, needed for build re-ingestion).
+ // The query-time arrayAgg function reads and validates the same tag before
decoding the payload.
+ //
+ // IMPORTANT: tags are an on-disk format. Never change or reuse an existing
value's tag; only append new values with
+ // new tags. Reordering the enum is safe because the tag, not the ordinal,
is persisted.
+ public enum ElementType {
+ INT(0), LONG(1), FLOAT(2), DOUBLE(3), BIG_DECIMAL(4), STRING(5), BYTES(6);
+
+ private final byte _tag;
+
+ ElementType(int tag) {
+ _tag = (byte) tag;
+ }
+
+ public byte getTag() {
+ return _tag;
+ }
+
+ private static final ElementType[] BY_TAG;
+
+ static {
+ ElementType[] values = values();
+ int maxTag = 0;
+ for (ElementType value : values) {
+ maxTag = Math.max(maxTag, value._tag);
+ }
+ BY_TAG = new ElementType[maxTag + 1];
+ for (ElementType value : values) {
+ BY_TAG[value._tag] = value;
+ }
+ }
+
+ public static ElementType fromTag(byte tag) {
+ if (tag < 0 || tag >= BY_TAG.length || BY_TAG[tag] == null) {
+ throw new IllegalStateException("Invalid arrayAgg element type tag: "
+ tag);
+ }
+ return BY_TAG[tag];
+ }
+ }
+
+ @Nullable
+ private ElementType _elementType;
+ private int _maxByteSize;
+
+ @Override
+ public AggregationFunctionType getAggregationType() {
+ return AggregationFunctionType.ARRAYAGG;
+ }
+
+ @Override
+ public DataType getAggregatedValueType() {
+ return AGGREGATED_VALUE_TYPE;
+ }
+
+ @Override
+ public ObjectSet<Object> getInitialAggregatedValue(@Nullable Object
rawValue) {
+ ObjectSet<Object> set = new ObjectOpenHashSet<>();
+ if (rawValue != null) {
+ addRawValue(set, rawValue);
+ }
+ updateMaxByteSize(set);
+ return set;
+ }
+
+ @Override
+ public ObjectSet<Object> applyRawValue(ObjectSet<Object> value, Object
rawValue) {
+ if (rawValue != null) {
+ addRawValue(value, rawValue);
+ updateMaxByteSize(value);
+ }
+ return value;
+ }
+
+ @Override
+ public ObjectSet<Object> applyAggregatedValue(ObjectSet<Object> value,
ObjectSet<Object> aggregatedValue) {
+ value.addAll(aggregatedValue);
+ updateMaxByteSize(value);
+ return value;
+ }
+
+ @Override
+ public ObjectSet<Object> cloneAggregatedValue(ObjectSet<Object> value) {
+ return new ObjectOpenHashSet<>(value);
+ }
+
+ @Override
+ public boolean isAggregatedValueFixedSize() {
+ return false;
+ }
+
+ @Override
+ public int getMaxAggregatedValueByteSize() {
+ return _maxByteSize;
+ }
+
+ @Override
+ public byte[] serializeAggregatedValue(ObjectSet<Object> value) {
+ // Layout: [1-byte element type tag][int size][payload]. The tag makes the
cell self-describing so it can be
+ // deserialized (during build re-ingestion) without aggregator instance
state. An all-empty aggregator that never
+ // saw a value defaults to LONG; the payload is empty either way so the
tag choice is immaterial for empty cells.
+ ElementType elementType = _elementType != null ? _elementType :
ElementType.LONG;
+ switch (elementType) {
+ case INT:
+ case LONG:
+ case FLOAT:
+ case DOUBLE:
+ return serializeFixedWidth(value, elementType);
+ case BIG_DECIMAL:
+ case STRING:
+ case BYTES:
+ return serializeVariableWidth(value, elementType);
+ default:
+ throw new IllegalStateException("Unsupported element type: " +
elementType);
+ }
+ }
+
+ @Override
+ public ObjectSet<Object> deserializeAggregatedValue(byte[] bytes) {
+ ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
+ ElementType elementType = ElementType.fromTag(byteBuffer.get());
+ int size = byteBuffer.getInt();
+ ObjectSet<Object> set = new ObjectOpenHashSet<>(size);
+ for (int i = 0; i < size; i++) {
+ switch (elementType) {
+ case INT:
+ set.add(byteBuffer.getInt());
+ break;
+ case LONG:
+ set.add(byteBuffer.getLong());
+ break;
+ case FLOAT:
+ set.add(byteBuffer.getFloat());
+ break;
+ case DOUBLE:
+ set.add(byteBuffer.getDouble());
+ break;
+ case BIG_DECIMAL: {
+ byte[] valueBytes = readLengthPrefixed(byteBuffer);
+ set.add(BigDecimalUtils.deserialize(valueBytes));
+ break;
+ }
+ case STRING: {
+ byte[] valueBytes = readLengthPrefixed(byteBuffer);
+ set.add(new String(valueBytes, StandardCharsets.UTF_8));
+ break;
+ }
+ case BYTES: {
+ byte[] valueBytes = readLengthPrefixed(byteBuffer);
+ set.add(new ByteArray(valueBytes));
+ break;
+ }
+ default:
+ throw new IllegalStateException("Unsupported element type: " +
elementType);
+ }
+ }
+ return set;
+ }
+
+ private void addRawValue(ObjectSet<Object> set, Object rawValue) {
+ if (_elementType == null) {
+ _elementType = inferElementType(rawValue);
+ }
+ // Normalize BYTES to ByteArray so hashCode/equals work as a set key;
other types are already value types.
+ if (rawValue instanceof byte[]) {
+ set.add(new ByteArray((byte[]) rawValue));
+ } else {
+ set.add(rawValue);
+ }
+ }
+
+ private static ElementType inferElementType(Object rawValue) {
+ if (rawValue instanceof Integer) {
+ return ElementType.INT;
+ } else if (rawValue instanceof Long) {
+ return ElementType.LONG;
+ } else if (rawValue instanceof Float) {
+ return ElementType.FLOAT;
+ } else if (rawValue instanceof Double) {
+ return ElementType.DOUBLE;
+ } else if (rawValue instanceof BigDecimal) {
+ return ElementType.BIG_DECIMAL;
+ } else if (rawValue instanceof String) {
+ return ElementType.STRING;
+ } else if (rawValue instanceof byte[] || rawValue instanceof ByteArray) {
+ return ElementType.BYTES;
+ } else {
+ throw new IllegalStateException(
+ "Unsupported data type for arrayAgg star-tree aggregation: " +
rawValue.getClass().getSimpleName());
+ }
+ }
+
+ private void updateMaxByteSize(ObjectSet<Object> value) {
+ _maxByteSize = Math.max(_maxByteSize,
serializeAggregatedValue(value).length);
Review Comment:
Fixed in d12afc0. The aggregator no longer re-serializes the set to track
size: fixed-width cell sizes are computed O(1) from the set size, and
variable-width sizes are tracked incrementally — the set is a private
`ObjectOpenHashSet` subclass carrying a running payload total, updated only
when an element is actually added (duplicates are free). STRING/BIG_DECIMAL
elements are encoded once when first added rather than on every size update.
`maxByteSizeStaysExactThroughDuplicatesAndMerges` and
`maxByteSizeStaysExactForFixedWidthDuplicates` assert the tracked max equals
the exact serialized length through duplicates, merges and clones.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
