diff --git a/be/src/exec/operator/materialization_opertor.cpp b/be/src/exec/operator/materialization_opertor.cpp index 629df2ad52baab..695efb5ef78df4 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,57 @@ 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 { + ROW_VERSION version = ROW_VERSION::FILE_LOCAL_ROW_ID; + int64_t backend_id = 0; + uint32_t file_id = 0; + uint64_t row_id = 0; +}; + +Status decode_global_row_location_v2(const StringRef& encoded, ROW_VERSION version, + MaterializationRowLocation* decoded) { + if (encoded.size != sizeof(GlobalRowLoacationV2)) { + return Status::InternalError( + "invalid global row location size for version {}: actual={}, expected={}", + static_cast(version), encoded.size, sizeof(GlobalRowLoacationV2)); + } + GlobalRowLoacationV2 location(GlobalRowLoacationV2::VERSION, 0, 0, 0); + std::memcpy(&location, encoded.data, sizeof(location)); + decoded->version = version; + decoded->backend_id = location.backend_id; + switch (version) { + case ROW_VERSION::FILE_LOCAL_ROW_ID: + decoded->file_id = location.file_local.file_id; + decoded->row_id = location.file_local.row_id; + return Status::OK(); + case ROW_VERSION::LANCE_DATASET_ROW_ID: + decoded->file_id = location.lance_file_id; + decoded->row_id = location.lance_row_id; + return Status::OK(); + } + return Status::NotSupported("unsupported V2 global row location version: {}", + static_cast(version)); +} + +Status decode_materialization_row_location(const StringRef& encoded, + MaterializationRowLocation* decoded) { + if (encoded.size < sizeof(uint8_t)) { + return Status::InternalError("global row location is empty"); + } + + const auto version_value = static_cast(encoded.data[0]); + const auto version = static_cast(version_value); + // Keep size validation inside each version-family decoder. A future version can therefore + // use a different encoded size without being rejected by the current 24-byte V2 contract. + switch (version) { + case ROW_VERSION::FILE_LOCAL_ROW_ID: + case ROW_VERSION::LANCE_DATASET_ROW_ID: + return decode_global_row_location_v2(encoded, version, decoded); + } + return Status::NotSupported("unsupported global row location version: {}, encoded_size={}", + static_cast(version_value), encoded.size); +} + 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,19 +498,28 @@ 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( "MaterializationSinkOperatorX failed to find rpc_struct, backend_id={}", row_location.backend_id); } - rpc_struct->second.request.mutable_request_block_descs(i)->add_row_id( - row_location.row_id); - rpc_struct->second.request.mutable_request_block_descs(i)->add_file_id( - row_location.file_id); + auto* request_block_desc = + rpc_struct->second.request.mutable_request_block_descs(i); + const auto row_location_version = static_cast(row_location.version); + if (request_block_desc->row_id_size() == 0) { + request_block_desc->set_row_location_version(row_location_version); + } else if (request_block_desc->row_location_version() != row_location_version) { + return Status::InternalError( + "mixed row location versions in one materialization request: " + "actual={}, expected={}", + row_location_version, request_block_desc->row_location_version()); + } + request_block_desc->add_row_id(row_location.row_id); + request_block_desc->add_file_id(row_location.file_id); block_order[j] = row_location.backend_id; // Count rows per backend diff --git a/be/src/exec/rowid_fetcher.cpp b/be/src/exec/rowid_fetcher.cpp index cf0440d6b738a1..bddf2ec4ad0453 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 @@ -545,15 +547,38 @@ Status RowIdStorageReader::read_by_rowids(const PMultiGetRequestV2& request, const auto& request_block_desc = request.request_block_descs(i); PMultiGetBlockV2* pblock = response->add_blocks(); if (request_block_desc.row_id_size() >= 1) { - // Since this block belongs to the same table, we only need to take the first type for judgment. - auto first_file_id = request_block_desc.file_id(0); + if (request_block_desc.file_id_size() != request_block_desc.row_id_size()) { + return Status::InvalidArgument( + "Row-id fetch has mismatched file_id and row_id counts: " + "file_ids={}, row_ids={}", + request_block_desc.file_id_size(), request_block_desc.row_id_size()); + } + const auto row_location_version = request_block_desc.row_location_version(); + switch (row_location_version) { + case static_cast(ROW_VERSION::FILE_LOCAL_ROW_ID): + case static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID): + break; + default: + return Status::NotSupported("unsupported row location version in fetch: {}", + row_location_version); + } + // Rows in this request block belong to the same relation. The first mapping is + // sufficient to select the internal/external fetch path; individual file IDs + // are still used below to group the actual reads. + const auto first_file_id = request_block_desc.file_id(0); auto first_file_mapping = id_file_map->get_file_mapping(first_file_id); if (!first_file_mapping) { return Status::InternalError( - "Backend:{} file_mapping not found, query_id: {}, file_id: {}", + "Backend:{} file mapping not found, query_id: {}, file_id: {}", BackendOptions::get_localhost(), print_id(request.query_id()), first_file_id); } + if (first_file_mapping->type != FileMappingType::EXTERNAL && + row_location_version == + static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID)) { + return Status::InvalidArgument( + "Lance dataset row IDs require an external file mapping"); + } file_type_counts[first_file_mapping->type] += request_block_desc.row_id_size(); // prepare slots to build block @@ -658,7 +683,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 +696,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 +747,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 +821,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 +864,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(); } @@ -804,9 +888,24 @@ Status RowIdStorageReader::read_batch_external_row( int plan_node_id = external_info.plan_node_id; const auto& first_scan_range_desc = external_info.scan_range_desc; - DCHECK(id_file_map->get_external_scan_params().contains(plan_node_id)); - const auto* old_scan_params = &(id_file_map->get_external_scan_params().at(plan_node_id)); - rpc_scan_params = *old_scan_params; + const auto& external_scan_params = id_file_map->get_external_scan_params(); + const auto scan_params = external_scan_params.find(plan_node_id); + if (scan_params == external_scan_params.end()) { + return Status::InternalError("External scan params not found for plan node id {}", + plan_node_id); + } + rpc_scan_params = scan_params->second; + if (request_block_desc.row_location_version() == + static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID)) { + const auto format_type = first_scan_range_desc.__isset.format_type + ? first_scan_range_desc.format_type + : rpc_scan_params.format_type; + if (format_type != TFileFormatType::FORMAT_LANCE) { + return Status::InvalidArgument( + "Lance dataset row IDs cannot be fetched with file format {}", + static_cast(format_type)); + } + } rpc_scan_params.required_slots.clear(); rpc_scan_params.column_idxs.clear(); @@ -850,9 +949,9 @@ 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. - std::map, std::shared_ptr>> + // 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>> scan_rows; // Block corresponding to the order of `scan_rows` map. @@ -864,7 +963,21 @@ 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,14 +998,12 @@ 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}}; - scan_rows.emplace(scan_range_hash, std::make_pair(tmp, file_mapping)); + std::multimap rows {{request_block_desc.row_id(j), j}}; + scan_rows.emplace(scan_range_hash, std::make_pair(std::move(rows), file_mapping)); } } @@ -917,6 +1028,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}; @@ -926,14 +1038,30 @@ 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()) { + 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; + } + } + // 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))); @@ -944,6 +1072,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 142c35ac4d42bb..868c8715a89f47 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,6 +56,7 @@ 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"; size_t vector_element_width(TVectorElementType::type type) { @@ -68,22 +74,6 @@ size_t vector_element_width(TVectorElementType::type type) { return 0; } -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: @@ -270,8 +260,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,17 +302,11 @@ 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; - _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"); + 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", std::to_string(vector.query_vector.dimension)); } @@ -337,11 +319,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); } @@ -358,7 +358,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; @@ -369,33 +368,13 @@ Status LanceTableReader::prepare_split(const SplitReadOptions& options) { if (current_split_pruned()) { 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)); - _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)); - _search_split_prepared = _vector_search; return Status::OK(); } @@ -438,7 +417,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; } @@ -461,6 +440,62 @@ 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; @@ -475,32 +510,10 @@ 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 { + // 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) { @@ -508,17 +521,18 @@ Status LanceTableReader::_validate_external_search_request() const { "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) { + if (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"); + if (!request.__isset.search_query) { + return Status::InvalidArgument("external search request requires search_query"); } - const bool has_vector = request.query.__isset.vector; - const bool has_full_text = request.query.__isset.full_text; + 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"); } @@ -526,7 +540,7 @@ Status LanceTableReader::_validate_external_search_request() const { return Status::NotSupported("Lance Format V2 reader does not yet support full-text search"); } - const auto& vector = request.query.vector; + 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"); @@ -567,8 +581,8 @@ Status LanceTableReader::_validate_external_search_request() const { return Status::InvalidArgument("Lance vector search top_k + offset exceeds uint32 range"); } - if (request.__isset.filter) { - const auto& filter = request.filter; + 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"); @@ -588,8 +602,8 @@ Status LanceTableReader::_validate_external_search_request() 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 && options.nprobes <= 0) { return Status::InvalidArgument("Lance nprobes must be positive"); } @@ -603,6 +617,18 @@ Status LanceTableReader::_validate_external_search_request() const { return Status::OK(); } +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); @@ -623,10 +649,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); @@ -634,8 +666,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 = @@ -645,6 +678,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; @@ -655,8 +692,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 +730,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 +747,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 +829,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 +849,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"); } @@ -832,7 +873,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); @@ -853,7 +895,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 GlobalRowLoacationV2 location(ROW_VERSION::LANCE_DATASET_ROW_ID, 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()); @@ -861,6 +948,16 @@ 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 73448e4ed99e48..ed2654a32917d4 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,14 +83,18 @@ 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 _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); @@ -92,10 +103,10 @@ 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; - bool _search_split_prepared = false; bool _eof = false; }; diff --git a/be/src/storage/utils.h b/be/src/storage/utils.h index a7c8046968e1b3..53b636bd2fc57b 100644 --- a/be/src/storage/utils.h +++ b/be/src/storage/utils.h @@ -23,6 +23,8 @@ #include #include +#include +#include #include #include #include @@ -238,15 +240,89 @@ struct GlobalRowLoacation { } }; +// Wire-protocol values: never reorder or reuse an existing value. A new value may use a new +// encoded structure and size, provided its decoder keeps supporting all older values. +enum class ROW_VERSION : uint8_t { + // The row ID is a uint32 ordinal local to the FileMapping identified by file_id. + FILE_LOCAL_ROW_ID = 0, + // The row ID is an opaque uint64 ID in a fixed Lance dataset snapshot. + LANCE_DATASET_ROW_ID = 1, +}; + +/* + * A serialized global row location has a fixed size of 24 bytes. The version determines how the + * bytes at offsets 4..7 and 16..23 must be interpreted: + * + * FILE_LOCAL_ROW_ID (version = 0), used by Doris, Parquet, and ORC: + * + * byte offset 0 1..7 8..15 16..19 20..23 + * +--------+----------------+---------------+-----------+-----------+ + * | ver=0 | reserved | backend_id | file_id | row_id | + * +--------+----------------+---------------+-----------+-----------+ + * uint8 7 bytes int64 uint32 uint32 + * + * row_id is an ordinal local to the FileMapping selected by file_id. + * + * LANCE_DATASET_ROW_ID (version = 1), used by Lance: + * + * byte offset 0 1..3 4..7 8..15 16..23 + * +--------+----------+-------------+---------------+----------------+ + * | ver=1 | reserved | file_id | backend_id | lance_row_id | + * +--------+----------+-------------+---------------+----------------+ + * uint8 3 bytes uint32 int64 uint64 + * + * lance_row_id is an opaque row ID in the fixed dataset snapshot recorded by the FileMapping. + * + * The first union reuses four bytes that are padding in version 0 as Lance's file_id in version 1. + * The second union reuses the original {uint32 file_id, uint32 row_id} payload as one uint64 Lance + * row ID. Therefore, always check version before reading either union. + */ struct GlobalRowLoacationV2 { + static constexpr uint8_t VERSION = static_cast(ROW_VERSION::FILE_LOCAL_ROW_ID); + + struct FileLocalRowId { + uint32_t file_id; + uint32_t row_id; + }; + GlobalRowLoacationV2(uint8_t ver, uint64_t bid, uint32_t fid, uint32_t rid) - : version(ver), backend_id(bid), file_id(fid), row_id(rid) {} + : version(ver), + reserved_for_file_local(0), + backend_id(bid), + file_local {.file_id = fid, .row_id = rid} {} + GlobalRowLoacationV2(ROW_VERSION ver, uint64_t bid, uint32_t fid, uint64_t rid) + : version(static_cast(ver)), + lance_file_id(fid), + backend_id(bid), + lance_row_id(rid) {} + uint8_t version; + std::array reserved_before_file_id {}; + union { + // version 0: offsets 4..7 remain reserved, preserving the original V2 layout. + uint32_t reserved_for_file_local; + // version 1: offsets 4..7 identify the FileMapping for lance_row_id. + uint32_t lance_file_id; + }; int64_t backend_id; - uint32_t file_id; - uint32_t row_id; - - auto operator<=>(const GlobalRowLoacationV2&) const = default; + union { + // version 0: file_id is at offset 16 and its uint32 row ordinal is at offset 20. + FileLocalRowId file_local; + // version 1: offsets 16..23 are one opaque uint64 Lance row ID. + uint64_t lance_row_id; + }; }; +static_assert(sizeof(GlobalRowLoacationV2) == 24); +static_assert(sizeof(GlobalRowLoacationV2::FileLocalRowId) == 8); +static_assert(offsetof(GlobalRowLoacationV2, version) == 0); +static_assert(offsetof(GlobalRowLoacationV2, reserved_before_file_id) == 1); +static_assert(offsetof(GlobalRowLoacationV2, reserved_for_file_local) == 4); +static_assert(offsetof(GlobalRowLoacationV2, lance_file_id) == 4); +static_assert(offsetof(GlobalRowLoacationV2, backend_id) == 8); +static_assert(offsetof(GlobalRowLoacationV2, file_local) == 16); +static_assert(offsetof(GlobalRowLoacationV2::FileLocalRowId, file_id) == 0); +static_assert(offsetof(GlobalRowLoacationV2::FileLocalRowId, row_id) == 4); +static_assert(offsetof(GlobalRowLoacationV2, lance_row_id) == 16); + } // 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..ee0b067824cb8e 100644 --- a/be/test/exec/operator/materialization_shared_state_test.cpp +++ b/be/test/exec/operator/materialization_shared_state_test.cpp @@ -17,6 +17,9 @@ #include +#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 +32,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 +128,73 @@ 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; + GlobalRowLoacationV2 location(ROW_VERSION::LANCE_DATASET_ROW_ID, _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); + ASSERT_EQ(request_block.file_id_size(), 1); + EXPECT_EQ(request_block.file_id(0), 7); + EXPECT_EQ(request_block.row_location_version(), + static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID)); + EXPECT_EQ(_shared_state->block_order_results[0][0], _backend_id1); +} + +TEST_F(MaterializationSharedStateTest, TestRejectUnknownRowLocationVersionWithFutureSize) { + Columns columns; + auto rowid_col = _string_type->create_column(); + auto* col_data = reinterpret_cast(rowid_col.get()); + + std::array future_location {}; + future_location[0] = 99; + col_data->insert_data(future_location.data(), future_location.size()); + columns.push_back(std::move(rowid_col)); + + const Status st = _shared_state->create_muiltget_result(columns, false, false); + EXPECT_FALSE(st.ok()); + EXPECT_NE(st.to_string().find("unsupported global row location version: 99"), + std::string::npos); + EXPECT_NE(st.to_string().find("encoded_size=32"), std::string::npos); +} + +TEST_F(MaterializationSharedStateTest, TestRejectKnownRowLocationVersionWithInvalidSize) { + Columns columns; + auto rowid_col = _string_type->create_column(); + auto* col_data = reinterpret_cast(rowid_col.get()); + + std::array invalid_location {}; + invalid_location[0] = static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID); + col_data->insert_data(invalid_location.data(), invalid_location.size()); + columns.push_back(std::move(rowid_col)); + + const Status st = _shared_state->create_muiltget_result(columns, false, false); + EXPECT_FALSE(st.ok()); + EXPECT_NE(st.to_string().find("invalid global row location size for version 1"), + std::string::npos); + EXPECT_NE(st.to_string().find("actual=25, expected=24"), std::string::npos); } TEST_F(MaterializationSharedStateTest, TestMergeMultiResponse) { diff --git a/be/test/format/orc/orc_read_lines.cpp b/be/test/format/orc/orc_read_lines.cpp index 600ae451e31c79..ae2526ded8ce6d 100644 --- a/be/test/format/orc/orc_read_lines.cpp +++ b/be/test/format/orc/orc_read_lines.cpp @@ -166,8 +166,8 @@ static void read_orc_line(int64_t line, std::string block_dump, for (auto i = 0; i < row_id_string_column.size(); i++) { GlobalRowLoacationV2 info = *((GlobalRowLoacationV2*)row_id_string_column.get_data_at(i).data); - EXPECT_EQ(info.file_id, 10); - EXPECT_EQ(info.row_id, line); + EXPECT_EQ(info.file_local.file_id, 10); + EXPECT_EQ(info.file_local.row_id, line); EXPECT_EQ(info.backend_id, BackendOptions::get_backend_id()); EXPECT_EQ(info.version, IdManager::ID_VERSION); } diff --git a/be/test/format/parquet/parquet_read_lines.cpp b/be/test/format/parquet/parquet_read_lines.cpp index 090b6aa909dfbd..9a19e1a9b36c73 100644 --- a/be/test/format/parquet/parquet_read_lines.cpp +++ b/be/test/format/parquet/parquet_read_lines.cpp @@ -180,8 +180,8 @@ static void read_parquet_lines(std::vector numeric_types, for (auto i = 0; i < row_id_string_column.size(); i++) { GlobalRowLoacationV2 info = *((GlobalRowLoacationV2*)row_id_string_column.get_data_at(i).data); - EXPECT_EQ(info.file_id, 10); - EXPECT_EQ(info.row_id, read_lines_tmp.front()); + EXPECT_EQ(info.file_local.file_id, 10); + EXPECT_EQ(info.file_local.row_id, read_lines_tmp.front()); read_lines_tmp.pop_front(); EXPECT_EQ(info.backend_id, BackendOptions::get_backend_id()); EXPECT_EQ(info.version, IdManager::ID_VERSION); diff --git a/be/test/format_v2/orc/orc_reader_test.cpp b/be/test/format_v2/orc/orc_reader_test.cpp index 6ed1393ea9a698..82936cf2b633bb 100644 --- a/be/test/format_v2/orc/orc_reader_test.cpp +++ b/be/test/format_v2/orc/orc_reader_test.cpp @@ -5159,8 +5159,8 @@ TEST_F(NewOrcReaderTest, ReadGlobalRowIdVirtualColumn) { std::memcpy(&location, rowid.data, sizeof(location)); EXPECT_EQ(location.version, context.version); EXPECT_EQ(location.backend_id, context.backend_id); - EXPECT_EQ(location.file_id, context.file_id); - EXPECT_EQ(location.row_id, row); + EXPECT_EQ(location.file_local.file_id, context.file_id); + EXPECT_EQ(location.file_local.row_id, row); } } diff --git a/be/test/format_v2/parquet/parquet_reader_test.cpp b/be/test/format_v2/parquet/parquet_reader_test.cpp index 8f463713f9810c..a683d1834d817e 100644 --- a/be/test/format_v2/parquet/parquet_reader_test.cpp +++ b/be/test/format_v2/parquet/parquet_reader_test.cpp @@ -3631,8 +3631,8 @@ TEST_F(NewParquetReaderTest, GlobalRowIdSchemaAndSelectionUseFileRowPosition) { const auto location = decode_rowid(rowids, row); EXPECT_EQ(location.version, context.version); EXPECT_EQ(location.backend_id, context.backend_id); - EXPECT_EQ(location.file_id, context.file_id); - EXPECT_EQ(location.row_id, static_cast(row + 2)); + EXPECT_EQ(location.file_local.file_id, context.file_id); + EXPECT_EQ(location.file_local.row_id, static_cast(row + 2)); } } diff --git a/be/test/format_v2/parquet/parquet_scan_test.cpp b/be/test/format_v2/parquet/parquet_scan_test.cpp index e6a51bbf3a3810..7705d543fd3807 100644 --- a/be/test/format_v2/parquet/parquet_scan_test.cpp +++ b/be/test/format_v2/parquet/parquet_scan_test.cpp @@ -2519,8 +2519,8 @@ TEST_F(ParquetScanTest, GlobalRowIdUsesFileLocalPositionForScanRange) { const auto location = decode_rowid(rowid_column, row); EXPECT_EQ(location.version, context.version); EXPECT_EQ(location.backend_id, context.backend_id); - EXPECT_EQ(location.file_id, context.file_id); - row_ids.push_back(location.row_id); + EXPECT_EQ(location.file_local.file_id, context.file_id); + row_ids.push_back(location.file_local.row_id); } } diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 42e84c211ec8ac..60136789c4b214 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,8 @@ 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 +178,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) { + 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))); + make_lance_range(dataset_uri, fixture.version, std::move(fragment_ids)), + global_rowid_context); } TFileScanRangeParams make_float32_vector_search_params( @@ -203,18 +211,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,12 +230,22 @@ 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; +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, @@ -252,6 +270,16 @@ std::vector> read_vector_search_rows(LanceTableReader* return rows; } +GlobalRowLoacationV2 decode_lance_row_id(const ColumnString& column, size_t row) { + const auto encoded = column.get_data_at(row); + EXPECT_EQ(sizeof(GlobalRowLoacationV2), encoded.size); + GlobalRowLoacationV2 location(ROW_VERSION::LANCE_DATASET_ROW_ID, 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"; @@ -269,11 +297,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 +337,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 +348,189 @@ 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(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(static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID), + location.version); + EXPECT_EQ(context.backend_id, location.backend_id); + EXPECT_EQ(context.file_id, location.lance_file_id); + row_ids.emplace(logical_row_ids.get_data()[row], location.lance_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(static_cast(ROW_VERSION::LANCE_DATASET_ROW_ID), location.version); + EXPECT_EQ(context.backend_id, location.backend_id); + EXPECT_EQ(context.file_id, location.lance_file_id); + native_row_ids.emplace(location.lance_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/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..4438b7af577d96 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)); } } @@ -132,7 +133,8 @@ protected void convertPredicate() { if (isExternalSearch()) { // The TVF "filter" property is already serialized in externalSearchRequest and is // evaluated by Lance before vector search. Outer WHERE conjuncts have different - // semantics: Doris must keep them here and evaluate them after Lance returns TopK. + // semantics: keep them as Doris scan residuals. Each fragment first returns its Lance + // ANN candidates, then Doris evaluates these conjuncts before the local/global TopN. } else { LancePredicateConverter.ConversionResult result = new LancePredicateConverter(plannedMetadata.getSchema()).convert(conjuncts); @@ -152,42 +154,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 +193,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 +239,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 +249,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 +271,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/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/main/java/org/apache/doris/nereids/rules/RuleSet.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java index 580fb79f88fa6c..908ce4365418fe 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java @@ -130,6 +130,7 @@ import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughRepeat; import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughSetOperation; import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughSort; +import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughVectorSearchTopN; import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughWindow; import org.apache.doris.nereids.rules.rewrite.PushDownJoinOtherCondition; import org.apache.doris.nereids.rules.rewrite.PushDownLimitDistinctThroughJoin; @@ -172,6 +173,7 @@ public class RuleSet { new CreatePartitionTopNFromWindow(), new PushDownFilterThroughProject(), new PushDownFilterThroughSort(), + new PushDownFilterThroughVectorSearchTopN(), new PushDownJoinOtherCondition(), new PushDownFilterThroughJoin(), new PushDownExpressionsInHashCondition(), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index bc0944a35500a0..43383be8ad7bdb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -222,6 +222,7 @@ public enum RuleType { PUSH_DOWN_ALIAS_INTO_UNION_ALL(RuleTypeClass.REWRITE), PUSH_DOWN_FILTER_THROUGH_SET_OPERATION(RuleTypeClass.REWRITE), PUSH_DOWN_FILTER_THROUGH_SORT(RuleTypeClass.REWRITE), + PUSH_DOWN_FILTER_THROUGH_VECTOR_SEARCH_TOPN(RuleTypeClass.REWRITE), PUSH_DOWN_FILTER_THROUGH_GENERATE(RuleTypeClass.REWRITE), PUSH_DOWN_FILTER_THROUGH_CTE(RuleTypeClass.REWRITE), 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..090c07f7c5409d 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/nereids/rules/rewrite/PushDownFilterThroughVectorSearchTopN.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownFilterThroughVectorSearchTopN.java new file mode 100644 index 00000000000000..f6c7feed6dbb7a --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownFilterThroughVectorSearchTopN.java @@ -0,0 +1,60 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; + +/** + * Move an outer vector_search WHERE predicate below its Doris merge TopN. + * + *

The TopN immediately above a vector_search TVF is added by {@code BindExpression} to merge + * the candidates returned by all Lance fragment scans. The SQL WHERE predicate must therefore be + * evaluated below this TopN so it can become a residual conjunct on the Doris Lance scan node: + * + * 

+ * Filter                         TopN
+ *   TopN            ->            Filter
+ *     vector_search                 vector_search
+ *
+ * + *

This remains a postfilter relative to Lance nearest(): every fragment first returns its ANN + * candidates, and Doris filters those candidates before the local/global TopN. It is deliberately + * not converted into the Lance prefilter carried by the TVF's {@code filter} property. + */ +public class PushDownFilterThroughVectorSearchTopN extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalFilter(logicalTopN(logicalTVFRelation())) + .then(filter -> { + LogicalTopN topN = filter.child(); + if (!VectorSearchTableValuedFunction.NAME.equals( + topN.child().getFunction().getName())) { + return null; + } + LogicalFilter scanFilter = new LogicalFilter<>( + filter.getConjuncts(), topN.child()); + return topN.withChildren(scanFilter); + }).toRule(RuleType.PUSH_DOWN_FILTER_THROUGH_VECTOR_SEARCH_TOPN); + } +} 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..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 @@ -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,22 @@ 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().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 + "'"); @@ -283,6 +299,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/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/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..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,11 +18,20 @@ 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; +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 +51,49 @@ 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")); + } + + @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/gensrc/proto/internal_service.proto b/gensrc/proto/internal_service.proto index 96ed6f1395ff08..aa8fbaa147802c 100644 --- a/gensrc/proto/internal_service.proto +++ b/gensrc/proto/internal_service.proto @@ -841,9 +841,14 @@ message PRequestBlockDesc { repeated PSlotDescriptor slots = 2; repeated ColumnPB column_descs = 3; repeated uint32 file_id = 4; - repeated uint32 row_id = 5; + // The wire field is uint64 for every version. FILE_LOCAL_ROW_ID values are zero-extended + // from uint32, while LANCE_DATASET_ROW_ID uses the full uint64 range. + repeated uint64 row_id = 5; optional PTupleDescriptor desc = 6; repeated uint32 column_idxs = 7; + // Numeric value of ROW_VERSION. Keep this as uint32 so a receiver can detect and reject a + // future version that is unknown to its current binary instead of treating it as version 0. + optional uint32 row_location_version = 8 [default = 0]; } message PTopNLazyMaterializationFileCacheStats { 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..2d03f30d5eb42e 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 @@ -38,6 +38,7 @@ row_id bigint No false \N -- !post_search_filter -- 2 odd 16.0 +514 odd 4210704.0 -- !boundary_single_probe -- 256 item-0256 0.0 @@ -60,4 +61,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/data/external_table_p0/lance/test_lance_vector_search_two_phase.out b/regression-test/data/external_table_p0/lance/test_lance_vector_search_two_phase.out new file mode 100644 index 00000000000000..30532377adfa3e --- /dev/null +++ b/regression-test/data/external_table_p0/lance/test_lance_vector_search_two_phase.out @@ -0,0 +1,23 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !two_phase_execution -- +2 odd item-0002 16.0 +3 even item-0003 64.0 +4 odd item-0004 144.0 +5 even item-0005 256.0 +6 odd item-0006 400.0 + +-- !one_phase_execution -- +2 odd item-0002 16.0 +3 even item-0003 64.0 +4 odd item-0004 144.0 +5 even item-0005 256.0 +6 odd item-0006 400.0 + +-- !prefilter_execution -- +2 odd item-0002 16.0 +4 odd item-0004 144.0 +6 odd item-0006 400.0 + +-- !postfilter_execution -- +2 odd item-0002 16.0 +514 odd item-0514 4210704.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..39ea89dd3d6d29 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" } @@ -175,8 +175,9 @@ suite("test_lance_vector_search", "p0,external") { ORDER BY _distance, row_id """ - // An outer WHERE remains a Doris post-search predicate. Search first selects row_id 1 - // and 2; filtering for odd retains only row_id 2. + // An outer WHERE is a Doris scan postfilter. Each fragment first returns two ANN + // candidates, and Doris filters those candidates before the global TopN. This keeps + // row_id 2 from the first fragment and row_id 514 from the second fragment. qt_post_search_filter """ SELECT row_id, category, _distance FROM ${indexedTopTwo} diff --git a/regression-test/suites/external_table_p0/lance/test_lance_vector_search_two_phase.groovy b/regression-test/suites/external_table_p0/lance/test_lance_vector_search_two_phase.groovy new file mode 100644 index 00000000000000..c78bf520b252a5 --- /dev/null +++ b/regression-test/suites/external_table_p0/lance/test_lance_vector_search_two_phase.groovy @@ -0,0 +1,178 @@ +// 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. + +suite("test_lance_vector_search_two_phase", "p0,external") { + String enabled = context.config.otherConfigs.get("enableIcebergTest") + if (enabled == null || !enabled.equalsIgnoreCase("true")) { + logger.info("disable Lance vector search two-phase test because the Iceberg MinIO environment is disabled.") + return + } + + String externalEnvIp = context.config.otherConfigs.get("externalEnvIp") + String minioPort = context.config.otherConfigs.get("iceberg_minio_port") + String catalogName = "test_lance_vector_search_two_phase" + String tableName = "${catalogName}.doris.vs_ivf_pq_f32" + String headQuery = "[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]" + + // This frozen fixture has two fragments. With top_k=5 and offset=1, every fragment + // contributes at most six phase-one candidates and the global TopN keeps rows 2..6. + String indexedWithOffset = """vector_search( + "table"="${tableName}", + "column"="embedding", + "query_vector"="${headQuery}", + "top_k"="5", + "offset"="1", + "metric"="l2", + "nprobes"="4", + "refine_factor"="10", + "use_index"="true")""" + String resultQuery = """ + SELECT row_id, category, label, _distance + FROM ${indexedWithOffset} + ORDER BY _distance + """ + + sql "DROP CATALOG IF EXISTS `${catalogName}`" + try { + sql """ + CREATE CATALOG `${catalogName}` PROPERTIES ( + "type" = "lance", + "lance.catalog.type" = "filesystem", + "warehouse" = "s3://warehouse/lance", + "s3.endpoint" = "http://${externalEnvIp}:${minioPort}", + "s3.access_key" = "admin", + "s3.secret_key" = "password", + "s3.region" = "us-east-1", + "use_path_style" = "true" + ) + """ + sql "SET enable_file_scanner_v2 = true" + // The same SQL is executed with two different materialization settings below. Disable + // both caches so the second execution cannot reuse the first execution's result. + sql "SET enable_sql_cache = false" + sql "SET enable_query_cache = false" + + // Plan contract with TopN lazy materialization enabled: + // phase 1: Lance reads embedding for ANN, but outputs only _distance and the hidden row ID + // global TopN: OFFSET 1 / LIMIT 5 + // phase 2: fetch row_id, category and label by the hidden row ID + sql "SET topn_lazy_materialization_threshold = 1024" + explain { + sql "verbose ${resultQuery}" + check { explainString -> + assertTrue(explainString.contains("VMaterializeNode")) + assertTrue(explainString.contains("VTOP-N")) + assertTrue(explainString.contains( + "row_ids: [__DORIS_GLOBAL_ROWID_COL__vector_search]")) + assertTrue(explainString.contains("isTopMaterializeNode: true")) + assertTrue(explainString.contains("limit: 5")) + assertTrue(explainString.contains("offset: 1")) + assertTrue(explainString.contains("externalSearchType=VECTOR")) + assertTrue(explainString.contains("lanceTopK=5")) + assertTrue(explainString.contains("lanceOffset=1")) + assertTrue(explainString.contains("lanceSearchFragments=2")) + String lazyColumns = explainString.readLines() + .find { line -> line.contains("column_descs_lists") } + assertTrue(lazyColumns != null, "missing phase-two column_descs_lists") + assertTrue(lazyColumns.contains("row_id")) + assertTrue(lazyColumns.contains("category")) + assertTrue(lazyColumns.contains("label")) + assertFalse(lazyColumns.contains("embedding")) + return true + } + } + + qt_two_phase_execution "${resultQuery}" + + // Turning the threshold off removes both the Materialization node and the hidden row ID. + // The user-visible result must remain identical to the two-phase result. + sql "SET topn_lazy_materialization_threshold = -1" + explain { + sql "verbose ${resultQuery}" + notContains "VMaterializeNode" + notContains "__DORIS_GLOBAL_ROWID_COL__vector_search" + contains "externalSearchType=VECTOR" + contains "lanceSearchFragments=2" + } + qt_one_phase_execution "${resultQuery}" + + sql "SET topn_lazy_materialization_threshold = 1024" + + // TVF filter is a Lance prefilter. Lance reads category while selecting ANN candidates + // and removes ineligible rows before nearest(), so the three nearest category='odd' rows + // are row IDs 2, 4 and 6. It does not become a Doris residual predicate. + String prefilterQuery = """ + SELECT row_id, category, label, _distance + FROM vector_search( + "table"="${tableName}", + "column"="embedding", + "query_vector"="${headQuery}", + "top_k"="3", + "filter"="category = 'odd'", + "metric"="l2", + "nprobes"="4", + "refine_factor"="10", + "use_index"="true") + ORDER BY _distance + """ + explain { + sql "verbose ${prefilterQuery}" + contains "VMaterializeNode" + notContains "predicates:" + } + qt_prefilter_execution "${prefilterQuery}" + + // Outer WHERE is a Doris scan postfilter, not a Lance prefilter. Each of the two fragments + // first returns three ANN candidates: rows 1..3 and rows 513..515. Doris evaluates + // category='odd' in the phase-one scan, leaving rows 2 and 514 before the global TopN. + // Since Doris does not ask Lance for replacement candidates, fewer than top_k rows remain. + String postfilterQuery = """ + SELECT row_id, category, label, _distance + FROM vector_search( + "table"="${tableName}", + "column"="embedding", + "query_vector"="${headQuery}", + "top_k"="3", + "metric"="l2", + "nprobes"="4", + "refine_factor"="10", + "use_index"="true") + WHERE category = 'odd' + ORDER BY _distance + """ + explain { + sql "verbose ${postfilterQuery}" + check { explainString -> + assertTrue(explainString.contains("VMaterializeNode")) + assertFalse(explainString.contains("VSELECT")) + assertTrue(explainString.contains("predicates:")) + assertTrue(explainString.contains("category")) + String lazyColumns = explainString.readLines() + .find { line -> line.contains("column_descs_lists") } + assertTrue(lazyColumns != null, "missing phase-two column_descs_lists") + assertTrue(lazyColumns.contains("row_id")) + assertTrue(lazyColumns.contains("label")) + assertFalse(lazyColumns.contains("category"), + "postfilter column category must be read in phase one") + return true + } + } + qt_postfilter_execution "${postfilterQuery}" + } finally { + // sql "DROP CATALOG IF EXISTS `${catalogName}`" + } +} diff --git a/thirdparty/download-thirdparty.sh b/thirdparty/download-thirdparty.sh index 87eae93eaf4c82..7fa675da700227 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 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-doris.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-doris.patch b/thirdparty/patches/lance-c-0.1.6-doris.patch new file mode 100644 index 00000000000000..cc49293d9dedf8 --- /dev/null +++ b/thirdparty/patches/lance-c-0.1.6-doris.patch @@ -0,0 +1,632 @@ +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 +@@ -752,0 +753,25 @@ int32_t lance_dataset_take( ++/** ++ * 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 ++); ++ +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 +@@ -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, ++ 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(); ++ } ++ } ++ +diff --git a/src/dataset.rs b/src/dataset.rs +index 91a735f..364be09 100644 +--- a/src/dataset.rs ++++ b/src/dataset.rs +@@ -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 ++/// `_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::invalid_input("dataset must not be NULL")); ++ } ++ if out.is_null() { ++ 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::invalid_input(format!( ++ "row_ids must not be NULL when num_row_ids = {num_row_ids}" ++ ))); ++ } ++ ++ 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 = projection_from_columns(&snap, col_names.as_deref())?; ++ ++ 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) ++} ++ +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 2babda6..9ee5af9 100644 +--- a/tests/c_api_test.rs ++++ b/tests/c_api_test.rs +@@ -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}; +@@ -393,0 +395,94 @@ fn test_dataset_take() { ++#[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) }; ++} ++ ++#[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}; ++ ++ 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(); ++ 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() ++ }), ++ ) ++ .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); ++ 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); ++} ++ +@@ -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()); ++ ++ 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 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 +@@ -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); ++ ++ 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(); ++} ++ +@@ -695,0 +725 @@ int main(int argc, char** argv) { ++ test_dataset_take_rows(uri);