github-actions[bot] commented on code in PR #68988: URL: https://github.com/apache/doris/pull/68988#discussion_r4322110266
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoPaimonScan.java: ########## @@ -0,0 +1,651 @@ +// 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.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); + + @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 Paimon table. + PaimonExternalTable paimonTable = (PaimonExternalTable) scan.getTable(); + CoreOptions coreOptions; + try { + Table basePaimonTable = paimonTable.getPaimonTable(Optional.empty()); Review Comment: [P1] Validate the index against the relation's bound table `getPaimonTable(Optional.empty())` fetches the cached main table, while `PaimonSource` and vector split planning use this relation's retained branch/snapshot. For example: ```text TopN(dis ASC LIMIT 1) Project(id, dis := l2_distance_approximate(vec, [0.25,0,...])) Scan(t@branch(b)) ``` Main can have an L2 index while branch `b` has an independently configured inner-product index. With branch vectors `[1,0,...]` and `[2,0,...]` in one bucket, this preflight accepts main's L2 metric, but Rust reads the branch schema and selects the second vector with score `0.5`. BE then applies the requested L2 inverse and publishes distance `1`; ordinary L2 selects the first vector at `0.75`. Full recall still gives this wrong result. Resolve all index/metric/eligibility checks from the scan's bound table before replacing the expression, and cover differing branch index options. ########## fe/fe-core/src/main/java/org/apache/doris/datasource/FileQueryScanNode.java: ########## @@ -379,6 +380,12 @@ private void setColumnPositionMapping() if (colName.startsWith(Column.GLOBAL_ROWID_COL)) { continue; } + // Reader-produced Paimon PK-vector search distance: a synthetic output column + // filled by the BE reader by name, not present in the table's file schema, + // so it has no Parquet/ORC column position to map. + if (PaimonVectorSearch.SEARCH_DISTANCE_COLUMN.equalsIgnoreCase(colName)) { Review Comment: [P2] Limit this mapping exception to an actual synthetic slot This method serves every external catalog, so the name alone also excludes a real Hive column. For a Hive text/CSV table `(id INT, __paimon_search_score_to_dis STRING)`, selecting both columns leaves two required file slots but emits only `column_idxs=[0]`. With FileScannerV2, `DelimitedTextReader::_init_reader` then rejects the index-count mismatch; Hive positional ORC/Parquet mapping likewise lacks the second ordinal. Preserve the mapping for physical same-named columns and identify the actual Paimon vector-produced slot. This is a separate generic FE/native-scan instance from the already-discussed ordinary Rust `_build_read_columns` exclusion. ########## fe/fe-core/src/main/java/org/apache/doris/datasource/paimon/source/PaimonScanNode.java: ########## @@ -982,72 +1029,174 @@ && splitHasOrcFile((DataSplit) split) // rust. Gate until the rust leaf has a Variant Arrow decoder. boolean projectedVariant = desc.getSlots().stream() .anyMatch(slot -> PaimonUtil.containsVariant(slot.getType())); - // Scheme capability gate: the pinned paimon-rust storage - // dispatcher (io/storage.rs) selects the FileIO parser from the - // table location's URI scheme, and libpaimon_c.a compiles in - // separate COS, OBS, GCS and Azdls parsers besides the OSS and S3 - // ones. Doris normalizes every object store's credentials into - // the AWS_* / use_path_style aliases (see the *Properties storage - // classes), which the BE rust bridge translates only into the - // fs.oss.* and s3.* key families — a cosn:// / obs:// / gs:// / - // abfs:// warehouse would reach its scheme's parser without the - // key family it reads (fs.cosn.userinfo.*, fs.obs.*, gcs.*, - // azure.*) and fail the open instead of using JNI. Only the - // schemes whose property translation is implemented and - // open-tested (s3 / s3a / oss, via RUST_VERIFIED_LOCATION_SCHEMES) - // plus the credential-free hdfs and local-filesystem parsers stay - // rust-eligible; every other scheme falls back to JNI. A null - // location also routes to JNI: the rust path needs the - // paimon_table that only a real location can provide (BE rejects - // a split without it). - boolean schemeCapabilityVerified = isRustVerifiedLocationScheme(source.getTableLocation()); - // An hdfs:// location is scheme-verified only together with the credential-free - // backend shape: the backend storage properties that ship to BE also carry an - // HDFS catalog's authentication (kerberos principal / keytab, proxy user, HA - // nameservice config), none of which the pinned rust HDFS parser reads — the - // scan would open as the BE process's ambient identity instead of the - // catalog's configured one and fail the access JNI honors. See - // isRustVerifiedHdfsBackend. - boolean hdfsBackendVerified = !isHdfsLocationScheme(source.getTableLocation()) - || isRustVerifiedHdfsBackend(backendStorageProperties); - // With merge-on-read=true the whole table already routes to JNI (dvMergeOnRead); - // for the remaining partial-update/aggregation DV tables, a split that is + // With merge-on-read=true the whole table already routes to JNI + // (rustEligibility.dvMergeOnRead); for the remaining + // partial-update/aggregation DV tables, a split that is // not fully materialized (uncompacted level-0 data, or retractions not // known to be applied — even a split with no deletion file attached yet) // fails the rust is_fully_materialized_pk_dv guard, so it falls back per // split instead of turning into a BE-open failure. nativeSplit and // !fallbackRead only guard the cast — those splits already route to JNI. - boolean splitDvNotMaterialized = puAggDeletionVectors && !dvMergeOnRead + boolean splitDvNotMaterialized = rustEligibility.puAggDeletionVectors + && !rustEligibility.dvMergeOnRead && nativeSplit && !fallbackRead && !isFullyMaterializedPkDvSplit((DataSplit) split); + boolean vectorSearchSplit = paimonSplit.getVectorPayloadBytes() != null; + // Session, shared table / storage / backend, and scan-shape gates. The + // shared PaimonRustEligibility computation is also the + // PushDownVectorTopNIntoPaimonScan rewrite's pre-flight + // (PaimonVectorSearch.vectorScanEligible): the rewrite is one-way, so both + // sites must agree on the same gates, and the scan node stays authoritative + // — it evaluates the relation-option-processed table, so a t@options + // override the raw-table pre-flight cannot see still rejects here (loudly, + // as the vector path has no JNI fallback). !isIncremental and the + // split-half of the fallback-read gate are the same backstops: an + // incremental scan has no PK-vector semantics (the VectorScan builder has + // no incremental mode), and a FallbackSplit shape cannot ride the vector + // payload either — both fail loudly here rather than silently returning + // snapshot rows for a query that asked for changes. Each gate's value is + // also what the fallback logging below reports. boolean canUseRust = sessionVariable.isEnablePaimonRustReader() - && sessionVariable.enableFileScannerV2 && nativeSplit && !fallbackRead - && !isIncremental && providerModeTranslatable && !queryAuthTable - && !restTokenTable - && !dvMergeOnRead && !splitDvNotMaterialized - && !orcLtzSchema && !projectedVariant && !externalFileSplit - && !deduplicateIgnoreDelete && !rustUnsupportedMergeOption - && schemeCapabilityVerified && hdfsBackendVerified - && paimonFileStoreTable != null; - if (canUseRust) { - // Branch tables can retain persisted modes without a selector. Validate the - // transported options too: the pinned Rust builder only accepts default here. - String scanMode = PaimonScanParams.withoutTimeTravelSelectors(paimonFileStoreTable.schema()) - .options().get(CoreOptions.SCAN_MODE.key()); - canUseRust = scanMode == null || "default".equalsIgnoreCase(scanMode); - } - if (canUseRust) { - // Every candidate can receive this range; older BEs cannot decode reader type 3. - canUseRust = !backendPolicy.getBackends().isEmpty() - && backendPolicy.getBackends().stream().allMatch(Backend::isPaimonRustReaderSupported); - } - if (canUseRust) { - if (rustReaderCapabilities == null) { - rustReaderCapabilities = new PaimonRustReaderCapabilities(paimonFileStoreTable, desc); + && sessionVariable.enableFileScannerV2 + && rustEligibility.eligible() + && !fallbackRead && !isIncremental + && !orcLtzSchema && !projectedVariant; + // The remaining gates judge the DataSplit the BATCH reader opens: the + // native DataSplit shape, unmaterialized partial-update / aggregation + // deletion vectors, external-path files, and the + // PaimonRustReaderCapabilities whitelist. A PK-vector search split is + // never read as a DataSplit (it rides serialized in vector_payload and + // executes paimon-rust's eager Top-K chain, which applies deletion + // vectors and the merge inside the search itself), so it does not + // consult them. + // + // The nested-projection bound is NOT DataSplit-specific: the ABI + // projects only root names and the Arrow struct SerDe binds children by + // ordinal, so a pruned complex slot (nested pruning, on by default) + // would return the wrong child on the vector path exactly as it would + // on the batch path. The rewrite pre-flights this and declines; this + // is the loud backstop for shapes the raw-table pre-flight cannot see. + if (vectorSearchSplit) { + canUseRust = canUseRust && rustEligibility.vectorEligible() + && PaimonRustReaderCapabilities.hasFullNestedProjection(desc); + } + if (!vectorSearchSplit && canUseRust) { Review Comment: [P1] Retain historical scalar compatibility checks for vector reads `canRead` also protects historical value conversions, which remain relevant when vector search materializes projected columns. On a PK-vector Parquet table, write `ts TIMESTAMP(3) = '2300-01-01 00:00:00.123'`, then widen `ts` to `TIMESTAMP(9)` while retaining the old files/index. `SELECT ts, l2_distance_approximate(vec,q) AS d FROM t ORDER BY d LIMIT 1` now bypasses that check. Rust's materialization uses `DataFileReader` -> `reconcile_column` -> Arrow's safe millisecond-to-nanosecond cast; the year-2300 multiplication overflows and returns NULL. The ordinary path explicitly rejects this evolution in `PaimonRustReaderCapabilities` and uses JNI, which preserves the timestamp (also documented by `testPersistedTimestampEvolutionIntoNanosecondRange`). Check current/historical schema compatibility before committing ANN, or provide equivalent materialization. This is a scalar full-projection case, separate from the existing pruned-STRUCT thread. ########## be/src/format_v2/table/paimon_rust_table_reader.cpp: ########## @@ -509,6 +626,316 @@ Status PaimonRustTableReader::_open_split_reader(const TFileRangeDesc& range) { return Status::OK(); } +Status PaimonRustTableReader::_open_vector_split_reader(const TFileRangeDesc& range) { + // Validate every __isset / value-range up front (mirrors + // LanceTableReader::_validate_external_search_request). The build path below + // then uses the fields unchecked, matching lance's split of validate-once / + // build-without-recheck. + RETURN_IF_ERROR(_validate_external_search_request()); + const auto& request = _scan_params->external_search_request; + const auto& vector = request.search_query.vector_search; + + // 0. Distance metric. It decides which score transform the reader has to undo + // (see _convert_score_to_distance), so defaulting would silently return wrong + // distances. __isset was checked by _validate_external_search_request; the + // value is rejected here if unsupported. + ScoreTransform transform; + RETURN_IF_ERROR(parse_score_transform(vector.metric, &transform)); + _score_transform = transform; + + // 1. The per-split payload is the serialized BucketVectorSearchSplit (raw + // thrift binary, NOT base64; magic "PKVSPLIT", v1). Grab it first so we fail + // fast before any filesystem IO. + const auto& payload = _current_range.table_format_params.paimon_params.vector_payload; + if (payload.payload_type != TPaimonVectorPayloadType::BUCKET_SPLIT_BYTES) { + return Status::InternalError("paimon-rust vector search unsupported payload type: {}", + static_cast<int>(payload.payload_type)); + } + const std::string& split_bytes = payload.bytes; + if (split_bytes.empty()) { + return Status::InternalError("paimon-rust vector search split payload is empty"); + } + + // 2. Resolve and open (or reuse) the table (same as the normal read path); + // options carry storage credentials to the FileIO. + RETURN_IF_ERROR(_open_paimon_table(range)); + + // 3. Build and configure the vector-search builder. No read.batch-size here, + // unlike the batch path (_open_split_reader): paimon-rust's PK vector readers + // construct their DataFileReaders without with_batch_size, so a + // read.batch-size option — from the schema, t@options or these search + // options — is silently ignored on the vector path, which batches at the + // parquet-rust default. Known difference; the vector stages are + // single-column streamed into bounded Top-K heaps and the output is + // capped by top_k, so no action needed until paimon-rust wires it. + paimon_result_vector_search_builder vb_res = + paimon_table_new_vector_search_builder(_handles->table.get()); + if (vb_res.error != nullptr) { + return Status::InternalError("paimon-rust new vector search builder failed: {}", + consume_error(vb_res.error)); + } + vector_search_builder_ptr vb(vb_res.builder); + + // 3a. Target vector column. + if (paimon_error* err = + paimon_vector_search_builder_with_vector_column(vb.get(), vector.column.c_str())) { + return Status::InternalError("paimon-rust set vector column failed: {}", + consume_error(err)); + } + // 3b. Query vector. The wire form is a TSearchVector carrying a FLOAT32, + // little-endian binary blob; decode it into the f32 array the C API wants. + // v1 only supports FLOAT32 (paimon-rust's index backends operate on f32); any + // other element type was already rejected by validate. Dimension and + // byte-size were validated there too, so decode unchecked here. + const auto& search_vec = vector.query_vector; + const auto dimension = static_cast<size_t>(search_vec.dimension); + std::vector<float> query_vector(dimension); + const auto* bytes = search_vec.values.data(); + for (size_t i = 0; i < dimension; ++i) { + query_vector[i] = std::bit_cast<float>(LittleEndian::Load32( + reinterpret_cast<const uint8_t*>(bytes + i * sizeof(uint32_t)))); + } + if (paimon_error* err = paimon_vector_search_builder_with_query_vector( + vb.get(), query_vector.data(), query_vector.size())) { + return Status::InternalError("paimon-rust set query vector failed: {}", consume_error(err)); + } + // 3c. Retrieval limit (user_limit + user_offset, folded on FE). A negative + // limit would wrap when cast to uintptr_t, so guard it. + if (vector.top_k < 0) { + return Status::InternalError("paimon-rust vector search limit is negative: {}", + vector.top_k); + } + if (paimon_error* err = paimon_vector_search_builder_with_limit( + vb.get(), static_cast<uintptr_t>(vector.top_k))) { + return Status::InternalError("paimon-rust set limit failed: {}", consume_error(err)); + } + // 3d. Index/search options (refine_factor, ivf.nprobe, "<algo>.metric", ...). + if (request.__isset.paimon_options && request.paimon_options.__isset.options && + !request.paimon_options.options.empty()) { + std::vector<paimon_option> c_search_options; + c_search_options.reserve(request.paimon_options.options.size()); + for (const auto& kv : request.paimon_options.options) { + c_search_options.emplace_back(kv.first.c_str(), kv.second.c_str()); + } + if (paimon_error* err = paimon_vector_search_builder_with_options( + vb.get(), c_search_options.data(), c_search_options.size())) { + return Status::InternalError("paimon-rust set search options failed: {}", + consume_error(err)); + } + } + // 3e. Materialized user columns. The score field is always appended by + // paimon-rust itself and must NOT be listed here. Empty projection = only the + // score column (index-only scan). Log the user columns: an empty list is + // legitimate, so it is logged as such rather than skipped. + auto read_columns = _build_read_columns(); + std::vector<const char*> projection; + projection.reserve(read_columns.size() + 1); + for (const auto& col : read_columns) { + projection.push_back(col.c_str()); + } + projection.push_back(nullptr); + LOG(INFO) << "paimon-rust vector search: projection=[" << join(read_columns, ", ") << "] (" + << read_columns.size() << " column(s), distance column '" + << kPaimonSearchDistanceColumn << "' produced by the reader)"; + if (paimon_error* err = + paimon_vector_search_builder_with_projection(vb.get(), projection.data())) { + return Status::InternalError("paimon-rust set vector projection failed: {}", + consume_error(err)); + } + + // 3f. Optional residual scalar filter from the FE data-half push-down. The + // FE-side DV/merge-on-read interception is intentionally not done (see + // PushDownVectorTopNIntoPaimonScan): the Rust candidate search stage applies + // the filter together with deletion vectors, and a table shape Rust cannot + // handle correctly is Rust's responsibility to reject. + if (!_conjuncts.empty()) { + std::vector<std::string> names; + std::vector<DataTypePtr> types; + names.reserve(_projected_columns.size()); + types.reserve(_projected_columns.size()); + for (const auto& col : _projected_columns) { + if (col.is_partition_key || col.name == kPaimonSearchDistanceColumn) { + continue; + } + names.push_back(col.name); + types.push_back(col.type); + } + PaimonRustPredicateConverter converter(names, types, _handles->table.get()); + paimon_predicate* predicate = converter.build(_conjuncts); + if (predicate != nullptr) { + // with_filter consumes the predicate on every path. + if (paimon_error* err = paimon_vector_search_builder_with_filter(vb.get(), predicate)) { + return Status::InternalError("paimon-rust set vector filter failed: {}", + consume_error(err)); + } + LOG(INFO) << "paimon-rust vector search: applied residual filter"; + } + } + + // 4. Execute the search over the single FE-planned bucket split, via the + // eager Top-K chain (paimon-rust #822). Search, refine, local Top-K and + // materialization all happen in paimon_vector_read_read; the returned arrow + // stream is best-first and carries the __paimon_search_score column. The + // intermediate handles are locals: the migration contract (see + // doris-vector-search-api-migration.md) guarantees the record batch reader + // is independent once paimon_vector_read_read returns, so they are released + // when this function exits. + // + // split: the bytes only need to live until deserialize returns. + // scan: only needs to live until the plan is constructed. + // builder: only needs to live until scan and read are created. + // read/plan: only need to live until paimon_vector_read_read returns. + paimon_result_bucket_vector_search_split split_result = + paimon_bucket_vector_search_split_deserialize( + reinterpret_cast<const uint8_t*>(split_bytes.data()), split_bytes.size()); + if (split_result.error != nullptr) { + return Status::InternalError("paimon-rust deserialize vector split failed: {}", + consume_error(split_result.error)); + } + bucket_vector_search_split_ptr split(split_result.split); + + paimon_result_vector_scan scan_result = paimon_vector_search_builder_new_scan(vb.get()); + if (scan_result.error != nullptr) { + return Status::InternalError("paimon-rust new vector scan failed: {}", + consume_error(scan_result.error)); + } + vector_scan_ptr scan(scan_result.scan); + + paimon_result_vector_read read_result = paimon_vector_search_builder_new_read(vb.get()); + if (read_result.error != nullptr) { + return Status::InternalError("paimon-rust new vector read failed: {}", + consume_error(read_result.error)); + } + vector_read_ptr read(read_result.read); + + // Must use plan_from_bucket_splits: it plans from the FE-assigned split, + // whereas paimon_vector_scan_plan / execute_read replan from the table and + // would ignore the split this scanner was given. + const paimon_bucket_vector_search_split* splits[] = {split.get()}; + paimon_result_vector_plan plan_result = + paimon_vector_scan_plan_from_bucket_splits(scan.get(), splits, 1); + if (plan_result.error != nullptr) { + return Status::InternalError("paimon-rust build vector plan failed: {}", + consume_error(plan_result.error)); + } + vector_plan_ptr plan(plan_result.plan); + + paimon_result_record_batch_reader rdr_res = paimon_vector_read_read(read.get(), plan.get()); Review Comment: [P2] Bound the eager materialization before entering Rust This call finishes materialization before returning the Arrow reader. In the pinned Rust SDK, `materialize_positions` retains every projected Top-K batch, then `reorder_and_strip_position` copies those rows into one output batch. Unlike the ordinary path, this path ignores the scanner's adaptive batch size and has no reservation/accounting bridge for those Rust buffers. A one-bucket query with `LIMIT 1000000` and a projected 4-KiB payload can therefore retain roughly 4 GiB plus another 4 GiB reorder copy during `prepare_split`, even with a tiny probe batch and query memory budget. Raw malloc interposition forwards to jemalloc; it does not perform Doris `Allocator` tracking. Bound and account for this eager work before the call, or use bounded materialization; splitting the returned batch afterward cannot prevent the allocation spike. ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoPaimonScan.java: ########## @@ -0,0 +1,651 @@ +// 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.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); + + @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; Review Comment: [P2] Preserve validation of NULL vectors stored in the table The RHS check above does not cover a live NULL in the indexed column: ```text TopN(dis ASC LIMIT 1) Project(id, dis := l2_distance_approximate(vec, valid_q)) Scan(rows: vec=NULL; vec=[1,0,...]) ``` Paimon permits nullable VECTOR columns and NULL stored values. The ordinary `FunctionArrayDistance::_validate_no_nulls` throws `First argument ... cannot be null` before calculating distances. The Rust vector reader instead skips outer-NULL vectors, so replacing the function makes this query successfully return the finite row, even with full recall. Both supported metrics have this difference. Preserve the stored-vector validation or decline the rewrite when its error contract cannot be guaranteed, and add an enabled/disabled NULL-data regression. This is distinct from the earlier NULL coordinate in the query literal: here the query vector is valid. ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoPaimonScan.java: ########## @@ -0,0 +1,651 @@ +// 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.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); + + @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 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; + } + + // 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. + try { + Table table = paimonTable.getPaimonTable(Optional.empty()); + if ((optionalFilter.isPresent() + && filterHasDataConjunct(optionalFilter.get(), table) + && !PaimonScanNode.vectorSearchSupportsDataPredicate(table)) + || table.rowType().getFieldNames().stream().anyMatch( + name -> name.equalsIgnoreCase(PaimonVectorSearch.SEARCH_DISTANCE_COLUMN)) + || !complexColumnsRideAsWholeColumns(project, orderKeyAlias)) { + return null; + } + } catch (Throwable t) { + return null; + } + + // Pre-flight the projection-dependent bounds the vector split would hit at + // the scan node: a PK-vector split has no JNI fallback (it would return an + // empty distance column), so a query whose projected columns would set + // projectedVariant -- or an ORC table with TIMESTAMP_LTZ in its schema, the + // orc-ltz-shape bound -- must not be rewritten at all. Without the rewrite + // those queries run through the ordinary paths, where they are supported. + try { + if (projectedColumnsContainVariant(project, optionalFilter) + || orcSchemaContainsTimestampLtz(paimonTable)) { + return null; + } + } catch (Throwable t) { + 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, + backendStorageProperties, + Env.getCurrentSystemInfo().getAllClusterBackends(true))) { + return null; + } + } catch (Throwable t) { + return null; + } + // NOTE: the FE-side data-predicate interception is intentionally MINIMAL. + // The data predicate is handed in full to the BE via the vector search + // builder's residual filter, where the Rust candidate search stage applies + // it together with deletion vectors. The two planner-side exceptions above + // (unridable conjunct shapes, and the SDK's no-deletion-vector table mode) + // exist because those failures would surface only AFTER the one-way rewrite + // replaced the distance expression. Any other table shape Rust cannot + // handle correctly is Rust's responsibility to reject or answer correctly, + // not the planner's. + + // paimon-rust's vector index operates on f32, so carry the query as floats + // rather than List<Double>: it avoids a double→float narrowing on the wire + // and lets PaimonScanNode encode straight into a FLOAT32 TSearchVector. + List<Literal> vectorLiterals = ((ArrayLiteral) right).getValue(); + float[] annQueryVector = new float[vectorLiterals.size()]; + for (int i = 0; i < vectorLiterals.size(); i++) { + annQueryVector[i] = (float) vectorLiterals.get(i).getDouble(); + } + + // Build the reader-produced distance column, matched by name downstream (BE fills + // any block column named __paimon_search_score_to_dis, bridging it from + // paimon-rust's __paimon_search_score Arrow field). It holds the metric-native + // distance, not the score: the BE reader converts it in place while filling the + // block, so nothing has to be computed here. It is backed by a synthetic hidden + // Column so the SlotDescriptor keeps a non-null column through + // FileQueryScanNode.initSchemaParams; the column is excluded from the + // Parquet/ORC column-position mapping on the scan node. + SlotReference distanceSlot = SlotReference.fromColumn( + StatementScopeIdGenerator.newExprId(), scan.getTable(), + PaimonVectorSearch.getDistanceColumn(), scan.qualified()); + + // Reuse the index-produced value everywhere the same distance expression + // appears, mirroring PushDownVectorTopNIntoOlapScan: + // - orderKeyAlias -> matches the `dis` projection at its Alias root, so that + // item is rebuilt as `dis := <distance slot>` in one step. + // - orderKeyExpr -> matches the bare `dist_fn(vec, q)` wherever else it + // occurs (filter conjuncts, other projections), so those reuse the index + // value instead of recomputing from the raw vector. + // Both keys map to the same slot, so `dis` stays numerically identical to Doris + // semantics on every path. Note ExpressionUtils.replace rewrites top-down and + // short-circuits on a hit, and Alias.withChildren preserves the ExprId/name, so + // the alias identity the TopN references is never disturbed. + Map<Expression, Expression> replaceMap = Maps.newHashMap(); + replaceMap.put(orderKeyAlias, distanceSlot); + replaceMap.put(orderKeyExpr, distanceSlot); + + Plan plan = scan.withVectorTopN( Review Comment: [P2] Keep nonmovable project expressions before truncation The rewrite accepts a Project containing `NoneMovableFunction`, such as: ```text TopN(dis ASC LIMIT 1) Project(c := assert_true(id > 0, 'bad id'), dis := l2_distance_approximate(vec,q)) Scan(rows: id=-1/far; id=1/near) ``` On the ordinary path the Project evaluates before TopN and raises for the far row. `withVectorTopN` truncates the scan to the near row first, so the assertion never sees `id=-1` and the query succeeds. The scalar assertion passes `complexColumnsRideAsWholeColumns`, and there is no nonmovable-expression guard. Decline this rewrite for project expressions whose evaluation must remain below TopN, and cover the assertion case. This concerns a SELECT expression, separate from the existing WHERE residual-filter thread. -- 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]
