xiangfu0 commented on code in PR #20073:
URL: https://github.com/apache/pinot/pull/20073#discussion_r4407103935
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/reader/VectorIndexReader.java:
##########
@@ -31,6 +33,17 @@ public interface VectorIndexReader extends IndexReader {
/// @return bitmap of top k closest vectors
ImmutableRoaringBitmap getDocIds(float[] vector, int topK);
+ /// Returns the distance of each of the top k closest vectors from the given
vector, computed by the index, keyed by
+ /// doc id. Distances use the column's configured distance function and
match the corresponding distance function
Review Comment:
**MAJOR (API contract):** this javadoc says distances use the column's
*configured* function. But `IndexVectorDistanceTransformFunction` always treats
them as cosine and falls back to `cosineDistance`. A future EUCLIDEAN or IVF
implementation that follows this contract would make VECTOR_DISTANCE return L2
values under a cosine label.
The generic SQL name has the same problem today: on EUCLIDEAN or
INNER_PRODUCT columns it silently ranks by cosine. Could you either make the
contract and the name cosine-specific, or dispatch on or validate the column's
distance function? A SQL name is hard to change after release.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/HnswVectorIndexReader.java:
##########
@@ -64,11 +68,24 @@ public class HnswVectorIndexReader implements
FilterAwareVectorIndexReader, EfSe
private final ThreadLocal<Integer> _efSearchOverride = new ThreadLocal<>();
private final ThreadLocal<Boolean> _useRelativeDistanceOverride = new
ThreadLocal<>();
private final ThreadLocal<Boolean> _useBoundedQueueOverride = new
ThreadLocal<>();
+ /// Lucene returns cosine similarity scores, which convert exactly to
`cosineDistance`; other functions use scaled
+ /// scores, so distances are only served for COSINE.
+ private final boolean _servesDistances;
+ /// Distances from the most recent searches, so a projection right after the
filter reuses the filter's search.
+ private final Map<Integer, RecentSearch> _recentSearches =
Collections.synchronizedMap(
Review Comment:
**MAJOR (heap retention):** this cache lives for the life of the segment.
`searchAndRemember` fills it on every `VECTOR_SIMILARITY` search of a COSINE
index, even when no query uses VECTOR_DISTANCE. Measured locally: about 212 KB
per segment once its 64 entries fill (768-dim, k=10). That is roughly 200 MB
per 1,000 segments, and more at 1536 dims or with rerank candidate pools. The
cap counts entries, not bytes, and `close()` never clears it.
Could the distances be scoped to the query instead? For example, the filter
operator could hand them to the transform, or the reader could remember a
search only when the query projects VECTOR_DISTANCE.
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IndexVectorDistanceTransformFunction.java:
##########
@@ -0,0 +1,148 @@
+/**
+ * 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 it.unimi.dsi.fastutil.ints.Int2FloatMap;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.function.scalar.VectorFunctions;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.operator.blocks.ValueBlock;
+import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.index.reader.VectorIndexReader;
+
+
+/// Implements `VECTOR_DISTANCE(column, ARRAY[...], k)`: the cosine distance
of each document from the query vector,
+/// served by the column's vector index for the `k` documents closest to the
query vector.
+///
+/// Paired with `VECTOR_SIMILARITY(column, ARRAY[...], k)` on the same vector
and `k`, it returns the values of
+/// `cosineDistance(column, ARRAY[...])` (up to the index's single-precision
arithmetic) without reading the vectors
+/// of the matched documents from the forward index: the distances come from
the index search the filter already ran.
+/// A document the index search did not return (e.g. when the column's index
does not serve distances) falls back to
+/// computing the distance from the forward index.
+///
+/// Not thread-safe: like other transform functions, an instance serves one
segment on one thread.
+public class IndexVectorDistanceTransformFunction extends
BaseTransformFunction {
+ public static final String FUNCTION_NAME = "vectorDistance";
+
+ private TransformFunction _columnTransformFunction;
+ @Nullable
+ private VectorIndexReader _vectorIndexReader;
+ private float[] _queryVector;
+ private int _topK;
+ @Nullable
+ private Int2FloatMap _distances;
+ private boolean _distancesLoaded;
+
+ @Override
+ public String getName() {
+ return FUNCTION_NAME;
+ }
+
+ @Override
+ public void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap) {
+ super.init(arguments, columnContextMap);
+ if (arguments.size() != 3) {
+ throw new IllegalArgumentException(
+ "Expected 3 arguments for VECTOR_DISTANCE: VECTOR_DISTANCE(column,
ARRAY[...], topK)");
+ }
+ _columnTransformFunction = arguments.get(0);
+ if (!(_columnTransformFunction instanceof IdentifierTransformFunction)
+ || _columnTransformFunction.getResultMetadata().isSingleValue()) {
+ throw new IllegalArgumentException("The first argument of
VECTOR_DISTANCE must be a multi-valued column");
+ }
+ String columnName = ((IdentifierTransformFunction)
_columnTransformFunction).getColumnName();
+ DataSource dataSource = columnContextMap.get(columnName).getDataSource();
+ _vectorIndexReader = dataSource != null ? dataSource.getVectorIndex() :
null;
+
+ TransformFunction vectorArgument = arguments.get(1);
+ if (!(vectorArgument instanceof ArrayLiteralTransformFunction)) {
+ throw new IllegalArgumentException("The second argument of
VECTOR_DISTANCE must be an array literal");
+ }
+ _queryVector = toFloatArray((ArrayLiteralTransformFunction)
vectorArgument);
+
+ TransformFunction topKArgument = arguments.get(2);
+ if (!(topKArgument instanceof LiteralTransformFunction)) {
+ throw new IllegalArgumentException("The third argument of
VECTOR_DISTANCE must be an integer literal");
+ }
+ _topK = ((LiteralTransformFunction) topKArgument).getIntLiteral();
+ if (_topK <= 0) {
+ throw new IllegalArgumentException("The third argument of
VECTOR_DISTANCE must be positive, got: " + _topK);
+ }
+ }
+
+ /// Float array literals are rare in SQL (`ARRAY[0.1, ...]` parses as
DOUBLE); narrow like VECTOR_SIMILARITY does.
+ private static float[] toFloatArray(ArrayLiteralTransformFunction literal) {
+ float[] floats = literal.getFloatArrayLiteral();
+ if (floats != null) {
+ return floats;
+ }
+ double[] doubles = literal.getDoubleArrayLiteral();
+ if (doubles == null) {
+ throw new IllegalArgumentException("The second argument of
VECTOR_DISTANCE must be a numeric array literal");
Review Comment:
**MINOR:** an integer array literal such as `ARRAY[1, 0, 0]` becomes an INT
array literal, so neither the float nor the double getter is set and this
throws. `VECTOR_SIMILARITY` accepts INT and LONG arrays
(`RequestContextUtils.getVectorValue`), so the same vector works in the filter
but fails here.
--
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]
