pzhdfy commented on code in PR #69493: URL: https://github.com/apache/doris/pull/69493#discussion_r4358676778
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoPaimonScan.java: ########## @@ -0,0 +1,756 @@ +// 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.analysis.TableScanParams; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.Type; +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.PaimonUtil; +import org.apache.doris.datasource.paimon.source.PaimonScanNode; +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.ComparisonPredicate; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.InPredicate; +import org.apache.doris.nereids.trees.expressions.IsNull; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Not; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.NoneMovableFunction; +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.expressions.literal.NullLiteral; +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.ImmutableSet; +import com.google.common.collect.Maps; +import org.apache.paimon.CoreOptions; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.Table; +import org.apache.paimon.types.ArrayType; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.DataTypeRoot; +import org.apache.paimon.types.MapType; +import org.apache.paimon.types.MultisetType; +import org.apache.paimon.types.RowType; + +import java.util.HashSet; +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 { + + /** Doris column types the BE paimon-rust predicate converter can push; see + * PaimonRustPredicateConverter._is_supported_slot_type on the BE side. */ + private static final Set<Type> VECTOR_FILTER_SUPPORTED_COLUMN_TYPES = ImmutableSet.of( + Type.BOOLEAN, Type.TINYINT, Type.SMALLINT, Type.INT, Type.BIGINT, + Type.STRING, Type.DATE, Type.DATEV2); + + /** The BE rust predicate converter encodes the IN option list under a 16-bit + * child count; a bigger list is dropped whole into the post-fetch residual. */ + private static final int MAX_IN_PREDICATE_CHILDREN = 0xFFFF; + + @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.) + // + // The same holds for a pinned file-creation-time read state: the ordinary + // path enforces `scan.file-creation-time-millis` by dropping manifest + // entries older than the bound, but the vector search plans its own splits + // and never consults that bound, so old files' nearest neighbors could leak + // into the result. Snapshot/tag/version selectors are fine -- they are + // resolved into the table handle before the vector search -- so only the + // file-level bound forces the decline. + // + // `scan.creation-time-millis` must decline as well: when a snapshot exists + // past the bound it resolves into an ordinary snapshot pin (safe), but when + // none does — the exact read state the option exists for — PaimonScanParams + // resolves it into the SAME internal file-creation bound (resolveOptions + // falls through to resolveFileCreationTime), which the vector path again + // never consults. Which case applies depends on the snapshot manager's + // state, so the raw option key is declined unconditionally. + if (scan.getScanParams().isPresent()) { + TableScanParams scanParams = scan.getScanParams().get(); + if (scanParams.incrementalRead()) { + return null; + } + Map<String, String> rawParams = scanParams.getMapParams(); + if (rawParams.containsKey(CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS.key()) + || rawParams.containsKey(CoreOptions.SCAN_CREATION_TIME_MILLIS.key())) { + return null; + } + } + + // A bounded vector top-N is only correct when every filter conjunct is + // applied inside the search, before the limit is exhausted. The BE + // paimon-rust predicate converter skips conjuncts it cannot represent (and + // the synthetic distance column is not in its field registry at all), and + // the engine re-applies those residuals on the returned block -- after the + // top-K rows were already selected. A skipped residual can therefore drop + // rows out of a bounded result: with candidate distances 0, 2 and 3, a + // `WHERE dis > 1 ORDER BY dis LIMIT 1` would fetch only the distance-0 row, + // filter it out and return nothing instead of the distance-2 row. So the + // rewrite only fires when each conjunct is provably of a form the converter + // pushes: a simple comparison / IN / null-check on a table column whose + // Doris type is in the converter's supported set (see + // PaimonRustPredicateConverter._is_supported_slot_type on BE). Anything + // else -- an expression predicate, a predicate on the distance itself, an + // unsupported column type -- declines the rewrite and the query runs as an + // ordinary scan with the distance computed in Doris, which is always + // correct. + if (optionalFilter.isPresent() && !conjunctsRideIntoVectorSearch(optionalFilter.get())) { + 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; + } + + // The indexed column must be a direct slot reference. A Cast chain over the + // column is NOT stripped: a legal cast can change vector values (for example + // an ARRAY<FLOAT> -> ARRAY<INT> -> ARRAY<FLOAT> round trip loses precision), + // and searching the original vectors while the query means the casted ones + // changes both the distance and the row order. Doris's implicit coercions do + // not produce a cast here for a matching ARRAY<FLOAT> column, so anything + // wrapped in a Cast simply declines the rewrite. + Expression left; + Expression right; + if (l2Dist) { + left = ((L2DistanceApproximate) orderKeyExpr).left(); + right = ((L2DistanceApproximate) orderKeyExpr).right(); + } else { + left = ((InnerProductApproximate) orderKeyExpr).left(); + right = ((InnerProductApproximate) orderKeyExpr).right(); + } + if (!(left instanceof SlotReference && right.isConstant() && right instanceof ArrayLiteral)) { + return null; + } + // A typed NULL inside the query vector ([CAST(NULL AS FLOAT), 0.0, ...]) is a + // NullLiteral whose getDouble() returns 0, so extracting the coordinates + // would silently search against a zero vector. The ordinary path raises + // instead (FunctionArrayDistance::_validate_no_nulls: "Second argument ... + // cannot have null"), and the rewrite must not replace that error with rows + // for a different input, so any null element declines the rewrite. Both + // supported metrics share this. + if (((ArrayLiteral) right).getValue().stream().anyMatch(NullLiteral.class::isInstance)) { + 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 relation's bound Paimon + // table: a branch table carries its own index options and can differ from + // the main table (see boundPaimonTable). + PaimonExternalTable paimonTable = (PaimonExternalTable) scan.getTable(); + CoreOptions coreOptions; + try { + Table basePaimonTable = boundPaimonTable(paimonTable, scan); + 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; + } + + // A data predicate additionally requires the SDK's "physical row read" + // table mode: the predicate itself converts fine, but the pinned + // paimon-rust PkVectorSearchParams::resolve rejects every DATA predicate + // unless deletion vectors are enabled and merge-on-read is disabled — + // so the failure lands AFTER this one-way rewrite replaced the distance + // expression, turning a query that the ordinary path answers into a + // user-facing error. Partition-only conjuncts are enforced by the search + // planning itself and stay fine, so decline only on a non-partition one. + // + // The synthetic distance slot is named after a legal paimon user column + // name (__paimon_search_score_to_dis, see the BE reader). A table that + // actually HAS a column of that name would collide with the synthetic + // slot's block-column matching; decline and run the ordinary path. + // + // A complex-typed column (STRUCT/MAP/nested ARRAY) must ride as a whole + // column: nested pruning rewrites a slot whose fields are accessed — + // Project(struct_element(payload, 'b')) materializes the scan tuple as + // ROW<b> — but the vector reader projects only root names and returns + // the full Arrow value, and the struct SerDe then binds children by + // ORDINAL: the wrong child comes back. Filter conjuncts cannot reach + // this (the whitelist admits only scalar columns); the order-key alias + // is replaced by the synthetic slot and cannot reach it either; so only + // a Project output expression over a complex column has to decline. + // + // A NoneMovableFunction in the projection (assert_true) must keep + // evaluating below the TopN — see projectContainsNonMovableExpression. + // + // A historically evolved schema must stay off the vector reader: its + // materialization decodes projected values through Arrow safe casts, + // which can silently NULL out-of-range values (a TIMESTAMP(3) widened + // to TIMESTAMP(9) reading a year-2300 value) where the ordinary path + // preserves them — see PaimonRustReaderCapabilities.hasRustIncompatibleSchemaHistory. + try { + Table table = boundPaimonTable(paimonTable, scan); + if ((optionalFilter.isPresent() + && filterHasDataConjunct(optionalFilter.get(), table) + && !PaimonScanNode.vectorSearchSupportsDataPredicate(table)) + || table.rowType().getFieldNames().stream().anyMatch( + name -> name.equalsIgnoreCase(PaimonVectorSearch.SEARCH_DISTANCE_COLUMN)) Review Comment: Confirmed. `__paimon_search_score` and `_PKEY_VECTOR_POSITION` join the synthetic-distance collision check on the table's row-type field names, so a user column named either declines before the one-way rewrite can hand rust a table it refuses to open. Covered by `testReservedMaterializerNamesNotPushedDown`. -- 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]
