pzhdfy commented on code in PR #68988: URL: https://github.com/apache/doris/pull/68988#discussion_r4258238783
########## 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(); + } + if (!(left instanceof SlotReference && right.isConstant() && right instanceof ArrayLiteral)) { + return null; + } + SlotReference vectorSlot = (SlotReference) left; + if (!vectorSlot.getOriginalColumn().isPresent() || !vectorSlot.getOriginalTable().isPresent()) { + return null; + } + String annColumnName = vectorSlot.getOriginalColumn().get().getName(); + + // Validate the PK-vector index and metric on the Paimon table. + PaimonExternalTable paimonTable = (PaimonExternalTable) scan.getTable(); + CoreOptions coreOptions; + try { + Table basePaimonTable = paimonTable.getPaimonTable(Optional.empty()); + coreOptions = CoreOptions.fromMap(basePaimonTable.options()); + } catch (Throwable t) { + return null; + } + if (!coreOptions.primaryKeyVectorIndexEnabled()) { + return null; + } + boolean columnIndexed = coreOptions.primaryKeyVectorIndexColumns().stream() + .anyMatch(c -> c.equalsIgnoreCase(annColumnName)); + if (!columnIndexed) { + return null; + } + // v1 supports only l2 / inner_product; the function must match the index metric. + String metric = coreOptions.primaryKeyVectorDistanceMetric(annColumnName); + String expectedMetric = l2Dist ? "l2" : "inner_product"; + if (!expectedMetric.equals(metric)) { + return null; + } + // The thrift-level enum the AnnTopNInfo will carry. expectedMetric is kept as a + // string for the comparison above (coreOptions returns a string), but the + // AnnTopNInfo / wire contract uses TVectorMetric, so map once here. DOT_PRODUCT + // is the thrift enum that paimon-rust's "inner_product" score transform maps to + // (see the BE reader's parse_score_transform). + TVectorMetric expectedMetricEnum = l2Dist ? TVectorMetric.L2 : TVectorMetric.DOT_PRODUCT; + + // Pre-flight the rust eligibility the scan node will gate on. The rewrite is + // one-way: a PK-vector split has no JNI fallback (the JNI reader ignores the + // vector payload and would return a non-ANN scan with an empty distance + // column), so the scan node can only fail loudly — a table whose gate would + // fail (credential provider shape, REST-token IO, query-auth, merge-option + // matrix, deletion-vector merge-on-read, scan-mode, storage scheme or HDFS + // credentials, no rust-capable backend) must not be rewritten at all. Without + // the rewrite the query still runs, computing the distance in Doris. The + // projection-dependent bounds (ORC LTZ, projected variant) need the final scan + // tuple and stay scan-node-side, and the scan node re-evaluates on the + // relation-option-processed table, so it stays authoritative for t@options + // overrides this raw-table pre-flight cannot see. + try { + Table baseTable = paimonTable.getPaimonTable(Optional.empty()); + PaimonExternalCatalog catalog = (PaimonExternalCatalog) paimonTable.getCatalog(); + Map<StorageProperties.Type, StorageProperties> storagePropertiesMap = + VendedCredentialsFactory.getStoragePropertiesMapWithVendedCredentials( + catalog.getCatalogProperty().getMetastoreProperties(), + catalog.getCatalogProperty().getStoragePropertiesMap(), + baseTable); + Map<String, String> backendStorageProperties = + CredentialUtils.getBackendPropertiesFromStorageMap(storagePropertiesMap); + String tableLocation = baseTable instanceof FileStoreTable + ? ((FileStoreTable) baseTable).location().toString() + : baseTable.options().get("path"); + if (!PaimonVectorSearch.vectorScanEligible(baseTable, tableLocation, Review Comment: Right — the rewrite was one-way (a PK-vector split has no JNI fallback), so a projectedVariant query would hit the scan-node guard as an error instead of falling back. We also found the same shape on the orc-ltz bound: an ORC table with TIMESTAMP_LTZ in its schema sets canUseRust=false at the scan node and throws the same way. ▎ ▎ Fixed in 4790f5a: both bounds are pre-flighted before the expression is replaced — the projected columns are checked for Variant (nested included, mirroring the scan node's PaimonUtil.containsVariant) and the table for a non-parquet schema containing TIMESTAMP_LTZ. Either hit declines the rewrite, keeping the query on an ordinary supported path. The Variant-projection decline is covered by a UT case (testVariantProjectionNotPushedDown); a combined ANN + VARIANT e2e case needs a Spark 4.0 ThriftServer for VARIANT DDL, which our local 3.5 setup cannot build — the UT covers the rule decision until then. -- 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]
