pzhdfy commented on code in PR #68988: URL: https://github.com/apache/doris/pull/68988#discussion_r4258234501
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoPaimonScan.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.doris.nereids.rules.rewrite; + +import org.apache.doris.catalog.Env; +import org.apache.doris.datasource.credentials.CredentialUtils; +import org.apache.doris.datasource.credentials.VendedCredentialsFactory; +import org.apache.doris.datasource.paimon.PaimonExternalCatalog; +import org.apache.doris.datasource.paimon.PaimonExternalTable; +import org.apache.doris.datasource.paimon.source.PaimonVectorSearch; +import org.apache.doris.datasource.property.storage.StorageProperties; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.AnnTopNInfo; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.thrift.TVectorMetric; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Maps; +import org.apache.paimon.CoreOptions; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.Table; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * Push a primary-key vector (ANN) top-N down into the Paimon scan (index-only). + * + * <p>Recognizes {@code TopN[order by dis] -> Project[dis := dist_fn(vec, q)] -> [Filter] -> PaimonScan} + * where {@code dist_fn} is {@code l2_distance_approximate} (ASC) or + * {@code inner_product_approximate} (DESC) and the Paimon table has a PK-vector + * index on {@code vec} with a matching distance metric. + * + * <p>The rewrite mirrors the internal-table ANN path (see + * {@code PushDownVectorTopNIntoOlapScan}): the scan is told to produce a + * reader-generated {@code __paimon_search_score_to_dis} column via the PK-vector index, and + * every occurrence of {@code dist_fn(vec, q)} is replaced by that column's slot. The + * order key slot keeps its identity, so the surrounding TopN merges/limits unchanged. + * The distance always comes from the index, never from the raw vector. + * + * <p>No arithmetic is emitted here even though paimon-rust reports a score rather than + * a distance: the BE reader undoes the metric-dependent score transform in place while + * filling the block, exactly as the internal table sqrt's faiss's squared distance inside + * the index layer. So the column arrives already holding what {@code dist_fn} would + * return, and the plan stays a bare slot reference -- which also keeps {@code dis} + * non-nullable (an expression over the score would not be, since {@code Divide} and + * {@code Sqrt} are always nullable) and makes predicates like + * {@code dist_fn(vec, q) < 10} correct without any direction flip. + * + * <p>Being index-only is just an optimization here, not a precondition: if nothing + * else references the raw vector column it is left unreferenced and dropped by column + * pruning (pure index-only read); if the query still needs it ({@code SELECT vec} or a + * predicate on {@code vec}), it stays a normal projected column and the BE vector + * reader materializes it alongside the score. Either way ANN is applied -- the rewrite + * never falls back merely because the vector is read. + */ +public class PushDownVectorTopNIntoPaimonScan implements RewriteRuleFactory { + + @Override + public List<Rule> buildRules() { + return ImmutableList.of( + logicalTopN(logicalProject(logicalFileScan())).when(t -> t.getOrderKeys().size() == 1).then(topN -> { + LogicalProject<LogicalFileScan> project = topN.child(); + LogicalFileScan scan = project.child(); + return pushDown(topN, project, scan, Optional.empty()); + }).toRule(RuleType.PUSH_DOWN_VECTOR_TOPN_INTO_PAIMON_SCAN), + logicalTopN(logicalProject(logicalFilter(logicalFileScan()))) + .when(t -> t.getOrderKeys().size() == 1).then(topN -> { + LogicalProject<LogicalFilter<LogicalFileScan>> project = topN.child(); + LogicalFilter<LogicalFileScan> filter = project.child(); + LogicalFileScan scan = filter.child(); + return pushDown(topN, project, scan, Optional.of(filter)); + }).toRule(RuleType.PUSH_DOWN_VECTOR_TOPN_INTO_PAIMON_SCAN) + ); + } + + private Plan pushDown( + LogicalTopN<?> topN, + LogicalProject<?> project, + LogicalFileScan scan, + Optional<LogicalFilter<?>> optionalFilter) { + // Only Paimon external tables carry PK-vector indexes. + if (!(scan.getTable() instanceof PaimonExternalTable)) { + return null; + } + + // The index-only search path is executed exclusively by the BE paimon-rust reader, + // which only FileScannerV2 can select (the V1 FileScanner explicitly rejects + // PAIMON_RUST). If the rust reader or the V2 scanner is disabled, the split would + // fall through to the JNI reader, which ignores the vector payload and would + // return a non-ANN scan with an empty distance column. Fall back to a normal + // scan then. + ConnectContext ctx = ConnectContext.get(); + if (ctx == null || ctx.getSessionVariable() == null + || !ctx.getSessionVariable().isEnablePaimonRustReader() + || !ctx.getSessionVariable().enableFileScannerV2) { + return null; + } + + // An incremental scan (t@incr / BETWEEN) must not be rewritten into a + // PK-vector search: the VectorScan builder has no incremental mode, so a + // vector payload would silently return snapshot rows for a query that asked + // for changes. Keeping the query un-rewritten runs it as the ordinary + // incremental scan the user requested, with the distance computed in Doris. + // (The scan node's vector gate keeps !isIncremental as the authoritative + // backstop for any path that produces a vector split regardless.) + if (scan.getScanParams().isPresent() && scan.getScanParams().get().incrementalRead()) { + return null; + } + + // The order key must be a SlotReference produced by a Project alias. + Expression orderKey = topN.getOrderKeys().get(0).getExpr(); + if (!(orderKey instanceof SlotReference)) { + return null; + } + SlotReference keySlot = (SlotReference) orderKey; + Alias orderKeyAlias = null; + Expression orderKeyExpr = null; + for (NamedExpression projection : project.getProjects()) { + if (projection.toSlot().equals(keySlot) && projection instanceof Alias) { + orderKeyAlias = (Alias) projection; + orderKeyExpr = orderKeyAlias.child(); + break; + } + } + if (orderKeyExpr == null) { + return null; + } + + // l2 must sort ascending (smaller distance = closer); inner_product descending. + boolean isAsc = topN.getOrderKeys().get(0).isAsc(); + boolean l2Dist = orderKeyExpr instanceof L2DistanceApproximate; + boolean innerProduct = orderKeyExpr instanceof InnerProductApproximate; + if (!(l2Dist && isAsc) && !(innerProduct && !isAsc)) { + return null; + } + + Expression left; + Expression right; + if (l2Dist) { + left = ((L2DistanceApproximate) orderKeyExpr).left(); + right = ((L2DistanceApproximate) orderKeyExpr).right(); + } else { + left = ((InnerProductApproximate) orderKeyExpr).left(); + right = ((InnerProductApproximate) orderKeyExpr).right(); + } + while (left instanceof Cast) { + left = ((Cast) left).child(); Review Comment: Agreed — stripping the cast chain searched the original vectors while the query meant the casted ones, changing both the ranking and the reported distance. The lossy ARRAY<FLOAT> -> ARRAY<INT> -> ARRAY<FLOAT> round trip is not simplified away, so it was reachable. ▎ ▎ Fixed in 4790f5a: the rewrite no longer strips casts at all; it accepts a direct indexed column slot only. The lossy-array-cast negative rewrite case is added to PushDownVectorTopNIntoPaimonScanTest. -- 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]
