This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch release-1.6.0-rc in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 404e0b55c610cd06b9ada403f6d1c3da3ecc7af8 Author: Gonzalo Ortiz <[email protected]> AuthorDate: Thu Oct 8 16:06:15 2026 +0200 Store QueryContext shared values in a per-query map (DATA-3823) Backport of the QueryContext part of apache/pinot#19759 (1ea817199a). QueryContext kept its shared values (getOrComputeSharedValue) in a ClassValue (MemoizedClassAssociation). A ClassValue stores its values on the Class object, in a WeakHashMap keyed by the ClassValue. When a stored value strongly references its QueryContext, the weak key stays reachable, so the entry and the whole QueryContext are never collected. StarTree FilterAheadQueryState and AiQueryRuntime store such values, so every server query leaked its QueryContext. QueryContext now owns a ConcurrentHashMap<Class<?>, ConcurrentHashMap> and the shared values are collected with it. Only the type of the private _sharedValues field changes, so the change is binary compatible. MemoizedClassAssociation stays in place to keep the hotfix minimal. QueryContextTest stores a shared value that references its QueryContext and verifies that GC collects the QueryContext. It fails without the fix. Co-Authored-By: Claude Opus 5.5 <[email protected]> --- .../core/query/request/context/QueryContext.java | 8 +-- .../query/request/context/QueryContextTest.java | 61 ++++++++++++++++++++++ 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java index 1371fd3f3c7..41c12e6adce 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java @@ -40,7 +40,6 @@ import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; import org.apache.pinot.core.query.aggregation.function.AggregationFunction; import org.apache.pinot.core.query.aggregation.function.AggregationFunctionFactory; import org.apache.pinot.core.util.GroupByUtils; -import org.apache.pinot.core.util.MemoizedClassAssociation; import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.data.Schema; @@ -80,7 +79,9 @@ public class QueryContext { private final Map<ExpressionContext, ExpressionContext> _expressionOverrideHints; private final ExplainMode _explain; - private final Function<Class<?>, Map<?, ?>> _sharedValues = MemoizedClassAssociation.of(ConcurrentHashMap::new); + // Values that the segments of the query share, by value type. The query context owns the map, so the values can be + // collected with it. A ClassValue would keep each value in its type's cache until later queries clear it. + private final ConcurrentHashMap<Class<?>, ConcurrentHashMap<?, ?>> _sharedValues = new ConcurrentHashMap<>(); // Pre-calculate the aggregation functions and columns for the query so that it can be shared across all the segments private AggregationFunction[] _aggregationFunctions; @@ -560,7 +561,8 @@ public class QueryContext { /// @param <V> the value type /// @return the shared value public <K, V> V getOrComputeSharedValue(Class<V> type, K key, Function<K, V> mapper) { - return ((ConcurrentHashMap<K, V>) _sharedValues.apply(type)).computeIfAbsent(key, mapper); + return ((ConcurrentHashMap<K, V>) _sharedValues.computeIfAbsent(type, k -> new ConcurrentHashMap<>())) + .computeIfAbsent(key, mapper); } public int getEffectiveSegmentGroupTrimSize() { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/QueryContextTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/QueryContextTest.java new file mode 100644 index 00000000000..df4f95b3888 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/request/context/QueryContextTest.java @@ -0,0 +1,61 @@ +/** + * 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.query.request.context; + +import java.lang.ref.WeakReference; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; + + +public class QueryContextTest { + + @Test + public void testSharedValueReferencingQueryContextDoesNotLeak() + throws InterruptedException { + WeakReference<QueryContext> queryContextRef = createQueryContextWithSelfReferencingSharedValue(); + + // A shared value that references its query context must not keep the query context reachable + for (int i = 0; i < 100 && queryContextRef.get() != null; i++) { + System.gc(); + Thread.sleep(10); + } + assertNull(queryContextRef.get(), "QueryContext is still reachable after GC"); + } + + private static WeakReference<QueryContext> createQueryContextWithSelfReferencingSharedValue() { + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable"); + SelfReferencingValue value = + queryContext.getOrComputeSharedValue(SelfReferencingValue.class, "key", k -> new SelfReferencingValue( + queryContext)); + assertSame(queryContext.getOrComputeSharedValue(SelfReferencingValue.class, "key", k -> null), value); + return new WeakReference<>(queryContext); + } + + private static final class SelfReferencingValue { + @SuppressWarnings("unused") + private final QueryContext _queryContext; + + SelfReferencingValue(QueryContext queryContext) { + _queryContext = queryContext; + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
