[feature](lance) Push down SQL vector Top-N and BM25 search

be/src/exprs/function/array/function_array_distance.cpp

@@ -63,6 +63,10 @@ void register_function_array_distance(SimpleFunctionFactory& factory) {
     factory.register_function<FunctionArrayDistance<InnerProduct>>();
     factory.register_function<FunctionArrayDistance<L2DistanceApproximate>>();
     factory.register_function<FunctionArrayDistance<InnerProductApproximate>>();
+    factory.register_function<FunctionArrayDistance<CosineDistanceApproximate>>();
+    factory.register_function<FunctionArrayDistance<CosineSimilarityApproximate>>();
+    factory.register_function<FunctionArrayDistance<HammingDistance, TYPE_BIGINT>>();
+    factory.register_function<FunctionArrayDistance<HammingDistanceApproximate, TYPE_BIGINT>>();
 }
 
 } // namespace doris

be/src/exprs/function/array/function_array_distance.h

@@ -21,6 +21,8 @@
 #include <faiss/utils/distances.h>
 #include <gen_cpp/Types_types.h>
 
+#include <bit>
+
 #include "common/exception.h"
 #include "common/status.h"
 #include "core/assert_cast.h"
@@ -86,15 +88,50 @@ class InnerProductApproximate : public InnerProduct {
     static constexpr auto name = "inner_product_approximate";
 };
 
