From 31f4d092378eaa1904862928599489813cac347b Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Fri, 14 Aug 2026 16:12:53 +0800 Subject: [PATCH 1/4] [opt](lance)support distribute plan in lance vector_search --- be/src/format_v2/table/lance_reader.cpp | 194 ++---------------- be/src/format_v2/table/lance_reader.h | 3 - be/test/format_v2/table/lance_reader_test.cpp | 66 ++++-- .../datasource/lance/LanceFragmentInfo.java | 56 +++++ .../datasource/lance/LanceMetadataLoader.java | 6 +- .../datasource/lance/LanceTableMetadata.java | 33 +-- .../lance/source/LanceScanNode.java | 110 +++++----- .../datasource/lance/source/LanceSplit.java | 56 ++--- .../datasource/tvf/source/TVFScanNode.java | 12 +- .../rules/analysis/BindExpression.java | 29 ++- .../ExternalFileTableValuedFunction.java | 9 +- .../FileTableValuedFunction.java | 4 +- .../VectorSearchTableValuedFunction.java | 66 ++++-- .../datasource/lance/LanceSnapshotTest.java | 2 +- .../lance/source/LanceScanNodeTest.java | 158 +++++++++++++- .../tvf/source/TVFScanNodeTest.java | 10 +- .../VectorSearchTableValuedFunctionTest.java | 39 ++++ gensrc/thrift/PlanNodes.thrift | 23 ++- .../lance/test_lance_vector_search.out | 1 - .../lance/test_lance_vector_search.groovy | 2 +- thirdparty/download-thirdparty.sh | 13 ++ .../patches/lance-c-0.1.6-fragment-knn.patch | 48 +++++ 22 files changed, 564 insertions(+), 376 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceFragmentInfo.java create mode 100644 thirdparty/patches/lance-c-0.1.6-fragment-knn.patch diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index 142c35ac4d42bb..fc388de14a4be5 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -53,21 +53,6 @@ struct LanceBatchDeleter { constexpr std::string_view DISTANCE_COLUMN = "_distance"; constexpr std::string_view ARROW_EXTENSION_NAME = "ARROW:extension:name"; -size_t vector_element_width(TVectorElementType::type type) { - switch (type) { - case TVectorElementType::FLOAT16: - return sizeof(uint16_t); - case TVectorElementType::FLOAT32: - return sizeof(float); - case TVectorElementType::FLOAT64: - return sizeof(double); - case TVectorElementType::UINT8: - case TVectorElementType::INT8: - return sizeof(uint8_t); - } - return 0; -} - std::string vector_metric_name(TVectorMetric::type metric) { switch (metric) { case TVectorMetric::DEFAULT: @@ -270,8 +255,6 @@ Status LanceTableReader::fetch_schema(const TFileRangeDesc& range, if (column_names == nullptr || column_types == nullptr) { return Status::InvalidArgument("Lance schema output must not be null"); } - RETURN_IF_ERROR(_validate_range(range)); - const auto& params = range.table_format_params.lance_params; const auto storage_options = _storage_options(&scan_params); std::vector storage_option_ptrs; @@ -314,10 +297,9 @@ Status LanceTableReader::init(TableReadOptions&& options) { _ctz = _runtime_state->timezone_obj(); _vector_search = _scan_params->__isset.external_search_request; - _search_split_prepared = false; if (_vector_search) { - RETURN_IF_ERROR(_validate_external_search_request()); - const auto& vector = _scan_params->external_search_request.query.vector; + const auto& vector = + _scan_params->external_search_request.search_query.vector_search; _scanner_profile->add_info_string("ExternalSearchType", "VECTOR"); _scanner_profile->add_info_string("LanceVectorColumn", vector.column); _scanner_profile->add_info_string("LanceTopK", std::to_string(vector.top_k)); @@ -358,7 +340,6 @@ Status LanceTableReader::init(TableReadOptions&& options) { } Status LanceTableReader::prepare_split(const SplitReadOptions& options) { - RETURN_IF_ERROR(_validate_range(options.current_range)); _close_scanner(); _eof = false; @@ -370,21 +351,6 @@ Status LanceTableReader::prepare_split(const SplitReadOptions& options) { return Status::OK(); } - if (_vector_search) { - const auto& lance_params = options.current_range.table_format_params.lance_params; - if (lance_params.version <= 0) { - return Status::InvalidArgument( - "Lance vector search requires a fixed positive dataset version"); - } - if (lance_params.__isset.fragment_ids) { - return Status::InvalidArgument("Lance vector search split must not set fragment ids"); - } - if (_search_split_prepared) { - return Status::InvalidArgument( - "Lance vector search supports exactly one whole-dataset split"); - } - } - const auto key = _dataset_key(options.current_range); if (_dataset == nullptr) { RETURN_IF_ERROR(_open_dataset(key)); @@ -395,7 +361,6 @@ Status LanceTableReader::prepare_split(const SplitReadOptions& options) { } RETURN_IF_ERROR(_open_scanner(options.current_range)); - _search_split_prepared = _vector_search; return Status::OK(); } @@ -475,134 +440,6 @@ Status LanceTableReader::close() { return TableReader::close(); } -Status LanceTableReader::_validate_range(const TFileRangeDesc& range) const { - if (!range.__isset.table_format_params || !range.table_format_params.__isset.lance_params) { - return Status::InvalidArgument("Lance split requires lance_params in table format params"); - } - const auto& params = range.table_format_params.lance_params; - if (!params.__isset.dataset_uri || params.dataset_uri.empty()) { - return Status::InvalidArgument("Lance split requires a non-empty dataset_uri"); - } - if (!params.__isset.version || params.version < 0) { - return Status::InvalidArgument("Lance split requires a non-negative dataset version"); - } - std::unordered_set unique_ids; - for (const auto fragment_id : params.fragment_ids) { - if (fragment_id < 0) { - return Status::InvalidArgument("Lance fragment id must be non-negative: {}", - fragment_id); - } - if (!unique_ids.emplace(fragment_id).second) { - return Status::InvalidArgument("Lance split contains duplicate fragment id: {}", - fragment_id); - } - } - return Status::OK(); -} - -Status LanceTableReader::_validate_external_search_request() const { - DORIS_CHECK(_scan_params != nullptr); - DORIS_CHECK(_scan_params->__isset.external_search_request); - if (_scan_params->__isset.lance_substrait_filter) { - return Status::InvalidArgument( - "Lance vector search cannot combine its pre-search filter with " - "lance_substrait_filter"); - } - const auto& request = _scan_params->external_search_request; - if (request.__isset.schema_version && request.schema_version != 1) { - return Status::NotSupported("unsupported external search schema version: {}", - request.schema_version); - } - if (!request.__isset.query) { - return Status::InvalidArgument("external search request requires query"); - } - - const bool has_vector = request.query.__isset.vector; - const bool has_full_text = request.query.__isset.full_text; - if (has_vector == has_full_text) { - return Status::InvalidArgument("external search query must set exactly one search kind"); - } - if (has_full_text) { - return Status::NotSupported("Lance Format V2 reader does not yet support full-text search"); - } - - const auto& vector = request.query.vector; - if (!vector.__isset.column || vector.column.empty() || - vector.column.find('\0') != std::string::npos) { - return Status::InvalidArgument("Lance vector search requires a non-empty column"); - } - if (!vector.__isset.query_vector) { - return Status::InvalidArgument("Lance vector search requires a query vector"); - } - const auto& query_vector = vector.query_vector; - if (!query_vector.__isset.element_type || !query_vector.__isset.dimension || - !query_vector.__isset.values) { - return Status::InvalidArgument( - "Lance query vector requires element_type, dimension, and values"); - } - if (query_vector.dimension <= 0) { - return Status::InvalidArgument("Lance query vector dimension must be positive: {}", - query_vector.dimension); - } - const auto element_width = vector_element_width(query_vector.element_type); - if (element_width == 0) { - return Status::NotSupported("unsupported Lance query vector element type: {}", - static_cast(query_vector.element_type)); - } - const auto dimension = static_cast(query_vector.dimension); - if (dimension > std::numeric_limits::max() / element_width || - query_vector.values.size() != dimension * element_width) { - return Status::InvalidArgument( - "Lance query vector byte size {} does not match dimension {} and element width {}", - query_vector.values.size(), dimension, element_width); - } - if (!vector.__isset.top_k || vector.top_k <= 0) { - return Status::InvalidArgument("Lance vector search top_k must be positive"); - } - if (!vector.__isset.offset || vector.offset < 0) { - return Status::InvalidArgument("Lance vector search offset must be non-negative"); - } - constexpr auto UINT32_MAX_VALUE = static_cast(std::numeric_limits::max()); - if (vector.offset > UINT32_MAX_VALUE || vector.top_k > UINT32_MAX_VALUE - vector.offset) { - return Status::InvalidArgument("Lance vector search top_k + offset exceeds uint32 range"); - } - - if (request.__isset.filter) { - const auto& filter = request.filter; - if (!filter.__isset.format || !filter.__isset.payload || filter.payload.empty()) { - return Status::InvalidArgument( - "external search filter requires format and non-empty payload"); - } - switch (filter.format) { - case TSearchFilterFormat::SQL: - if (filter.payload.find('\0') != std::string::npos) { - return Status::InvalidArgument( - "Lance SQL search filter contains an embedded NUL byte"); - } - break; - case TSearchFilterFormat::SUBSTRAIT: - break; - default: - return Status::NotSupported("unsupported external search filter format: {}", - static_cast(filter.format)); - } - } - - if (request.__isset.lance_options) { - const auto& options = request.lance_options; - if (options.__isset.nprobes && options.nprobes <= 0) { - return Status::InvalidArgument("Lance nprobes must be positive"); - } - if (options.__isset.refine_factor && options.refine_factor <= 0) { - return Status::InvalidArgument("Lance refine_factor must be positive"); - } - if (options.__isset.ef && options.ef <= 0) { - return Status::InvalidArgument("Lance ef must be positive"); - } - } - return Status::OK(); -} - Status LanceTableReader::_open_dataset(const DatasetKey& key) { std::vector storage_option_ptrs; storage_option_ptrs.reserve(key.storage_options.size() + 1); @@ -634,8 +471,9 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { const char* sql_filter = nullptr; if (_vector_search) { const auto& request = _scan_params->external_search_request; - if (request.__isset.filter && request.filter.format == TSearchFilterFormat::SQL) { - sql_filter = request.filter.payload.c_str(); + if (request.__isset.search_filter && + request.search_filter.format == TSearchFilterFormat::SQL) { + sql_filter = request.search_filter.payload.c_str(); } } LanceScanner* scanner = @@ -655,8 +493,9 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { } if (_vector_search) { const auto& request = _scan_params->external_search_request; - if (request.__isset.filter && request.filter.format == TSearchFilterFormat::SUBSTRAIT) { - const auto& filter = request.filter.payload; + if (request.__isset.search_filter && + request.search_filter.format == TSearchFilterFormat::SUBSTRAIT) { + const auto& filter = request.search_filter.payload; if (lance_scanner_set_substrait_filter(scanner, reinterpret_cast(filter.data()), filter.size()) != 0) { @@ -692,6 +531,12 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { } } if (_vector_search) { + // Distributed vector search always restricts each scanner to an explicit fragment set. + // Tell Lance that this fragment scan is the input to nearest() before installing the + // query. The same prefilter path also applies the TVF search filter, when present. + if (lance_scanner_set_prefilter(scanner, true) != 0) { + return _lance_error("enable Lance vector prefilter"); + } RETURN_IF_ERROR(_configure_vector_search(scanner)); } _scanner = scanner_guard.release(); @@ -703,10 +548,10 @@ Status LanceTableReader::_configure_vector_search(LanceScanner* scanner) const { DORIS_CHECK(scanner != nullptr); DORIS_CHECK(_scan_params != nullptr); const auto& request = _scan_params->external_search_request; - const auto& vector = request.query.vector; + const auto& vector = request.search_query.vector_search; const auto& query = vector.query_vector; - const auto* bytes = query.values.data(); const auto dimension = static_cast(query.dimension); + const auto* bytes = query.values.data(); const auto candidate_k = static_cast(vector.top_k + vector.offset); const auto set_nearest = [&](const void* values, LanceDataType type) -> Status { @@ -785,8 +630,8 @@ Status LanceTableReader::_configure_vector_search(LanceScanner* scanner) const { } } - if (request.__isset.lance_options) { - const auto& options = request.lance_options; + if (request.__isset.vector_search_options) { + const auto& options = request.vector_search_options; if (options.__isset.nprobes && lance_scanner_set_nprobes(scanner, static_cast(options.nprobes)) != 0) { return _lance_error("set Lance vector nprobes"); @@ -805,9 +650,6 @@ Status LanceTableReader::_configure_vector_search(LanceScanner* scanner) const { return _lance_error("set Lance vector use_index"); } } - if (request.__isset.filter && lance_scanner_set_prefilter(scanner, true) != 0) { - return _lance_error("enable Lance vector prefilter"); - } if (lance_scanner_set_offset(scanner, vector.offset) != 0) { return _lance_error("set Lance vector offset"); } diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 73448e4ed99e48..5f73db2af3e7ff 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -76,8 +76,6 @@ class LanceTableReader final : public TableReader { bool operator==(const DatasetKey&) const = default; }; - Status _validate_range(const TFileRangeDesc& range) const; - Status _validate_external_search_request() const; Status _open_dataset(const DatasetKey& key); Status _open_scanner(const TFileRangeDesc& range); Status _configure_vector_search(LanceScanner* scanner) const; @@ -95,7 +93,6 @@ class LanceTableReader final : public TableReader { cctz::time_zone _ctz; size_t _scanner_batch_size = 0; bool _vector_search = false; - bool _search_split_prepared = false; bool _eof = false; }; diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 42e84c211ec8ac..16a316423fa2fd 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -203,18 +203,18 @@ TFileScanRangeParams make_float32_vector_search_params( vector_params.__set_metric(TVectorMetric::L2); TExternalSearchQuery query; - query.__set_vector(std::move(vector_params)); + query.__set_vector_search(std::move(vector_params)); TExternalSearchRequest request; request.__set_schema_version(1); - request.__set_query(std::move(query)); - TLanceVectorSearchOptions lance_options; - lance_options.__set_use_index(false); - request.__set_lance_options(std::move(lance_options)); + request.__set_search_query(std::move(query)); + TVectorSearchOptions vector_search_options; + vector_search_options.__set_use_index(false); + request.__set_vector_search_options(std::move(vector_search_options)); if (filter.has_value()) { TSearchFilter search_filter; search_filter.__set_format(TSearchFilterFormat::SQL); search_filter.__set_payload(*filter); - request.__set_filter(std::move(search_filter)); + request.__set_search_filter(std::move(search_filter)); } TFileScanRangeParams scan_params; @@ -222,14 +222,6 @@ TFileScanRangeParams make_float32_vector_search_params( return scan_params; } -TFileRangeDesc make_whole_dataset_lance_range(const std::filesystem::path& dataset_uri, - int64_t version) { - auto range = make_lance_range(dataset_uri, version, {}); - range.table_format_params.lance_params.fragment_ids.clear(); - range.table_format_params.lance_params.__isset.fragment_ids = false; - return range; -} - std::vector> read_vector_search_rows(LanceTableReader* reader, Block* block) { std::vector> rows; @@ -269,11 +261,14 @@ TEST(LanceTableReaderVectorSearchTest, SearchesWholeSnapshotWithOffsetAndDistanc state.set_query_options(query_options); RuntimeProfile profile("lance_vector_search_fixture"); auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 2, 1); + auto& search_options = scan_params.external_search_request.vector_search_options; + search_options.__set_nprobes(4); + search_options.__set_refine_factor(2); + search_options.__set_ef(16); LanceTableReader reader; ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); - ASSERT_TRUE(prepare_range(&reader, make_whole_dataset_lance_range(dataset_uri, fixture.version)) - .ok()); + ASSERT_TRUE(prepare_fixture(&reader, dataset_uri, fixture, fixture.fragment_ids).ok()); Block block; add_output_columns(&block, columns); @@ -306,8 +301,7 @@ TEST(LanceTableReaderVectorSearchTest, AppliesSearchFilterBeforeTopK) { LanceTableReader reader; ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); - ASSERT_TRUE(prepare_range(&reader, make_whole_dataset_lance_range(dataset_uri, fixture.version)) - .ok()); + ASSERT_TRUE(prepare_fixture(&reader, dataset_uri, fixture, fixture.fragment_ids).ok()); Block block; add_output_columns(&block, columns); @@ -318,6 +312,42 @@ TEST(LanceTableReaderVectorSearchTest, AppliesSearchFilterBeforeTopK) { EXPECT_TRUE(reader.close().ok()); } +TEST(LanceTableReaderVectorSearchTest, SearchesMultipleFragmentSplits) { + const std::filesystem::path dataset_uri = + "./be/test/format_v2/table/lance/data/all_types.lance"; + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(dataset_uri, &fixture).ok()); + ASSERT_GT(fixture.fragment_ids.size(), 1U); + + const Columns columns { + projected_column("row_id", TYPE_BIGINT, false), + projected_column("_distance", TYPE_FLOAT, true), + }; + TQueryOptions query_options; + query_options.__set_batch_size(2); + TQueryGlobals query_globals; + RuntimeState state(query_globals); + state.set_query_options(query_options); + RuntimeProfile profile("lance_vector_search_fragment_splits_fixture"); + auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 4, 0); + + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); + std::vector row_ids; + for (const auto fragment_id : fixture.fragment_ids) { + ASSERT_TRUE(prepare_fixture(&reader, dataset_uri, fixture, {fragment_id}).ok()); + Block block; + add_output_columns(&block, columns); + const auto rows = read_vector_search_rows(&reader, &block); + for (const auto& row : rows) { + row_ids.emplace_back(row.first); + } + } + std::ranges::sort(row_ids); + EXPECT_EQ((std::vector {1, 2, 3, 4}), row_ids); + EXPECT_TRUE(reader.close().ok()); +} + TEST(LanceTableReaderFilterTest, PushesFilterOnNonProjectedColumn) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/all_types.lance"; diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceFragmentInfo.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceFragmentInfo.java new file mode 100644 index 00000000000000..213bf4bf12c72e --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceFragmentInfo.java @@ -0,0 +1,56 @@ +// 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.datasource.lance; + +/** Immutable row-count metadata for one Lance fragment in a fixed dataset snapshot. */ +public final class LanceFragmentInfo { + private final long id; + private final long rowCount; + private final long physicalRows; + + public LanceFragmentInfo(long id, long rowCount, long physicalRows) { + this.id = id; + this.rowCount = rowCount; + this.physicalRows = physicalRows; + } + + /** + * Returns the fragment id stored by Lance as an unsigned 32-bit value. + * + *

Doris represents it as a {@code long} so ids with the high bit set do not become + * negative while crossing Java and Thrift boundaries. + */ + public long getId() { + return id; + } + + /** Returns the logical row count after deletion vectors have been applied. */ + public long getRowCount() { + return rowCount; + } + + /** + * Returns the number of physical rows stored before deletions. + * + *

The BE legacy reader reads and merges physical batches before applying the deletion + * vector, so split scheduling uses this value rather than {@link #getRowCount()}. + */ + public long getPhysicalRows() { + return physicalRows; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceMetadataLoader.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceMetadataLoader.java index a2bc6d9a60de37..c800930fa060da 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceMetadataLoader.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceMetadataLoader.java @@ -83,10 +83,10 @@ private static LanceTableMetadata loadInternal(String datasetUri, Map fragments = new ArrayList<>(); + List fragments = new ArrayList<>(); for (Fragment fragment : dataset.getFragments()) { - fragments.add(new LanceTableMetadata.LanceFragmentInfo( - fragment.getId(), fragment.metadata().getNumRows(), + fragments.add(new LanceFragmentInfo( + Integer.toUnsignedLong(fragment.getId()), fragment.metadata().getNumRows(), fragment.metadata().getPhysicalRows())); } return new LanceTableMetadata(datasetUri, resolvedVersion, dataset.getSchema(), fragments, diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceTableMetadata.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceTableMetadata.java index 674b5b79d15be3..a964165f403712 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceTableMetadata.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/LanceTableMetadata.java @@ -19,6 +19,7 @@ import org.apache.arrow.vector.types.pojo.Schema; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -37,7 +38,7 @@ public LanceTableMetadata(String datasetUri, long version, Schema schema, this.datasetUri = datasetUri; this.version = version; this.schema = schema; - this.fragments = Collections.unmodifiableList(fragments); + this.fragments = Collections.unmodifiableList(new ArrayList<>(fragments)); this.backendStorageOptions = Collections.unmodifiableMap(new HashMap<>(backendStorageOptions)); } @@ -64,34 +65,4 @@ public Map getBackendStorageOptions() { public long getRowCount() { return fragments.stream().mapToLong(LanceFragmentInfo::getRowCount).sum(); } - - public static class LanceFragmentInfo { - private final long id; - private final long rowCount; - private final long physicalRows; - - public LanceFragmentInfo(long id, long rowCount, long physicalRows) { - this.id = id; - this.rowCount = rowCount; - this.physicalRows = physicalRows; - } - - public long getId() { - return id; - } - - /** Logical rows after deletions, used for row-count statistics. */ - public long getRowCount() { - return rowCount; - } - - /** - * Physical rows on disk before deletions. The pinned BE legacy reader reads and merges - * physical batches before applying the deletion vector, so scan work scales with this - * value rather than the post-deletion row count. Used for split scheduling weight. - */ - public long getPhysicalRows() { - return physicalRows; - } - } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceScanNode.java index 4c433a30bce145..7faaecc569caad 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceScanNode.java @@ -26,6 +26,7 @@ import org.apache.doris.datasource.TableFormatType; import org.apache.doris.datasource.lance.LanceExternalCatalog; import org.apache.doris.datasource.lance.LanceExternalTable; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.lance.LanceTableMetadata; import org.apache.doris.datasource.mvcc.MvccSnapshot; import org.apache.doris.planner.PlanNodeId; @@ -56,9 +57,9 @@ * *

These modes share dataset metadata, storage properties, and BE scan-range serialization. * Keeping them in one node prevents those common parts from drifting apart. The search request is - * also an explicit mode marker: ordinary scans are split by fragment, while the first version of - * vector search deliberately sends one whole-snapshot split to one scanner so Lance can compute a - * global TopK result. + * also an explicit mode marker. Both ordinary scans and vector searches are split by fragment. + * Each search split produces local candidates; a Doris TopN above this scan merges them into the + * requested snapshot-wide result. */ public class LanceScanNode extends FileQueryScanNode { private LanceExternalTable lanceTable; @@ -110,7 +111,7 @@ protected void doInitialize() throws UserException { if (isExternalSearch()) { // Search output comes from the FunctionGenTable because it adds generated columns such // as _distance. The real Lance table is still retained for storage and metadata access. - params.setExternalSearchRequest(externalSearchRequest.deepCopy()); + params.setExternalSearchRequest(createFragmentSearchRequest(externalSearchRequest)); } } @@ -152,42 +153,34 @@ public void createScanRangeLocations() throws UserException { @Override public List getSplits(int numBackends) throws UserException { - if (isExternalSearch()) { - plannedVersion = plannedMetadata.getVersion(); - plannedFragments = plannedMetadata.getFragments().size(); - - // Do not attach fragment IDs. A vector index is a dataset-wide structure and one - // scanner must see every fragment visible in this pinned snapshot to produce global - // TopK. Fragment-level parallel search and result merging are intentionally deferred. - return Collections.singletonList(LanceSplit.wholeDatasetAtVersion( - plannedMetadata.getDatasetUri(), plannedMetadata.getVersion(), - plannedMetadata.getRowCount())); - } else { - LanceTableMetadata metadata = plannedMetadata; - plannedVersion = metadata.getVersion(); - plannedFragments = metadata.getFragments().size(); - Set fragmentIds = new HashSet<>(); - long targetRows = 1; - for (LanceTableMetadata.LanceFragmentInfo fragment : metadata.getFragments()) { - if (!fragmentIds.add(fragment.getId())) { - throw new UserException("Duplicate Lance fragment id " + fragment.getId() - + " at dataset version " + metadata.getVersion()); - } - targetRows = Math.max(targetRows, Math.max(fragment.getPhysicalRows(), 1)); + LanceTableMetadata metadata = plannedMetadata; + plannedVersion = metadata.getVersion(); + plannedFragments = metadata.getFragments().size(); + if (isExternalSearch() && plannedVersion <= 0) { + throw new UserException( + "Lance vector search requires a fixed positive dataset version"); + } + Set fragmentIds = new HashSet<>(); + long targetRows = 1; + for (LanceFragmentInfo fragment : metadata.getFragments()) { + if (!fragmentIds.add(fragment.getId())) { + throw new UserException("Duplicate Lance fragment id " + fragment.getId() + + " at dataset version " + metadata.getVersion()); } + targetRows = Math.max(targetRows, Math.max(fragment.getPhysicalRows(), 1)); + } - // Use the largest fragment as one standard split so smaller fragments keep - // their relative row-count weight during backend assignment. Physical rows drive - // the weight because the BE legacy reader scans physical batches before deletions. - List splits = new ArrayList<>(plannedFragments); - for (LanceTableMetadata.LanceFragmentInfo fragment : metadata.getFragments()) { - LanceSplit split = new LanceSplit(metadata.getDatasetUri(), metadata.getVersion(), - fragment.getId(), fragment.getPhysicalRows()); - split.setTargetSplitSize(targetRows); - splits.add(split); - } - return splits; + // Keep one fragment per split. Use the largest fragment's physical row count as the + // normalization baseline for split weights, so backend scheduling reflects the relative + // amount of physical data each fragment scans, including rows covered by deletion metadata. + List splits = new ArrayList<>(plannedFragments); + for (LanceFragmentInfo fragment : metadata.getFragments()) { + LanceSplit split = new LanceSplit(metadata.getDatasetUri(), metadata.getVersion(), + fragment.getId(), fragment.getPhysicalRows()); + split.setTargetSplitSize(targetRows); + splits.add(split); } + return splits; } @Override @@ -199,25 +192,14 @@ protected void setScanParams(TFileRangeDesc rangeDesc, Split split) { TLanceFileDesc lanceParams = new TLanceFileDesc(); lanceParams.setDatasetUri(lanceSplit.getDatasetUri()); lanceParams.setVersion(lanceSplit.getVersion()); - if (isExternalSearch()) { - if (lanceSplit.hasFragmentId()) { - throw new IllegalArgumentException( - "Lance external search split must cover the whole dataset"); - } - // Leaving fragment_ids unset instructs lance-c to scan/search all fragments in the - // selected dataset version. - } else { - if (!lanceSplit.hasFragmentId()) { - throw new IllegalArgumentException( - "Ordinary Lance scan split must contain one fragment"); - } - lanceParams.setFragmentIds(Collections.singletonList(lanceSplit.getFragmentId())); - // Push LIMIT into each fragment scanner only when it is safe to truncate a single - // fragment early. See canPushDownLimit(). Each scanner still returns at most `limit` - // rows and the upper LIMIT operator enforces the final bound across fragments. - if (canPushDownLimit()) { - lanceParams.setLimit(getLimit()); - } + if (lanceSplit.getFragmentIds().size() != 1) { + throw new IllegalArgumentException("Lance scan split must contain one fragment"); + } + lanceParams.setFragmentIds(lanceSplit.getFragmentIds()); + // Push LIMIT into each ordinary fragment scanner only when it is safe to truncate that + // fragment early. Vector search uses its own per-fragment candidate bound. + if (!isExternalSearch() && canPushDownLimit()) { + lanceParams.setLimit(getLimit()); } TTableFormatFileDesc tableFormatParams = new TTableFormatFileDesc(); @@ -256,7 +238,7 @@ protected Map getLocationProperties() { public String getNodeExplainString(String prefix, TExplainLevel detailLevel) { StringBuilder result = new StringBuilder(super.getNodeExplainString(prefix, detailLevel)); if (isExternalSearch()) { - TVectorSearchParams vector = externalSearchRequest.getQuery().getVector(); + TVectorSearchParams vector = externalSearchRequest.getSearchQuery().getVectorSearch(); result.append(prefix).append("externalSearchType=VECTOR\n"); result.append(prefix).append("lanceVectorColumn=").append(vector.getColumn()).append("\n"); result.append(prefix).append("lanceTopK=").append(vector.getTopK()).append("\n"); @@ -266,7 +248,8 @@ public String getNodeExplainString(String prefix, TExplainLevel detailLevel) { .append("\n"); result.append(prefix).append("lanceVersion=") .append(plannedMetadata.getVersion()).append("\n"); - result.append(prefix).append("lanceSearchScanners=1\n"); + result.append(prefix).append("lanceSearchFragments=") + .append(plannedFragments).append("\n"); } else { result.append(prefix).append("lanceCatalogType=") .append(((LanceExternalCatalog) lanceTable.getCatalog()).getLanceCatalogType()).append("\n"); @@ -287,6 +270,17 @@ private boolean isExternalSearch() { return externalSearchRequest != null; } + static TExternalSearchRequest createFragmentSearchRequest(TExternalSearchRequest searchRequest) { + TExternalSearchRequest fragmentRequest = searchRequest.deepCopy(); + TVectorSearchParams vector = fragmentRequest.getSearchQuery().getVectorSearch(); + // Every fragment must retain enough rows for the later global OFFSET/LIMIT. Applying the + // logical offset independently inside each fragment could discard rows that belong to the + // snapshot-wide result. + vector.setTopK(vector.getTopK() + vector.getOffset()); + vector.setOffset(0); + return fragmentRequest; + } + private static String metricName(TVectorMetric metric) { switch (metric) { case L2: diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceSplit.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceSplit.java index 22060e6eed96f8..b8712769aaee2b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceSplit.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/lance/source/LanceSplit.java @@ -21,49 +21,55 @@ import org.apache.doris.datasource.FileSplit; import org.apache.doris.datasource.TableFormatType; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; /** * A Lance scan split. Catalog and S3 scans normally use one fixed-version fragment per split. - * Backend-local TVFs use one whole-dataset latest-version split, while vector search uses one - * whole-dataset fixed-version split. + * Vector search also uses one fixed-version fragment per split. Backend-local TVFs use one + * whole-dataset latest-version split. */ public class LanceSplit extends FileSplit { private final String datasetUri; private final long version; - private final long fragmentId; - private final boolean hasFragmentId; + private final List fragmentIds; public LanceSplit(String datasetUri, long version, long fragmentId, long rowCount) { - super(LocationPath.of(datasetUri), 0, 0, 0, 0, null, Collections.emptyList()); + this(datasetUri, version, Collections.singletonList(fragmentId), rowCount); + } + + private LanceSplit(String datasetUri, long version, List fragmentIds, long rowCount) { + super(LocationPath.of(requireDatasetUri(datasetUri)), 0, 0, 0, 0, null, + Collections.emptyList()); + if (version < 0) { + throw new IllegalArgumentException("Lance dataset version must be non-negative"); + } + for (Long fragmentId : fragmentIds) { + if (fragmentId == null || fragmentId < 0) { + throw new IllegalArgumentException("Lance fragment id must be non-negative"); + } + } this.datasetUri = datasetUri; this.version = version; - this.fragmentId = fragmentId; - this.hasFragmentId = true; + this.fragmentIds = Collections.unmodifiableList(new ArrayList<>(fragmentIds)); this.tableFormatType = TableFormatType.LANCE; this.selfSplitWeight = Math.max(rowCount, 1); } private LanceSplit(String datasetUri, long version, long rowCount) { - super(LocationPath.of(datasetUri), 0, 0, 0, 0, null, Collections.emptyList()); - this.datasetUri = datasetUri; - this.version = version; - this.fragmentId = -1; - this.hasFragmentId = false; - this.tableFormatType = TableFormatType.LANCE; - this.selfSplitWeight = Math.max(rowCount, 1); + this(datasetUri, version, Collections.emptyList(), rowCount); } public static LanceSplit wholeDatasetAtLatest(String datasetUri) { return new LanceSplit(datasetUri, 0, 1); } - public static LanceSplit wholeDatasetAtVersion(String datasetUri, long version, long rowCount) { - if (version <= 0) { - throw new IllegalArgumentException( - "A fixed Lance dataset version must be positive: " + version); + private static String requireDatasetUri(String datasetUri) { + if (datasetUri == null || datasetUri.trim().isEmpty()) { + throw new IllegalArgumentException("Lance dataset URI must not be empty"); } - return new LanceSplit(datasetUri, version, rowCount); + return datasetUri; } public String getDatasetUri() { @@ -74,18 +80,18 @@ public long getVersion() { return version; } - public long getFragmentId() { - return fragmentId; + public List getFragmentIds() { + return fragmentIds; } - public boolean hasFragmentId() { - return hasFragmentId; + public boolean hasFragmentIds() { + return !fragmentIds.isEmpty(); } @Override public String getConsistentHashString() { - return hasFragmentId - ? datasetUri + "#" + version + "#" + fragmentId + return hasFragmentIds() + ? datasetUri + "#" + version + "#" + fragmentIds : datasetUri + "#" + (version == 0 ? "latest" : version) + "#all"; } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/tvf/source/TVFScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/tvf/source/TVFScanNode.java index 3254b5235036a1..462a201e6e90e3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/tvf/source/TVFScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/tvf/source/TVFScanNode.java @@ -31,7 +31,7 @@ import org.apache.doris.datasource.FileSplit.FileSplitCreator; import org.apache.doris.datasource.FileSplitter; import org.apache.doris.datasource.TableFormatType; -import org.apache.doris.datasource.lance.LanceTableMetadata; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.lance.source.LanceSplit; import org.apache.doris.planner.PlanNodeId; import org.apache.doris.planner.ScanContext; @@ -193,15 +193,15 @@ private List getLanceSplits() throws UserException { throw new UserException( "S3 Lance TVF metadata was not initialized with a fixed dataset version"); } - List fragments = tableValuedFunction.getLanceFragments(); + List fragments = tableValuedFunction.getLanceFragments(); // Mirror LanceScanNode: use the largest fragment as one standard split so smaller fragments // keep their relative physical-row weight, keeping the catalog and S3/file TVF paths in sync. long targetRows = 1; - for (LanceTableMetadata.LanceFragmentInfo fragment : fragments) { + for (LanceFragmentInfo fragment : fragments) { targetRows = Math.max(targetRows, Math.max(fragment.getPhysicalRows(), 1)); } List splits = new ArrayList<>(fragments.size()); - for (LanceTableMetadata.LanceFragmentInfo fragment : fragments) { + for (LanceFragmentInfo fragment : fragments) { LanceSplit split = new LanceSplit(tableValuedFunction.getFilePath(), version, fragment.getId(), fragment.getPhysicalRows()); split.setTargetSplitSize(targetRows); @@ -239,8 +239,8 @@ protected void setScanParams(TFileRangeDesc rangeDesc, Split split) { TLanceFileDesc lanceParams = new TLanceFileDesc(); lanceParams.setDatasetUri(lanceSplit.getDatasetUri()); lanceParams.setVersion(lanceSplit.getVersion()); - if (lanceSplit.hasFragmentId()) { - lanceParams.setFragmentIds(Collections.singletonList(lanceSplit.getFragmentId())); + if (lanceSplit.hasFragmentIds()) { + lanceParams.setFragmentIds(lanceSplit.getFragmentIds()); } TTableFormatFileDesc tableFormatFileDesc = new TTableFormatFileDesc(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java index 2cf8dbbe93be40..85d5b964c3d792 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java @@ -69,6 +69,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt; import org.apache.doris.nereids.trees.expressions.functions.scalar.GroupingScalarFunction; import org.apache.doris.nereids.trees.expressions.functions.table.TableValuedFunction; +import org.apache.doris.nereids.trees.expressions.functions.table.VectorSearch; import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; import org.apache.doris.nereids.trees.plans.JoinType; @@ -101,6 +102,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalSort; import org.apache.doris.nereids.trees.plans.logical.LogicalSubQueryAlias; import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; import org.apache.doris.nereids.trees.plans.logical.LogicalUsingJoin; import org.apache.doris.nereids.trees.plans.visitor.InferPlanOutputAlias; import org.apache.doris.nereids.types.BooleanType; @@ -114,6 +116,7 @@ import org.apache.doris.nereids.util.TypeCoercionUtils; import org.apache.doris.nereids.util.Utils; import org.apache.doris.qe.SqlModeHelper; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; import com.google.common.base.Joiner; import com.google.common.base.Preconditions; @@ -1762,7 +1765,7 @@ private Plan bindSortWithoutSetOperation(MatchingContext> ctx) return new LogicalSort<>(boundOrderKeys.build(), sort.child()); } - private LogicalTVFRelation bindTableValuedFunction(MatchingContext ctx) { + private Plan bindTableValuedFunction(MatchingContext ctx) { UnboundTVFRelation unboundTVFRelation = ctx.root; StatementContext statementContext = ctx.statementContext; Env env = statementContext.getConnectContext().getEnv(); @@ -1780,8 +1783,28 @@ private LogicalTVFRelation bindTableValuedFunction(MatchingContext slot.getName().equalsIgnoreCase( + VectorSearchTableValuedFunction.DISTANCE_COLUMN)) + .findFirst() + .orElseThrow(() -> new AnalysisException("vector_search() output is missing '" + + VectorSearchTableValuedFunction.DISTANCE_COLUMN + "'")); + OrderKey distanceAscending = new OrderKey(distance, true, false); + return new LogicalTopN<>(ImmutableList.of(distanceAscending), + vectorSearch.getTopK(), vectorSearch.getOffset(), relation); } private void checkIfOutputAliasNameDuplicatedForGroupBy(Collection expressions, diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/ExternalFileTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/ExternalFileTableValuedFunction.java index 8e06e8c3e3d3a3..5a61c3c3b16e51 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/ExternalFileTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/ExternalFileTableValuedFunction.java @@ -45,6 +45,7 @@ import org.apache.doris.common.util.S3Util; import org.apache.doris.common.util.Util; import org.apache.doris.datasource.TableFormatType; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.lance.LanceTableMetadata; import org.apache.doris.datasource.lance.LanceTypeConverter; import org.apache.doris.datasource.property.fileformat.CsvFileFormatProperties; @@ -138,7 +139,7 @@ public abstract class ExternalFileTableValuedFunction extends TableValuedFunctio public FileFormatProperties fileFormatProperties; private long tableId; private long lanceDatasetVersion = -1; - private List lanceFragments = Collections.emptyList(); + private List lanceFragments = Collections.emptyList(); public abstract TFileType getTFileType(); @@ -162,7 +163,7 @@ public long getLanceDatasetVersion() { return lanceDatasetVersion; } - public List getLanceFragments() { + public List getLanceFragments() { return lanceFragments; } @@ -358,10 +359,10 @@ protected void setLanceTableMetadata(LanceTableMetadata metadata) throws Analysi throw new AnalysisException("Lance returned an invalid dataset version: " + metadata.getVersion()); } - List fragments = + List fragments = new ArrayList<>(metadata.getFragments().size()); Set uniqueIds = new HashSet<>(); - for (LanceTableMetadata.LanceFragmentInfo fragment : metadata.getFragments()) { + for (LanceFragmentInfo fragment : metadata.getFragments()) { long fragmentId = fragment.getId(); if (fragmentId < 0) { throw new AnalysisException("Lance returned an invalid fragment id: " + fragmentId); diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FileTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FileTableValuedFunction.java index 7d4b1471436c61..6ec9f12a81cb3f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FileTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FileTableValuedFunction.java @@ -22,7 +22,7 @@ import org.apache.doris.catalog.Column; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.UserException; -import org.apache.doris.datasource.lance.LanceTableMetadata; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.property.storage.AbstractS3CompatibleProperties; import org.apache.doris.datasource.property.storage.AzureProperties; import org.apache.doris.datasource.property.storage.HdfsCompatibleProperties; @@ -101,7 +101,7 @@ public long getLanceDatasetVersion() { } @Override - public List getLanceFragments() { + public List getLanceFragments() { return delegateTvf.getLanceFragments(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java index 50ffc053d4c753..4893284acd1017 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java @@ -44,11 +44,11 @@ import org.apache.doris.qe.SessionVariable; import org.apache.doris.thrift.TExternalSearchQuery; import org.apache.doris.thrift.TExternalSearchRequest; -import org.apache.doris.thrift.TLanceVectorSearchOptions; import org.apache.doris.thrift.TSearchFilter; import org.apache.doris.thrift.TSearchFilterFormat; import org.apache.doris.thrift.TSearchVector; import org.apache.doris.thrift.TVectorMetric; +import org.apache.doris.thrift.TVectorSearchOptions; import org.apache.doris.thrift.TVectorSearchParams; import com.google.common.annotations.VisibleForTesting; @@ -62,8 +62,9 @@ import java.util.Map; import java.util.Set; import java.util.TreeMap; +import java.util.TreeSet; -/** Relation TVF for a whole-snapshot Lance vector search. */ +/** Relation TVF for a fixed-snapshot Lance vector search. */ public class VectorSearchTableValuedFunction extends TableValuedFunctionIf { public static final String NAME = "vector_search"; public static final String DISTANCE_COLUMN = "_distance"; @@ -128,38 +129,34 @@ public VectorSearchTableValuedFunction(Map properties) searchRequest = new TExternalSearchRequest() .setSchemaVersion(1) - .setQuery(TExternalSearchQuery.vector(vectorParams)); + .setSearchQuery(TExternalSearchQuery.vector_search(vectorParams)); if (params.containsKey(FILTER)) { - String filter = params.get(FILTER); - if (filter == null || filter.trim().isEmpty()) { - throw new AnalysisException("'filter' must not be empty"); - } - searchRequest.setFilter(new TSearchFilter() + searchRequest.setSearchFilter(new TSearchFilter() .setFormat(TSearchFilterFormat.SQL) - .setPayload(filter.getBytes(StandardCharsets.UTF_8))); + .setPayload(validateAndEncodeSqlFilter(params.get(FILTER)))); } - TLanceVectorSearchOptions lanceOptions = new TLanceVectorSearchOptions(); - boolean hasLanceOptions = false; + TVectorSearchOptions vectorSearchOptions = new TVectorSearchOptions(); + boolean hasVectorSearchOptions = false; if (params.containsKey(NPROBES)) { - lanceOptions.setNprobes(parsePositiveInt(params.get(NPROBES), NPROBES)); - hasLanceOptions = true; + vectorSearchOptions.setNprobes(parsePositiveInt(params.get(NPROBES), NPROBES)); + hasVectorSearchOptions = true; } if (params.containsKey(REFINE_FACTOR)) { - lanceOptions.setRefineFactor( + vectorSearchOptions.setRefineFactor( parsePositiveInt(params.get(REFINE_FACTOR), REFINE_FACTOR)); - hasLanceOptions = true; + hasVectorSearchOptions = true; } if (params.containsKey(EF)) { - lanceOptions.setEf(parsePositiveInt(params.get(EF), EF)); - hasLanceOptions = true; + vectorSearchOptions.setEf(parsePositiveInt(params.get(EF), EF)); + hasVectorSearchOptions = true; } if (params.containsKey(USE_INDEX)) { - lanceOptions.setUseIndex(parseBoolean(params.get(USE_INDEX), USE_INDEX)); - hasLanceOptions = true; + vectorSearchOptions.setUseIndex(parseBoolean(params.get(USE_INDEX), USE_INDEX)); + hasVectorSearchOptions = true; } - if (hasLanceOptions) { - searchRequest.setLanceOptions(lanceOptions); + if (hasVectorSearchOptions) { + searchRequest.setVectorSearchOptions(vectorSearchOptions); } columns = buildOutputColumns(metadata); } @@ -176,6 +173,14 @@ public TExternalSearchRequest getSearchRequest() { return searchRequest.deepCopy(); } + public long getTopK() { + return searchRequest.getSearchQuery().getVectorSearch().getTopK(); + } + + public long getOffset() { + return searchRequest.getSearchQuery().getVectorSearch().getOffset(); + } + @Override public String getTableName() { return "VectorSearchTableValuedFunction<" + sourceTableName + ">"; @@ -257,11 +262,17 @@ private static LanceExternalTable findLanceExternalTable(TableName tableName) return (LanceExternalTable) table; } - private static List buildOutputColumns(LanceTableMetadata metadata) + @VisibleForTesting + static List buildOutputColumns(LanceTableMetadata metadata) throws AnalysisException { List result = new ArrayList<>(metadata.getSchema().getFields().size() + 1); + Set fieldNames = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); int position = 0; for (Field field : metadata.getSchema().getFields()) { + if (!fieldNames.add(field.getName())) { + throw new AnalysisException("Duplicate Lance schema column under " + + "case-insensitive matching: '" + field.getName() + "'"); + } if (field.getName().equalsIgnoreCase(DISTANCE_COLUMN)) { throw new AnalysisException("Lance table already contains reserved vector search " + "column '" + DISTANCE_COLUMN + "'"); @@ -283,6 +294,17 @@ private static List buildOutputColumns(LanceTableMetadata metadata) return result; } + @VisibleForTesting + static byte[] validateAndEncodeSqlFilter(String filter) throws AnalysisException { + if (filter == null || filter.trim().isEmpty()) { + throw new AnalysisException("'filter' must not be empty"); + } + if (filter.indexOf('\0') >= 0) { + throw new AnalysisException("'filter' must not contain an embedded NUL byte"); + } + return filter.getBytes(StandardCharsets.UTF_8); + } + private static long parseLong(String value, String property, long min, long max) throws AnalysisException { try { diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/LanceSnapshotTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/LanceSnapshotTest.java index a48afa669ad771..9f4d7e45e4b0ab 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/LanceSnapshotTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/LanceSnapshotTest.java @@ -88,7 +88,7 @@ public void testBoundSnapshotCarriesItsOwnSchema() { private static LanceTableMetadata metadata(long version, Field field) { return new LanceTableMetadata("s3://bucket/table.lance", version, new Schema(Collections.singletonList(field)), - Collections.singletonList(new LanceTableMetadata.LanceFragmentInfo(version, 1, 1)), + Collections.singletonList(new LanceFragmentInfo(version, 1, 1)), Collections.singletonMap("s3.endpoint", "http://minio:9000")); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/source/LanceScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/source/LanceScanNodeTest.java index 51fa7cad4b5114..4edd2c6134e521 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/source/LanceScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/lance/source/LanceScanNodeTest.java @@ -19,11 +19,17 @@ import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.analysis.TupleId; +import org.apache.doris.common.UserException; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.lance.LanceTableMetadata; import org.apache.doris.planner.PlanNodeId; import org.apache.doris.planner.ScanContext; import org.apache.doris.qe.SessionVariable; import org.apache.doris.spi.Split; +import org.apache.doris.thrift.TExternalSearchQuery; +import org.apache.doris.thrift.TExternalSearchRequest; +import org.apache.doris.thrift.TFileRangeDesc; +import org.apache.doris.thrift.TVectorSearchParams; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.Assert; @@ -42,9 +48,9 @@ public void testFragmentRowsDetermineSplitWeights() throws Exception { 42, new Schema(Collections.emptyList()), Arrays.asList( - new LanceTableMetadata.LanceFragmentInfo(7, 1000, 1000), - new LanceTableMetadata.LanceFragmentInfo(11, 250, 250), - new LanceTableMetadata.LanceFragmentInfo(13, 0, 0)), + new LanceFragmentInfo(7, 1000, 1000), + new LanceFragmentInfo(11, 250, 250), + new LanceFragmentInfo(13, 0, 0)), Collections.emptyMap()); LanceScanNode node = newNode(); setMetadata(node, metadata); @@ -68,8 +74,8 @@ public void testDeletionHeavyFragmentKeepsPhysicalScanWeight() throws Exception 42, new Schema(Collections.emptyList()), Arrays.asList( - new LanceTableMetadata.LanceFragmentInfo(7, 1000, 1000), - new LanceTableMetadata.LanceFragmentInfo(11, 10, 1000)), + new LanceFragmentInfo(7, 1000, 1000), + new LanceFragmentInfo(11, 10, 1000)), Collections.emptyMap()); LanceScanNode node = newNode(); setMetadata(node, metadata); @@ -81,6 +87,128 @@ public void testDeletionHeavyFragmentKeepsPhysicalScanWeight() throws Exception assertSplit(splits.get(1), 11, 1000, 100); } + @Test + public void testExternalSearchUsesFragmentSplits() throws Exception { + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", + 42, + new Schema(Collections.emptyList()), + Arrays.asList( + new LanceFragmentInfo(7, 1000, 1000), + new LanceFragmentInfo(11, 250, 250)), + Collections.emptyMap()); + TExternalSearchRequest request = vectorSearchRequest(5, 2); + LanceScanNode node = new LanceScanNode( + new PlanNodeId(0), new TupleDescriptor(new TupleId(0)), null, + metadata, request, new SessionVariable()); + + List splits = node.getSplits(2); + + Assert.assertEquals(2, splits.size()); + assertSplit(splits.get(0), 7, 1000, 100); + assertSplit(splits.get(1), 11, 1000, 25); + TFileRangeDesc range = new TFileRangeDesc(); + node.setScanParams(range, splits.get(1)); + Assert.assertEquals(Collections.singletonList(11L), range.getTableFormatParams() + .getLanceParams().getFragmentIds()); + Assert.assertEquals(42L, range.getTableFormatParams().getLanceParams().getVersion()); + } + + @Test + public void testExternalSearchRejectsNonPositiveSnapshotVersionInFrontend() { + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", + 0, + new Schema(Collections.emptyList()), + Collections.singletonList(new LanceFragmentInfo(7, 1000, 1000)), + Collections.emptyMap()); + LanceScanNode node = new LanceScanNode( + new PlanNodeId(0), new TupleDescriptor(new TupleId(0)), null, + metadata, vectorSearchRequest(5, 0), new SessionVariable()); + + UserException exception = Assert.assertThrows(UserException.class, + () -> node.getSplits(1)); + + Assert.assertTrue(exception.getMessage().contains("fixed positive dataset version")); + } + + @Test + public void testExternalSearchUsesFragmentRowsForSplitWeights() throws Exception { + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", + 42, + new Schema(Collections.emptyList()), + Arrays.asList( + new LanceFragmentInfo(7, 1000, 1000), + new LanceFragmentInfo(11, 250, 250), + new LanceFragmentInfo(13, 800, 800), + new LanceFragmentInfo(17, 100, 100)), + Collections.emptyMap()); + LanceScanNode node = new LanceScanNode( + new PlanNodeId(0), new TupleDescriptor(new TupleId(0)), null, + metadata, vectorSearchRequest(5, 0), new SessionVariable()); + + List splits = node.getSplits(3); + + Assert.assertEquals(4, splits.size()); + assertSplit(splits.get(0), 7, 1000, 100); + assertSplit(splits.get(1), 11, 1000, 25); + assertSplit(splits.get(2), 13, 1000, 80); + assertSplit(splits.get(3), 17, 1000, 10); + TFileRangeDesc range = new TFileRangeDesc(); + node.setScanParams(range, splits.get(0)); + Assert.assertEquals(Collections.singletonList(7L), + range.getTableFormatParams().getLanceParams().getFragmentIds()); + } + + @Test + public void testExternalSearchUsesOneSplitPerFragmentRegardlessOfBackendCount() throws Exception { + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", + 42, + new Schema(Collections.emptyList()), + Arrays.asList( + new LanceFragmentInfo(1, 8, 8), + new LanceFragmentInfo(2, 7, 7), + new LanceFragmentInfo(3, 6, 6), + new LanceFragmentInfo(4, 5, 5)), + Collections.emptyMap()); + LanceScanNode node = new LanceScanNode( + new PlanNodeId(0), new TupleDescriptor(new TupleId(0)), null, + metadata, vectorSearchRequest(5, 0), new SessionVariable()); + + List splits = node.getSplits(2); + + Assert.assertEquals(4, splits.size()); + assertSplit(splits.get(0), 1, 8, 100); + assertSplit(splits.get(1), 2, 8, 88); + assertSplit(splits.get(2), 3, 8, 75); + assertSplit(splits.get(3), 4, 8, 63); + } + + @Test + public void testFragmentSearchRetainsTopKPlusOffsetCandidates() { + TExternalSearchRequest logicalRequest = vectorSearchRequest(5, 2); + + TExternalSearchRequest fragmentRequest = + LanceScanNode.createFragmentSearchRequest(logicalRequest); + + Assert.assertEquals(7, fragmentRequest.getSearchQuery().getVectorSearch().getTopK()); + Assert.assertEquals(0, fragmentRequest.getSearchQuery().getVectorSearch().getOffset()); + Assert.assertEquals(5, logicalRequest.getSearchQuery().getVectorSearch().getTopK()); + Assert.assertEquals(2, logicalRequest.getSearchQuery().getVectorSearch().getOffset()); + } + + @Test + public void testLanceSplitRejectsInvalidRangeFieldsInFrontend() { + assertInvalidSplit(() -> new LanceSplit("", 42, 1, 1), + "Lance dataset URI must not be empty"); + assertInvalidSplit(() -> new LanceSplit("s3://bucket/table.lance", -1, 1, 1), + "Lance dataset version must be non-negative"); + assertInvalidSplit(() -> new LanceSplit("s3://bucket/table.lance", 42, -1, 1), + "Lance fragment id must be non-negative"); + } + private static LanceScanNode newNode() { return new LanceScanNode( new PlanNodeId(0), @@ -98,8 +226,26 @@ private static void setMetadata(LanceScanNode node, LanceTableMetadata metadata) private static void assertSplit(Split split, long fragmentId, long targetRows, long weight) { LanceSplit lanceSplit = (LanceSplit) split; - Assert.assertEquals(fragmentId, lanceSplit.getFragmentId()); + Assert.assertEquals(Collections.singletonList(fragmentId), lanceSplit.getFragmentIds()); Assert.assertEquals(targetRows, lanceSplit.getTargetSplitSize().longValue()); Assert.assertEquals(weight, lanceSplit.getSplitWeight().getRawValue()); } + + private static void assertInvalidSplit(Runnable action, String expectedMessage) { + try { + action.run(); + Assert.fail("Expected IllegalArgumentException"); + } catch (IllegalArgumentException e) { + Assert.assertEquals(expectedMessage, e.getMessage()); + } + } + + private static TExternalSearchRequest vectorSearchRequest(long topK, long offset) { + TVectorSearchParams vector = new TVectorSearchParams() + .setColumn("vector") + .setTopK(topK) + .setOffset(offset); + return new TExternalSearchRequest() + .setSearchQuery(TExternalSearchQuery.vector_search(vector)); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/tvf/source/TVFScanNodeTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/tvf/source/TVFScanNodeTest.java index 03d489e58a0e5e..f7d2ea8a702c52 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/tvf/source/TVFScanNodeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/tvf/source/TVFScanNodeTest.java @@ -23,7 +23,7 @@ import org.apache.doris.catalog.FunctionGenTable; import org.apache.doris.datasource.FileQueryScanNode; import org.apache.doris.datasource.FileSplitter; -import org.apache.doris.datasource.lance.LanceTableMetadata; +import org.apache.doris.datasource.lance.LanceFragmentInfo; import org.apache.doris.datasource.lance.source.LanceSplit; import org.apache.doris.planner.PlanNodeId; import org.apache.doris.planner.ScanContext; @@ -129,8 +129,8 @@ public void testLanceSplitsPinVersionAndFragments() throws Exception { Mockito.when(tvf.getFilePath()).thenReturn("s3://bucket/table.lance"); Mockito.when(tvf.getLanceDatasetVersion()).thenReturn(42L); Mockito.when(tvf.getLanceFragments()).thenReturn(Arrays.asList( - new LanceTableMetadata.LanceFragmentInfo(7, 1000, 1000), - new LanceTableMetadata.LanceFragmentInfo(11, 250, 250))); + new LanceFragmentInfo(7, 1000, 1000), + new LanceFragmentInfo(11, 250, 250))); desc.setTable(table); TVFScanNode node = new TVFScanNode(new PlanNodeId(0), desc, false, sv, ScanContext.EMPTY); @@ -140,7 +140,7 @@ public void testLanceSplitsPinVersionAndFragments() throws Exception { LanceSplit first = (LanceSplit) splits.get(0); Assert.assertEquals("s3://bucket/table.lance", first.getDatasetUri()); Assert.assertEquals(42L, first.getVersion()); - Assert.assertEquals(7L, first.getFragmentId()); + Assert.assertEquals(Collections.singletonList(7L), first.getFragmentIds()); // The S3/file TVF path must weight fragments by physical rows, in sync with LanceScanNode. Assert.assertEquals(100L, first.getSplitWeight().getRawValue()); Assert.assertEquals(25L, ((LanceSplit) splits.get(1)).getSplitWeight().getRawValue()); @@ -175,7 +175,7 @@ public void testLocalLanceUsesOneLatestWholeDatasetSplit() throws Exception { LanceSplit split = (LanceSplit) splits.get(0); Assert.assertEquals("/data/table.lance", split.getDatasetUri()); Assert.assertEquals(0L, split.getVersion()); - Assert.assertFalse(split.hasFragmentId()); + Assert.assertFalse(split.hasFragmentIds()); TFileRangeDesc range = new TFileRangeDesc(); node.setScanParams(range, split); diff --git a/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java index 32786124a8bff7..18985cbfbe3a84 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java @@ -19,10 +19,18 @@ import org.apache.doris.analysis.TableName; import org.apache.doris.common.AnalysisException; +import org.apache.doris.datasource.lance.LanceTableMetadata; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; import org.junit.Assert; import org.junit.Test; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; + public class VectorSearchTableValuedFunctionTest { @Test public void testParseQuotedMultiLevelNamespace() throws AnalysisException { @@ -42,4 +50,35 @@ public void testRejectAmbiguousUnquotedMultiLevelNamespace() { Assert.assertTrue(exception.getMessage().contains("catalog.database.table")); } + + @Test + public void testValidateAndEncodeSqlFilterInFrontend() throws Exception { + Assert.assertArrayEquals("category = 'book'".getBytes(StandardCharsets.UTF_8), + VectorSearchTableValuedFunction.validateAndEncodeSqlFilter( + "category = 'book'")); + + AnalysisException empty = Assert.assertThrows(AnalysisException.class, + () -> VectorSearchTableValuedFunction.validateAndEncodeSqlFilter(" ")); + Assert.assertTrue(empty.getMessage().contains("must not be empty")); + + AnalysisException nul = Assert.assertThrows(AnalysisException.class, + () -> VectorSearchTableValuedFunction.validateAndEncodeSqlFilter( + "category = 'book'\0 OR true")); + Assert.assertTrue(nul.getMessage().contains("NUL")); + } + + @Test + public void testRejectCaseInsensitiveDuplicateOutputColumnsInFrontend() { + Schema schema = new Schema(Arrays.asList( + Field.nullable("Category", ArrowType.Utf8.INSTANCE), + Field.nullable("category", ArrowType.Utf8.INSTANCE))); + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", 42, schema, + Collections.emptyList(), Collections.emptyMap()); + + AnalysisException duplicate = Assert.assertThrows(AnalysisException.class, + () -> VectorSearchTableValuedFunction.buildOutputColumns(metadata)); + + Assert.assertTrue(duplicate.getMessage().contains("case-insensitive")); + } } diff --git a/gensrc/thrift/PlanNodes.thrift b/gensrc/thrift/PlanNodes.thrift index d3f0583c2206a1..6c22a6182670f5 100644 --- a/gensrc/thrift/PlanNodes.thrift +++ b/gensrc/thrift/PlanNodes.thrift @@ -496,9 +496,9 @@ struct TSearchFilter { 2: optional binary payload } -// Lance-only vector search tuning. Logical query fields stay in TVectorSearchParams so another -// provider, such as Paimon, can reuse the same request without depending on Lance options. -struct TLanceVectorSearchOptions { +// Optional vector search tuning. Logical query fields stay in TVectorSearchParams so providers +// can reuse the same request and interpret the supported tuning fields independently. +struct TVectorSearchOptions { 1: optional i32 nprobes 2: optional i32 refine_factor 3: optional i32 ef @@ -508,17 +508,17 @@ struct TLanceVectorSearchOptions { // The active union field identifies the logical search kind. A future hybrid field can contain both // vector and full-text subqueries plus its fusion parameters without changing either existing field. union TExternalSearchQuery { - 1: TVectorSearchParams vector - 2: TFullTextSearchParams full_text + 1: TVectorSearchParams vector_search + 2: TFullTextSearchParams full_text_search } // A provider-independent logical search request. Physical target information remains in the // provider FileDesc (for example, dataset_uri/version/fragment_ids in TLanceFileDesc). struct TExternalSearchRequest { 1: optional i32 schema_version = 1 - 2: optional TExternalSearchQuery query - 3: optional TSearchFilter filter - 4: optional TLanceVectorSearchOptions lance_options + 2: optional TExternalSearchQuery search_query + 3: optional TSearchFilter search_filter + 4: optional TVectorSearchOptions vector_search_options } // A catalog/S3 range reads fragments from a fixed snapshot. A local TVF range uses version zero @@ -631,10 +631,11 @@ struct TFileScanRangeParams { 35: optional string serialized_table_cache_key // Serialized Substrait ExtendedExpression executed by the native Lance scanner. Set at // ScanNode level so it is not serialized once per fragment split. - 36: optional binary lance_substrait_filter + 37: optional binary lance_substrait_filter // Provider-independent search request. Set at ScanNode level so all ranges use the same logical - // query. The first implementation uses one whole-dataset range for Lance vector search. - 37: optional TExternalSearchRequest external_search_request + // query. Lance vector search uses one range per fragment and Doris merges the split-local + // candidates. + 38: optional TExternalSearchRequest external_search_request } struct TFileRangeDesc { diff --git a/regression-test/data/external_table_p0/lance/test_lance_vector_search.out b/regression-test/data/external_table_p0/lance/test_lance_vector_search.out index 0a98ba5b3a717a..9dca6511c639ce 100644 --- a/regression-test/data/external_table_p0/lance/test_lance_vector_search.out +++ b/regression-test/data/external_table_p0/lance/test_lance_vector_search.out @@ -60,4 +60,3 @@ row_id bigint No false \N 259 item-0259 144.0 252 item-0252 256.0 260 item-0260 256.0 - diff --git a/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy b/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy index bde006738bec9c..55c48650ad1388 100644 --- a/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy +++ b/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy @@ -243,4 +243,4 @@ suite("test_lance_vector_search", "p0,external") { } finally { // sql """DROP CATALOG IF EXISTS `${catalogName}`""" } -} +} \ No newline at end of file diff --git a/thirdparty/download-thirdparty.sh b/thirdparty/download-thirdparty.sh index 87eae93eaf4c82..715155c6044182 100755 --- a/thirdparty/download-thirdparty.sh +++ b/thirdparty/download-thirdparty.sh @@ -774,6 +774,19 @@ if [[ " ${TP_ARCHIVES[*]} " =~ " PAIMON_CPP " ]]; then echo "Finished patching ${PAIMON_CPP_SOURCE}" fi +# Patch lance-c to allow nearest-neighbor search over an explicitly selected fragment set. +if [[ " ${TP_ARCHIVES[*]} " =~ " LANCE_C " ]]; then + if [[ "${LANCE_C_SOURCE}" == "lance-c-0.1.6" ]]; then + cd "${TP_SOURCE_DIR}/${LANCE_C_SOURCE}" + if [[ ! -f "${PATCHED_MARK}" ]]; then + patch -p1 <"${TP_PATCH_DIR}/lance-c-0.1.6-fragment-knn.patch" + touch "${PATCHED_MARK}" + fi + cd - + fi + echo "Finished patching ${LANCE_C_SOURCE}" +fi + if [[ " ${TP_ARCHIVES[*]} " =~ " CCTZ " ]] ; then cd $TP_SOURCE_DIR/$CCTZ_SOURCE if [[ ! -f "$PATCHED_MARK" ]] ; then diff --git a/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch b/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch new file mode 100644 index 00000000000000..ef38ea0c5289aa --- /dev/null +++ b/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch @@ -0,0 +1,48 @@ +diff --git a/src/scanner.rs b/src/scanner.rs +--- a/src/scanner.rs ++++ b/src/scanner.rs +@@ -158,6 +158,12 @@ impl LanceScanner { + location: snafu::location!(), + }); + } ++ // Lance validates fragment-scoped nearest searches when nearest() is ++ // configured. Such searches are supported when the fragment scan is ++ // the input to a prefilter, so this flag must be set first. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } + if let Some(n) = &self.nearest { + scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; + if let Some(np) = self.nprobes { +@@ -175,9 +181,6 @@ impl LanceScanner { + if let Some(ui) = self.use_index { + scanner.use_index(ui); + } +- if self.prefilter { +- scanner.prefilter(true); +- } + if let Some(segments) = &self.index_segments { + scanner.with_index_segments(segments.clone())?; + } +@@ -219,6 +222,11 @@ impl LanceScanner { + location: snafu::location!(), + }); + } ++ // nearest() checks the current prefilter setting before accepting a ++ // fragment-scoped search. Enable it before installing the query. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } + if let Some(n) = &self.nearest { + scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; + if let Some(np) = self.nprobes { +@@ -236,9 +244,6 @@ impl LanceScanner { + if let Some(ui) = self.use_index { + scanner.use_index(ui); + } +- if self.prefilter { +- scanner.prefilter(true); +- } + if let Some(segments) = &self.index_segments { + scanner.with_index_segments(segments.clone())?; + } From 27bc405c64e0898e9b100fda8379fc9733eb43a0 Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Mon, 17 Aug 2026 17:09:08 +0800 Subject: [PATCH 2/4] fetch rowid --- .../exec/operator/materialization_opertor.cpp | 49 +- be/src/exec/rowid_fetcher.cpp | 138 ++++- be/src/exec/rowid_fetcher.h | 13 +- be/src/format_v2/table/lance_reader.cpp | 205 ++++++-- be/src/format_v2/table/lance_reader.h | 15 +- be/src/storage/utils.h | 23 + .../materialization_shared_state_test.cpp | 34 +- be/test/format_v2/table/lance_reader_test.cpp | 175 ++++++- .../materialize/MaterializeProbeVisitor.java | 14 + .../MaterializeProbeVisitorTest.java | 57 +++ gensrc/proto/internal_service.proto | 3 +- thirdparty/download-thirdparty.sh | 4 +- thirdparty/patches/lance-c-0.1.6-doris.patch | 473 ++++++++++++++++++ .../patches/lance-c-0.1.6-fragment-knn.patch | 48 -- 14 files changed, 1128 insertions(+), 123 deletions(-) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitorTest.java create mode 100644 thirdparty/patches/lance-c-0.1.6-doris.patch delete mode 100644 thirdparty/patches/lance-c-0.1.6-fragment-knn.patch diff --git a/be/src/exec/operator/materialization_opertor.cpp b/be/src/exec/operator/materialization_opertor.cpp index 629df2ad52baab..03fc98098dd92f 100644 --- a/be/src/exec/operator/materialization_opertor.cpp +++ b/be/src/exec/operator/materialization_opertor.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -66,6 +67,48 @@ constexpr const char* TOPN_LAZY_MAT_PHASE2_PER_BACKEND_REMOTE_IO_TIME = constexpr const char* TOPN_LAZY_MAT_PHASE2_PER_BACKEND_WRITE_CACHE_IO_TIME = "TopNLazyMaterializationSecondPhasePerBackendWriteCacheIOTime"; +struct MaterializationRowLocation { + int64_t backend_id = 0; + uint32_t file_id = 0; + uint64_t row_id = 0; +}; + +Status decode_materialization_row_location(const StringRef& encoded, + MaterializationRowLocation* decoded) { + if (encoded.size == 0) { + return Status::InternalError("global row location is empty"); + } + + const auto version = static_cast(encoded.data[0]); + if (version == GlobalRowLoacationV2::VERSION) { + if (encoded.size != sizeof(GlobalRowLoacationV2)) { + return Status::InternalError( + "invalid global row location V2 size: actual={}, expected={}", encoded.size, + sizeof(GlobalRowLoacationV2)); + } + GlobalRowLoacationV2 location(0, 0, 0, 0); + std::memcpy(&location, encoded.data, sizeof(location)); + decoded->backend_id = location.backend_id; + decoded->file_id = location.file_id; + decoded->row_id = location.row_id; + return Status::OK(); + } + if (version == GlobalRowLocationV3::VERSION) { + if (encoded.size != sizeof(GlobalRowLocationV3)) { + return Status::InternalError( + "invalid global row location V3 size: actual={}, expected={}", encoded.size, + sizeof(GlobalRowLocationV3)); + } + GlobalRowLocationV3 location(0, 0, 0); + std::memcpy(&location, encoded.data, sizeof(location)); + decoded->backend_id = location.backend_id; + decoded->file_id = location.file_id; + decoded->row_id = location.row_id; + return Status::OK(); + } + return Status::InternalError("unsupported global row location version: {}", version); +} + void update_counter(RuntimeProfile* profile, const std::string& name, TUnit::type unit, int64_t value) { COUNTER_UPDATE(ADD_COUNTER_WITH_LEVEL(profile, name, unit, 2), value); @@ -446,9 +489,9 @@ Status MaterializationSharedState::create_muiltget_result(const Columns& columns for (int j = 0; j < rows; ++j) { if (!null_map || !null_map[j]) { - DCHECK(column_rowid->get_data_at(j).size == sizeof(GlobalRowLoacationV2)); - GlobalRowLoacationV2 row_location = - *((GlobalRowLoacationV2*)column_rowid->get_data_at(j).data); + MaterializationRowLocation row_location; + RETURN_IF_ERROR(decode_materialization_row_location( + column_rowid->get_data_at(j), &row_location)); auto rpc_struct = rpc_struct_map.find(row_location.backend_id); if (UNLIKELY(rpc_struct == rpc_struct_map.end())) { return Status::InternalError( diff --git a/be/src/exec/rowid_fetcher.cpp b/be/src/exec/rowid_fetcher.cpp index cf0440d6b738a1..7b853c029c2404 100644 --- a/be/src/exec/rowid_fetcher.cpp +++ b/be/src/exec/rowid_fetcher.cpp @@ -30,6 +30,7 @@ #include #include +#include #include #include #include @@ -58,6 +59,7 @@ #include "exec/scan/file_scanner.h" #include "format/orc/vorc_reader.h" #include "format/parquet/vparquet_reader.h" +#include "format_v2/table/lance_reader.h" #include "io/io_common.h" #include "runtime/descriptors.h" #include "runtime/exec_env.h" // ExecEnv @@ -658,7 +660,11 @@ Status RowIdStorageReader::read_batch_doris_format_row( auto max_k = 0; for (int j = 0; j < request_block_desc.row_id_size();) { auto file_id = request_block_desc.file_id(j); - row_ids.emplace_back(request_block_desc.row_id(j)); + const auto row_id = request_block_desc.row_id(j); + if (row_id > std::numeric_limits::max()) { + return Status::InvalidArgument("internal row id exceeds uint32 range: {}", row_id); + } + row_ids.emplace_back(static_cast(row_id)); auto file_mapping = id_file_map->get_file_mapping(file_id); if (!file_mapping) { return Status::InternalError( @@ -667,7 +673,12 @@ Status RowIdStorageReader::read_batch_doris_format_row( } for (k = 1; j + k < request_block_desc.row_id_size(); ++k) { if (request_block_desc.file_id(j + k) == file_id) { - row_ids.emplace_back(request_block_desc.row_id(j + k)); + const auto next_row_id = request_block_desc.row_id(j + k); + if (next_row_id > std::numeric_limits::max()) { + return Status::InvalidArgument("internal row id exceeds uint32 range: {}", + next_row_id); + } + row_ids.emplace_back(static_cast(next_row_id)); } else { break; } @@ -713,22 +724,60 @@ const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRowsRead const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseSegmentsRead = "TopNLazyMaterializationSecondPhaseSegmentsRead"; +Status RowIdStorageReader::read_lance_rows_by_row_ids( + const TFileRangeDesc& scan_range_desc, const std::vector& row_ids, + const std::vector& slots, RuntimeState* runtime_state, + RuntimeProfile* runtime_profile, const TFileScanRangeParams& scan_params, Block* block, + RowIdStorageReader::ExternalFetchStatistics* fetch_statistics) { + // Unlike Parquet/ORC, a Lance row ID is a native uint64 ID in a fixed dataset snapshot, + // rather than a row ordinal interpreted by a reader for one physical file range. Phase-two + // fetch therefore opens that dataset snapshot and uses dataset-level take_rows directly. It + // does not create a FileScanner or rescan the fragment split that produced the row ID. + std::vector projected_columns; + projected_columns.reserve(slots.size()); + for (const auto& slot : slots) { + format::ColumnDefinition column; + column.identifier = Field::create_field(slot.col_name()); + column.name = slot.col_name(); + column.type = slot.get_data_type_ptr(); + projected_columns.emplace_back(std::move(column)); + } + + format::lance::LanceTableReader reader; + auto lance_scan_params = scan_params; + RETURN_IF_ERROR(scope_timer_run( + [&]() { + return reader.init({ + .projected_columns = projected_columns, + .conjuncts = {}, + .format = format::FileFormat::LANCE, + .scan_params = &lance_scan_params, + .io_ctx = nullptr, + .runtime_state = runtime_state, + .scanner_profile = runtime_profile, + }); + }, + &fetch_statistics->init_reader_ms)); + RETURN_IF_ERROR(scope_timer_run( + [&]() { return reader.read_by_row_ids(scan_range_desc, row_ids, block); }, + &fetch_statistics->get_block_ms)); + return reader.close(); +} + Status RowIdStorageReader::read_external_row_from_file_mapping( - size_t idx, const std::multimap& row_ids, + size_t idx, const std::multimap& row_ids, const std::shared_ptr& file_mapping, const std::vector& slots, const TUniqueId& query_id, const std::shared_ptr& runtime_state, std::vector& scan_blocks, std::vector>& row_id_block_idx, std::vector& fetch_statistics, const TFileScanRangeParams& rpc_scan_params, const std::unordered_map& colname_to_slot_id, - std::atomic& producer_count, size_t scan_rows_count, - std::counting_semaphore<>& semaphore, std::condition_variable& cv, std::mutex& mtx, TupleDescriptor& tuple_desc) { SCOPED_ATTACH_TASK(ExecEnv::GetInstance()->rowid_storage_reader_tracker()); signal::set_signal_task_id(query_id); - std::list read_ids; - //Generate an ordered list with the help of the orderliness of the map. + std::vector read_ids; + // Generate an ordered, deduplicated list with the help of the multimap ordering. for (const auto& [row_id, result_block_idx] : row_ids) { if (read_ids.empty() || read_ids.back() != row_id) { read_ids.emplace_back(row_id); @@ -749,14 +798,31 @@ Status RowIdStorageReader::read_external_row_from_file_mapping( std::unique_ptr sub_runtime_profile = std::make_unique("ExternalRowIDFetcher"); - { + const auto format_type = scan_range_desc.__isset.format_type ? scan_range_desc.format_type + : rpc_scan_params.format_type; + if (format_type == TFileFormatType::FORMAT_LANCE) { + RETURN_IF_ERROR(read_lance_rows_by_row_ids( + scan_range_desc, read_ids, slots, runtime_state.get(), sub_runtime_profile.get(), + rpc_scan_params, &scan_blocks[idx], &fetch_statistics[idx])); + } else { + // Parquet/ORC row IDs are consumed as row ordinals within the exact physical file range + // recorded by phase one. Keep using FileScanner so the format reader can resolve those + // ordinals against that range; unlike Lance, ranges cannot be merged at dataset level. + std::list legacy_read_ids; + for (const auto row_id : read_ids) { + if (row_id > static_cast(std::numeric_limits::max())) { + return Status::InvalidArgument("legacy external row id exceeds int64 range: {}", + row_id); + } + legacy_read_ids.emplace_back(static_cast(row_id)); + } std::unique_ptr vfile_scanner_ptr = FileScanner::create_unique(runtime_state.get(), sub_runtime_profile.get(), &rpc_scan_params, &colname_to_slot_id, &tuple_desc); RETURN_IF_ERROR(vfile_scanner_ptr->prepare_for_read_lines(scan_range_desc)); RETURN_IF_ERROR(vfile_scanner_ptr->read_lines_from_range( - scan_range_desc, read_ids, &scan_blocks[idx], external_info, + scan_range_desc, legacy_read_ids, &scan_blocks[idx], external_info, &fetch_statistics[idx].init_reader_ms, &fetch_statistics[idx].get_block_ms)); } @@ -775,11 +841,6 @@ Status RowIdStorageReader::read_external_row_from_file_mapping( file_read_times_counter->value(), file_read_times_counter->type()); } - semaphore.release(); - if (++producer_count == scan_rows_count) { - std::lock_guard lock(mtx); - cv.notify_one(); - } return Status::OK(); } @@ -850,9 +911,10 @@ Status RowIdStorageReader::read_batch_external_row( } // Hash(TFileRangeDesc) => { all the rows that need to be read and their positions in the result block. } + file mapping - // std::multimap : The reason for using multimap is: may need the same row of data multiple times. + // The multimap retains duplicate row IDs because the same source row can appear more than once + // in the materialized result. std::map, std::shared_ptr>> + std::pair, std::shared_ptr>> scan_rows; // Block corresponding to the order of `scan_rows` map. @@ -864,7 +926,22 @@ Status RowIdStorageReader::read_batch_external_row( // Count the time/bytes it takes to read each TFileRangeDesc. (for profile) std::vector fetch_statistics; - auto hash_file_range = [](const TFileRangeDesc& file_range_desc) { + auto hash_file_range = [&rpc_scan_params](const ExternalFileMappingInfo& external_info) { + const auto& file_range_desc = external_info.scan_range_desc; + const auto format_type = file_range_desc.__isset.format_type + ? file_range_desc.format_type + : rpc_scan_params.format_type; + if (format_type == TFileFormatType::FORMAT_LANCE) { + // Parquet and ORC row IDs are offsets within a physical file range, so their fetch + // path must keep each path/start_offset pair separate. Lance row IDs instead belong + // to a fixed dataset snapshot. Although phase one registers one file mapping per + // fragment split, phase two uses dataset-level take_rows and can fetch row IDs from + // all of those fragments together. Group the mappings by scan node, snapshot, and + // dataset URI; the plan node keeps independent scans of the same snapshot isolated. + const auto& lance_params = file_range_desc.table_format_params.lance_params; + return fmt::format("lance:{}:{}:{}", external_info.plan_node_id, + lance_params.version, lance_params.dataset_uri); + } std::string value; value.resize(file_range_desc.path.size() + sizeof(file_range_desc.start_offset)); auto* ptr = value.data(); @@ -885,13 +962,11 @@ Status RowIdStorageReader::read_batch_external_row( } const auto& external_info = file_mapping->get_external_file_info(); - const auto& scan_range_desc = external_info.scan_range_desc; - - auto scan_range_hash = hash_file_range(scan_range_desc); + const auto& scan_range_hash = hash_file_range(external_info); if (scan_rows.contains(scan_range_hash)) { scan_rows.at(scan_range_hash).first.emplace(request_block_desc.row_id(j), j); } else { - std::multimap tmp {{request_block_desc.row_id(j), j}}; + std::multimap tmp {{request_block_desc.row_id(j), j}}; scan_rows.emplace(scan_range_hash, std::make_pair(tmp, file_mapping)); } } @@ -917,6 +992,7 @@ Status RowIdStorageReader::read_batch_external_row( std::atomic producer_count {0}; std::condition_variable cv; std::mutex mtx; + Status scan_status = Status::OK(); //semaphore: Limit the number of scan tasks submitted at one time std::counting_semaphore semaphore {max_file_scanners}; @@ -927,13 +1003,26 @@ Status RowIdStorageReader::read_batch_external_row( RETURN_IF_ERROR(remote_scan_sched->submit_scan_task( SimplifiedScanTask( [&, idx, scan_info]() -> Status { + Defer complete_task {[&]() { + semaphore.release(); + if (++producer_count == scan_rows.size()) { + std::lock_guard lock(mtx); + cv.notify_one(); + } + }}; const auto& [row_ids, file_mapping] = scan_info; - return read_external_row_from_file_mapping( + auto status = read_external_row_from_file_mapping( idx, row_ids, file_mapping, slots, query_id, runtime_state, scan_blocks, row_id_block_idx, fetch_statistics, rpc_scan_params, - colname_to_slot_id, producer_count, - scan_rows.size(), semaphore, cv, mtx, tuple_desc); + colname_to_slot_id, tuple_desc); + if (!status.ok()) { + std::lock_guard lock(mtx); + if (scan_status.ok()) { + scan_status = status; + } + } + return status; }, nullptr, nullptr), fmt::format("{}-read_batch_external_row-{}", print_id(query_id), idx))); @@ -944,6 +1033,7 @@ Status RowIdStorageReader::read_batch_external_row( std::unique_lock lock(mtx); cv.wait(lock, [&] { return producer_count == scan_rows.size(); }); } + RETURN_IF_ERROR(scan_status); return Status::OK(); }, &scan_running_time)); diff --git a/be/src/exec/rowid_fetcher.h b/be/src/exec/rowid_fetcher.h index 790f9cf17e7e4e..41acea97c1da6d 100644 --- a/be/src/exec/rowid_fetcher.h +++ b/be/src/exec/rowid_fetcher.h @@ -36,6 +36,7 @@ namespace doris { class DorisNodesInfo; +class RuntimeProfile; class RuntimeState; class TupleDescriptor; namespace io { @@ -142,8 +143,16 @@ class RowIdStorageReader { Block& result_block, PRuntimeProfileTree* pprofile, int64_t* init_reader_avg_ms, int64_t* get_block_avg_ms, size_t* scan_range_cnt); + static Status read_lance_rows_by_row_ids(const TFileRangeDesc& scan_range_desc, + const std::vector& row_ids, + const std::vector& slots, + RuntimeState* runtime_state, + RuntimeProfile* runtime_profile, + const TFileScanRangeParams& scan_params, Block* block, + ExternalFetchStatistics* fetch_statistics); + static Status read_external_row_from_file_mapping( - size_t idx, const std::multimap& row_ids, + size_t idx, const std::multimap& row_ids, const std::shared_ptr& file_mapping, const std::vector& slots, const TUniqueId& query_id, const std::shared_ptr& runtime_state, std::vector& scan_blocks, @@ -151,8 +160,6 @@ class RowIdStorageReader { std::vector& fetch_statistics, const TFileScanRangeParams& rpc_scan_params, const std::unordered_map& colname_to_slot_id, - std::atomic& producer_count, size_t scan_rows_count, - std::counting_semaphore<>& semaphore, std::condition_variable& cv, std::mutex& mtx, TupleDescriptor& tuple_desc); struct ExternalFetchStatistics { diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index fc388de14a4be5..f5f3932a7a28ba 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -17,6 +17,7 @@ #include "format_v2/table/lance_reader.h" +#include #include #include #include @@ -28,12 +29,16 @@ #include #include +#include "common/consts.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" #include "core/data_type/data_type_array.h" #include "core/data_type/data_type_factory.hpp" #include "core/data_type/data_type_map.h" #include "core/data_type/data_type_nothing.h" #include "core/data_type/data_type_struct.h" #include "exec/common/endian.h" +#include "storage/utils.h" namespace doris::format::lance { namespace { @@ -51,24 +56,9 @@ struct LanceBatchDeleter { }; constexpr std::string_view DISTANCE_COLUMN = "_distance"; +constexpr std::string_view ROW_ID_COLUMN = "_rowid"; constexpr std::string_view ARROW_EXTENSION_NAME = "ARROW:extension:name"; -std::string vector_metric_name(TVectorMetric::type metric) { - switch (metric) { - case TVectorMetric::DEFAULT: - return "default"; - case TVectorMetric::L2: - return "l2"; - case TVectorMetric::COSINE: - return "cosine"; - case TVectorMetric::DOT_PRODUCT: - return "dot"; - case TVectorMetric::HAMMING: - return "hamming"; - } - return "unknown"; -} - int arrow_time_precision(arrow::TimeUnit::type unit) { switch (unit) { case arrow::TimeUnit::SECOND: @@ -300,13 +290,8 @@ Status LanceTableReader::init(TableReadOptions&& options) { if (_vector_search) { const auto& vector = _scan_params->external_search_request.search_query.vector_search; - _scanner_profile->add_info_string("ExternalSearchType", "VECTOR"); - _scanner_profile->add_info_string("LanceVectorColumn", vector.column); - _scanner_profile->add_info_string("LanceTopK", std::to_string(vector.top_k)); - _scanner_profile->add_info_string("LanceOffset", std::to_string(vector.offset)); - _scanner_profile->add_info_string("LanceMetric", vector.__isset.metric - ? vector_metric_name(vector.metric) - : "default"); + _scanner_profile->add_info_string("LanceFragmentTopK", std::to_string(vector.top_k)); + _scanner_profile->add_info_string("LanceFragmentOffset", std::to_string(vector.offset)); _scanner_profile->add_info_string("LanceVectorDimension", std::to_string(vector.query_vector.dimension)); } @@ -319,11 +304,29 @@ Status LanceTableReader::init(TableReadOptions&& options) { _output_name_to_idx.clear(); _output_name_to_idx.reserve(_projected_columns.size()); + _global_rowid_output_idx.reset(); for (size_t idx = 0; idx < _projected_columns.size(); ++idx) { const auto& column = _projected_columns[idx]; if (column.type == nullptr) { return Status::InvalidArgument("Lance projected column '{}' has no type", column.name); } + if (column.name.starts_with(BeConsts::GLOBAL_ROWID_COL)) { + if (!_vector_search) { + return Status::NotSupported( + "Lance global row id is currently supported only for vector search"); + } + if (_global_rowid_output_idx.has_value()) { + return Status::InvalidArgument( + "duplicate Lance global row id projected column: {}", column.name); + } + if (remove_nullable(column.type)->get_primitive_type() != TYPE_STRING) { + return Status::InvalidArgument( + "Lance global row id column '{}' must have Doris STRING type, but was {}", + column.name, column.type->get_name()); + } + _global_rowid_output_idx = idx; + continue; + } if (!_output_name_to_idx.emplace(column.name, idx).second) { return Status::InvalidArgument("duplicate Lance projected column: {}", column.name); } @@ -350,16 +353,12 @@ Status LanceTableReader::prepare_split(const SplitReadOptions& options) { if (current_split_pruned()) { return Status::OK(); } - - const auto key = _dataset_key(options.current_range); - if (_dataset == nullptr) { - RETURN_IF_ERROR(_open_dataset(key)); - _opened_dataset_key = key; - } else if (!_opened_dataset_key.has_value() || *_opened_dataset_key != key) { + if (_global_rowid_output_idx.has_value() && !_global_rowid_context.has_value()) { return Status::InvalidArgument( - "Lance reader cannot mix dataset snapshots or storage options in one scan"); + "Lance global row id requested without global row id context"); } + RETURN_IF_ERROR(_ensure_dataset_open(options.current_range)); RETURN_IF_ERROR(_open_scanner(options.current_range)); return Status::OK(); } @@ -403,7 +402,7 @@ Status LanceTableReader::get_block(Block* block, bool* eos) { std::unique_ptr batch(raw_batch); size_t rows = 0; - RETURN_IF_ERROR(_fill_block_from_arrow(batch.get(), block, &rows)); + RETURN_IF_ERROR(_fill_block_from_lance_batch(batch.get(), block, &rows)); _record_scan_rows(rows); raw_rows += rows; } @@ -426,6 +425,63 @@ Status LanceTableReader::get_block(Block* block, bool* eos) { } } +Status LanceTableReader::read_by_row_ids(const TFileRangeDesc& range, + const std::vector& row_ids, Block* block) { + DORIS_CHECK(block != nullptr); + DORIS_CHECK(block->columns() == _projected_columns.size()); + if (row_ids.empty()) { + return Status::OK(); + } + + RETURN_IF_ERROR(_ensure_dataset_open(range)); + std::vector columns; + columns.reserve(_projected_columns.size() + 1); + for (const auto& column : _projected_columns) { + columns.emplace_back(column.name.c_str()); + } + columns.emplace_back(nullptr); + + ArrowArrayStream stream {}; + if (lance_dataset_take_rows(_dataset, row_ids.data(), row_ids.size(), columns.data(), + &stream) != 0) { + if (stream.release != nullptr) { + stream.release(&stream); + } + return _lance_error("take Lance rows by row id"); + } + auto imported_reader = arrow::ImportRecordBatchReader(&stream); + if (!imported_reader.ok()) { + if (stream.release != nullptr) { + stream.release(&stream); + } + return Status::InternalError("import Lance take-rows stream failed: {}", + imported_reader.status().message()); + } + + size_t fetched_rows = 0; + auto batch_reader = std::move(imported_reader).ValueUnsafe(); + while (true) { + std::shared_ptr record_batch; + const auto read_status = batch_reader->ReadNext(&record_batch); + if (!read_status.ok()) { + return Status::InternalError("read Lance take-rows batch failed: {}", + read_status.message()); + } + if (record_batch == nullptr) { + break; + } + size_t rows = 0; + RETURN_IF_ERROR(_fill_block_from_record_batch(record_batch, block, &rows)); + fetched_rows += rows; + } + if (fetched_rows != row_ids.size()) { + return Status::InternalError( + "Lance row-id fetch returned {} rows for {} requested row ids", + fetched_rows, row_ids.size()); + } + return Status::OK(); +} + Status LanceTableReader::abort_split() { _close_scanner(); _eof = true; @@ -440,6 +496,18 @@ Status LanceTableReader::close() { return TableReader::close(); } +Status LanceTableReader::_ensure_dataset_open(const TFileRangeDesc& range) { + const auto key = _dataset_key(range); + if (_dataset == nullptr) { + RETURN_IF_ERROR(_open_dataset(key)); + _opened_dataset_key = key; + } else if (!_opened_dataset_key.has_value() || *_opened_dataset_key != key) { + return Status::InvalidArgument( + "Lance reader cannot mix dataset snapshots or storage options"); + } + return Status::OK(); +} + Status LanceTableReader::_open_dataset(const DatasetKey& key) { std::vector storage_option_ptrs; storage_option_ptrs.reserve(key.storage_options.size() + 1); @@ -460,10 +528,16 @@ Status LanceTableReader::_open_dataset(const DatasetKey& key) { Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { std::vector columns; columns.reserve(_projected_columns.size() + 1); - for (const auto& column : _projected_columns) { + for (size_t idx = 0; idx < _projected_columns.size(); ++idx) { + if (_global_rowid_output_idx == idx) { + continue; + } + const auto& column = _projected_columns[idx]; columns.emplace_back(column.name.c_str()); } - if (_vector_search && _projected_columns.empty()) { + if (_vector_search && columns.empty()) { + // Keep an explicit empty user projection from becoming `nullptr`, which means all dataset + // columns to lance-c. nearest() already returns this optional system column. columns.emplace_back(DISTANCE_COLUMN.data()); } columns.emplace_back(nullptr); @@ -483,6 +557,10 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { } std::unique_ptr scanner_guard(scanner); + if (_global_rowid_output_idx.has_value() && lance_scanner_with_row_id(scanner, true) != 0) { + return _lance_error("enable Lance row id output"); + } + if (_scan_params->__isset.lance_substrait_filter && !_scan_params->lance_substrait_filter.empty()) { const auto& filter = _scan_params->lance_substrait_filter; @@ -674,7 +752,8 @@ void LanceTableReader::_close_dataset() { } } -Status LanceTableReader::_fill_block_from_arrow(LanceBatch* batch, Block* block, size_t* rows) { +Status LanceTableReader::_fill_block_from_lance_batch(LanceBatch* batch, Block* block, + size_t* rows) { DORIS_CHECK(batch != nullptr); DORIS_CHECK(block != nullptr); DORIS_CHECK(rows != nullptr); @@ -695,7 +774,52 @@ Status LanceTableReader::_fill_block_from_arrow(LanceBatch* batch, Block* block, result.status().message()); } - const auto record_batch = std::move(result).ValueUnsafe(); + return _fill_block_from_record_batch(std::move(result).ValueUnsafe(), block, rows); +} + +Status LanceTableReader::_append_global_row_ids(const std::shared_ptr& row_ids, + MutableColumnPtr& output_column) const { + DORIS_CHECK(row_ids != nullptr); + DORIS_CHECK(_global_rowid_context.has_value()); + if (row_ids->type_id() != arrow::Type::UINT64) { + return Status::InternalError("Lance row id column must be Arrow UINT64, but was {}", + row_ids->type()->ToString()); + } + + ColumnString* data_column = nullptr; + ColumnUInt8::Container* null_map = nullptr; + if (auto* nullable = check_and_get_column(*output_column)) { + data_column = check_and_get_column(nullable->get_nested_column()); + null_map = &nullable->get_null_map_data(); + } else { + data_column = check_and_get_column(*output_column); + } + if (data_column == nullptr) { + return Status::InternalError("Lance global row id output column must be STRING"); + } + + const auto typed_row_ids = std::static_pointer_cast(row_ids); + if (typed_row_ids->null_count() != 0) { + return Status::InternalError("Lance returned null row id"); + } + const auto row_count = static_cast(typed_row_ids->length()); + if (null_map != nullptr) { + null_map->resize_fill(null_map->size() + row_count, 0); + } + const auto& context = *_global_rowid_context; + for (size_t row = 0; row < row_count; ++row) { + const GlobalRowLocationV3 location(context.backend_id, context.file_id, + typed_row_ids->Value(row)); + data_column->insert_data(reinterpret_cast(&location), sizeof(location)); + } + return Status::OK(); +} + +Status LanceTableReader::_fill_block_from_record_batch( + const std::shared_ptr& record_batch, Block* block, size_t* rows) { + DORIS_CHECK(record_batch != nullptr); + DORIS_CHECK(block != nullptr); + DORIS_CHECK(rows != nullptr); const auto row_count = static_cast(record_batch->num_rows()); std::unordered_set materialized_columns; materialized_columns.reserve(record_batch->num_columns()); @@ -703,6 +827,17 @@ Status LanceTableReader::_fill_block_from_arrow(LanceBatch* batch, Block* block, auto& columns = columns_guard.mutable_columns(); for (int arrow_idx = 0; arrow_idx < record_batch->num_columns(); ++arrow_idx) { const auto& field = record_batch->schema()->field(arrow_idx); + if (field->name() == ROW_ID_COLUMN && _global_rowid_output_idx.has_value()) { + const auto output_idx = *_global_rowid_output_idx; + const auto& output_name = _projected_columns[output_idx].name; + if (!materialized_columns.emplace(output_name).second) { + return Status::InternalError("Lance returned duplicate column '{}'", + ROW_ID_COLUMN); + } + RETURN_IF_ERROR( + _append_global_row_ids(record_batch->column(arrow_idx), columns[output_idx])); + continue; + } const auto output_it = _output_name_to_idx.find(field->name()); if (output_it == _output_name_to_idx.end()) { if (_vector_search && field->name() == DISTANCE_COLUMN) { diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 5f73db2af3e7ff..09e413eb208f3c 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -37,6 +37,8 @@ struct LanceDataset; struct LanceScanner; namespace arrow { +class Array; +class RecordBatch; class Schema; } // namespace arrow @@ -65,6 +67,11 @@ class LanceTableReader final : public TableReader { Status init(TableReadOptions&& options) override; Status prepare_split(const SplitReadOptions& options) override; Status get_block(Block* block, bool* eos) override; + // Fetch top-level projected columns by native Lance row IDs from one fixed dataset snapshot. + // Input order and duplicates are preserved by lance-c. Missing rows are rejected because row + // IDs produced by phase one must still exist in the same snapshot during materialization. + Status read_by_row_ids(const TFileRangeDesc& range, const std::vector& row_ids, + Block* block); Status abort_split() override; Status close() override; @@ -76,12 +83,17 @@ class LanceTableReader final : public TableReader { bool operator==(const DatasetKey&) const = default; }; + Status _ensure_dataset_open(const TFileRangeDesc& range); Status _open_dataset(const DatasetKey& key); Status _open_scanner(const TFileRangeDesc& range); Status _configure_vector_search(LanceScanner* scanner) const; void _close_scanner(); void _close_dataset(); - Status _fill_block_from_arrow(LanceBatch* batch, Block* block, size_t* rows); + Status _fill_block_from_lance_batch(LanceBatch* batch, Block* block, size_t* rows); + Status _fill_block_from_record_batch(const std::shared_ptr& record_batch, + Block* block, size_t* rows); + Status _append_global_row_ids(const std::shared_ptr& row_ids, + MutableColumnPtr& output_column) const; static std::vector _storage_options(const TFileScanRangeParams* scan_params); DatasetKey _dataset_key(const TFileRangeDesc& range) const; static Status _lance_error(std::string_view operation); @@ -90,6 +102,7 @@ class LanceTableReader final : public TableReader { LanceScanner* _scanner = nullptr; std::optional _opened_dataset_key; std::unordered_map _output_name_to_idx; + std::optional _global_rowid_output_idx; cctz::time_zone _ctz; size_t _scanner_batch_size = 0; bool _vector_search = false; diff --git a/be/src/storage/utils.h b/be/src/storage/utils.h index a7c8046968e1b3..7fb4f9a3834f15 100644 --- a/be/src/storage/utils.h +++ b/be/src/storage/utils.h @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -239,6 +240,8 @@ struct GlobalRowLoacation { }; struct GlobalRowLoacationV2 { + static constexpr uint8_t VERSION = 0; + GlobalRowLoacationV2(uint8_t ver, uint64_t bid, uint32_t fid, uint32_t rid) : version(ver), backend_id(bid), file_id(fid), row_id(rid) {} uint8_t version; @@ -249,4 +252,24 @@ struct GlobalRowLoacationV2 { auto operator<=>(const GlobalRowLoacationV2&) const = default; }; +// Global row location for data sources whose native row ID is wider than the +// uint32_t row ordinal carried by GlobalRowLoacationV2. Keep V2 unchanged for +// existing internal, Parquet, and ORC readers. +struct GlobalRowLocationV3 { + static constexpr uint8_t VERSION = 1; + + GlobalRowLocationV3(int64_t bid, uint32_t fid, uint64_t rid) + : version(VERSION), backend_id(bid), file_id(fid), row_id(rid) {} + uint8_t version; + std::array reserved_before_backend_id {}; + int64_t backend_id; + uint32_t file_id; + std::array reserved_before_row_id {}; + uint64_t row_id; + + auto operator<=>(const GlobalRowLocationV3&) const = default; +}; + +static_assert(sizeof(GlobalRowLocationV3) == 32); + } // namespace doris diff --git a/be/test/exec/operator/materialization_shared_state_test.cpp b/be/test/exec/operator/materialization_shared_state_test.cpp index 3632e479a2393c..56ae03723d857b 100644 --- a/be/test/exec/operator/materialization_shared_state_test.cpp +++ b/be/test/exec/operator/materialization_shared_state_test.cpp @@ -17,6 +17,8 @@ #include +#include + #include "core/column/column_vector.h" #include "core/data_type/data_type_number.h" #include "core/data_type/data_type_string.h" @@ -29,7 +31,7 @@ namespace doris { namespace { -void add_request_row(PRequestBlockDesc* request_block_desc, uint32_t row_id, uint32_t file_id) { +void add_request_row(PRequestBlockDesc* request_block_desc, uint64_t row_id, uint32_t file_id) { request_block_desc->add_row_id(row_id); request_block_desc->add_file_id(file_id); } @@ -125,6 +127,36 @@ TEST_F(MaterializationSharedStateTest, TestCreateMultiGetResult) { // Verify block_order_results EXPECT_EQ(_shared_state->block_order_results.size(), columns.size()); EXPECT_EQ(_shared_state->eos, true); + const auto& backend1_request = + _shared_state->rpc_struct_map[_backend_id1].request.request_block_descs(0); + ASSERT_EQ(backend1_request.row_id_size(), 1); + EXPECT_EQ(backend1_request.row_id(0), 1); + const auto& backend2_request = + _shared_state->rpc_struct_map[_backend_id2].request.request_block_descs(0); + ASSERT_EQ(backend2_request.row_id_size(), 1); + EXPECT_EQ(backend2_request.row_id(0), 2); +} + +TEST_F(MaterializationSharedStateTest, TestCreateMultiGetResultWithUint64RowId) { + Columns columns; + auto rowid_col = _string_type->create_column(); + auto* col_data = reinterpret_cast(rowid_col.get()); + + constexpr uint64_t large_row_id = + static_cast(std::numeric_limits::max()) + 17; + GlobalRowLocationV3 location(_backend_id1, 7, large_row_id); + col_data->insert_data(reinterpret_cast(&location), sizeof(location)); + columns.push_back(std::move(rowid_col)); + + ASSERT_TRUE(_shared_state->create_muiltget_result(columns, false, false).ok()); + + const auto& request = _shared_state->rpc_struct_map[_backend_id1].request; + ASSERT_EQ(request.request_block_descs_size(), 1); + const auto& request_block = request.request_block_descs(0); + ASSERT_EQ(request_block.row_id_size(), 1); + EXPECT_EQ(request_block.row_id(0), large_row_id); + EXPECT_EQ(request_block.file_id(0), 7); + EXPECT_EQ(_shared_state->block_order_results[0][0], _backend_id1); } TEST_F(MaterializationSharedStateTest, TestMergeMultiResponse) { diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 16a316423fa2fd..ee2105b0742285 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -28,16 +28,20 @@ #include #include #include +#include #include #include #include +#include #include #include +#include #include #include #include #include +#include "common/consts.h" #include "core/assert_cast.h" #include "core/block/block.h" #include "core/column/column_array.h" @@ -56,6 +60,7 @@ #include "exprs/vexpr.h" #include "runtime/runtime_profile.h" #include "runtime/runtime_state.h" +#include "storage/utils.h" #include "util/defer_op.h" #include "util/timezone_utils.h" #include "util/url_coding.h" @@ -163,7 +168,9 @@ Status init_reader(LanceTableReader* reader, const Columns& projected_columns, }); } -Status prepare_range(LanceTableReader* reader, TFileRangeDesc range) { +Status prepare_range( + LanceTableReader* reader, TFileRangeDesc range, + std::optional global_rowid_context = std::nullopt) { return reader->prepare_split({.partition_values = {}, .conjuncts = std::nullopt, .partition_prune_conjuncts = {}, @@ -172,13 +179,15 @@ Status prepare_range(LanceTableReader* reader, TFileRangeDesc range) { .cache = nullptr, .current_range = std::move(range), .current_split_format = FileFormat::LANCE, - .global_rowid_context = std::nullopt}); + .global_rowid_context = global_rowid_context}); } Status prepare_fixture(LanceTableReader* reader, const std::filesystem::path& dataset_uri, - const LanceFixtureInfo& fixture, std::vector fragment_ids) { - return prepare_range(reader, - make_lance_range(dataset_uri, fixture.version, std::move(fragment_ids))); + const LanceFixtureInfo& fixture, std::vector fragment_ids, + std::optional global_rowid_context = std::nullopt) { + return prepare_range( + reader, make_lance_range(dataset_uri, fixture.version, std::move(fragment_ids)), + global_rowid_context); } TFileScanRangeParams make_float32_vector_search_params( @@ -244,6 +253,16 @@ std::vector> read_vector_search_rows(LanceTableReader* return rows; } +GlobalRowLocationV3 decode_lance_row_id(const ColumnString& column, size_t row) { + const auto encoded = column.get_data_at(row); + EXPECT_EQ(sizeof(GlobalRowLocationV3), encoded.size); + GlobalRowLocationV3 location(0, 0, 0); + if (encoded.size == sizeof(location)) { + std::memcpy(&location, encoded.data, sizeof(location)); + } + return location; +} + TEST(LanceTableReaderVectorSearchTest, SearchesWholeSnapshotWithOffsetAndDistance) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/all_types.lance"; @@ -348,6 +367,152 @@ TEST(LanceTableReaderVectorSearchTest, SearchesMultipleFragmentSplits) { EXPECT_TRUE(reader.close().ok()); } +TEST(LanceTableReaderVectorSearchTest, ReturnsStableGlobalRowIdsAndFetchesPayload) { + const std::filesystem::path dataset_uri = + "./be/test/format_v2/table/lance/data/all_types.lance"; + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(dataset_uri, &fixture).ok()); + + const auto global_rowid_name = BeConsts::GLOBAL_ROWID_COL + std::string("topn_fetch_lance"); + const Columns columns { + projected_column("row_id", TYPE_BIGINT, false), + projected_column(global_rowid_name, TYPE_STRING, false), + }; + TQueryOptions query_options; + query_options.__set_batch_size(1); + TQueryGlobals query_globals; + RuntimeState state(query_globals); + state.set_query_options(query_options); + RuntimeProfile profile("lance_vector_search_global_rowid_fixture"); + auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 4, 0); + const GlobalRowIdContext context {.backend_id = 123456789, .file_id = 42}; + + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); + const auto scan_row_ids = [&]() { + std::map row_ids; + const auto prepare_status = + prepare_fixture(&reader, dataset_uri, fixture, fixture.fragment_ids, context); + EXPECT_TRUE(prepare_status.ok()) << prepare_status.to_string(); + if (!prepare_status.ok()) { + return row_ids; + } + Block block; + add_output_columns(&block, columns); + bool eos = false; + while (!eos) { + const auto read_status = reader.get_block(&block, &eos); + EXPECT_TRUE(read_status.ok()) << read_status.to_string(); + if (!read_status.ok()) { + break; + } + if (eos) { + continue; + } + const auto& logical_row_ids = + assert_cast(*block.get_by_position(0).column); + const auto& global_row_ids = + assert_cast(*block.get_by_position(1).column); + for (size_t row = 0; row < block.rows(); ++row) { + const auto location = decode_lance_row_id(global_row_ids, row); + EXPECT_EQ(GlobalRowLocationV3::VERSION, location.version); + EXPECT_EQ(context.backend_id, location.backend_id); + EXPECT_EQ(context.file_id, location.file_id); + row_ids.emplace(logical_row_ids.get_data()[row], location.row_id); + } + } + return row_ids; + }; + + const auto first_scan = scan_row_ids(); + const auto second_scan = scan_row_ids(); + EXPECT_EQ(4U, first_scan.size()); + EXPECT_EQ(first_scan, second_scan); + EXPECT_TRUE(reader.close().ok()); + + ASSERT_TRUE(first_scan.contains(2)); + ASSERT_TRUE(first_scan.contains(4)); + const std::vector fetch_row_ids { + first_scan.at(4), first_scan.at(2), first_scan.at(4)}; + const Columns payload_columns { + projected_column("row_id", TYPE_BIGINT, false), + projected_column("label", TYPE_STRING, true), + }; + RuntimeProfile fetch_profile("lance_vector_search_rowid_fetch_fixture"); + LanceTableReader payload_reader; + ASSERT_TRUE(init_reader(&payload_reader, payload_columns, &state, &fetch_profile, &scan_params) + .ok()); + Block payload_block; + add_output_columns(&payload_block, payload_columns); + ASSERT_TRUE(payload_reader + .read_by_row_ids(make_lance_range(dataset_uri, fixture.version, + fixture.fragment_ids), + fetch_row_ids, &payload_block) + .ok()); + + ASSERT_EQ(3U, payload_block.rows()); + const auto& logical_row_ids = + assert_cast(*payload_block.get_by_position(0).column); + const auto& labels = + assert_cast(*payload_block.get_by_position(1).column); + const auto& label_values = assert_cast(labels.get_nested_column()); + EXPECT_EQ((std::vector {4, 2, 4}), + std::vector(logical_row_ids.get_data().begin(), + logical_row_ids.get_data().end())); + ASSERT_EQ(3U, labels.size()); + EXPECT_EQ(0, labels.get_null_map_data()[0]); + EXPECT_EQ(0, labels.get_null_map_data()[1]); + EXPECT_EQ(0, labels.get_null_map_data()[2]); + EXPECT_EQ("extra", label_values.get_data_at(0).to_string()); + EXPECT_EQ("unit-x", label_values.get_data_at(1).to_string()); + EXPECT_EQ("extra", label_values.get_data_at(2).to_string()); + EXPECT_TRUE(payload_reader.close().ok()); +} + +TEST(LanceTableReaderVectorSearchTest, ReadsOnlyGlobalRowIdVirtualColumn) { + const std::filesystem::path dataset_uri = + "./be/test/format_v2/table/lance/data/all_types.lance"; + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(dataset_uri, &fixture).ok()); + + const auto global_rowid_name = BeConsts::GLOBAL_ROWID_COL + std::string("topn_fetch_lance"); + const Columns columns {projected_column(global_rowid_name, TYPE_STRING, false)}; + TQueryOptions query_options; + query_options.__set_batch_size(2); + TQueryGlobals query_globals; + RuntimeState state(query_globals); + state.set_query_options(query_options); + RuntimeProfile profile("lance_vector_search_only_global_rowid_fixture"); + auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 2, 0); + const GlobalRowIdContext context {.backend_id = 13579, .file_id = 24}; + + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); + ASSERT_TRUE(prepare_fixture(&reader, dataset_uri, fixture, fixture.fragment_ids, context).ok()); + + Block block; + add_output_columns(&block, columns); + std::set native_row_ids; + bool eos = false; + while (!eos) { + ASSERT_TRUE(reader.get_block(&block, &eos).ok()); + if (eos) { + continue; + } + const auto& global_row_ids = + assert_cast(*block.get_by_position(0).column); + for (size_t row = 0; row < block.rows(); ++row) { + const auto location = decode_lance_row_id(global_row_ids, row); + EXPECT_EQ(GlobalRowLocationV3::VERSION, location.version); + EXPECT_EQ(context.backend_id, location.backend_id); + EXPECT_EQ(context.file_id, location.file_id); + native_row_ids.emplace(location.row_id); + } + } + EXPECT_EQ(2U, native_row_ids.size()); + EXPECT_TRUE(reader.close().ok()); +} + TEST(LanceTableReaderFilterTest, PushesFilterOnNonProjectedColumn) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/all_types.lance"; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitor.java index b6b4b31d8b6b34..8336105a004819 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitor.java @@ -38,6 +38,7 @@ import org.apache.doris.nereids.trees.plans.physical.PhysicalTVFRelation; import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanVisitor; import org.apache.doris.qe.SessionVariable; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; import com.google.common.collect.ImmutableSet; import org.apache.logging.log4j.LogManager; @@ -126,6 +127,10 @@ boolean checkRelationTableSupportedType(PhysicalCatalogRelation relation) { } boolean checkTVFRelationTableSupportedType(PhysicalTVFRelation tvfRelation) { + if (isVectorSearch(tvfRelation)) { + return true; + } + Map properties = tvfRelation.getFunction().getTVFProperties().getMap(); String functionName = tvfRelation.getFunction().getName(); @@ -139,6 +144,10 @@ boolean checkTVFRelationTableSupportedType(PhysicalTVFRelation tvfRelation) { return false; } + private boolean isVectorSearch(PhysicalTVFRelation tvfRelation) { + return VectorSearchTableValuedFunction.NAME.equals(tvfRelation.getFunction().getName()); + } + @Override public Optional visitPhysicalOlapScan(PhysicalOlapScan scan, ProbeContext context) { if (scan.getSelectedIndexId() != scan.getTable().getBaseIndexId()) { @@ -174,6 +183,11 @@ public Optional visitPhysicalCatalogRelation( @Override public Optional visitPhysicalTVFRelation( PhysicalTVFRelation tvfRelation, ProbeContext context) { + // The first Lance implementation fetches top-level columns by row ID. Keep nested + // sub-column projections in the search phase until take_rows supports access paths. + if (isVectorSearch(tvfRelation) && context.slot.hasSubColPath()) { + return Optional.empty(); + } if (checkTVFRelationTableSupportedType(tvfRelation) && tvfRelation.getOutput().contains(context.slot) && !tvfRelation.getOperativeSlots().contains(context.slot)) { // lazy materialize slot must be a passive slot diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitorTest.java new file mode 100644 index 00000000000000..de143cea0a3ced --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/MaterializeProbeVisitorTest.java @@ -0,0 +1,57 @@ +// 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.processor.post.materialize; + +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.table.VectorSearch; +import org.apache.doris.nereids.trees.plans.physical.PhysicalTVFRelation; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +class MaterializeProbeVisitorTest { + + @Test + void testVectorSearchSupportsLazyMaterialization() { + MaterializeProbeVisitor visitor = new MaterializeProbeVisitor(); + PhysicalTVFRelation relation = mockVectorSearchRelation(); + + Assertions.assertTrue(visitor.checkTVFRelationTableSupportedType(relation)); + } + + @Test + void testVectorSearchKeepsNestedSubColumnInSearchPhase() { + MaterializeProbeVisitor visitor = new MaterializeProbeVisitor(); + PhysicalTVFRelation relation = mockVectorSearchRelation(); + SlotReference nestedSlot = Mockito.mock(SlotReference.class); + Mockito.when(nestedSlot.hasSubColPath()).thenReturn(true); + + Assertions.assertFalse(visitor.visitPhysicalTVFRelation( + relation, new MaterializeProbeVisitor.ProbeContext(nestedSlot)).isPresent()); + } + + private PhysicalTVFRelation mockVectorSearchRelation() { + PhysicalTVFRelation relation = Mockito.mock(PhysicalTVFRelation.class); + VectorSearch function = Mockito.mock(VectorSearch.class); + Mockito.when(function.getName()).thenReturn(VectorSearchTableValuedFunction.NAME); + Mockito.when(relation.getFunction()).thenReturn(function); + return relation; + } +} diff --git a/gensrc/proto/internal_service.proto b/gensrc/proto/internal_service.proto index 96ed6f1395ff08..fac4c2f31706d6 100644 --- a/gensrc/proto/internal_service.proto +++ b/gensrc/proto/internal_service.proto @@ -841,7 +841,8 @@ message PRequestBlockDesc { repeated PSlotDescriptor slots = 2; repeated ColumnPB column_descs = 3; repeated uint32 file_id = 4; - repeated uint32 row_id = 5; + // V2 row locations use uint32 values; V3 preserves native uint64 external row IDs. + repeated uint64 row_id = 5; optional PTupleDescriptor desc = 6; repeated uint32 column_idxs = 7; } diff --git a/thirdparty/download-thirdparty.sh b/thirdparty/download-thirdparty.sh index 715155c6044182..7fa675da700227 100755 --- a/thirdparty/download-thirdparty.sh +++ b/thirdparty/download-thirdparty.sh @@ -774,12 +774,12 @@ if [[ " ${TP_ARCHIVES[*]} " =~ " PAIMON_CPP " ]]; then echo "Finished patching ${PAIMON_CPP_SOURCE}" fi -# Patch lance-c to allow nearest-neighbor search over an explicitly selected fragment set. +# Patch lance-c for fragment-scoped nearest-neighbor search and row-ID-based fetching. if [[ " ${TP_ARCHIVES[*]} " =~ " LANCE_C " ]]; then if [[ "${LANCE_C_SOURCE}" == "lance-c-0.1.6" ]]; then cd "${TP_SOURCE_DIR}/${LANCE_C_SOURCE}" if [[ ! -f "${PATCHED_MARK}" ]]; then - patch -p1 <"${TP_PATCH_DIR}/lance-c-0.1.6-fragment-knn.patch" + patch -p1 <"${TP_PATCH_DIR}/lance-c-0.1.6-doris.patch" touch "${PATCHED_MARK}" fi cd - diff --git a/thirdparty/patches/lance-c-0.1.6-doris.patch b/thirdparty/patches/lance-c-0.1.6-doris.patch new file mode 100644 index 00000000000000..7e225dc637feff --- /dev/null +++ b/thirdparty/patches/lance-c-0.1.6-doris.patch @@ -0,0 +1,473 @@ +diff --git a/include/lance/lance.h b/include/lance/lance.h +index 1de72fa..a0a8b5e 100644 +--- a/include/lance/lance.h ++++ b/include/lance/lance.h +@@ -750,6 +750,31 @@ int32_t lance_dataset_take( + struct ArrowArrayStream* out + ); + ++/** ++ * Take rows by dataset row IDs. ++ * ++ * Row IDs are values from the `_rowid` scanner column, not zero-based row ++ * offsets. They must belong to the same dataset snapshot used for this read. ++ * Missing or deleted row IDs may be omitted from the result. For found rows, ++ * input order and duplicates are preserved. ++ * ++ * @param dataset Open dataset snapshot. ++ * @param row_ids Array of dataset row IDs. May be NULL only when ++ * `num_row_ids` is zero. ++ * @param num_row_ids Length of `row_ids`. ++ * @param columns NULL-terminated column names, or NULL for all. The ++ * system column `_rowid` may be requested explicitly. ++ * @param out Pointer to caller-allocated ArrowArrayStream. ++ * @return 0 on success, -1 on error. ++ */ ++int32_t lance_dataset_take_rows( ++ const LanceDataset* dataset, ++ const uint64_t* row_ids, ++ size_t num_row_ids, ++ const char* const* columns, ++ struct ArrowArrayStream* out ++); ++ + /* ─── Scanner builder ─── */ + + /** +diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp +index 40aa9e3..d54dc5a 100644 +--- a/include/lance/lance.hpp ++++ b/include/lance/lance.hpp +@@ -652,6 +652,30 @@ public: + } + } + ++ /// Take rows by dataset row IDs. Results exported as ArrowArrayStream. ++ void take_rows(const uint64_t* row_ids, size_t num_row_ids, ++ const std::vector& columns, ++ ArrowArrayStream* out) const { ++ std::vector col_ptrs; ++ for (auto& c : columns) col_ptrs.push_back(c.c_str()); ++ col_ptrs.push_back(nullptr); ++ const char* const* cols_ptr = columns.empty() ? nullptr : col_ptrs.data(); ++ ++ if (lance_dataset_take_rows( ++ handle_.get(), row_ids, num_row_ids, cols_ptr, out) != 0) { ++ check_error(); ++ } ++ } ++ ++ /// Take all columns by dataset row IDs. ++ void take_rows(const uint64_t* row_ids, size_t num_row_ids, ++ ArrowArrayStream* out) const { ++ if (lance_dataset_take_rows( ++ handle_.get(), row_ids, num_row_ids, nullptr, out) != 0) { ++ check_error(); ++ } ++ } ++ + /// Create a Scanner builder for this dataset. + Scanner scan() const; + +diff --git a/src/dataset.rs b/src/dataset.rs +index 91a735f..a7b9e13 100644 +--- a/src/dataset.rs ++++ b/src/dataset.rs +@@ -261,6 +261,85 @@ unsafe fn dataset_take_inner( + Ok(0) + } + ++/// Take rows by dataset row IDs, returning results as an ArrowArrayStream. ++/// ++/// - `row_ids`: array of dataset row IDs, such as values returned in the ++/// `_rowid` scanner column ++/// - `num_row_ids`: length of the row ID array ++/// - `columns`: NULL-terminated column name array, or NULL for all columns ++/// - `out`: pointer to a stack-allocated `ArrowArrayStream` ++/// ++/// `row_ids` may be NULL only when `num_row_ids` is zero. Row IDs must belong ++/// to the same dataset snapshot used for this read. Missing or deleted row IDs ++/// may be omitted from the result by the upstream Lance implementation. ++/// ++/// Returns 0 on success, -1 on error. ++#[unsafe(no_mangle)] ++pub unsafe extern "C" fn lance_dataset_take_rows( ++ dataset: *const LanceDataset, ++ row_ids: *const u64, ++ num_row_ids: usize, ++ columns: *const *const c_char, ++ out: *mut FFI_ArrowArrayStream, ++) -> i32 { ++ ffi_try!( ++ unsafe { dataset_take_rows_inner(dataset, row_ids, num_row_ids, columns, out) }, ++ neg ++ ) ++} ++ ++unsafe fn dataset_take_rows_inner( ++ dataset: *const LanceDataset, ++ row_ids: *const u64, ++ num_row_ids: usize, ++ columns: *const *const c_char, ++ out: *mut FFI_ArrowArrayStream, ++) -> Result { ++ if dataset.is_null() { ++ return Err(lance_core::Error::InvalidInput { ++ source: "dataset must not be NULL".into(), ++ location: snafu::location!(), ++ }); ++ } ++ if out.is_null() { ++ return Err(lance_core::Error::InvalidInput { ++ source: "out must not be NULL".into(), ++ location: snafu::location!(), ++ }); ++ } ++ if num_row_ids > 0 && row_ids.is_null() { ++ return Err(lance_core::Error::InvalidInput { ++ source: format!("row_ids must not be NULL when num_row_ids = {num_row_ids}").into(), ++ location: snafu::location!(), ++ }); ++ } ++ ++ let ds = unsafe { &*dataset }; ++ let row_id_slice = if num_row_ids == 0 { ++ &[] ++ } else { ++ unsafe { std::slice::from_raw_parts(row_ids, num_row_ids) } ++ }; ++ let col_names = unsafe { helpers::parse_c_string_array(columns)? }; ++ ++ let snap = ds.snapshot(); ++ let projection = match &col_names { ++ Some(cols) => lance::dataset::ProjectionRequest::from_columns(cols.iter(), snap.schema()), ++ None => lance::dataset::ProjectionRequest::from_schema(snap.schema().clone()), ++ }; ++ ++ let batch = block_on(snap.take_rows(row_id_slice, projection))?; ++ ++ // Match lance_dataset_take: export the single RecordBatch as an Arrow stream. ++ let schema = batch.schema(); ++ let reader = arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema); ++ let ffi_stream = FFI_ArrowArrayStream::new(Box::new(reader)); ++ unsafe { ++ std::ptr::write_unaligned(out, ffi_stream); ++ } ++ Ok(0) ++} ++ + // --------------------------------------------------------------------------- + // Fragment enumeration + // --------------------------------------------------------------------------- +diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs +index f184bbe..0818c6e 100644 +--- a/tests/c_api_test.rs ++++ b/tests/c_api_test.rs +@@ -15,7 +15,7 @@ use arrow::ffi::from_ffi; + use arrow::ffi_stream::ArrowArrayStreamReader; + use arrow::ffi_stream::FFI_ArrowArrayStream; + use arrow::record_batch::RecordBatchReader; +-use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray}; ++use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray, UInt64Array}; + use arrow_schema::{DataType, Field, Schema}; + use lance::Dataset; + use lance_c::*; +@@ -391,6 +391,50 @@ fn test_dataset_take() { + unsafe { lance_dataset_close(ds) }; + } + ++#[test] ++fn test_dataset_take_rows_empty_and_null_validation() { ++ let (_tmp, uri) = create_test_dataset(); ++ let c_uri = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut empty_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, ptr::null(), 0, ptr::null(), &mut empty_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut empty_stream) }.unwrap(); ++ assert_eq!( ++ reader.map(|batch| batch.unwrap().num_rows()).sum::(), ++ 0 ++ ); ++ ++ let mut invalid_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, ptr::null(), 1, ptr::null(), &mut invalid_stream) }, ++ -1 ++ ); ++ let message = unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()) }.to_string_lossy(); ++ assert!( ++ message.contains("row_ids must not be NULL when num_row_ids = 1"), ++ "unexpected error: {message}" ++ ); ++ ++ let row_id = 0_u64; ++ assert_eq!( ++ unsafe { ++ lance_dataset_take_rows(ptr::null(), &row_id, 1, ptr::null(), &mut invalid_stream) ++ }, ++ -1 ++ ); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, ptr::null(), ptr::null_mut()) }, ++ -1 ++ ); ++ ++ unsafe { lance_dataset_close(ds) }; ++} ++ + // --------------------------------------------------------------------------- + // Error handling tests + // --------------------------------------------------------------------------- +@@ -2500,6 +2544,7 @@ fn create_vector_dataset(num_rows: i32, dim: i32) -> (tempfile::TempDir, String) + fn create_multi_fragment_vector_dataset( + rows_per_fragment: i32, + dim: i32, ++ enable_stable_row_ids: bool, + ) -> (tempfile::TempDir, String) { + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + +@@ -2544,7 +2589,10 @@ fn create_multi_fragment_vector_dataset( + Dataset::write( + arrow::record_batch::RecordBatchIterator::new(vec![Ok(first)], schema.clone()), + &uri, +- None, ++ Some(lance::dataset::WriteParams { ++ enable_stable_row_ids, ++ ..Default::default() ++ }), + ) + .await + .unwrap(); +@@ -2553,6 +2601,7 @@ fn create_multi_fragment_vector_dataset( + &uri, + Some(lance::dataset::WriteParams { + mode: lance::dataset::WriteMode::Append, ++ enable_stable_row_ids, + ..Default::default() + }), + ) +@@ -2789,6 +2838,105 @@ fn test_scanner_nearest_brute_force() { + unsafe { lance_dataset_close(ds) }; + } + ++fn assert_dataset_take_rows_from_multi_fragment_ann_result(enable_stable_row_ids: bool) { ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, enable_stable_row_ids); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ // A non-NULL array whose first element is NULL is an explicit empty ++ // projection. The ANN result should therefore contain only _distance and ++ // the explicitly requested _rowid system column. ++ let no_columns: [*const c_char; 1] = [ptr::null()]; ++ let scanner = unsafe { lance_scanner_new(ds, no_columns.as_ptr(), ptr::null()) }; ++ assert!(!scanner.is_null()); ++ assert_eq!(unsafe { lance_scanner_with_row_id(scanner, true) }, 0); ++ ++ let column = c_str("embedding"); ++ let query = [40.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 1, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let mut ann_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_scanner_to_arrow_stream(scanner, &mut ann_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut ann_stream) }.unwrap(); ++ let ann_batches = reader.map(|batch| batch.unwrap()).collect::>(); ++ assert_eq!( ++ ann_batches.iter().map(RecordBatch::num_rows).sum::(), ++ 1 ++ ); ++ assert_eq!(ann_batches[0].num_columns(), 2); ++ ++ let distance = ann_batches[0] ++ .column_by_name("_distance") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .value(0); ++ assert_eq!(distance, 0.0); ++ let row_id = ann_batches[0] ++ .column_by_name("_rowid") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .value(0); ++ if !enable_stable_row_ids { ++ assert_ne!( ++ row_id >> 32, ++ 0, ++ "expected an address-style row ID from the second fragment" ++ ); ++ } ++ ++ let id_column = c_str("id"); ++ let columns = [id_column.as_ptr(), ptr::null()]; ++ let mut take_stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, columns.as_ptr(), &mut take_stream) }, ++ 0 ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut take_stream) }.unwrap(); ++ let batches = reader.map(|batch| batch.unwrap()).collect::>(); ++ assert_eq!(batches.len(), 1); ++ let ids = batches[0] ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap(); ++ assert_eq!(ids.values(), &[40]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ ++#[test] ++fn test_dataset_take_rows_from_multi_fragment_ann_result() { ++ assert_dataset_take_rows_from_multi_fragment_ann_result(false); ++} ++ ++#[test] ++fn test_dataset_take_rows_from_multi_fragment_ann_result_with_stable_row_ids() { ++ assert_dataset_take_rows_from_multi_fragment_ann_result(true); ++} ++ + #[test] + fn test_scanner_nearest_with_ivf_pq_index() { + let (_tmp, uri) = create_vector_dataset(512, 16); +@@ -2923,7 +3071,7 @@ fn test_scanner_nearest_filter_postfilter() { + + #[test] + fn test_scanner_nearest_prefilter_with_fragment_ids_next() { +- let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8); ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); + let uri_c = c_str(&uri); + let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); +@@ -2986,7 +3134,7 @@ fn test_scanner_nearest_prefilter_with_fragment_ids_next() { + + #[test] + fn test_scanner_nearest_prefilter_with_fragment_ids_arrow_stream() { +- let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8); ++ let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); + let uri_c = c_str(&uri); + let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); +diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp +index 43491f2..444fbdc 100644 +--- a/tests/cpp/test_cpp_api.cpp ++++ b/tests/cpp/test_cpp_api.cpp +@@ -122,6 +122,35 @@ static void test_dataset_take(const std::string& uri) { + PASS(); + } + ++static void test_dataset_take_rows(const std::string& uri) { ++ TEST(test_dataset_take_rows); ++ ++ auto ds = lance::Dataset::open(uri); ++ ++ // The smoke fixture has one fragment, so its first row IDs are 0, 1, 2. ++ uint64_t row_ids[] = {0, 1, 2}; ++ ArrowArrayStream stream; ++ memset(&stream, 0, sizeof(stream)); ++ ds.take_rows(row_ids, 3, &stream); ++ ++ uint64_t total = 0; ++ while (true) { ++ ArrowArray arr; ++ memset(&arr, 0, sizeof(arr)); ++ int rc = stream.get_next(&stream, &arr); ++ assert(rc == 0); ++ if (!arr.release) break; ++ total += (uint64_t)arr.length; ++ arr.release(&arr); ++ } ++ ++ assert(total == 3); ++ printf("rows=%llu... ", (unsigned long long)total); ++ ++ if (stream.release) stream.release(&stream); ++ PASS(); ++} ++ + static void test_raii_cleanup(const std::string& uri) { + TEST(test_raii_cleanup); + +@@ -693,6 +722,7 @@ int main(int argc, char** argv) { + test_dataset_schema(uri); + test_scanner_fluent(uri); + test_dataset_take(uri); ++ test_dataset_take_rows(uri); + test_raii_cleanup(uri); + test_versions(uri); + test_restore_to_current(uri); +diff --git a/src/scanner.rs b/src/scanner.rs +--- a/src/scanner.rs ++++ b/src/scanner.rs +@@ -158,6 +158,12 @@ impl LanceScanner { + location: snafu::location!(), + }); + } ++ // Lance validates fragment-scoped nearest searches when nearest() is ++ // configured. Such searches are supported when the fragment scan is ++ // the input to a prefilter, so this flag must be set first. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } + if let Some(n) = &self.nearest { + scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; + if let Some(np) = self.nprobes { +@@ -175,9 +181,6 @@ impl LanceScanner { + if let Some(ui) = self.use_index { + scanner.use_index(ui); + } +- if self.prefilter { +- scanner.prefilter(true); +- } + if let Some(segments) = &self.index_segments { + scanner.with_index_segments(segments.clone())?; + } +@@ -219,6 +222,11 @@ impl LanceScanner { + location: snafu::location!(), + }); + } ++ // nearest() checks the current prefilter setting before accepting a ++ // fragment-scoped search. Enable it before installing the query. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } + if let Some(n) = &self.nearest { + scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; + if let Some(np) = self.nprobes { +@@ -236,9 +244,6 @@ impl LanceScanner { + if let Some(ui) = self.use_index { + scanner.use_index(ui); + } +- if self.prefilter { +- scanner.prefilter(true); +- } + if let Some(segments) = &self.index_segments { + scanner.with_index_segments(segments.clone())?; + } diff --git a/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch b/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch deleted file mode 100644 index ef38ea0c5289aa..00000000000000 --- a/thirdparty/patches/lance-c-0.1.6-fragment-knn.patch +++ /dev/null @@ -1,48 +0,0 @@ -diff --git a/src/scanner.rs b/src/scanner.rs ---- a/src/scanner.rs -+++ b/src/scanner.rs -@@ -158,6 +158,12 @@ impl LanceScanner { - location: snafu::location!(), - }); - } -+ // Lance validates fragment-scoped nearest searches when nearest() is -+ // configured. Such searches are supported when the fragment scan is -+ // the input to a prefilter, so this flag must be set first. -+ if self.prefilter { -+ scanner.prefilter(true); -+ } - if let Some(n) = &self.nearest { - scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; - if let Some(np) = self.nprobes { -@@ -175,9 +181,6 @@ impl LanceScanner { - if let Some(ui) = self.use_index { - scanner.use_index(ui); - } -- if self.prefilter { -- scanner.prefilter(true); -- } - if let Some(segments) = &self.index_segments { - scanner.with_index_segments(segments.clone())?; - } -@@ -219,6 +222,11 @@ impl LanceScanner { - location: snafu::location!(), - }); - } -+ // nearest() checks the current prefilter setting before accepting a -+ // fragment-scoped search. Enable it before installing the query. -+ if self.prefilter { -+ scanner.prefilter(true); -+ } - if let Some(n) = &self.nearest { - scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; - if let Some(np) = self.nprobes { -@@ -236,9 +244,6 @@ impl LanceScanner { - if let Some(ui) = self.use_index { - scanner.use_index(ui); - } -- if self.prefilter { -- scanner.prefilter(true); -- } - if let Some(segments) = &self.index_segments { - scanner.with_index_segments(segments.clone())?; - } From 984d5e5a4057fd7ad5ac7bfc99a563aa2d19b0a7 Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Mon, 17 Aug 2026 20:12:34 +0800 Subject: [PATCH 3/4] update --- .../exec/operator/materialization_opertor.cpp | 4 +- be/src/exec/rowid_fetcher.cpp | 12 +- be/src/format_v2/table/lance_reader.cpp | 138 ++++- be/src/format_v2/table/lance_reader.h | 1 + be/test/format_v2/table/lance_reader_test.cpp | 35 +- .../VectorSearchTableValuedFunction.java | 5 + .../VectorSearchTableValuedFunctionTest.java | 15 + .../lance/test_lance_vector_search.groovy | 8 +- thirdparty/patches/lance-c-0.1.6-doris.patch | 479 ++++++++++++------ 9 files changed, 506 insertions(+), 191 deletions(-) diff --git a/be/src/exec/operator/materialization_opertor.cpp b/be/src/exec/operator/materialization_opertor.cpp index 03fc98098dd92f..c1b10e7abc7ffb 100644 --- a/be/src/exec/operator/materialization_opertor.cpp +++ b/be/src/exec/operator/materialization_opertor.cpp @@ -490,8 +490,8 @@ Status MaterializationSharedState::create_muiltget_result(const Columns& columns for (int j = 0; j < rows; ++j) { if (!null_map || !null_map[j]) { MaterializationRowLocation row_location; - RETURN_IF_ERROR(decode_materialization_row_location( - column_rowid->get_data_at(j), &row_location)); + RETURN_IF_ERROR(decode_materialization_row_location(column_rowid->get_data_at(j), + &row_location)); auto rpc_struct = rpc_struct_map.find(row_location.backend_id); if (UNLIKELY(rpc_struct == rpc_struct_map.end())) { return Status::InternalError( diff --git a/be/src/exec/rowid_fetcher.cpp b/be/src/exec/rowid_fetcher.cpp index 7b853c029c2404..32cbacc6f46234 100644 --- a/be/src/exec/rowid_fetcher.cpp +++ b/be/src/exec/rowid_fetcher.cpp @@ -913,8 +913,7 @@ Status RowIdStorageReader::read_batch_external_row( // Hash(TFileRangeDesc) => { all the rows that need to be read and their positions in the result block. } + file mapping // The multimap retains duplicate row IDs because the same source row can appear more than once // in the materialized result. - std::map, std::shared_ptr>> + std::map, std::shared_ptr>> scan_rows; // Block corresponding to the order of `scan_rows` map. @@ -928,9 +927,8 @@ Status RowIdStorageReader::read_batch_external_row( auto hash_file_range = [&rpc_scan_params](const ExternalFileMappingInfo& external_info) { const auto& file_range_desc = external_info.scan_range_desc; - const auto format_type = file_range_desc.__isset.format_type - ? file_range_desc.format_type - : rpc_scan_params.format_type; + const auto format_type = file_range_desc.__isset.format_type ? file_range_desc.format_type + : rpc_scan_params.format_type; if (format_type == TFileFormatType::FORMAT_LANCE) { // Parquet and ORC row IDs are offsets within a physical file range, so their fetch // path must keep each path/start_offset pair separate. Lance row IDs instead belong @@ -939,8 +937,8 @@ Status RowIdStorageReader::read_batch_external_row( // all of those fragments together. Group the mappings by scan node, snapshot, and // dataset URI; the plan node keeps independent scans of the same snapshot isolated. const auto& lance_params = file_range_desc.table_format_params.lance_params; - return fmt::format("lance:{}:{}:{}", external_info.plan_node_id, - lance_params.version, lance_params.dataset_uri); + return fmt::format("lance:{}:{}:{}", external_info.plan_node_id, lance_params.version, + lance_params.dataset_uri); } std::string value; value.resize(file_range_desc.path.size() + sizeof(file_range_desc.start_offset)); diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index f5f3932a7a28ba..2ae35de613d0b6 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -59,6 +59,21 @@ constexpr std::string_view DISTANCE_COLUMN = "_distance"; constexpr std::string_view ROW_ID_COLUMN = "_rowid"; constexpr std::string_view ARROW_EXTENSION_NAME = "ARROW:extension:name"; +size_t vector_element_width(TVectorElementType::type type) { + switch (type) { + case TVectorElementType::FLOAT16: + return sizeof(uint16_t); + case TVectorElementType::FLOAT32: + return sizeof(float); + case TVectorElementType::FLOAT64: + return sizeof(double); + case TVectorElementType::UINT8: + case TVectorElementType::INT8: + return sizeof(uint8_t); + } + return 0; +} + int arrow_time_precision(arrow::TimeUnit::type unit) { switch (unit) { case arrow::TimeUnit::SECOND: @@ -288,8 +303,8 @@ Status LanceTableReader::init(TableReadOptions&& options) { _ctz = _runtime_state->timezone_obj(); _vector_search = _scan_params->__isset.external_search_request; if (_vector_search) { - const auto& vector = - _scan_params->external_search_request.search_query.vector_search; + RETURN_IF_ERROR(_validate_external_search_request()); + const auto& vector = _scan_params->external_search_request.search_query.vector_search; _scanner_profile->add_info_string("LanceFragmentTopK", std::to_string(vector.top_k)); _scanner_profile->add_info_string("LanceFragmentOffset", std::to_string(vector.offset)); _scanner_profile->add_info_string("LanceVectorDimension", @@ -316,8 +331,8 @@ Status LanceTableReader::init(TableReadOptions&& options) { "Lance global row id is currently supported only for vector search"); } if (_global_rowid_output_idx.has_value()) { - return Status::InvalidArgument( - "duplicate Lance global row id projected column: {}", column.name); + return Status::InvalidArgument("duplicate Lance global row id projected column: {}", + column.name); } if (remove_nullable(column.type)->get_primitive_type() != TYPE_STRING) { return Status::InvalidArgument( @@ -475,9 +490,8 @@ Status LanceTableReader::read_by_row_ids(const TFileRangeDesc& range, fetched_rows += rows; } if (fetched_rows != row_ids.size()) { - return Status::InternalError( - "Lance row-id fetch returned {} rows for {} requested row ids", - fetched_rows, row_ids.size()); + return Status::InternalError("Lance row-id fetch returned {} rows for {} requested row ids", + fetched_rows, row_ids.size()); } return Status::OK(); } @@ -496,6 +510,113 @@ Status LanceTableReader::close() { return TableReader::close(); } +Status LanceTableReader::_validate_external_search_request() const { + // FE validates requests produced by vector_search(), but this reader consumes a deserialized + // Thrift boundary. Recheck structural invariants and values used for allocation, pointer + // arithmetic, C-string calls, and narrowing conversions before accessing them below. + DORIS_CHECK(_scan_params != nullptr); + DORIS_CHECK(_scan_params->__isset.external_search_request); + if (_scan_params->__isset.lance_substrait_filter) { + return Status::InvalidArgument( + "Lance vector search cannot combine its pre-search filter with " + "lance_substrait_filter"); + } + + const auto& request = _scan_params->external_search_request; + if (request.schema_version != 1) { + return Status::NotSupported("unsupported external search schema version: {}", + request.schema_version); + } + if (!request.__isset.search_query) { + return Status::InvalidArgument("external search request requires search_query"); + } + + const bool has_vector = request.search_query.__isset.vector_search; + const bool has_full_text = request.search_query.__isset.full_text_search; + if (has_vector == has_full_text) { + return Status::InvalidArgument("external search query must set exactly one search kind"); + } + if (has_full_text) { + return Status::NotSupported("Lance Format V2 reader does not yet support full-text search"); + } + + const auto& vector = request.search_query.vector_search; + if (!vector.__isset.column || vector.column.empty() || + vector.column.find('\0') != std::string::npos) { + return Status::InvalidArgument("Lance vector search requires a non-empty column"); + } + if (!vector.__isset.query_vector) { + return Status::InvalidArgument("Lance vector search requires a query vector"); + } + const auto& query_vector = vector.query_vector; + if (!query_vector.__isset.element_type || !query_vector.__isset.dimension || + !query_vector.__isset.values) { + return Status::InvalidArgument( + "Lance query vector requires element_type, dimension, and values"); + } + if (query_vector.dimension <= 0) { + return Status::InvalidArgument("Lance query vector dimension must be positive: {}", + query_vector.dimension); + } + const auto element_width = vector_element_width(query_vector.element_type); + if (element_width == 0) { + return Status::NotSupported("unsupported Lance query vector element type: {}", + static_cast(query_vector.element_type)); + } + const auto dimension = static_cast(query_vector.dimension); + if (dimension > std::numeric_limits::max() / element_width || + query_vector.values.size() != dimension * element_width) { + return Status::InvalidArgument( + "Lance query vector byte size {} does not match dimension {} and element width {}", + query_vector.values.size(), dimension, element_width); + } + if (!vector.__isset.top_k || vector.top_k <= 0) { + return Status::InvalidArgument("Lance vector search top_k must be positive"); + } + if (!vector.__isset.offset || vector.offset < 0) { + return Status::InvalidArgument("Lance vector search offset must be non-negative"); + } + constexpr auto UINT32_MAX_VALUE = static_cast(std::numeric_limits::max()); + if (vector.offset > UINT32_MAX_VALUE || vector.top_k > UINT32_MAX_VALUE - vector.offset) { + return Status::InvalidArgument("Lance vector search top_k + offset exceeds uint32 range"); + } + + if (request.__isset.search_filter) { + const auto& filter = request.search_filter; + if (!filter.__isset.format || !filter.__isset.payload || filter.payload.empty()) { + return Status::InvalidArgument( + "external search filter requires format and non-empty payload"); + } + switch (filter.format) { + case TSearchFilterFormat::SQL: + if (filter.payload.find('\0') != std::string::npos) { + return Status::InvalidArgument( + "Lance SQL search filter contains an embedded NUL byte"); + } + break; + case TSearchFilterFormat::SUBSTRAIT: + break; + default: + return Status::NotSupported("unsupported external search filter format: {}", + static_cast(filter.format)); + } + } + + if (request.__isset.vector_search_options) { + const auto& options = request.vector_search_options; + if (options.__isset.nprobes && options.nprobes <= 0) { + return Status::InvalidArgument("Lance nprobes must be positive"); + } + if (options.__isset.refine_factor && options.refine_factor <= 0) { + return Status::InvalidArgument("Lance refine_factor must be positive"); + } + if (options.__isset.ef && options.ef <= 0) { + return Status::InvalidArgument("Lance ef must be positive"); + } + } + return Status::OK(); +} + Status LanceTableReader::_ensure_dataset_open(const TFileRangeDesc& range) { const auto key = _dataset_key(range); if (_dataset == nullptr) { @@ -831,8 +952,7 @@ Status LanceTableReader::_fill_block_from_record_batch( const auto output_idx = *_global_rowid_output_idx; const auto& output_name = _projected_columns[output_idx].name; if (!materialized_columns.emplace(output_name).second) { - return Status::InternalError("Lance returned duplicate column '{}'", - ROW_ID_COLUMN); + return Status::InternalError("Lance returned duplicate column '{}'", ROW_ID_COLUMN); } RETURN_IF_ERROR( _append_global_row_ids(record_batch->column(arrow_idx), columns[output_idx])); diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 09e413eb208f3c..ed2654a32917d4 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -83,6 +83,7 @@ class LanceTableReader final : public TableReader { bool operator==(const DatasetKey&) const = default; }; + Status _validate_external_search_request() const; Status _ensure_dataset_open(const TFileRangeDesc& range); Status _open_dataset(const DatasetKey& key); Status _open_scanner(const TFileRangeDesc& range); diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index ee2105b0742285..08f46f2496a038 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -168,9 +168,8 @@ Status init_reader(LanceTableReader* reader, const Columns& projected_columns, }); } -Status prepare_range( - LanceTableReader* reader, TFileRangeDesc range, - std::optional global_rowid_context = std::nullopt) { +Status prepare_range(LanceTableReader* reader, TFileRangeDesc range, + std::optional global_rowid_context = std::nullopt) { return reader->prepare_split({.partition_values = {}, .conjuncts = std::nullopt, .partition_prune_conjuncts = {}, @@ -185,9 +184,9 @@ Status prepare_range( Status prepare_fixture(LanceTableReader* reader, const std::filesystem::path& dataset_uri, const LanceFixtureInfo& fixture, std::vector fragment_ids, std::optional global_rowid_context = std::nullopt) { - return prepare_range( - reader, make_lance_range(dataset_uri, fixture.version, std::move(fragment_ids)), - global_rowid_context); + return prepare_range(reader, + make_lance_range(dataset_uri, fixture.version, std::move(fragment_ids)), + global_rowid_context); } TFileScanRangeParams make_float32_vector_search_params( @@ -231,6 +230,24 @@ TFileScanRangeParams make_float32_vector_search_params( return scan_params; } +TEST(LanceTableReaderVectorSearchTest, RejectsMalformedVectorPayloadBeforeReadingIt) { + const Columns columns { + projected_column("row_id", TYPE_BIGINT, false), + projected_column("_distance", TYPE_FLOAT, true), + }; + TQueryGlobals query_globals; + RuntimeState state(query_globals); + RuntimeProfile profile("lance_vector_search_invalid_request"); + auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 2, 0); + scan_params.external_search_request.search_query.vector_search.query_vector.__set_dimension(4); + + LanceTableReader reader; + const auto status = init_reader(&reader, columns, &state, &profile, &scan_params); + + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("query vector byte size"), std::string::npos); +} + std::vector> read_vector_search_rows(LanceTableReader* reader, Block* block) { std::vector> rows; @@ -432,8 +449,8 @@ TEST(LanceTableReaderVectorSearchTest, ReturnsStableGlobalRowIdsAndFetchesPayloa ASSERT_TRUE(first_scan.contains(2)); ASSERT_TRUE(first_scan.contains(4)); - const std::vector fetch_row_ids { - first_scan.at(4), first_scan.at(2), first_scan.at(4)}; + const std::vector fetch_row_ids {first_scan.at(4), first_scan.at(2), + first_scan.at(4)}; const Columns payload_columns { projected_column("row_id", TYPE_BIGINT, false), projected_column("label", TYPE_STRING, true), @@ -446,7 +463,7 @@ TEST(LanceTableReaderVectorSearchTest, ReturnsStableGlobalRowIdsAndFetchesPayloa add_output_columns(&payload_block, payload_columns); ASSERT_TRUE(payload_reader .read_by_row_ids(make_lance_range(dataset_uri, fixture.version, - fixture.fragment_ids), + fixture.fragment_ids), fetch_row_ids, &payload_block) .ok()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java index 4893284acd1017..a0b900b0479511 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunction.java @@ -273,6 +273,11 @@ static List buildOutputColumns(LanceTableMetadata metadata) throw new AnalysisException("Duplicate Lance schema column under " + "case-insensitive matching: '" + field.getName() + "'"); } + if (field.getName().startsWith(Column.GLOBAL_ROWID_COL)) { + throw new AnalysisException("Lance table contains column '" + field.getName() + + "' using reserved Doris internal column prefix '" + + Column.GLOBAL_ROWID_COL + "'"); + } if (field.getName().equalsIgnoreCase(DISTANCE_COLUMN)) { throw new AnalysisException("Lance table already contains reserved vector search " + "column '" + DISTANCE_COLUMN + "'"); diff --git a/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java index 18985cbfbe3a84..cba2c6383043e6 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/tablefunction/VectorSearchTableValuedFunctionTest.java @@ -18,6 +18,7 @@ package org.apache.doris.tablefunction; import org.apache.doris.analysis.TableName; +import org.apache.doris.catalog.Column; import org.apache.doris.common.AnalysisException; import org.apache.doris.datasource.lance.LanceTableMetadata; @@ -81,4 +82,18 @@ public void testRejectCaseInsensitiveDuplicateOutputColumnsInFrontend() { Assert.assertTrue(duplicate.getMessage().contains("case-insensitive")); } + + @Test + public void testRejectReservedGlobalRowIdPrefixInFrontend() { + Schema schema = new Schema(Collections.singletonList( + Field.nullable(Column.GLOBAL_ROWID_COL + "payload", ArrowType.Utf8.INSTANCE))); + LanceTableMetadata metadata = new LanceTableMetadata( + "s3://bucket/table.lance", 42, schema, + Collections.emptyList(), Collections.emptyMap()); + + AnalysisException reserved = Assert.assertThrows(AnalysisException.class, + () -> VectorSearchTableValuedFunction.buildOutputColumns(metadata)); + + Assert.assertTrue(reserved.getMessage().contains(Column.GLOBAL_ROWID_COL)); + } } diff --git a/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy b/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy index 55c48650ad1388..0bf2c4807a1e09 100644 --- a/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy +++ b/regression-test/suites/external_table_p0/lance/test_lance_vector_search.groovy @@ -57,8 +57,8 @@ suite("test_lance_vector_search", "p0,external") { * * An outer WHERE is deliberately different from the TVF filter property: it is evaluated by * Doris after Lance returns Top-K and can therefore reduce the final result below top_k. - * The current implementation pins one Lance dataset version and searches its entire snapshot - * with one scanner. Multi-scanner search plus global Top-K merging remains future work. + * The current implementation pins one Lance dataset version, searches every visible fragment + * independently, and lets Doris merge the fragment-local candidates with a global Top-N. * * Fixture: doris.vs_ivf_pq_f32 is generated offline by * docker/thirdparties/docker-compose/iceberg/scripts/lance_build_preinstalled_catalog.py. @@ -131,7 +131,7 @@ suite("test_lance_vector_search", "p0,external") { contains "lanceTopK=5" contains "lanceOffset=0" contains "lanceMetric=l2" - contains "lanceSearchScanners=1" + contains "lanceSearchFragments=2" // The raw query vector must not be echoed into the plan output. notContains "[0,1,2,3" } @@ -243,4 +243,4 @@ suite("test_lance_vector_search", "p0,external") { } finally { // sql """DROP CATALOG IF EXISTS `${catalogName}`""" } -} \ No newline at end of file +} diff --git a/thirdparty/patches/lance-c-0.1.6-doris.patch b/thirdparty/patches/lance-c-0.1.6-doris.patch index 7e225dc637feff..cc49293d9dedf8 100644 --- a/thirdparty/patches/lance-c-0.1.6-doris.patch +++ b/thirdparty/patches/lance-c-0.1.6-doris.patch @@ -2,10 +2,7 @@ diff --git a/include/lance/lance.h b/include/lance/lance.h index 1de72fa..a0a8b5e 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h -@@ -750,6 +750,31 @@ int32_t lance_dataset_take( - struct ArrowArrayStream* out - ); - +@@ -752,0 +753,25 @@ int32_t lance_dataset_take( +/** + * Take rows by dataset row IDs. + * @@ -31,17 +28,11 @@ index 1de72fa..a0a8b5e 100644 + struct ArrowArrayStream* out +); + - /* ─── Scanner builder ─── */ - - /** diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 40aa9e3..d54dc5a 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp -@@ -652,6 +652,30 @@ public: - } - } - +@@ -654,0 +655,24 @@ public: + /// Take rows by dataset row IDs. Results exported as ArrowArrayStream. + void take_rows(const uint64_t* row_ids, size_t num_row_ids, + const std::vector& columns, @@ -66,17 +57,38 @@ index 40aa9e3..d54dc5a 100644 + } + } + - /// Create a Scanner builder for this dataset. - Scanner scan() const; - diff --git a/src/dataset.rs b/src/dataset.rs -index 91a735f..a7b9e13 100644 +index 91a735f..364be09 100644 --- a/src/dataset.rs +++ b/src/dataset.rs -@@ -261,6 +261,85 @@ unsafe fn dataset_take_inner( - Ok(0) - } - +@@ -44,0 +45,20 @@ impl LanceDataset { ++fn projection_from_columns( ++ dataset: &Dataset, ++ columns: Option<&[String]>, ++) -> Result { ++ match columns { ++ Some(columns) => { ++ let schema = dataset ++ .schema() ++ .project_preserve_system_columns(columns) ++ .map_err(|err| { ++ lance_core::Error::invalid_input(format!("invalid columns {columns:?}: {err}")) ++ })?; ++ Ok(lance::dataset::ProjectionRequest::from_schema(schema)) ++ } ++ None => Ok(lance::dataset::ProjectionRequest::from_schema( ++ dataset.schema().clone(), ++ )), ++ } ++} ++ +@@ -247,4 +267 @@ unsafe fn dataset_take_inner( +- let projection = match &col_names { +- Some(cols) => lance::dataset::ProjectionRequest::from_columns(cols.iter(), snap.schema()), +- None => lance::dataset::ProjectionRequest::from_schema(snap.schema().clone()), +- }; ++ let projection = projection_from_columns(&snap, col_names.as_deref())?; +@@ -263,0 +281,69 @@ unsafe fn dataset_take_inner( +/// Take rows by dataset row IDs, returning results as an ArrowArrayStream. +/// +/// - `row_ids`: array of dataset row IDs, such as values returned in the @@ -112,22 +124,15 @@ index 91a735f..a7b9e13 100644 + out: *mut FFI_ArrowArrayStream, +) -> Result { + if dataset.is_null() { -+ return Err(lance_core::Error::InvalidInput { -+ source: "dataset must not be NULL".into(), -+ location: snafu::location!(), -+ }); ++ return Err(lance_core::Error::invalid_input("dataset must not be NULL")); + } + if out.is_null() { -+ return Err(lance_core::Error::InvalidInput { -+ source: "out must not be NULL".into(), -+ location: snafu::location!(), -+ }); ++ return Err(lance_core::Error::invalid_input("out must not be NULL")); + } + if num_row_ids > 0 && row_ids.is_null() { -+ return Err(lance_core::Error::InvalidInput { -+ source: format!("row_ids must not be NULL when num_row_ids = {num_row_ids}").into(), -+ location: snafu::location!(), -+ }); ++ return Err(lance_core::Error::invalid_input(format!( ++ "row_ids must not be NULL when num_row_ids = {num_row_ids}" ++ ))); + } + + let ds = unsafe { &*dataset }; @@ -139,10 +144,7 @@ index 91a735f..a7b9e13 100644 + let col_names = unsafe { helpers::parse_c_string_array(columns)? }; + + let snap = ds.snapshot(); -+ let projection = match &col_names { -+ Some(cols) => lance::dataset::ProjectionRequest::from_columns(cols.iter(), snap.schema()), -+ None => lance::dataset::ProjectionRequest::from_schema(snap.schema().clone()), -+ }; ++ let projection = projection_from_columns(&snap, col_names.as_deref())?; + + let batch = block_on(snap.take_rows(row_id_slice, projection))?; + @@ -156,26 +158,41 @@ index 91a735f..a7b9e13 100644 + Ok(0) +} + - // --------------------------------------------------------------------------- - // Fragment enumeration - // --------------------------------------------------------------------------- +diff --git a/src/scanner.rs b/src/scanner.rs +index c1c2f67..de646be 100644 +--- a/src/scanner.rs ++++ b/src/scanner.rs +@@ -160,0 +161,6 @@ impl LanceScanner { ++ // Lance validates fragment-scoped nearest searches when nearest() is ++ // configured. Such searches are supported when the fragment scan is ++ // the input to a prefilter, so this flag must be set first. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } +@@ -178,3 +183,0 @@ impl LanceScanner { +- if self.prefilter { +- scanner.prefilter(true); +- } +@@ -221,0 +225,5 @@ impl LanceScanner { ++ // nearest() checks the current prefilter setting before accepting a ++ // fragment-scoped search. Enable it before installing the query. ++ if self.prefilter { ++ scanner.prefilter(true); ++ } +@@ -239,3 +246,0 @@ impl LanceScanner { +- if self.prefilter { +- scanner.prefilter(true); +- } diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs -index f184bbe..0818c6e 100644 +index 2babda6..9ee5af9 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs -@@ -15,7 +15,7 @@ use arrow::ffi::from_ffi; - use arrow::ffi_stream::ArrowArrayStreamReader; - use arrow::ffi_stream::FFI_ArrowArrayStream; - use arrow::record_batch::RecordBatchReader; +@@ -9,0 +10 @@ use std::ffi::{CString, c_char}; ++use std::process::Command; +@@ -18 +19 @@ use arrow::record_batch::RecordBatchReader; -use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray}; +use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray, UInt64Array}; - use arrow_schema::{DataType, Field, Schema}; - use lance::Dataset; - use lance_c::*; -@@ -391,6 +391,50 @@ fn test_dataset_take() { - unsafe { lance_dataset_close(ds) }; - } - +@@ -393,0 +395,94 @@ fn test_dataset_take() { +#[test] +fn test_dataset_take_rows_empty_and_null_validation() { + let (_tmp, uri) = create_test_dataset(); @@ -220,41 +237,131 @@ index f184bbe..0818c6e 100644 + unsafe { lance_dataset_close(ds) }; +} + - // --------------------------------------------------------------------------- - // Error handling tests - // --------------------------------------------------------------------------- -@@ -2500,6 +2544,7 @@ fn create_vector_dataset(num_rows: i32, dim: i32) -> (tempfile::TempDir, String) - fn create_multi_fragment_vector_dataset( - rows_per_fragment: i32, - dim: i32, ++#[test] ++fn test_dataset_take_rows_invalid_column() { ++ const CHILD_ENV: &str = "LANCE_C_TEST_TAKE_ROWS_INVALID_COLUMN_CHILD"; ++ ++ if std::env::var_os(CHILD_ENV).is_some() { ++ let (_tmp, uri) = create_test_dataset(); ++ let c_uri = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let row_id = 0_u64; ++ let invalid_column = c_str("unknown_column"); ++ let columns = [invalid_column.as_ptr(), ptr::null()]; ++ let mut stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_dataset_take_rows(ds, &row_id, 1, columns.as_ptr(), &mut stream) }, ++ -1 ++ ); ++ assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); ++ ++ let message_ptr = lance_last_error_message(); ++ assert!(!message_ptr.is_null()); ++ let message = unsafe { std::ffi::CStr::from_ptr(message_ptr) } ++ .to_string_lossy() ++ .into_owned(); ++ unsafe { lance_free_string(message_ptr) }; ++ assert!( ++ message.contains("unknown_column"), ++ "unexpected error: {message}" ++ ); ++ ++ unsafe { lance_dataset_close(ds) }; ++ return; ++ } ++ ++ let output = Command::new(std::env::current_exe().unwrap()) ++ .arg("--exact") ++ .arg("test_dataset_take_rows_invalid_column") ++ .arg("--nocapture") ++ .env(CHILD_ENV, "1") ++ .output() ++ .unwrap(); ++ assert!( ++ output.status.success(), ++ "invalid-column subprocess failed\nstdout:\n{}\nstderr:\n{}", ++ String::from_utf8_lossy(&output.stdout), ++ String::from_utf8_lossy(&output.stderr) ++ ); ++} ++ +@@ -2497,0 +2593,73 @@ fn create_vector_dataset(num_rows: i32, dim: i32) -> (tempfile::TempDir, String) ++/// Create two vector fragments with deterministic vectors. Every component of ++/// row `id` is `id as f32`, making nearest-neighbor expectations unambiguous. ++fn create_multi_fragment_vector_dataset( ++ rows_per_fragment: i32, ++ dim: i32, + enable_stable_row_ids: bool, - ) -> (tempfile::TempDir, String) { - use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; - -@@ -2544,7 +2589,10 @@ fn create_multi_fragment_vector_dataset( - Dataset::write( - arrow::record_batch::RecordBatchIterator::new(vec![Ok(first)], schema.clone()), - &uri, -- None, ++) -> (tempfile::TempDir, String) { ++ use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; ++ ++ let tmp = tempfile::tempdir().unwrap(); ++ let uri = tmp ++ .path() ++ .join("multi_fragment_vec_ds") ++ .to_str() ++ .unwrap() ++ .to_string(); ++ let schema = Arc::new(Schema::new(vec![ ++ Field::new("id", DataType::Int32, false), ++ Field::new( ++ "embedding", ++ DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), ++ false, ++ ), ++ ])); ++ ++ let make_batch = |first_id: i32| { ++ let ids: Vec = (first_id..first_id + rows_per_fragment).collect(); ++ let mut embeddings = FixedSizeListBuilder::new(Float32Builder::new(), dim); ++ for id in &ids { ++ for _ in 0..dim { ++ embeddings.values().append_value(*id as f32); ++ } ++ embeddings.append(true); ++ } ++ RecordBatch::try_new( ++ schema.clone(), ++ vec![ ++ Arc::new(Int32Array::from(ids)), ++ Arc::new(embeddings.finish()), ++ ], ++ ) ++ .unwrap() ++ }; ++ ++ let first = make_batch(0); ++ let second = make_batch(rows_per_fragment); ++ lance_c::runtime::block_on(async { ++ Dataset::write( ++ arrow::record_batch::RecordBatchIterator::new(vec![Ok(first)], schema.clone()), ++ &uri, + Some(lance::dataset::WriteParams { + enable_stable_row_ids, + ..Default::default() + }), - ) - .await - .unwrap(); -@@ -2553,6 +2601,7 @@ fn create_multi_fragment_vector_dataset( - &uri, - Some(lance::dataset::WriteParams { - mode: lance::dataset::WriteMode::Append, ++ ) ++ .await ++ .unwrap(); ++ Dataset::write( ++ arrow::record_batch::RecordBatchIterator::new(vec![Ok(second)], schema), ++ &uri, ++ Some(lance::dataset::WriteParams { ++ mode: lance::dataset::WriteMode::Append, + enable_stable_row_ids, - ..Default::default() - }), - ) -@@ -2789,6 +2838,105 @@ fn test_scanner_nearest_brute_force() { - unsafe { lance_dataset_close(ds) }; - } - ++ ..Default::default() ++ }), ++ ) ++ .await ++ .unwrap(); ++ }); ++ ++ (tmp, uri) ++} ++ +@@ -2723,0 +2892,99 @@ fn test_scanner_nearest_brute_force() { +fn assert_dataset_take_rows_from_multi_fragment_ann_result(enable_stable_row_ids: bool) { + let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, enable_stable_row_ids); + let uri_c = c_str(&uri); @@ -354,35 +461,144 @@ index f184bbe..0818c6e 100644 + assert_dataset_take_rows_from_multi_fragment_ann_result(true); +} + - #[test] - fn test_scanner_nearest_with_ivf_pq_index() { - let (_tmp, uri) = create_vector_dataset(512, 16); -@@ -2923,7 +3071,7 @@ fn test_scanner_nearest_filter_postfilter() { - - #[test] - fn test_scanner_nearest_prefilter_with_fragment_ids_next() { -- let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8); +@@ -2855,0 +3123,132 @@ fn test_scanner_nearest_filter_postfilter() { ++#[test] ++fn test_scanner_nearest_prefilter_with_fragment_ids_next() { + let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); - let uri_c = c_str(&uri); - let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; - assert!(!ds.is_null()); -@@ -2986,7 +3134,7 @@ fn test_scanner_nearest_prefilter_with_fragment_ids_next() { - - #[test] - fn test_scanner_nearest_prefilter_with_fragment_ids_arrow_stream() { -- let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut fragment_ids = vec![0; unsafe { lance_dataset_fragment_count(ds) } as usize]; ++ assert_eq!(fragment_ids.len(), 2); ++ assert_eq!( ++ unsafe { lance_dataset_fragment_ids(ds, fragment_ids.as_mut_ptr()) }, ++ 0 ++ ); ++ ++ let filter = c_str("id >= 40"); ++ let scanner = unsafe { lance_scanner_new(ds, ptr::null(), filter.as_ptr()) }; ++ assert_eq!( ++ unsafe { lance_scanner_set_fragment_ids(scanner, fragment_ids[1..].as_ptr(), 1) }, ++ 0 ++ ); ++ ++ // Match the Doris call order: nearest is configured before prefilter. ++ let column = c_str("embedding"); ++ let query = [40.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 5, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, true) }, 0); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let batches = scan_all_rows_from_scanner(scanner); ++ let mut ids = batches ++ .iter() ++ .flat_map(|batch| { ++ batch ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .values() ++ .iter() ++ .copied() ++ .collect::>() ++ }) ++ .collect::>(); ++ ids.sort_unstable(); ++ assert_eq!(ids, vec![40, 41, 42, 43, 44]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ ++#[test] ++fn test_scanner_nearest_prefilter_with_fragment_ids_arrow_stream() { + let (_tmp, uri) = create_multi_fragment_vector_dataset(32, 8, false); - let uri_c = c_str(&uri); - let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; - assert!(!ds.is_null()); ++ let uri_c = c_str(&uri); ++ let ds = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; ++ assert!(!ds.is_null()); ++ ++ let mut fragment_ids = vec![0; unsafe { lance_dataset_fragment_count(ds) } as usize]; ++ assert_eq!(fragment_ids.len(), 2); ++ assert_eq!( ++ unsafe { lance_dataset_fragment_ids(ds, fragment_ids.as_mut_ptr()) }, ++ 0 ++ ); ++ ++ let filter = c_str("id >= 60"); ++ let scanner = unsafe { lance_scanner_new(ds, ptr::null(), filter.as_ptr()) }; ++ assert_eq!( ++ unsafe { lance_scanner_set_fragment_ids(scanner, fragment_ids[1..].as_ptr(), 1) }, ++ 0 ++ ); ++ ++ let column = c_str("embedding"); ++ let query = [60.0_f32; 8]; ++ assert_eq!( ++ unsafe { ++ lance_scanner_nearest( ++ scanner, ++ column.as_ptr(), ++ query.as_ptr().cast(), ++ query.len(), ++ LanceDataType::Float32 as i32, ++ 10, ++ ) ++ }, ++ 0 ++ ); ++ assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, true) }, 0); ++ assert_eq!(unsafe { lance_scanner_set_use_index(scanner, false) }, 0); ++ ++ let mut stream = FFI_ArrowArrayStream::empty(); ++ assert_eq!( ++ unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, ++ 0, ++ "{}", ++ unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()) }.to_string_lossy() ++ ); ++ let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); ++ let mut ids = reader ++ .flat_map(|batch| { ++ let batch = batch.unwrap(); ++ batch ++ .column_by_name("id") ++ .unwrap() ++ .as_any() ++ .downcast_ref::() ++ .unwrap() ++ .values() ++ .iter() ++ .copied() ++ .collect::>() ++ }) ++ .collect::>(); ++ ids.sort_unstable(); ++ assert_eq!(ids, vec![60, 61, 62, 63]); ++ ++ unsafe { lance_scanner_close(scanner) }; ++ unsafe { lance_dataset_close(ds) }; ++} ++ diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 43491f2..444fbdc 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp -@@ -122,6 +122,35 @@ static void test_dataset_take(const std::string& uri) { - PASS(); - } - +@@ -124,0 +125,29 @@ static void test_dataset_take(const std::string& uri) { +static void test_dataset_take_rows(const std::string& uri) { + TEST(test_dataset_take_rows); + @@ -412,62 +628,5 @@ index 43491f2..444fbdc 100644 + PASS(); +} + - static void test_raii_cleanup(const std::string& uri) { - TEST(test_raii_cleanup); - -@@ -693,6 +722,7 @@ int main(int argc, char** argv) { - test_dataset_schema(uri); - test_scanner_fluent(uri); - test_dataset_take(uri); +@@ -695,0 +725 @@ int main(int argc, char** argv) { + test_dataset_take_rows(uri); - test_raii_cleanup(uri); - test_versions(uri); - test_restore_to_current(uri); -diff --git a/src/scanner.rs b/src/scanner.rs ---- a/src/scanner.rs -+++ b/src/scanner.rs -@@ -158,6 +158,12 @@ impl LanceScanner { - location: snafu::location!(), - }); - } -+ // Lance validates fragment-scoped nearest searches when nearest() is -+ // configured. Such searches are supported when the fragment scan is -+ // the input to a prefilter, so this flag must be set first. -+ if self.prefilter { -+ scanner.prefilter(true); -+ } - if let Some(n) = &self.nearest { - scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; - if let Some(np) = self.nprobes { -@@ -175,9 +181,6 @@ impl LanceScanner { - if let Some(ui) = self.use_index { - scanner.use_index(ui); - } -- if self.prefilter { -- scanner.prefilter(true); -- } - if let Some(segments) = &self.index_segments { - scanner.with_index_segments(segments.clone())?; - } -@@ -219,6 +222,11 @@ impl LanceScanner { - location: snafu::location!(), - }); - } -+ // nearest() checks the current prefilter setting before accepting a -+ // fragment-scoped search. Enable it before installing the query. -+ if self.prefilter { -+ scanner.prefilter(true); -+ } - if let Some(n) = &self.nearest { - scanner.nearest(&n.column, n.query.as_ref(), n.k as usize)?; - if let Some(np) = self.nprobes { -@@ -236,9 +244,6 @@ impl LanceScanner { - if let Some(ui) = self.use_index { - scanner.use_index(ui); - } -- if self.prefilter { -- scanner.prefilter(true); -- } - if let Some(segments) = &self.index_segments { - scanner.with_index_segments(segments.clone())?; - } From ab427b671012ba70817d627f9f2ad9c6bec6f14f Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 18 Aug 2026 10:55:56 +0800 Subject: [PATCH 4/4] update --- be/src/exec/rowid_fetcher.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/be/src/exec/rowid_fetcher.cpp b/be/src/exec/rowid_fetcher.cpp index 32cbacc6f46234..26f9f2a12bd81b 100644 --- a/be/src/exec/rowid_fetcher.cpp +++ b/be/src/exec/rowid_fetcher.cpp @@ -1000,7 +1000,7 @@ Status RowIdStorageReader::read_batch_external_row( semaphore.acquire(); RETURN_IF_ERROR(remote_scan_sched->submit_scan_task( SimplifiedScanTask( - [&, idx, scan_info]() -> Status { + [&, idx, scan_info]() -> bool { Defer complete_task {[&]() { semaphore.release(); if (++producer_count == scan_rows.size()) { @@ -1020,7 +1020,10 @@ Status RowIdStorageReader::read_batch_external_row( scan_status = status; } } - return status; + // The return value indicates whether this one-shot + // scheduler task has completed, not whether the fetch + // succeeded. The fetch status is propagated by scan_status. + return true; }, nullptr, nullptr), fmt::format("{}-read_batch_external_row-{}", print_id(query_id), idx)));