Skip to content

[opt](lance) Support distributed execution for vector_search #66841

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhangstar333 wants to merge 3 commits into
base: branch-4.1
Choose a base branch
from
Open

[opt](lance) Support distributed execution for vector_search #66841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
31f4d09
[opt](lance)support distribute plan in lance vector_search
zhangstar333 Aug 14, 2026
27bc405
fetch rowid
zhangstar333 Aug 17, 2026
984d5e5
update
zhangstar333 Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 46 additions & 3 deletions be/src/exec/operator/materialization_opertor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <fmt/format.h>
#include <gen_cpp/internal_service.pb.h>

#include <cstring>
#include <set>
#include <sstream>
#include <utility>
Expand Down Expand Up @@ -66,6 +67,48 @@ constexpr const char* TOPN_LAZY_MAT_PHASE2_PER_BACKEND_REMOTE_IO_TIME =
constexpr const char* TOPN_LAZY_MAT_PHASE2_PER_BACKEND_WRITE_CACHE_IO_TIME =
"TopNLazyMaterializationSecondPhasePerBackendWriteCacheIOTime";

struct MaterializationRowLocation {
int64_t backend_id = 0;
uint32_t file_id = 0;
uint64_t row_id = 0;
};

Status decode_materialization_row_location(const StringRef& encoded,
MaterializationRowLocation* decoded) {
if (encoded.size == 0) {
return Status::InternalError("global row location is empty");
}

const auto version = static_cast<uint8_t>(encoded.data[0]);
if (version == GlobalRowLoacationV2::VERSION) {
if (encoded.size != sizeof(GlobalRowLoacationV2)) {
return Status::InternalError(
"invalid global row location V2 size: actual={}, expected={}", encoded.size,
sizeof(GlobalRowLoacationV2));
}
GlobalRowLoacationV2 location(0, 0, 0, 0);
std::memcpy(&location, encoded.data, sizeof(location));
decoded->backend_id = location.backend_id;
decoded->file_id = location.file_id;
decoded->row_id = location.row_id;
return Status::OK();
}
if (version == GlobalRowLocationV3::VERSION) {
if (encoded.size != sizeof(GlobalRowLocationV3)) {
return Status::InternalError(
"invalid global row location V3 size: actual={}, expected={}", encoded.size,
sizeof(GlobalRowLocationV3));
}
GlobalRowLocationV3 location(0, 0, 0);
std::memcpy(&location, encoded.data, sizeof(location));
decoded->backend_id = location.backend_id;
decoded->file_id = location.file_id;
decoded->row_id = location.row_id;
return Status::OK();
}
return Status::InternalError("unsupported global row location version: {}", version);
}