-template <typename DistanceImpl>
+class CosineDistanceApproximate : public CosineDistance {
+public:
+    static constexpr auto name = "cosine_distance_approximate";
+};
+
+class CosineSimilarityApproximate : public CosineSimilarity {
+public:
+    static constexpr auto name = "cosine_similarity_approximate";
+};
+
+class HammingDistance {
+public:
+    static constexpr auto name = "hamming_distance";
+    static float distance(const int64_t* x, const int64_t* y, size_t d) {
+        uint64_t bits = 0;
+        for (size_t i = 0; i < d; ++i) {
+            // Validate the wide SQL integers before converting: truncation would silently
+            // turn invalid values such as 256 into a different binary vector.
+            if (x[i] < 0 || x[i] > 255 || y[i] < 0 || y[i] > 255) {
+                throw doris::Exception(ErrorCode::INVALID_ARGUMENT,
+                                       "Hamming vector elements must be integers in [0, 255]");
+            }
+            bits += std::popcount(static_cast<unsigned>(x[i]) ^ static_cast<unsigned>(y[i]));
+        }
+        return static_cast<float>(bits);
+    }
+};
+
+class HammingDistanceApproximate : public HammingDistance {
+public:
+    static constexpr auto name = "hamming_distance_approximate";
+};
+
+template <typename DistanceImpl, PrimitiveType ArgumentType = TYPE_FLOAT>
 class FunctionArrayDistance : public IFunction {
 public:
     using DataType = PrimitiveTypeTraits<TYPE_FLOAT>::DataType;
     using ColumnType = PrimitiveTypeTraits<TYPE_FLOAT>::ColumnType;
 
     static constexpr auto name = DistanceImpl::name;
     String get_name() const override { return name; }
-    static FunctionPtr create() { return std::make_shared<FunctionArrayDistance<DistanceImpl>>(); }
+    static FunctionPtr create() {
+        return std::make_shared<FunctionArrayDistance<DistanceImpl, ArgumentType>>();
+    }
     size_t get_number_of_arguments() const override { return 2; }
 
     DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
@@ -154,17 +191,17 @@ class FunctionArrayDistance : public IFunction {
         _validate_no_nulls(col2, "Second argument", get_name());
 
         // Create views — handles Const/Nullable unwrapping automatically
-        auto view1 = ColumnArrayView<TYPE_FLOAT>::create(col1);
-        auto view2 = ColumnArrayView<TYPE_FLOAT>::create(col2);
+        auto view1 = ColumnArrayView<ArgumentType>::create(col1);
+        auto view2 = ColumnArrayView<ArgumentType>::create(col2);
 
         auto dst = ColumnType::create(input_rows_count);
         auto& dst_data = dst->get_data();
 
         for (size_t row = 0; row < input_rows_count; ++row) {
             auto a1 = view1[row];
             auto a2 = view2[row];
-            const float* p1 = a1.get_data();
-            const float* p2 = a2.get_data();
+            const auto* p1 = a1.get_data();
+            const auto* p2 = a2.get_data();
             auto dim1 = a1.size();
             auto dim2 = a2.size();
 

be/src/format_v2/table/lance_reader.cpp

@@ -32,11 +32,13 @@
 #include "common/config.h"
 #include "common/logging.h"
 #include "exec/common/endian.h"
+#include "format/format_common.h"
 #include "format_v2/lance/lance_reader_helper.h"
 #include "format_v2/lance/lance_runtime_filter_helper.h"
 #include "format_v2/lance/lance_session_manager.h"
 #include "runtime/file_scan_profile.h"
 #include "runtime/runtime_state.h"
+#include "util/debug_points.h"
 
 namespace doris::format::lance {
 
@@ -591,11 +593,13 @@ Status LanceTableReader::_validate_external_search_request() const {
             return Status::InvalidArgument(
                     "external search filter requires format and non-empty payload");
         }
-        if (filter.format != TSearchFilterFormat::SQL) {
+        if (filter.format != TSearchFilterFormat::SQL &&
+            filter.format != TSearchFilterFormat::SUBSTRAIT) {
             return Status::NotSupported("unsupported external search filter format: {}",
                                         static_cast<int>(filter.format));
         }
-        if (filter.payload.find('\0') != std::string::npos) {
+        if (filter.format == TSearchFilterFormat::SQL &&
+            filter.payload.find('\0') != std::string::npos) {
             return Status::InvalidArgument("Lance SQL search filter contains an embedded NUL byte");
         }
     }
@@ -753,6 +757,14 @@ void LanceTableReader::_init_scanner_profile() {
     _runtime_filter_sql_translation_time =
             add_lance_counter(_scanner_profile, "LanceRuntimeFilterSqlTranslationTime",
                               TUnit::TIME_NS, LANCE_TIMING_PROFILE);
+    _row_id_prefilter_time = add_lance_counter(_scanner_profile, "LanceRowIdPrefilterTime",
+                                               TUnit::TIME_NS, LANCE_TIMING_PROFILE);
+    _row_id_prefilter_rows = add_lance_counter(_scanner_profile, "LanceRowIdPrefilterRows",
+                                               TUnit::UNIT, LANCE_READER_PROFILE);
+    _row_id_prefilter_cache_hits = add_lance_counter(
+            _scanner_profile, "LanceRowIdPrefilterCacheHits", TUnit::UNIT, LANCE_READER_PROFILE);
+    _row_id_prefilter_fallbacks = add_lance_counter(
+            _scanner_profile, "LanceRowIdPrefilterFallbacks", TUnit::UNIT, LANCE_READER_PROFILE);
     _scanner_next_time = add_lance_counter(_scanner_profile, "LanceScannerNextTime", TUnit::TIME_NS,
                                            LANCE_TIMING_PROFILE);
     _execution_iops =
@@ -851,6 +863,19 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) {
         return lance_error("enable Lance row id output");
     }
 
+    bool row_id_prefilter_applied = false;
+    RETURN_IF_ERROR(_apply_row_id_prefilter(scanner, &row_id_prefilter_applied));
+    if (_search_kind != SearchKind::NORMAL && !row_id_prefilter_applied) {
+        const auto& request = lance_scan_params.external_search_request;
+        if (request.__isset.search_filter &&
+            request.search_filter.format == TSearchFilterFormat::SUBSTRAIT &&
+            lance_scanner_set_substrait_filter(
+                    scanner, reinterpret_cast<const uint8_t*>(request.search_filter.payload.data()),
+                    request.search_filter.payload.size()) != 0) {
+            return lance_error("set Lance search Substrait filter");
+        }
+    }
+
     if (lance_scan_params.__isset.lance_substrait_filter &&
         lance_scanner_set_substrait_filter(
                 scanner,
@@ -891,6 +916,148 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) {
     return Status::OK();
 }
 
+Status LanceTableReader::_check_row_id_prefilter_stop() const {
+    RETURN_IF_CANCELLED(_runtime_state);
+    // Downstream early completion stops a scanner without cancelling the whole query.
+    if (_io_ctx != nullptr && _io_ctx->should_stop) {
+        return Status::EndOfFile("stop Lance row-id prefilter");
+    }
+    return Status::OK();
+}
+
+Status LanceTableReader::_read_row_id_prefilter(const std::string& filter, RowIdPrefilter* result) {
+    RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+    // A query-local cache shares this list between index-segment readers. Bound it to 512 KiB
+    // of IDs; broader predicates retain Lance's native filtering instead of growing a C++ list.
+    constexpr size_t MAX_ROW_IDS = 64 * 1024;
+    const char* columns[] = {nullptr};
+    std::unique_ptr<LanceScanner, LanceScannerDeleter> scanner(
+            lance_scanner_new(_dataset, columns, nullptr));
+    if (scanner == nullptr) {
+        return lance_error("create Lance row-id prefilter scanner");
+    }
+    if (lance_scanner_with_row_id(scanner.get(), true) != 0 ||
+        lance_scanner_set_substrait_filter(scanner.get(),
+                                           reinterpret_cast<const uint8_t*>(filter.data()),
+                                           filter.size()) != 0 ||
+        lance_scanner_set_limit(scanner.get(), MAX_ROW_IDS + 1) != 0 ||
+        lance_scanner_set_batch_size(scanner.get(), _runtime_state->batch_size()) != 0) {
+        return lance_error("configure Lance row-id prefilter scanner");
+    }
+    RETURN_IF_ERROR(_configure_scan_options(scanner.get()));
+    const auto collect = [](void* ctx, const LanceScanStatistics* statistics) noexcept {
+        try {
+            LanceTableReader::_collect_scan_statistics(ctx, statistics);
+        } catch (...) {
+            // Statistics are best-effort and must not unwind through the C ABI.
+        }
+    };
+    if (lance_scanner_set_statistics_callback(scanner.get(), collect, this) != 0) {
+        return lance_error("set Lance row-id prefilter statistics callback");
+    }
+    ArrowArrayStream stream {};
+    RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+    if (lance_scanner_to_arrow_stream(scanner.get(), &stream) != 0) {
+        if (stream.release != nullptr) stream.release(&stream);
+        return lance_error("open Lance row-id prefilter stream");
+    }
+    auto imported = arrow::ImportRecordBatchReader(&stream);
+    if (!imported.ok()) {
+        if (stream.release != nullptr) stream.release(&stream);
+        return Status::InternalError("import Lance row-id prefilter: {}",
+                                     imported.status().message());
+    }
+    auto reader = std::move(imported).ValueUnsafe();
+    DorisVector<uint64_t> row_ids;
+    while (true) {
+        RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+        std::shared_ptr<arrow::RecordBatch> batch;
+        auto status = reader->ReadNext(&batch);
+        DBUG_EXECUTE_IF("LanceTableReader.row_id_prefilter.after_read",
+                        DBUG_RUN_CALLBACK(batch == nullptr));
+        RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+        if (!status.ok()) {
+            return Status::InternalError("read Lance row-id prefilter: {}", status.message());
+        }
+        if (batch == nullptr) {
+            result->row_ids = std::move(row_ids);
+            return Status::OK();
+        }
+        auto column = batch->GetColumnByName("_rowid");
+        if (column == nullptr || column->type_id() != arrow::Type::UINT64 ||
+            column->null_count() != 0) {
+            return Status::InternalError("Lance row-id prefilter requires non-null UINT64 _rowid");
+        }
+        const auto count = static_cast<size_t>(column->length());
+        if (count > MAX_ROW_IDS - row_ids.size()) {
+            // A partial allowlist would drop valid rows. Discard it completely on fallback.
+            return Status::OK();
+        }
+        const auto& ids = static_cast<const arrow::UInt64Array&>(*column);
+        if (count > 0) {
+            row_ids.insert(row_ids.end(), ids.raw_values(), ids.raw_values() + count);
+        }
+    }
+}
+
+Status LanceTableReader::_apply_row_id_prefilter(LanceScanner* scanner, bool* applied) {
+    *applied = false;
+    if (_search_kind == SearchKind::NORMAL) {
+        return Status::OK();
+    }
+    const auto& request = _scan_params->lance_scan_params.external_search_request;
+    // Explicit TVF SQL filters retain their existing path. Automatic SQL rewrites carry a
+    // fully validated Substrait prefilter, whose row source can be shared across search splits.
+    if (!request.__isset.prefilter || !request.prefilter || !request.__isset.search_filter ||
+        request.search_filter.format != TSearchFilterFormat::SUBSTRAIT) {
+        return Status::OK();
+    }
+    RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+    SCOPED_TIMER(_row_id_prefilter_time);
+    const auto& filter = request.search_filter.payload;
+    const auto& uri = _opened_dataset_key->uri;
+    const auto key = "lance-row-id-prefilter:" + std::to_string(uri.size()) + ":" + uri + ":" +
+                     std::to_string(_opened_version) + ":" + filter;
+    RowIdPrefilter local;
+    const RowIdPrefilter* result = &local;
+    if (_runtime_filter_cache != nullptr) {
+        bool hit = false;
+        result = _runtime_filter_cache->get<RowIdPrefilter>(
+                key,
+                [&] {
+                    auto value = std::make_unique<RowIdPrefilter>();
+                    value->status = _read_row_id_prefilter(filter, value.get());
+                    if (!value->status.ok()) {
+                        // IOContext stop is reader-local. Never cache an interrupted list or EOF
+                        // that could terminate another active reader sharing this key.
+                        local.status = value->status;
+                        return static_cast<RowIdPrefilter*>(nullptr);
+                    }
+                    return value.release();
+                },
+                &hit);
+        if (hit) COUNTER_UPDATE(_row_id_prefilter_cache_hits, 1);
+    } else {
+        local.status = _read_row_id_prefilter(filter, &local);
+    }
+    RETURN_IF_ERROR(local.status);
+    RETURN_IF_ERROR(_check_row_id_prefilter_stop());
+    RETURN_IF_ERROR(result->status);
+    if (!result->row_ids.has_value()) {
+        COUNTER_UPDATE(_row_id_prefilter_fallbacks, 1);
+        return Status::OK();
+    }
+    const auto& ids = *result->row_ids;
+    // The dataset handle is the same pinned snapshot used to collect these native _rowids.
+    // In particular, an empty list must be installed: it means no rows, not no filtering.
+    if (lance_scanner_set_row_addr_allowlist(scanner, ids.data(), ids.size()) != 0) {
+        return lance_error("set Lance search row-id allowlist");
+    }
+    COUNTER_UPDATE(_row_id_prefilter_rows, ids.size());
+    *applied = true;
+    return Status::OK();
+}
+
 Status LanceTableReader::_configure_scan_options(LanceScanner* scanner) const {
     DORIS_CHECK(scanner != nullptr);
     // Doris runs multiple scanners concurrently. Limit each scanner's read-ahead;

be/src/format_v2/table/lance_reader.h

@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "common/status.h"
+#include "core/custom_allocator.h"
 #include "format_v2/lance/lance_record_batch_converter.h"
 #include "format_v2/table_reader.h"
 #include "runtime/runtime_profile.h"
@@ -82,6 +83,15 @@ class LanceTableReader final : public TableReader {
     std::string _snapshot_description() const;
     Status _prepare_fts_query_context();
     Status _open_scanner(const TFileRangeDesc& range);
+    struct RowIdPrefilter {
+        Status status;
+        // Empty selects no rows; nullopt keeps the provider's native filter for a broad result.
+        // Charge both collection and cache retention to the query's memory limit.
+        std::optional<DorisVector<uint64_t>> row_ids;
+    };
+    Status _check_row_id_prefilter_stop() const;
+    Status _read_row_id_prefilter(const std::string& filter, RowIdPrefilter* result);
+    Status _apply_row_id_prefilter(LanceScanner* scanner, bool* applied);
     void _init_scanner_profile();
     Status _configure_scan_options(LanceScanner* scanner) const;
     Status _configure_normal_scan(LanceScanner* scanner, const TLanceFileDesc& lance_params) const;
@@ -118,6 +128,10 @@ class LanceTableReader final : public TableReader {
     RuntimeProfile::Counter* _scanner_configure_time = nullptr;
     RuntimeProfile::Counter* _runtime_filter_sql_translation_time = nullptr;
     RuntimeProfile::Counter* _scanner_next_time = nullptr;
+    RuntimeProfile::Counter* _row_id_prefilter_time = nullptr;
+    RuntimeProfile::Counter* _row_id_prefilter_rows = nullptr;
+    RuntimeProfile::Counter* _row_id_prefilter_cache_hits = nullptr;
+    RuntimeProfile::Counter* _row_id_prefilter_fallbacks = nullptr;
     RuntimeProfile::Counter* _arrow_to_doris_block_time = nullptr;
     RuntimeProfile::Counter* _row_id_take_read_time = nullptr;
     RuntimeProfile::Counter* _row_id_fetch_total_time = nullptr;