void update_counter(RuntimeProfile* profile, const std::string& name, TUnit::type unit,
int64_t value) {
COUNTER_UPDATE(ADD_COUNTER_WITH_LEVEL(profile, name, unit, 2), value);
Expand Down Expand Up @@ -446,9 +489,9 @@ Status MaterializationSharedState::create_muiltget_result(const Columns& columns

for (int j = 0; j < rows; ++j) {
if (!null_map || !null_map[j]) {
DCHECK(column_rowid->get_data_at(j).size == sizeof(GlobalRowLoacationV2));
GlobalRowLoacationV2 row_location =
*((GlobalRowLoacationV2*)column_rowid->get_data_at(j).data);
MaterializationRowLocation row_location;
RETURN_IF_ERROR(decode_materialization_row_location(column_rowid->get_data_at(j),
&row_location));
auto rpc_struct = rpc_struct_map.find(row_location.backend_id);
if (UNLIKELY(rpc_struct == rpc_struct_map.end())) {
return Status::InternalError(
Expand Down
138 changes: 113 additions & 25 deletions be/src/exec/rowid_fetcher.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

#include <algorithm>
#include <cstdint>
#include <limits>
#include <memory>
#include <ostream>
#include <string>
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -658,7 +660,11 @@ Status RowIdStorageReader::read_batch_doris_format_row(
auto max_k = 0;
for (int j = 0; j < request_block_desc.row_id_size();) {
auto file_id = request_block_desc.file_id(j);
row_ids.emplace_back(request_block_desc.row_id(j));
const auto row_id = request_block_desc.row_id(j);
if (row_id > std::numeric_limits<uint32_t>::max()) {
return Status::InvalidArgument("internal row id exceeds uint32 range: {}", row_id);
}
row_ids.emplace_back(static_cast<uint32_t>(row_id));
auto file_mapping = id_file_map->get_file_mapping(file_id);
if (!file_mapping) {
return Status::InternalError(
Expand All @@ -667,7 +673,12 @@ Status RowIdStorageReader::read_batch_doris_format_row(
}
for (k = 1; j + k < request_block_desc.row_id_size(); ++k) {
if (request_block_desc.file_id(j + k) == file_id) {
row_ids.emplace_back(request_block_desc.row_id(j + k));
const auto next_row_id = request_block_desc.row_id(j + k);
if (next_row_id > std::numeric_limits<uint32_t>::max()) {
return Status::InvalidArgument("internal row id exceeds uint32 range: {}",
next_row_id);
}
row_ids.emplace_back(static_cast<uint32_t>(next_row_id));
} else {
break;
}
Expand Down Expand Up @@ -713,22 +724,60 @@ const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRowsRead
const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseSegmentsRead =
"TopNLazyMaterializationSecondPhaseSegmentsRead";

Status RowIdStorageReader::read_lance_rows_by_row_ids(
const TFileRangeDesc& scan_range_desc, const std::vector<uint64_t>& row_ids,
const std::vector<SlotDescriptor>& 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<format::ColumnDefinition> projected_columns;
projected_columns.reserve(slots.size());
for (const auto& slot : slots) {
format::ColumnDefinition column;
column.identifier = Field::create_field<TYPE_STRING>(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<segment_v2::rowid_t, size_t>& row_ids,
size_t idx, const std::multimap<uint64_t, size_t>& row_ids,
const std::shared_ptr<FileMapping>& file_mapping, const std::vector<SlotDescriptor>& slots,
const TUniqueId& query_id, const std::shared_ptr<RuntimeState>& runtime_state,
std::vector<Block>& scan_blocks, std::vector<std::pair<size_t, size_t>>& row_id_block_idx,
std::vector<RowIdStorageReader::ExternalFetchStatistics>& fetch_statistics,
const TFileScanRangeParams& rpc_scan_params,
const std::unordered_map<std::string, int>& colname_to_slot_id,
std::atomic<int>& 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<int64_t> read_ids;
//Generate an ordered list with the help of the orderliness of the map.
std::vector<uint64_t> 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);
Expand All @@ -749,14 +798,31 @@ Status RowIdStorageReader::read_external_row_from_file_mapping(

std::unique_ptr<RuntimeProfile> sub_runtime_profile =
std::make_unique<RuntimeProfile>("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<int64_t> legacy_read_ids;
for (const auto row_id : read_ids) {
if (row_id > static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
return Status::InvalidArgument("legacy external row id exceeds int64 range: {}",
row_id);
}
legacy_read_ids.emplace_back(static_cast<int64_t>(row_id));
}
std::unique_ptr<FileScanner> 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));
}

Expand All @@ -775,11 +841,6 @@ Status RowIdStorageReader::read_external_row_from_file_mapping(
file_read_times_counter->value(), file_read_times_counter->type());
}

semaphore.release();
if (++producer_count == scan_rows_count) {
std::lock_guard<std::mutex> lock(mtx);
cv.notify_one();
}
return Status::OK();
}

Expand Down Expand Up @@ -850,9 +911,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<segment_v2::rowid_t, size_t> : The reason for using multimap is: may need the same row of data multiple times.
std::map<std::string,
std::pair<std::multimap<segment_v2::rowid_t, size_t>, std::shared_ptr<FileMapping>>>
// The multimap retains duplicate row IDs because the same source row can appear more than once
// in the materialized result.
std::map<std::string, std::pair<std::multimap<uint64_t, size_t>, std::shared_ptr<FileMapping>>>
scan_rows;

// Block corresponding to the order of `scan_rows` map.
Expand All @@ -864,7 +925,21 @@ Status RowIdStorageReader::read_batch_external_row(
// Count the time/bytes it takes to read each TFileRangeDesc. (for profile)
std::vector<ExternalFetchStatistics> 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();
Expand All @@ -885,13 +960,11 @@ Status RowIdStorageReader::read_batch_external_row(
}

const auto& external_info = file_mapping->get_external_file_info();
const auto& scan_range_desc = external_info.scan_range_desc;

auto scan_range_hash = hash_file_range(scan_range_desc);
const auto& scan_range_hash = hash_file_range(external_info);
if (scan_rows.contains(scan_range_hash)) {
scan_rows.at(scan_range_hash).first.emplace(request_block_desc.row_id(j), j);
} else {
std::multimap<segment_v2::rowid_t, size_t> tmp {{request_block_desc.row_id(j), j}};
std::multimap<uint64_t, size_t> tmp {{request_block_desc.row_id(j), j}};
scan_rows.emplace(scan_range_hash, std::make_pair(tmp, file_mapping));
}
}
Expand All @@ -917,6 +990,7 @@ Status RowIdStorageReader::read_batch_external_row(
std::atomic<int> 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};
Expand All @@ -927,13 +1001,26 @@ Status RowIdStorageReader::read_batch_external_row(
RETURN_IF_ERROR(remote_scan_sched->submit_scan_task(
SimplifiedScanTask(
[&, idx, scan_info]() -> Status {
Defer complete_task {[&]() {
semaphore.release();
if (++producer_count == scan_rows.size()) {
std::lock_guard<std::mutex> 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<std::mutex> lock(mtx);
if (scan_status.ok()) {
scan_status = status;
}
}
return status;

@github-actions github-actions Bot Aug 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Mark this one-shot scheduler task complete even when the fetch fails. SimplifiedScanTask stores a std::function<bool()>, so returning an error Status here converts to false. ScannerSplitRunner then never sets its completion future, it is not auto-rescheduled, and the one-shot wrapper removes its task handle only for a true result. The new defer wakes the waiter, so the caller returns scan_status while the executor still retains an unfinished split/closure capturing this stack frame by reference; repeated read errors leak those scheduler tasks for the process lifetime. Record the error in scan_status, but return true from the scheduler callback after the single attempt.

},
nullptr, nullptr),
fmt::format("{}-read_batch_external_row-{}", print_id(query_id), idx)));
Expand All @@ -944,6 +1031,7 @@ Status RowIdStorageReader::read_batch_external_row(
std::unique_lock<std::mutex> lock(mtx);
cv.wait(lock, [&] { return producer_count == scan_rows.size(); });
}
RETURN_IF_ERROR(scan_status);
return Status::OK();
},
&scan_running_time));
Expand Down
13 changes: 10 additions & 3 deletions be/src/exec/rowid_fetcher.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
namespace doris {

class DorisNodesInfo;
class RuntimeProfile;
class RuntimeState;
class TupleDescriptor;
namespace io {
Expand Down Expand Up @@ -142,17 +143,23 @@ 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<uint64_t>& row_ids,
const std::vector<SlotDescriptor>& 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<segment_v2::rowid_t, size_t>& row_ids,
size_t idx, const std::multimap<uint64_t, size_t>& row_ids,
const std::shared_ptr<FileMapping>& file_mapping,
const std::vector<SlotDescriptor>& slots, const TUniqueId& query_id,
const std::shared_ptr<RuntimeState>& runtime_state, std::vector<Block>& scan_blocks,
std::vector<std::pair<size_t, size_t>>& row_id_block_idx,
std::vector<ExternalFetchStatistics>& fetch_statistics,
const TFileScanRangeParams& rpc_scan_params,
const std::unordered_map<std::string, int>& colname_to_slot_id,
std::atomic<int>& producer_count, size_t scan_rows_count,
std::counting_semaphore<>& semaphore, std::condition_variable& cv, std::mutex& mtx,
TupleDescriptor& tuple_desc);

struct ExternalFetchStatistics {
Expand Down
Loading
Loading