From 1871330d0af8322f7460e8b8ea8b22a89adda1c2 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 20 Oct 2026 19:37:16 +0800 Subject: [PATCH 1/8] [feature](lance) Push down SQL vector Top-N and BM25 search --- be/src/format_v2/table/lance_reader.cpp | 147 +++++++- be/src/format_v2/table/lance_reader.h | 11 + be/test/format_v2/table/lance_reader_test.cpp | 252 ++++++++++++- .../doris/nereids/jobs/executor/Rewriter.java | 2 + .../apache/doris/nereids/rules/RuleType.java | 1 + .../PushDownSearchTopNIntoLanceScan.java | 323 +++++++++++++++++ .../functions/table/LanceSearch.java | 43 +++ .../FullTextSearchTableValuedFunction.java | 24 +- ...anceExternalSearchTableValuedFunction.java | 19 + .../VectorSearchTableValuedFunction.java | 23 +- .../LanceSearchExpressionPushDownTest.java | 340 ++++++++++++++++++ ...st_lance_search_expression_pushdown.groovy | 113 ++++++ 12 files changed, 1290 insertions(+), 8 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/LanceSearch.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java create mode 100644 regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index e0e2884efe7a4a..739d61bc188aec 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -32,6 +32,7 @@ #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" @@ -591,11 +592,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(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 +756,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 +862,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(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 +915,125 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { return Status::OK(); } +Status LanceTableReader::_read_row_id_prefilter(const std::string& filter, RowIdPrefilter* result) { + // 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 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(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 {}; + 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(); + std::vector row_ids; + while (true) { + RETURN_IF_CANCELLED(_runtime_state); + std::shared_ptr batch; + auto status = reader->ReadNext(&batch); + 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(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(*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(); + } + 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( + key, + [&] { + auto value = std::make_unique(); + value->status = _read_row_id_prefilter(filter, value.get()); + 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(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; diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 205c082d12052d..3a548d95f1ac2d 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -82,6 +82,13 @@ 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. + std::optional> row_ids; + }; + 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 +125,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; diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 6531fc21de72e8..1c5fc552379d0a 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -77,6 +77,7 @@ #include "exprs/vexpr_context.h" #include "exprs/vliteral.h" #include "exprs/vslot_ref.h" +#include "format/format_common.h" #include "format_v2/lance/lance_reader_helper.h" #include "runtime/runtime_profile.h" #include "runtime/runtime_state.h" @@ -308,7 +309,8 @@ Status convert_record_batch_for_test(RuntimeState* runtime_state, const Columns& } Status prepare_range(LanceTableReader* reader, TFileRangeDesc range, - std::optional global_rowid_context = std::nullopt) { + std::optional global_rowid_context = std::nullopt, + ShardedKVCache* cache = nullptr) { // Assign after value initialization so adding optional split state cannot break this fixture's // designated initializer under -Wmissing-designated-field-initializers. SplitReadOptions options; @@ -316,6 +318,7 @@ Status prepare_range(LanceTableReader* reader, TFileRangeDesc range, options.current_range = std::move(range); options.current_split_format = FileFormat::LANCE; options.global_rowid_context = global_rowid_context; + options.cache = cache; return reader->prepare_split(options); } @@ -708,6 +711,63 @@ TEST(LanceTableReaderFullTextSearchTest, SearchesIndexedSnapshotWithOptionalScor EXPECT_TRUE(row_id_reader.close().ok()); } +TEST(LanceTableReaderFullTextSearchTest, RowIdPrefilterPreservesGlobalScoresAndEmptySet) { + const std::filesystem::path dataset_uri = + "./be/test/format_v2/table/lance/data/fts_indexed.lance"; + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(dataset_uri, &fixture).ok()); + auto range = make_full_text_search_range(dataset_uri, fixture, "body_fts"); + const Columns columns {projected_column("row_id", TYPE_BIGINT, false), + projected_column("_score", TYPE_FLOAT, true)}; + TQueryGlobals globals; + RuntimeState state(globals); + auto params = make_full_text_search_params("lance", 8, 0); + RuntimeProfile baseline_profile("lance_fts_baseline"); + LanceTableReader baseline; + ASSERT_TRUE(init_reader(&baseline, columns, &state, &baseline_profile, ¶ms).ok()); + ASSERT_TRUE(prepare_range(&baseline, range).ok()); + Block block; + add_output_columns(&block, columns); + const auto all_rows = read_full_text_search_rows(&baseline, &block); + ASSERT_EQ(4, all_rows.size()); + ASSERT_TRUE(baseline.close().ok()); + // FE-produced row_id >= 3 and row_id >= 100 envelopes for the complete FTS fixture schema. + const std::vector filters { + "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSERoPCAEaC2d0ZTphbnlfYW55GisKHBoaGgQKAhACIg" + "oaCBIGCgISACIAIgYaBAoCOAMaC2ZpbHRlcl9tYXNrIjsKBnJvd19pZAoFdGl0bGUKBGJvZHkKCGNhdGVnb3J5" + "EhoKBDoCEAIKBGICEAIKBGICEAIKBGICEAIYAjoPEEYqC2RvcmlzLWxhbmNl", + "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSERoPCAEaC2d0ZTphbnlfYW55GisKHBoaGgQKAhACIg" + "oaCBIGCgISACIAIgYaBAoCOGQaC2ZpbHRlcl9tYXNrIjsKBnJvd19pZAoFdGl0bGUKBGJvZHkKCGNhdGVnb3J5" + "EhoKBDoCEAIKBGICEAIKBGICEAIKBGICEAIYAjoPEEYqC2RvcmlzLWxhbmNl"}; + for (size_t iteration = 0; iteration < filters.size(); ++iteration) { + std::string payload; + ASSERT_TRUE(base64_decode(filters[iteration], &payload)); + TSearchFilter filter; + filter.__set_format(TSearchFilterFormat::SUBSTRAIT); + filter.__set_payload(payload); + params.lance_scan_params.external_search_request.__set_search_filter(filter); + params.lance_scan_params.external_search_request.__set_prefilter(true); + RuntimeProfile profile("lance_fts_row_id_prefilter"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, ¶ms).ok()); + ASSERT_TRUE(prepare_range(&reader, range).ok()); + Block output; + add_output_columns(&output, columns); + const auto rows = read_full_text_search_rows(&reader, &output); + EXPECT_EQ(iteration == 0 ? 2 : 0, rows.size()); + for (const auto& [id, score] : rows) { + auto found = std::ranges::find_if(all_rows, + [&](const auto& row) { return row.first == id; }); + ASSERT_NE(all_rows.end(), found); + EXPECT_GE(id, 3); + EXPECT_FLOAT_EQ(found->second, score); + } + EXPECT_EQ(iteration == 0 ? 6 : 0, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(0, profile.get_counter("LanceRowIdPrefilterFallbacks")->value()); + EXPECT_TRUE(reader.close().ok()); + } +} + TEST(LanceTableReaderFullTextSearchTest, SupportsMatchAndPhraseQueries) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/fts_indexed.lance"; @@ -1878,6 +1938,196 @@ TEST(LanceTableReaderVectorSearchTest, AppliesSearchFilterBeforeTopKWhenPrefilte EXPECT_TRUE(reader.close().ok()); } +TEST(LanceTableReaderVectorSearchTest, AppliesSubstraitSearchFilterBeforeTopK) { + 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 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); + auto scan_params = make_float32_vector_search_params({0.0F, 0.0F, 0.0F}, 1, 0, "row_id >= 3"); + const std::string substrait_filter_base64 = + "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSERoPCAEaC2d0ZTphbnlfYW55GisKHBoaGgQKAhACIg" + "oa" + "CBIGCgISACIAIgYaBAoCOAMaC2ZpbHRlcl9tYXNrIrkGCgZyb3dfaWQKCmJvb2xfdmFsdWUKDXRpbnlpbnRfdm" + "Fs" + "dWUKDnNtYWxsaW50X3ZhbHVlCglpbnRfdmFsdWUKDGJpZ2ludF92YWx1ZQoLZmxvYXRfdmFsdWUKDGRvdWJsZV" + "92" + "YWx1ZQoLZGVjaW1hbF8xXzAKC2RlY2ltYWxfOV8yCgxkZWNpbWFsXzEwXzAKDGRlY2ltYWxfMThfNAoMZGVjaW" + "1h" + "bF8xOV8wCg1kZWNpbWFsXzM4XzEwCiRfX3VubGlrZWx5X25hbWVfcGxhY2Vob2xkZXJfZG9yaXNfMTQKJF9fdW" + "5s" + "aWtlbHlfbmFtZV9wbGFjZWhvbGRlcl9kb3Jpc18xNQoKdGV4dF92YWx1ZQokX191bmxpa2VseV9uYW1lX3BsYW" + "Nl" + "aG9sZGVyX2RvcmlzXzE3CiRfX3VubGlrZWx5X25hbWVfcGxhY2Vob2xkZXJfZG9yaXNfMTgKJF9fdW5saWtlbH" + "lf" + "bmFtZV9wbGFjZWhvbGRlcl9kb3Jpc18xOQokX191bmxpa2VseV9uYW1lX3BsYWNlaG9sZGVyX2RvcmlzXzIwCi" + "Rf" + "X3VubGlrZWx5X25hbWVfcGxhY2Vob2xkZXJfZG9yaXNfMjEKJF9fdW5saWtlbHlfbmFtZV9wbGFjZWhvbGRlcl" + "9k" + "b3Jpc18yMgokX191bmxpa2VseV9uYW1lX3BsYWNlaG9sZGVyX2RvcmlzXzIzCiRfX3VubGlrZWx5X25hbWVfcG" + "xh" + "Y2Vob2xkZXJfZG9yaXNfMjQKBWxhYmVsCiRfX3VubGlrZWx5X25hbWVfcGxhY2Vob2xkZXJfZG9yaXNfMjYSxw" + "EK" + "BDoCEAIKBAoCEAEKBBICEAEKBBoCEAEKBCoCEAEKBDoCEAEKBFICEAEKBFoCEAEKB8IBBBABIAEKCcIBBggCEA" + "kg" + "AQoHwgEEEAogAQoJwgEGCAQQEiABCgfCAQQQEyABCgnCAQYIChAmIAEKBfIBAhgBCgXyAQIYAQoEYgIQAQoF8g" + "EC" + "GAEKBfIBAhgBCgXyAQIYAQoF8gECGAEKBfIBAhgBCgXyAQIYAQoF8gECGAEKBfIBAhgBCgRiAhABCgXyAQIYAR" + "gC" + "Og8QRioLZG9yaXMtbGFuY2U="; + std::string substrait_filter; + ASSERT_TRUE(base64_decode(substrait_filter_base64, &substrait_filter)); + TSearchFilter search_filter; + search_filter.__set_format(TSearchFilterFormat::SUBSTRAIT); + search_filter.__set_payload(substrait_filter); + scan_params.lance_scan_params.external_search_request.__set_search_filter(search_filter); + scan_params.lance_scan_params.external_search_request.__set_prefilter(true); + + ShardedKVCache cache(1); + for (int iteration = 0; iteration < 2; ++iteration) { + RuntimeProfile profile("lance_row_id_prefilter_cache"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); + ASSERT_TRUE( + prepare_range(&reader, + make_lance_range(dataset_uri, fixture.version, fixture.fragment_ids), + std::nullopt, &cache) + .ok()); + Block block; + add_output_columns(&block, columns); + const auto rows = read_vector_search_rows(&reader, &block); + ASSERT_EQ(1U, rows.size()); + EXPECT_EQ(4, rows[0].first); + EXPECT_FLOAT_EQ(8.25F, rows[0].second); + ASSERT_NE(nullptr, profile.get_counter("LanceRowIdPrefilterRows")); + EXPECT_EQ(2, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(iteration, profile.get_counter("LanceRowIdPrefilterCacheHits")->value()); + EXPECT_TRUE(reader.close().ok()); + } +} + +TEST(LanceTableReaderVectorSearchTest, BroadRowIdPrefilterFallsBackWithoutLosingNearestRow) { + const auto unique = std::chrono::steady_clock::now().time_since_epoch().count(); + const auto uri = std::filesystem::temp_directory_path() / + ("doris_lance_broad_prefilter_" + std::to_string(unique) + ".lance"); + Defer cleanup {[&] { + std::error_code error; + std::filesystem::remove_all(uri, error); + }}; + constexpr int64_t ROWS = 65538; + arrow::Int64Builder id_builder; + arrow::FloatBuilder vector_builder; + for (int64_t row = 0; row < ROWS; ++row) { + ASSERT_TRUE(id_builder.Append(row).ok()); + ASSERT_TRUE(vector_builder.AppendValues({static_cast(row), 0.0F, 0.0F}).ok()); + } + std::shared_ptr ids; + std::shared_ptr values; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(vector_builder.Finish(&values).ok()); + auto vectors = arrow::FixedSizeListArray::FromArrays(values, 3); + ASSERT_TRUE(vectors.ok()); + const auto schema = + arrow::schema({arrow::field("row_id", arrow::int64(), false), + arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3))}); + ASSERT_TRUE(write_lance_record_batch( + uri, arrow::RecordBatch::Make(schema, ROWS, + {ids, std::move(vectors).ValueUnsafe()})) + .ok()); + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(uri, &fixture).ok()); + auto params = + make_float32_vector_search_params({static_cast(ROWS - 1), 0.0F, 0.0F}, 1, 0); + // FE-generated row_id >= 0. A truncated allowlist would omit the nearest (last) row. + std::string payload; + ASSERT_TRUE(base64_decode( + "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSERoPCAEaC2d0ZTphbnlfYW55GisKHBoaGgQKAhACIg" + "oaCBIGCgISACIAIgYaBAoCOAAaC2ZpbHRlcl9tYXNrIj4KBnJvd19pZAojX191bmxpa2VseV9uYW1lX3BsYWNl" + "aG9sZGVyX2RvcmlzXzESDwoEOgIQAgoF8gECGAEYAjoPEEYqC2RvcmlzLWxhbmNl", + &payload)); + TSearchFilter filter; + filter.__set_format(TSearchFilterFormat::SUBSTRAIT); + filter.__set_payload(payload); + params.lance_scan_params.external_search_request.__set_search_filter(filter); + params.lance_scan_params.external_search_request.__set_prefilter(true); + const Columns columns {projected_column("row_id", TYPE_BIGINT, false), + projected_column("_distance", TYPE_FLOAT, true)}; + TQueryGlobals globals; + RuntimeState state(globals); + RuntimeProfile profile("lance_broad_prefilter"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, ¶ms).ok()); + ASSERT_TRUE(prepare_fixture(&reader, uri, fixture, fixture.fragment_ids).ok()); + Block block; + add_output_columns(&block, columns); + const auto result = read_vector_search_rows(&reader, &block); + ASSERT_EQ(1, result.size()); + EXPECT_EQ(ROWS - 1, result[0].first); + EXPECT_FLOAT_EQ(0.0F, result[0].second); + EXPECT_EQ(0, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(1, profile.get_counter("LanceRowIdPrefilterFallbacks")->value()); + EXPECT_TRUE(reader.close().ok()); +} + +TEST(LanceTableReaderVectorSearchTest, NativeApproximateSearchSkipsNullVectorsAndNaNDistances) { + const auto unique = std::chrono::steady_clock::now().time_since_epoch().count(); + const auto uri = std::filesystem::temp_directory_path() / + ("doris_lance_invalid_vectors_" + std::to_string(unique) + ".lance"); + Defer cleanup {[&] { + std::error_code error; + std::filesystem::remove_all(uri, error); + }}; + arrow::Int64Builder id_builder; + ASSERT_TRUE(id_builder.AppendValues({0, 1, 2}).ok()); + auto values = std::make_shared(); + arrow::FixedSizeListBuilder vector_builder(arrow::default_memory_pool(), values, 3); + ASSERT_TRUE(vector_builder.AppendNull().ok()); + ASSERT_TRUE(vector_builder.Append().ok()); + ASSERT_TRUE(values->AppendValues({std::numeric_limits::quiet_NaN(), 0.0F, 0.0F}).ok()); + ASSERT_TRUE(vector_builder.Append().ok()); + ASSERT_TRUE(values->AppendValues({1.0F, 0.0F, 0.0F}).ok()); + std::shared_ptr ids; + std::shared_ptr vectors; + ASSERT_TRUE(id_builder.Finish(&ids).ok()); + ASSERT_TRUE(vector_builder.Finish(&vectors).ok()); + const auto schema = arrow::schema({arrow::field("row_id", arrow::int64(), false), + arrow::field("embedding", vectors->type())}); + ASSERT_TRUE(write_lance_record_batch(uri, arrow::RecordBatch::Make(schema, 3, {ids, vectors})) + .ok()); + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(uri, &fixture).ok()); + const Columns columns {projected_column("row_id", TYPE_BIGINT, false), + projected_column("_distance", TYPE_FLOAT, true)}; + TQueryGlobals globals; + RuntimeState state(globals); + for (const auto metric : {TVectorMetric::L2, TVectorMetric::DOT_PRODUCT}) { + auto params = make_float32_vector_search_params({1.0F, 0.0F, 0.0F}, 10, 0); + auto& request = params.lance_scan_params.external_search_request; + request.__set_prefilter(true); + request.search_query.vector_search.__set_metric(metric); + RuntimeProfile profile("lance_invalid_vectors"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, ¶ms).ok()); + ASSERT_TRUE(prepare_fixture(&reader, uri, fixture, fixture.fragment_ids).ok()); + Block block; + add_output_columns(&block, columns); + const auto rows = read_vector_search_rows(&reader, &block); + ASSERT_EQ(1, rows.size()); + EXPECT_EQ(2, rows[0].first); + EXPECT_FLOAT_EQ(0.0F, rows[0].second); + EXPECT_TRUE(reader.close().ok()); + } +} + TEST(LanceTableReaderVectorSearchTest, SearchFilterDefaultsToPostfilter) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/all_types.lance"; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java index db96a03926ddc8..6cd36cbd0aaad8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java @@ -138,6 +138,7 @@ import org.apache.doris.nereids.rules.rewrite.PushDownMatchProjectionAsVirtualColumn; import org.apache.doris.nereids.rules.rewrite.PushDownProjectThroughLimit; import org.apache.doris.nereids.rules.rewrite.PushDownScoreTopNIntoOlapScan; +import org.apache.doris.nereids.rules.rewrite.PushDownSearchTopNIntoLanceScan; import org.apache.doris.nereids.rules.rewrite.PushDownTopNDistinctThroughJoin; import org.apache.doris.nereids.rules.rewrite.PushDownTopNDistinctThroughUnion; import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughJoin; @@ -790,6 +791,7 @@ public class Rewriter extends AbstractBatchJobExecutor { )), topDown(DistinctAggregateRewriter.INSTANCE), custom(RuleType.ELIMINATE_UNNECESSARY_PROJECT, EliminateUnnecessaryProject::new), + topDown(new PushDownSearchTopNIntoLanceScan()), topDown(new PushDownVectorTopNIntoOlapScan()), topDown(new PushDownVirtualColumnsIntoOlapScan()), topDown(new PushDownMatchProjectionAsVirtualColumn()), 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 052261dba9690d..1dc2f00d3147ac 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 @@ -342,6 +342,7 @@ public enum RuleType { PUSH_DOWN_MATCH_PROJECTION_AS_VIRTUAL_COLUMN(RuleTypeClass.REWRITE), PUSH_DOWN_SCORE_TOPN_INTO_OLAP_SCAN(RuleTypeClass.REWRITE), PUSH_DOWN_VECTOR_TOPN_INTO_OLAP_SCAN(RuleTypeClass.REWRITE), + PUSH_DOWN_SEARCH_TOPN_INTO_LANCE_SCAN(RuleTypeClass.REWRITE), CHECK_SCORE_USAGE(RuleTypeClass.REWRITE), OLAP_SCAN_TABLET_PRUNE(RuleTypeClass.REWRITE), PUSH_AGGREGATE_TO_OLAP_SCAN(RuleTypeClass.REWRITE), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java new file mode 100644 index 00000000000000..ce5e6710191092 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -0,0 +1,323 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.analysis.Expr; +import org.apache.doris.analysis.TableSnapshot; +import org.apache.doris.analysis.TupleDescriptor; +import org.apache.doris.datasource.lance.LanceExternalTable; +import org.apache.doris.datasource.lance.metadata.LanceRefSelector; +import org.apache.doris.datasource.lance.metadata.LanceTableMetadata; +import org.apache.doris.datasource.lance.source.LancePredicateConverter; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.glue.translator.ExpressionTranslator; +import org.apache.doris.nereids.glue.translator.PlanTranslatorContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.GreaterThan; +import org.apache.doris.nereids.trees.expressions.GreaterThanEqual; +import org.apache.doris.nereids.trees.expressions.LessThan; +import org.apache.doris.nereids.trees.expressions.LessThanEqual; +import org.apache.doris.nereids.trees.expressions.Match; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Properties; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.Subtract; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Sqrt; +import org.apache.doris.nereids.trees.expressions.functions.table.LanceSearch; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.literal.FloatLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.types.DoubleType; +import org.apache.doris.nereids.types.FloatType; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.tablefunction.FullTextSearchTableValuedFunction; +import org.apache.doris.tablefunction.TableValuedFunctionIf; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** Reuses the fixed-snapshot Lance search relations for ordinary SQL search expressions. */ +public class PushDownSearchTopNIntoLanceScan extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalTopN(logicalProject(any())) + .thenApply(ctx -> rewrite(ctx.root, ctx.cascadesContext)) + .toRule(RuleType.PUSH_DOWN_SEARCH_TOPN_INTO_LANCE_SCAN); + } + + private Plan rewrite(LogicalTopN> topN, CascadesContext context) { + LogicalProject project = topN.child(); + Plan input = project.child(); + Set predicates = ImmutableSet.of(); + if (input instanceof LogicalFilter) { + predicates = ((LogicalFilter) input).getConjuncts(); + input = input.child(0); + } + // A non-movable projection may raise an error on rows outside native Top-K. + // Keep its evaluation before pruning, as LogicalProject.pruneOutputs requires. + if (!(input instanceof LogicalFileScan) || !(((LogicalFileScan) input).getTable() + instanceof LanceExternalTable) || project.isDistinct() + || project.containsNoneMovableFunction()) { + return null; + } + LogicalFileScan scan = (LogicalFileScan) input; + if (scan.getTableSample().isPresent() || scan.getScanParams().isPresent() + || topN.getOrderKeys().size() != 1 || topN.getLimit() <= 0) { + return null; + } + Expression order = topN.getOrderKeys().get(0).getExpr(); + Expression expression = project.getProjects().stream() + .filter(p -> p.toSlot().equals(order) && p instanceof Alias) + .map(p -> ((Alias) p).child()).findFirst().orElse(null); + boolean l2 = expression instanceof L2DistanceApproximate; + boolean dot = expression instanceof InnerProductApproximate; + boolean score = expression instanceof Score; + if (!l2 && !dot && !score) { + return null; + } + // The explicit approximate functions select native Lance candidate semantics: NULL + // vectors and NaN distances are excluded. Exact functions retain scalar validation. + // Each split can prune only the nearest (or highest-scoring) candidates. Reversing + // that order or adding tie-breakers would require a different local candidate contract. + if (topN.getOrderKeys().get(0).isAsc() != l2) { + if (score) { + throw new AnalysisException("Lance BM25 Top-N requires ORDER BY score() DESC"); + } + return null; + } + if (topN.getOffset() < 0 || topN.getLimit() > 0xffff_ffffL - topN.getOffset()) { + throw new AnalysisException("Lance search requires limit + offset <= 4294967295"); + } + LanceExternalTable table = (LanceExternalTable) scan.getTable(); + LanceTableMetadata metadata = table.getMetadata(scan.getRelationSnapshot()); + if (!score && metadata.getSchema().getFields().stream().anyMatch( + field -> field.getName().equalsIgnoreCase(VectorSearchTableValuedFunction.DISTANCE_COLUMN))) { + // The TVF reserves _distance, but an ordinary table may legally store that column. + // Keep its existing scalar execution rather than turning a valid SQL query into an error. + return null; + } + Map properties = new LinkedHashMap<>(); + properties.put("table", scan.qualified().stream() + .map(part -> "`" + part.replace("`", "``") + "`").collect(Collectors.joining("."))); + // OFFSET belongs to the global Doris TopN, never independently to each split. + properties.put("top_k", Long.toString(topN.getLimit() + topN.getOffset())); + properties.put("version", Long.toString(metadata.getVersion())); + List filters = new ArrayList<>(predicates); + List scoreFilters = new ArrayList<>(); + if (score) { + List matches = predicates.stream().filter(Match.class::isInstance) + .map(Match.class::cast).collect(Collectors.toList()); + if (matches.size() != 1) { + throw new AnalysisException("Lance BM25 Top-N requires one top-level MATCH_ANY, MATCH_ALL " + + "or MATCH_PHRASE predicate on a single column"); + } + Match match = matches.get(0); + if (!(match.left() instanceof SlotReference) || !(match.right() instanceof StringLikeLiteral) + || match.getAnalyzer().isPresent()) { + throw new AnalysisException("Lance MATCH requires a string literal and the index's stored analyzer"); + } + properties.put("column", physicalName((SlotReference) match.left())); + properties.put("query", ((StringLikeLiteral) match.right()).getStringValue()); + switch (match.op()) { + case MATCH_ANY: + properties.put("operator", "or"); + break; + case MATCH_ALL: + properties.put("operator", "and"); + break; + case MATCH_PHRASE: + properties.put("query_type", "phrase"); + break; + default: + throw new AnalysisException("Unsupported Lance BM25 MATCH operator: " + match.op()); + } + filters.remove(match); + for (Expression predicate : filters) { + if (predicate.anyMatch(Score.class::isInstance)) { + if (!isScoreLowerBound(predicate)) { + throw new AnalysisException("Lance BM25 supports only score() > literal or score() >= literal"); + } + scoreFilters.add(predicate); + } + } + filters.removeAll(scoreFilters); + } else { + Expression column = expression.child(0); + if (column instanceof Cast && column.getDataType().equals(column.child(0).getDataType())) { + column = column.child(0); + } + Expression query = expression.child(1); + if (!(column instanceof SlotReference) || !(query instanceof ArrayLiteral)) { + return null; + } + Field field = metadata.getSchema().findField(physicalName((SlotReference) column)); + // These SQL functions consume ARRAY. Searching a wider physical vector + // directly would bypass Doris's element casts and can change distances/ranking. + if (field == null || !(field.getType() instanceof ArrowType.FixedSizeList) + || field.getChildren().size() != 1 + || !(field.getChildren().get(0).getType() instanceof ArrowType.FloatingPoint) + || ((ArrowType.FloatingPoint) field.getChildren().get(0).getType()).getPrecision() + != FloatingPointPrecision.SINGLE) { + return null; + } + properties.put("column", field.getName()); + properties.put("query_vector", query.toSql()); + properties.put("metric", l2 ? "l2" : "dot"); + } + Optional filter = convertFilters(filters, scan, metadata, context); + if (!filter.isPresent()) { + if (score) { + throw new AnalysisException("Lance BM25 WHERE conditions must be fully pushable before Top-N"); + } + return null; + } + metadata = withIndexes(table, metadata); + TableValuedFunctionIf function; + String name; + String resultName; + try { + if (score) { + function = FullTextSearchTableValuedFunction.forScan(properties, table, metadata, filter.get()); + name = FullTextSearchTableValuedFunction.NAME; + resultName = FullTextSearchTableValuedFunction.SCORE_COLUMN; + } else { + function = VectorSearchTableValuedFunction.forScan(properties, table, metadata, filter.get()); + name = VectorSearchTableValuedFunction.NAME; + resultName = VectorSearchTableValuedFunction.DISTANCE_COLUMN; + } + } catch (org.apache.doris.common.AnalysisException e) { + throw new AnalysisException(e.getMessage(), e); + } + LogicalTVFRelation relation = new LogicalTVFRelation(scan.getRelationId(), + new LanceSearch(name, new Properties(properties), function), ImmutableList.of()); + Map sourceSlots = new HashMap<>(); + scan.getOutput().forEach(slot -> sourceSlots.put(physicalName((SlotReference) slot), (SlotReference) slot)); + List outputs = new ArrayList<>(); + Slot result = null; + for (Slot slot : relation.getOutput()) { + if (slot.getName().equals(resultName)) { + result = slot; + } + SlotReference source = sourceSlots.get(slot.getName()); + outputs.add(source == null ? slot : ((SlotReference) slot) + .withExprId(source.getExprId()).withQualifier(source.getQualifier())); + } + if (result == null) { + throw new AnalysisException("Lance search result column is missing: " + resultName); + } + Expression value = score ? result : l2 + ? new Cast(new Sqrt(new Cast(result, DoubleType.INSTANCE)), FloatType.INSTANCE) + : new Subtract(new FloatLiteral(1.0f), result); + Map replacements = new HashMap<>(); + replacements.put(expression, value); + List projects = ExpressionUtils.replaceNamedExpressions(project.getProjects(), replacements); + List residuals = scoreFilters.stream().map(e -> ExpressionUtils.replace(e, replacements)) + .collect(Collectors.toList()); + Set required = new HashSet<>(); + projects.forEach(p -> required.addAll(p.getInputSlots())); + residuals.forEach(p -> required.addAll(p.getInputSlots())); + // This rewrite runs after column pruning. Do not reintroduce the embedding or wide + // payload columns just because they are part of the underlying search TVF's schema. + relation = relation.withCachedOutputs(outputs.stream().filter(required::contains).collect(Collectors.toList())); + Plan child = residuals.isEmpty() ? relation : new LogicalFilter<>(ImmutableSet.copyOf(residuals), relation); + context.getStatementContext().getSqlCacheContext().ifPresent(cache -> cache.setCannotProcessExpression(true)); + return topN.withChildren(project.withProjectsAndChild(projects, child)); + } + + private Optional convertFilters(List filters, LogicalFileScan scan, + LanceTableMetadata metadata, CascadesContext context) { + if (filters.isEmpty()) { + return Optional.of(new byte[0]); + } + PlanTranslatorContext translator = new PlanTranslatorContext(context); + TupleDescriptor tuple = translator.getDescTable().createTupleDescriptor(); + scan.getOutput().forEach(slot -> translator.createSlotDesc(tuple, (SlotReference) slot, scan.getTable())); + List expressions = filters.stream().map(e -> ExpressionTranslator.translate(e, translator)) + .collect(Collectors.toList()); + LancePredicateConverter.ConversionResult result = new LancePredicateConverter(metadata.getSchema()) + .convert(expressions); + return result.getResidualConjuncts().isEmpty() ? Optional.of(result.getSubstraitFilter()) : Optional.empty(); + } + + private LanceTableMetadata withIndexes(LanceExternalTable table, LanceTableMetadata metadata) { + if (metadata.getIndexMetadataState().canPlanIndexSegments()) { + return metadata; + } + Optional version = Optional.of(new TableSnapshot(Long.toString(metadata.getVersion()), + TableSnapshot.VersionType.VERSION)); + LanceRefSelector selector = metadata.getBranch().isPresent() + ? LanceRefSelector.branch(metadata.getBranch().get(), version) : LanceRefSelector.snapshot(version); + LanceTableMetadata loaded = table.loadMetadataForSearch(selector); + if (loaded.getVersion() != metadata.getVersion() || !loaded.getDatasetUri().equals(metadata.getDatasetUri()) + || !loaded.getSchema().equals(metadata.getSchema())) { + throw new AnalysisException("Lance search index metadata does not match the bound snapshot"); + } + return loaded; + } + + private static String physicalName(SlotReference slot) { + return slot.getOriginalColumn().map(org.apache.doris.catalog.Column::getName).orElse(slot.getName()); + } + + private boolean isScoreLowerBound(Expression predicate) { + if ((predicate instanceof GreaterThan || predicate instanceof GreaterThanEqual) + && unwrapScore(predicate.child(0)) instanceof Score && predicate.child(1) instanceof Literal) { + return true; + } + return (predicate instanceof LessThan || predicate instanceof LessThanEqual) + && predicate.child(0) instanceof Literal && unwrapScore(predicate.child(1)) instanceof Score; + } + + private Expression unwrapScore(Expression expression) { + while (expression instanceof Cast && (expression.getDataType().isFloatType() + || expression.getDataType().isDoubleType())) { + expression = expression.child(0); + } + return expression; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/LanceSearch.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/LanceSearch.java new file mode 100644 index 00000000000000..38f381c9118f42 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/table/LanceSearch.java @@ -0,0 +1,43 @@ +// 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.trees.expressions.functions.table; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Properties; +import org.apache.doris.nereids.types.coercion.AnyDataType; +import org.apache.doris.tablefunction.TableValuedFunctionIf; + +/** A search relation created by an optimizer rewrite of an already bound Lance scan. */ +public class LanceSearch extends TableValuedFunction { + private final TableValuedFunctionIf catalogFunction; + + public LanceSearch(String name, Properties properties, TableValuedFunctionIf catalogFunction) { + super(name, properties); + this.catalogFunction = catalogFunction; + } + + @Override + public FunctionSignature customSignature() { + return FunctionSignature.of(AnyDataType.INSTANCE_WITHOUT_INDEX, getArgumentsTypes()); + } + + @Override + protected TableValuedFunctionIf toCatalogFunction() { + return catalogFunction; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java index 1eeb6f00399d6b..439cf60144bb62 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java @@ -18,6 +18,7 @@ package org.apache.doris.tablefunction; import org.apache.doris.common.AnalysisException; +import org.apache.doris.datasource.lance.LanceExternalTable; import org.apache.doris.datasource.lance.metadata.LanceTableMetadata; import org.apache.doris.thrift.TExternalSearchQuery; import org.apache.doris.thrift.TExternalSearchRequest; @@ -55,11 +56,28 @@ public FullTextSearchTableValuedFunction(Map properties) super(prepare(properties)); } - private static PreparedSearch prepare(Map properties) + /** Builds a search from the metadata snapshot already bound to the source scan. */ + public static FullTextSearchTableValuedFunction forScan(Map properties, + LanceExternalTable table, LanceTableMetadata metadata, byte[] substraitFilter) throws AnalysisException { + return new FullTextSearchTableValuedFunction(prepare(properties, table, metadata), substraitFilter); + } + + private FullTextSearchTableValuedFunction(PreparedSearch prepared, byte[] substraitFilter) { + super(prepared, substraitFilter); + } + + private static PreparedSearch prepare(Map properties) throws AnalysisException { + return prepare(properties, null, null); + } + + private static PreparedSearch prepare(Map properties, + LanceExternalTable table, LanceTableMetadata metadata) throws AnalysisException { Map params = normalizeProperties(properties, PROPERTIES, NAME); - CommonSearch common = prepareCommon(params, NAME, - "FullTextSearchTableValuedFunction", "full-text search", true); + CommonSearch common = table == null + ? prepareCommon(params, NAME, "FullTextSearchTableValuedFunction", "full-text search", true) + : prepareCommonForScan(params, "FullTextSearchTableValuedFunction", + parseTableName(required(params, TABLE, NAME)), table, metadata); Field field = findStringField(common.metadata(), required(params, COLUMN, NAME)); int fieldId = requireLanceFieldId(common.metadata(), field, "full-text"); diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java index 1800d2b3cae3dd..08e211726a7a63 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java @@ -108,6 +108,16 @@ protected LanceExternalSearchTableValuedFunction(PreparedSearch prepared) { this.offset = common.offset; } + protected LanceExternalSearchTableValuedFunction(PreparedSearch prepared, byte[] substraitFilter) { + this(prepared); + // Automatic rewrites must apply the original WHERE before choosing search candidates. + searchRequest.setPrefilter(true); + if (substraitFilter.length > 0) { + searchRequest.setSearchFilter(new TSearchFilter().setFormat(TSearchFilterFormat.SUBSTRAIT) + .setPayload(substraitFilter)); + } + } + public final LanceExternalTable getSourceTable() { return sourceTable; } @@ -190,6 +200,15 @@ protected static CommonSearch prepareCommon(Map params, String f + " requires a fixed positive dataset version"); } + return prepareCommonForScan(params, displayName, sourceTableName, sourceTable, metadata); + } + + protected static CommonSearch prepareCommonForScan(Map params, String displayName, + TableName sourceTableName, LanceExternalTable sourceTable, LanceTableMetadata metadata) + throws AnalysisException { + if (metadata.getVersion() <= 0) { + throw new AnalysisException("Lance search requires a fixed positive dataset version"); + } long topK = parseLong(params.getOrDefault(TOP_K, "10"), TOP_K, 1, Long.MAX_VALUE); long offset = parseLong(params.getOrDefault(OFFSET, "0"), OFFSET, 0, Long.MAX_VALUE); if (offset > UINT32_MAX || topK > UINT32_MAX - offset) { 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 fbb42ed7c0c09a..87c871f9a03225 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 @@ -19,6 +19,7 @@ import org.apache.doris.catalog.Column; import org.apache.doris.common.AnalysisException; +import org.apache.doris.datasource.lance.LanceExternalTable; import org.apache.doris.datasource.lance.metadata.LanceTableMetadata; import org.apache.doris.thrift.TExternalSearchQuery; import org.apache.doris.thrift.TExternalSearchRequest; @@ -82,15 +83,33 @@ public VectorSearchTableValuedFunction(Map properties, boolean d super(prepare(properties, deferQueryVector)); } + /** Builds a search from a bound scan without reopening latest or resolving the table again. */ + public static VectorSearchTableValuedFunction forScan(Map properties, + LanceExternalTable table, LanceTableMetadata metadata, byte[] substraitFilter) + throws AnalysisException { + return new VectorSearchTableValuedFunction(prepare(properties, false, table, metadata), substraitFilter); + } + + private VectorSearchTableValuedFunction(PreparedSearch prepared, byte[] substraitFilter) { + super(prepared, substraitFilter); + } + private static PreparedSearch prepare(Map properties, boolean deferQueryVector) throws AnalysisException { + return prepare(properties, deferQueryVector, null, null); + } + + private static PreparedSearch prepare(Map properties, boolean deferQueryVector, + LanceExternalTable table, LanceTableMetadata metadata) throws AnalysisException { Map params = normalizeProperties(properties, PROPERTIES, NAME); TVectorSearchParams vectorParams = parseDistanceBounds(params); boolean hasDistanceRange = vectorParams.isSetDistanceLowerBound() || vectorParams.isSetDistanceUpperBound(); boolean useIndex = !params.containsKey(USE_INDEX) || parseBoolean(params.get(USE_INDEX), USE_INDEX); - CommonSearch common = prepareCommon(params, NAME, - "VectorSearchTableValuedFunction", "vector search", useIndex); + CommonSearch common = table == null + ? prepareCommon(params, NAME, "VectorSearchTableValuedFunction", "vector search", useIndex) + : prepareCommonForScan(params, "VectorSearchTableValuedFunction", + parseTableName(required(params, TABLE, NAME)), table, metadata); Field vectorField = requireSearchColumn( common.metadata().getSchema(), required(params, COLUMN, NAME), "vector"); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java new file mode 100644 index 00000000000000..a40b3d7dc9450a --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -0,0 +1,340 @@ +// 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.catalog.Column; +import org.apache.doris.catalog.Type; +import org.apache.doris.datasource.lance.LanceExternalTable; +import org.apache.doris.datasource.lance.metadata.LanceMvccSnapshot; +import org.apache.doris.datasource.lance.metadata.LanceRefSelector; +import org.apache.doris.datasource.lance.metadata.LanceTableAccess; +import org.apache.doris.datasource.lance.metadata.LanceTableMetadata; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.properties.OrderKey; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.GreaterThan; +import org.apache.doris.nereids.trees.expressions.LessThan; +import org.apache.doris.nereids.trees.expressions.MatchAll; +import org.apache.doris.nereids.trees.expressions.MatchAny; +import org.apache.doris.nereids.trees.expressions.MatchPhrase; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.Subtract; +import org.apache.doris.nereids.trees.expressions.functions.scalar.AssertTrue; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProduct; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.FloatLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan.SelectedPartitions; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.tablefunction.FullTextSearchTableValuedFunction; +import org.apache.doris.tablefunction.VectorSearchTableValuedFunction; +import org.apache.doris.thrift.TFtsMatchOperator; +import org.apache.doris.thrift.TFtsQueryType; +import org.apache.doris.thrift.TSearchFilterFormat; +import org.apache.doris.thrift.TVectorMetric; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.stream.Collectors; + +class LanceSearchExpressionPushDownTest { + private ConnectContext context; + private ConnectContext previous; + private LogicalFileScan scan; + private LanceExternalTable table; + private LanceTableMetadata metadata; + private SlotReference id; + private SlotReference body; + private SlotReference embedding; + + @BeforeEach + void setup() { + previous = ConnectContext.get(); + context = MemoTestUtils.createConnectContext(); + Schema schema = new Schema(ImmutableList.of( + Field.nullable("id", new ArrowType.Int(64, true)), + Field.nullable("body", ArrowType.Utf8.INSTANCE), + new Field("embedding", FieldType.nullable( + new ArrowType.FixedSizeList(2)), ImmutableList.of(Field.nullable("item", + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)))))); + metadata = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, schema, + Collections.emptyList(), ImmutableMap.of("id", 0, "body", 1, "embedding", 2), + Collections.emptyList()); + table = Mockito.mock(LanceExternalTable.class); + Mockito.when(table.getName()).thenReturn("items"); + Mockito.when(table.getDbName()).thenReturn("db"); + Mockito.when(table.initSelectedPartitions(Mockito.any())).thenReturn(SelectedPartitions.NOT_PRUNED); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(metadata); + List columns = ImmutableList.of(new Column("id", Type.BIGINT), + new Column("body", Type.STRING), new Column("embedding", + org.apache.doris.catalog.ArrayType.create(Type.FLOAT, true))); + Mockito.when(table.getFullSchema(Mockito.any())).thenReturn(columns); + List qualifier = ImmutableList.of("catalog", "db", "items"); + id = SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, columns.get(0), qualifier); + body = SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, columns.get(1), qualifier); + embedding = SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, columns.get(2), qualifier); + scan = new LogicalFileScan(new RelationId(1), table, ImmutableList.of("catalog", "db"), + Collections.emptyList(), Optional.empty(), Optional.empty(), Optional.empty(), + Optional.of(ImmutableList.of(id, body, embedding)), + Optional.of(new LanceMvccSnapshot(metadata))); + } + + @AfterEach + void cleanup() { + ConnectContext.remove(); + if (previous != null) { + previous.setThreadLocalInfo(); + } + } + + private ArrayLiteral vector() { + return new ArrayLiteral(ImmutableList.of(new FloatLiteral(1), new FloatLiteral(2))); + } + + private Plan rewrite(Expression expression, boolean ascending, Plan input) { + Alias result = new Alias(expression, "result"); + LogicalProject project = new LogicalProject<>(ImmutableList.of(id, result), input); + LogicalTopN topN = new LogicalTopN<>( + ImmutableList.of(new OrderKey(result.toSlot(), ascending, false)), 5, 2, project); + return PlanChecker.from(context, topN) + .applyTopDown(new PushDownSearchTopNIntoLanceScan()) + .getPlan(); + } + + private LogicalTVFRelation search(Plan plan) { + List searches = plan.collectToList(LogicalTVFRelation.class::isInstance); + Assertions.assertEquals(1, searches.size(), plan.treeString()); + return searches.get(0); + } + + @Test + void l2TopNUsesLanceAndRestoresEuclideanDistance() { + Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, scan); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertEquals(TVectorMetric.L2, function.getSearchRequest().getSearchQuery() + .getVectorSearch().getMetric()); + Assertions.assertEquals(42, function.getMetadata().getVersion()); + Assertions.assertEquals(7, function.getTopK()); + Assertions.assertEquals(0, function.getOffset()); + Assertions.assertTrue(plan.treeString().toLowerCase().contains("sqrt"), plan.treeString()); + Assertions.assertEquals(2, ((LogicalTopN) plan).getOffset()); + Assertions.assertEquals(ImmutableList.of("id", "result"), + plan.getOutput().stream().map(Slot::getName).collect(Collectors.toList())); + } + + @Test + void innerProductTopNUsesDotAndRestoresSimilarity() { + Plan plan = rewrite(new InnerProductApproximate(embedding, vector()), false, scan); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertEquals(TVectorMetric.DOT_PRODUCT, function.getSearchRequest().getSearchQuery() + .getVectorSearch().getMetric()); + Assertions.assertTrue(plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.anyMatch(Subtract.class::isInstance))), + plan.treeString()); + } + + @Test + void matchAndScoreUseLanceBm25() { + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral("apple"))), scan); + Plan plan = rewrite(new Score(), false, input); + FullTextSearchTableValuedFunction function = (FullTextSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertEquals(TFtsMatchOperator.OR, function.getSearchRequest().getSearchQuery() + .getFullTextSearch().getMatchOperator()); + Assertions.assertEquals("body", function.getSearchRequest().getSearchQuery() + .getFullTextSearch().getColumn()); + Assertions.assertFalse(plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.anyMatch(Score.class::isInstance))), plan.treeString()); + } + + @Test + void ordinaryFiltersRunBeforeVectorTopK() { + Plan input = new LogicalFilter<>(ImmutableSet.of(new EqualTo(id, new BigIntLiteral(9))), scan); + Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, input); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertTrue(function.getSearchRequest().isPrefilter()); + Assertions.assertEquals(TSearchFilterFormat.SUBSTRAIT, + function.getSearchRequest().getSearchFilter().getFormat()); + Assertions.assertTrue(function.getSearchRequest().getSearchFilter().getPayload().length > 0); + Assertions.assertTrue(plan.collectToList(LogicalFilter.class::isInstance).isEmpty(), plan.treeString()); + } + + @Test + void allAndPhraseMatchesUseTheirNativeLanceQueryTypes() { + for (boolean phrase : new boolean[] {false, true}) { + Expression match = phrase ? new MatchPhrase(body, new StringLiteral("apple banana")) + : new MatchAll(body, new StringLiteral("apple banana")); + Plan input = new LogicalFilter<>(ImmutableSet.of(match), scan); + FullTextSearchTableValuedFunction function = (FullTextSearchTableValuedFunction) + search(rewrite(new Score(), false, input)).getFunction().getCatalogFunction(); + Assertions.assertEquals(phrase ? TFtsQueryType.PHRASE + : TFtsQueryType.MATCH, + function.getSearchRequest().getSearchQuery().getFullTextSearch().getQueryType()); + if (!phrase) { + Assertions.assertEquals(TFtsMatchOperator.AND, + function.getSearchRequest().getSearchQuery().getFullTextSearch().getMatchOperator()); + } + } + } + + @Test + void scoreThresholdRemainsAboveSearchWithoutEvaluatingScoreFunction() { + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral("apple")), + new GreaterThan(new Score(), new FloatLiteral(0.5f)), + new EqualTo(id, new BigIntLiteral(9))), scan); + Plan plan = rewrite(new Score(), false, input); + FullTextSearchTableValuedFunction function = (FullTextSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertTrue(function.getSearchRequest().isPrefilter()); + Assertions.assertTrue(function.getSearchRequest().isSetSearchFilter()); + Assertions.assertEquals(1, plan.collectToList(LogicalFilter.class::isInstance).size()); + Assertions.assertFalse(plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.anyMatch(Score.class::isInstance))), plan.treeString()); + } + + @Test + void indexDiscoveryUsesTheBoundVersionInsteadOfLatest() { + LanceTableMetadata basic = LanceTableMetadata.createBasicSnapshot( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + metadata.getSchema(), Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(basic); + Mockito.when(table.loadMetadataForSearch(Mockito.any())).thenReturn(metadata); + Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, scan); + search(plan); + ArgumentCaptor selector = + ArgumentCaptor.forClass(LanceRefSelector.class); + Mockito.verify(table).loadMetadataForSearch(selector.capture()); + Assertions.assertEquals("42", selector.getValue().getSnapshot().get().getValue()); + Mockito.verify(table, Mockito.never()).loadBasicMetadata(Mockito.any()); + } + + @Test + void rejectScoreUpperBoundsAndAscendingBm25() { + Expression match = new MatchAny(body, new StringLiteral("apple")); + Assertions.assertThrows(AnalysisException.class, + () -> rewrite(new Score(), true, new LogicalFilter<>(ImmutableSet.of(match), scan))); + Plan input = new LogicalFilter<>(ImmutableSet.of(match, + new LessThan(new Score(), new FloatLiteral(1))), scan); + Assertions.assertThrows(AnalysisException.class, + () -> rewrite(new Score(), false, input)); + } + + @Test + void doNotReintroduceTheEmbeddingProjection() { + Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, scan); + Assertions.assertEquals(ImmutableList.of("id", "_distance"), search(plan).getOutput().stream() + .map(Slot::getName).collect(Collectors.toList())); + } + + @Test + void fullRewritePipelineUsesTheLanceSearchRule() { + Alias result = new Alias(new L2DistanceApproximate(embedding, vector()), "distance"); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(result.toSlot(), true, false)), + 5, 0, new LogicalProject<>(ImmutableList.of(id, result), scan)); + Plan plan = PlanChecker.from(context, topN).rewrite().getPlan(); + search(plan); + } + + @Test + void reservedResultColumnKeepsAnOtherwiseValidVectorQueryScalar() { + List fields = new ArrayList<>(metadata.getSchema().getFields()); + fields.add(Field.nullable("_distance", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); + LanceTableMetadata collision = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + new Schema(fields), Collections.emptyList(), + ImmutableMap.of("id", 0, "body", 1, "embedding", 2, "_distance", 3), Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(collision); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void nonMovableProjectionCannotBeSkippedByCandidatePruning() { + Alias result = new Alias(new L2DistanceApproximate(embedding, vector()), "distance"); + Alias checked = new Alias(new AssertTrue(new GreaterThan(id, new BigIntLiteral(0)), + new StringLiteral("bad")), "checked"); + LogicalProject project = new LogicalProject<>(ImmutableList.of(id, checked, result), scan); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(result.toSlot(), true, false)), + 1, 0, project); + Plan plan = PlanChecker.from(context, topN) + .applyTopDown(new PushDownSearchTopNIntoLanceScan()).getPlan(); + Assertions.assertTrue(plan.collectToList(LogicalTVFRelation.class::isInstance).isEmpty(), plan.treeString()); + Assertions.assertTrue(plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.containsType(AssertTrue.class))), plan.treeString()); + } + + @Test + void approximateFunctionsSelectNativeSearchOnNullableVectorSchemas() { + // Approximate SQL intentionally uses the provider's candidate domain, unlike exact SQL. + // The corresponding reader test covers exclusion of stored NULL vectors / NaN distances. + Assertions.assertTrue(metadata.getSchema().findField("embedding").isNullable()); + search(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan)); + search(rewrite(new InnerProductApproximate(embedding, vector()), false, scan)); + Assertions.assertTrue(rewrite(new L2Distance(embedding, vector()), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void exactDistanceAndReverseOrderRemainScalar() { + Assertions.assertTrue(rewrite(new InnerProduct(embedding, vector()), false, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + Assertions.assertTrue(rewrite(new L2Distance(embedding, vector()), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), false, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } +} diff --git a/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy new file mode 100644 index 00000000000000..99c6896c2cdf34 --- /dev/null +++ b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy @@ -0,0 +1,113 @@ +// 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_search_expression_pushdown", "p0,external") { + if (context.config.otherConfigs.get("enableIcebergTest") != "true") { + logger.info("Lance fixtures are disabled") + return + } + def catalog = "test_lance_search_expression_pushdown" + def endpoint = context.config.otherConfigs.get("externalEnvIp") + def port = context.config.otherConfigs.get("iceberg_minio_port") + sql "DROP CATALOG IF EXISTS ${catalog}" + try { + sql """CREATE CATALOG ${catalog} PROPERTIES ( + "type"="lance", "lance.catalog.type"="filesystem", "warehouse"="s3://warehouse/lance", + "s3.endpoint"="http://${endpoint}:${port}", "s3.access_key"="admin", + "s3.secret_key"="password", "s3.region"="us-east-1", "use_path_style"="true")""" + sql "SET enable_file_scanner_v2 = true" + def vector = "[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]" + def table = "${catalog}.doris.vs_ivf_flat_f32" + def autoVector = """SELECT row_id, l2_distance_approximate(embedding, ${vector}) AS distance + FROM ${table} WHERE category = 'odd' ORDER BY distance LIMIT 5 OFFSET 1""" + explain { + sql autoVector + contains "externalSearchType=VECTOR" + contains "lanceMetric=l2" + } + // Compare with the explicit TVF, including metric conversion and a global OFFSET. + def explicitVector = """SELECT row_id, CAST(sqrt(_distance) AS FLOAT) AS distance + FROM vector_search("table"="${table}", "column"="embedding", "query_vector"="${vector}", + "top_k"="5", "offset"="1", "metric"="l2", "filter"="category = 'odd'", "prefilter"="true") + ORDER BY distance""" + assertEquals(sql(explicitVector), sql(autoVector)) + // The scalar function returns Euclidean distance, while Lance returns squared L2. + for (row in sql(autoVector)) { + def actual = sql "SELECT l2_distance(embedding, ${vector}) FROM ${table} WHERE row_id = ${row[0]}" + assertTrue(Math.abs((row[1] as double) - (actual[0][0] as double)) < 0.0001) + } + def dotTable = "${catalog}.doris.vs_ivf_pq_f32_dot" + def autoDot = """SELECT row_id, inner_product_approximate(embedding, ${vector}) AS similarity + FROM ${dotTable} ORDER BY similarity DESC LIMIT 5""" + explain { + sql autoDot + contains "externalSearchType=VECTOR" + contains "lanceMetric=dot" + } + def explicitDot = """SELECT row_id, 1.0 - _distance AS similarity + FROM vector_search("table"="${dotTable}", "column"="embedding", "query_vector"="${vector}", + "top_k"="5", "metric"="dot") ORDER BY similarity DESC""" + def dotActual = sql(autoDot) + def dotExpected = sql(explicitDot) + assertEquals(dotExpected.collect { it[0] }, dotActual.collect { it[0] }) + for (int i = 0; i < dotActual.size(); ++i) { + assertTrue(Math.abs((dotExpected[i][1] as double) - (dotActual[i][1] as double)) < 0.001) + } + def textTable = "${catalog}.doris.full_text_search" + for (entry in [["MATCH_ANY", "match", "or"], ["MATCH_ALL", "match", "and"], + ["MATCH_PHRASE", "phrase", null]]) { + def autoText = """SELECT row_id, score() AS relevance FROM ${textTable} + WHERE body ${entry[0]} 'lance search' AND row_id >= 2 AND score() > CAST(0 AS FLOAT) + ORDER BY relevance DESC LIMIT 100""" + explain { + sql autoText + contains "externalSearchType=FULL_TEXT" + contains "lanceFullTextColumn=body" + } + def operator = entry[2] == null ? "" : ', "operator"="' + entry[2] + '"' + def explicitText = """SELECT row_id, _score FROM full_text_search( + "table"="${textTable}", "column"="body", "query"="lance search", "top_k"="100", + "query_type"="${entry[1]}", "filter"="row_id >= 2", "prefilter"="true"${operator}) + WHERE _score > 0 ORDER BY _score DESC""" + // Equal-score rows need not have the same physical order. + def actual = sql(autoText).sort { a, b -> a[0] <=> b[0] } + def expected = sql(explicitText).sort { a, b -> a[0] <=> b[0] } + assertEquals(expected, actual) + } + def empty = sql """SELECT row_id, score() AS relevance FROM ${textTable} + WHERE body MATCH_ANY 'lance' AND row_id < 0 ORDER BY relevance DESC LIMIT 10""" + assertEquals([], empty) + explain { + sql "SELECT row_id FROM ${table} ORDER BY l2_distance(embedding, ${vector}) LIMIT 5" + notContains "externalSearchType=VECTOR" + } + explain { + sql """SELECT row_id FROM ${table} + WHERE md5(category) = 'unsupported' + ORDER BY l2_distance_approximate(embedding, ${vector}) LIMIT 5""" + notContains "externalSearchType=VECTOR" + } + test { + sql """SELECT row_id, score() AS relevance FROM ${textTable} + WHERE body MATCH_ANY 'lance' AND md5(category) = 'unsupported' + ORDER BY relevance DESC LIMIT 5""" + exception "WHERE conditions must be fully pushable before Top-N" + } + } finally { + sql "DROP CATALOG IF EXISTS ${catalog}" + } +} From 16f8c096e7656516ac72c06b04abc63ea53e033e Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 21 Oct 2026 19:15:33 +0800 Subject: [PATCH 2/8] [fix](lance) Preserve bound search values and scalar fallbacks --- .../PushDownSearchTopNIntoLanceScan.java | 9 ++ .../FullTextSearchTableValuedFunction.java | 4 +- ...anceExternalSearchTableValuedFunction.java | 14 ++- .../VectorSearchTableValuedFunction.java | 2 +- .../LanceSearchExpressionPushDownTest.java | 105 +++++++++++++++++- 5 files changed, 128 insertions(+), 6 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java index ce5e6710191092..3fd82a005d6d22 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -131,6 +131,10 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con return null; } if (topN.getOffset() < 0 || topN.getLimit() > 0xffff_ffffL - topN.getOffset()) { + // UINT32 bounds the native candidate budget, not ordinary scalar SQL LIMIT/OFFSET. + if (!score) { + return null; + } throw new AnalysisException("Lance search requires limit + offset <= 4294967295"); } LanceExternalTable table = (LanceExternalTable) scan.getTable(); @@ -217,6 +221,11 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con return null; } metadata = withIndexes(table, metadata); + // SDK schema conversion can leave ordinary fragment scans usable without Lance field IDs. + // Missing index-planning metadata must not turn that scalar path into an analysis error. + if (!score && !metadata.getLanceFieldId(properties.get("column")).isPresent()) { + return null; + } TableValuedFunctionIf function; String name; String resultName; diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java index 439cf60144bb62..6160930903c43a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/FullTextSearchTableValuedFunction.java @@ -79,9 +79,9 @@ private static PreparedSearch prepare(Map properties, : prepareCommonForScan(params, "FullTextSearchTableValuedFunction", parseTableName(required(params, TABLE, NAME)), table, metadata); - Field field = findStringField(common.metadata(), required(params, COLUMN, NAME)); + Field field = findStringField(common.metadata(), required(params, COLUMN, NAME, table == null)); int fieldId = requireLanceFieldId(common.metadata(), field, "full-text"); - String query = required(params, QUERY, NAME); + String query = required(params, QUERY, NAME, table == null); if (query.indexOf('\0') >= 0) { throw new AnalysisException("'query' must not contain an embedded NUL byte"); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java index 08e211726a7a63..803cedddcef55b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/tablefunction/LanceExternalSearchTableValuedFunction.java @@ -173,12 +173,22 @@ protected static Map normalizeProperties(Map pro protected static String required(Map params, String key, String functionName) throws AnalysisException { + return required(params, key, functionName, true); + } + + protected static String required(Map params, String key, String functionName, boolean trim) + throws AnalysisException { String value = params.get(key); - if (value == null || value.trim().isEmpty()) { + // Bound SQL literals and physical names are data: whitespace can change raw-tokenizer + // matches or select a different column. The explicit TVF path keeps its normalization. + if (value != null && trim) { + value = value.trim(); + } + if (value == null || value.isEmpty()) { throw new AnalysisException( "Missing required " + functionName + "() property '" + key + "'"); } - return value.trim(); + return value; } protected static CommonSearch prepareCommon(Map params, String functionName, 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 87c871f9a03225..7574ab79ca185e 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 @@ -112,7 +112,7 @@ private static PreparedSearch prepare(Map properties, boolean de parseTableName(required(params, TABLE, NAME)), table, metadata); Field vectorField = requireSearchColumn( - common.metadata().getSchema(), required(params, COLUMN, NAME), "vector"); + common.metadata().getSchema(), required(params, COLUMN, NAME, table == null), "vector"); // Reject matrices at PREPARE as well as EXECUTE; lance-c only supports single-vector ranges. if (hasDistanceRange && vectorField.getType() instanceof ArrowType.List) { throw new AnalysisException("Distance bounds are not supported for Lance multi-vector search"); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java index a40b3d7dc9450a..ce3d0295419e92 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -43,6 +43,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Md5; import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; @@ -143,10 +144,14 @@ private ArrayLiteral vector() { } private Plan rewrite(Expression expression, boolean ascending, Plan input) { + return rewrite(expression, ascending, input, 5, 2); + } + + private Plan rewrite(Expression expression, boolean ascending, Plan input, long limit, long offset) { Alias result = new Alias(expression, "result"); LogicalProject project = new LogicalProject<>(ImmutableList.of(id, result), input); LogicalTopN topN = new LogicalTopN<>( - ImmutableList.of(new OrderKey(result.toSlot(), ascending, false)), 5, 2, project); + ImmutableList.of(new OrderKey(result.toSlot(), ascending, false)), limit, offset, project); return PlanChecker.from(context, topN) .applyTopDown(new PushDownSearchTopNIntoLanceScan()) .getPlan(); @@ -328,6 +333,104 @@ void approximateFunctionsSelectNativeSearchOnNullableVectorSchemas() { .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); } + @Test + void fullTextRewritePreservesBoundQueryWhitespace() { + for (String query : ImmutableList.of(" foo ", " ")) { + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral(query))), scan); + FullTextSearchTableValuedFunction function = (FullTextSearchTableValuedFunction) + search(rewrite(new Score(), false, input)).getFunction().getCatalogFunction(); + Assertions.assertEquals(query, function.getSearchRequest().getSearchQuery().getFullTextSearch().getQuery()); + } + } + + @Test + void boundSearchFactoriesPreservePhysicalColumnWhitespace() throws Exception { + List fields = new ArrayList<>(metadata.getSchema().getFields()); + fields.add(Field.nullable(" body ", ArrowType.Utf8.INSTANCE)); + Field vectorField = metadata.getSchema().findField("embedding"); + fields.add(new Field(" embedding ", vectorField.getFieldType(), vectorField.getChildren())); + LanceTableMetadata padded = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + new Schema(fields), Collections.emptyList(), + ImmutableMap.of("id", 0, "body", 1, "embedding", 2, " body ", 3, " embedding ", 4), + Collections.emptyList()); + VectorSearchTableValuedFunction vectorSearch = VectorSearchTableValuedFunction.forScan( + ImmutableMap.of("table", "catalog.db.items", "column", " embedding ", "query_vector", "[1,2]"), + table, padded, new byte[0]); + Assertions.assertEquals(" embedding ", + vectorSearch.getSearchRequest().getSearchQuery().getVectorSearch().getColumn()); + FullTextSearchTableValuedFunction textSearch = FullTextSearchTableValuedFunction.forScan( + ImmutableMap.of("table", "catalog.db.items", "column", " body ", "query", " foo "), + table, padded, new byte[0]); + Assertions.assertEquals(" body ", + textSearch.getSearchRequest().getSearchQuery().getFullTextSearch().getColumn()); + Assertions.assertEquals(" foo ", + textSearch.getSearchRequest().getSearchQuery().getFullTextSearch().getQuery()); + } + + @Test + void unavailableFieldIdsAfterEnrichmentKeepVectorQueriesScalar() { + LanceTableMetadata basic = LanceTableMetadata.createBasicSnapshot( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + metadata.getSchema(), Collections.emptyList()); + LanceTableMetadata unavailable = LanceTableMetadata.createSnapshotWithUnavailableFieldIds( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + metadata.getSchema(), Collections.emptyList(), Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(basic); + Mockito.when(table.loadMetadataForSearch(Mockito.any())).thenReturn(unavailable); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + Assertions.assertTrue(rewrite(new InnerProductApproximate(embedding, vector()), false, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void nativeCandidateOverflowKeepsVectorQueriesScalar() { + for (long[] bounds : new long[][] {{0x1_0000_0000L, 0}, {1, 0xffff_ffffL}, + {Long.MAX_VALUE, Long.MAX_VALUE}}) { + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan, + bounds[0], bounds[1]).collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + Assertions.assertTrue(rewrite(new InnerProductApproximate(embedding, vector()), false, scan, + bounds[0], bounds[1]).collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + + @Test + void nativeCandidateOverflowDoesNotRejectOtherScalarFallbacks() { + Plan filtered = new LogicalFilter<>(ImmutableSet.of( + new EqualTo(new Md5(body), new StringLiteral("unpushable"))), scan); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, filtered, + 0x1_0000_0000L, 0).collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + List fields = new ArrayList<>(metadata.getSchema().getFields()); + Field field = fields.get(2); + fields.set(2, new Field("embedding", FieldType.nullable(ArrowType.List.INSTANCE), field.getChildren())); + LanceTableMetadata variableList = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + new Schema(fields), Collections.emptyList(), ImmutableMap.of("id", 0, "body", 1, "embedding", 2), + Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(variableList); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan, + 0x1_0000_0000L, 0).collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void largestNativeCandidateCountStillPushesVectorSearch() { + Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, scan, 1, 0xffff_fffeL); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertEquals(0xffff_ffffL, function.getTopK()); + Assertions.assertEquals(0, function.getOffset()); + Assertions.assertEquals(0xffff_fffeL, ((LogicalTopN) plan).getOffset()); + } + + @Test + void bm25RetainsItsExplicitNativeCandidateLimitError() { + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral("apple"))), scan); + AnalysisException error = Assertions.assertThrows(AnalysisException.class, + () -> rewrite(new Score(), false, input, 0x1_0000_0000L, 0)); + Assertions.assertTrue(error.getMessage().contains("limit + offset <= 4294967295")); + } + @Test void exactDistanceAndReverseOrderRemainScalar() { Assertions.assertTrue(rewrite(new InnerProduct(embedding, vector()), false, scan) From 0aaee0a50f509e9c4e2dfb253a165db07d23984d Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 09:39:19 +0800 Subject: [PATCH 3/8] [fix](lance) Preserve vector search result nullability --- .../PushDownSearchTopNIntoLanceScan.java | 6 +++ .../LanceSearchExpressionPushDownTest.java | 37 +++++++++++++++++++ ...st_lance_search_expression_pushdown.groovy | 5 +++ 3 files changed, 48 insertions(+) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java index 3fd82a005d6d22..d6b465e94a5a37 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -45,6 +45,7 @@ import org.apache.doris.nereids.trees.expressions.Subtract; import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.NonNullable; import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; import org.apache.doris.nereids.trees.expressions.functions.scalar.Sqrt; import org.apache.doris.nereids.trees.expressions.functions.table.LanceSearch; @@ -262,6 +263,11 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con Expression value = score ? result : l2 ? new Cast(new Sqrt(new Cast(result, DoubleType.INSTANCE)), FloatType.INSTANCE) : new Subtract(new FloatLiteral(1.0f), result); + if (!score) { + // Preserve the approximate function's NOT NULL contract for existing Top-N and ancestor slots. + // The generic TVF result and sqrt are nullable; NonNullable also rejects unexpected runtime NULLs. + value = new NonNullable(value); + } Map replacements = new HashMap<>(); replacements.put(expression, value); List projects = ExpressionUtils.replaceNamedExpressions(project.getProjects(), replacements); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java index ce3d0295419e92..b052e1a749c0fe 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -27,6 +27,7 @@ import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.properties.OrderKey; import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.GreaterThan; @@ -57,6 +58,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.types.DoubleType; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.nereids.util.PlanChecker; import org.apache.doris.qe.ConnectContext; @@ -120,6 +122,7 @@ void setup() { List columns = ImmutableList.of(new Column("id", Type.BIGINT), new Column("body", Type.STRING), new Column("embedding", org.apache.doris.catalog.ArrayType.create(Type.FLOAT, true))); + columns.forEach(column -> column.setIsAllowNull(true)); Mockito.when(table.getFullSchema(Mockito.any())).thenReturn(columns); List qualifier = ImmutableList.of("catalog", "db", "items"); id = SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, columns.get(0), qualifier); @@ -294,6 +297,40 @@ void fullRewritePipelineUsesTheLanceSearchRule() { search(plan); } + @Test + void l2FullRewritePreservesNonNullResultsWithDebugValidation() { + assertNonNullVectorFullRewrite(new L2DistanceApproximate(embedding, vector()), true); + } + + @Test + void innerProductFullRewritePreservesNonNullResultsWithDebugValidation() { + assertNonNullVectorFullRewrite(new InnerProductApproximate(embedding, vector()), false); + } + + private void assertNonNullVectorFullRewrite(Expression expression, boolean ascending) { + for (boolean debug : new boolean[] {true, false}) { + context.getSessionVariable().feDebug = debug; + Alias result = new Alias(expression, "distance"); + Assertions.assertFalse(result.nullable()); + LogicalTopN topN = new LogicalTopN<>( + ImmutableList.of(new OrderKey(result.toSlot(), ascending, false)), 5, 0, + new LogicalProject<>(ImmutableList.of(id, result), scan)); + Alias outer = new Alias(new Cast(result.toSlot(), DoubleType.INSTANCE), "distance_as_double"); + Plan plan = PlanChecker.from(context, + new LogicalProject<>(ImmutableList.of(id, result.toSlot(), outer), topN)).rewrite().getPlan(); + LogicalTVFRelation relation = search(plan); + Assertions.assertTrue(relation.getOutput().stream() + .filter(slot -> slot.getName().equals("_distance")).findFirst().get().nullable()); + Assertions.assertEquals(result.getExprId(), plan.getOutput().get(1).getExprId()); + Assertions.assertFalse(plan.getOutput().get(1).nullable(), plan.treeString()); + Assertions.assertFalse(plan.getOutput().get(2).nullable(), plan.treeString()); + List topNs = plan.collectToList(LogicalTopN.class::isInstance); + Assertions.assertEquals(1, topNs.size(), plan.treeString()); + LogicalTopN rewrittenTopN = (LogicalTopN) topNs.get(0); + Assertions.assertFalse(rewrittenTopN.getOrderKeys().get(0).getExpr().nullable(), plan.treeString()); + } + } + @Test void reservedResultColumnKeepsAnOtherwiseValidVectorQueryScalar() { List fields = new ArrayList<>(metadata.getSchema().getFields()); diff --git a/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy index 99c6896c2cdf34..d19881e932e150 100644 --- a/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy +++ b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy @@ -23,6 +23,7 @@ suite("test_lance_search_expression_pushdown", "p0,external") { def catalog = "test_lance_search_expression_pushdown" def endpoint = context.config.otherConfigs.get("externalEnvIp") def port = context.config.otherConfigs.get("iceberg_minio_port") + def previousFeDebug = sql("SHOW VARIABLES LIKE 'fe_debug'")[0][1] sql "DROP CATALOG IF EXISTS ${catalog}" try { sql """CREATE CATALOG ${catalog} PROPERTIES ( @@ -30,6 +31,8 @@ suite("test_lance_search_expression_pushdown", "p0,external") { "s3.endpoint"="http://${endpoint}:${port}", "s3.access_key"="admin", "s3.secret_key"="password", "s3.region"="us-east-1", "use_path_style"="true")""" sql "SET enable_file_scanner_v2 = true" + // Validate that vector rewrites preserve the original functions' NOT NULL contract. + sql "SET fe_debug = true" def vector = "[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]" def table = "${catalog}.doris.vs_ivf_flat_f32" def autoVector = """SELECT row_id, l2_distance_approximate(embedding, ${vector}) AS distance @@ -67,6 +70,7 @@ suite("test_lance_search_expression_pushdown", "p0,external") { for (int i = 0; i < dotActual.size(); ++i) { assertTrue(Math.abs((dotExpected[i][1] as double) - (dotActual[i][1] as double)) < 0.001) } + sql "SET fe_debug = ${previousFeDebug}" def textTable = "${catalog}.doris.full_text_search" for (entry in [["MATCH_ANY", "match", "or"], ["MATCH_ALL", "match", "and"], ["MATCH_PHRASE", "phrase", null]]) { @@ -108,6 +112,7 @@ suite("test_lance_search_expression_pushdown", "p0,external") { exception "WHERE conditions must be fully pushable before Top-N" } } finally { + sql "SET fe_debug = ${previousFeDebug}" sql "DROP CATALOG IF EXISTS ${catalog}" } } From 53ce165cf7cc99b755234daf18900b1bf3006404 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 10:05:35 +0800 Subject: [PATCH 4/8] [fix](lance) Preserve scalar plans on native search validation failures --- .../PushDownSearchTopNIntoLanceScan.java | 5 + .../LanceSearchExpressionPushDownTest.java | 99 ++++++++++++++++++- 2 files changed, 99 insertions(+), 5 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java index d6b465e94a5a37..5fd578ab8399c4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -241,6 +241,11 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con resultName = VectorSearchTableValuedFunction.DISTANCE_COLUMN; } } catch (org.apache.doris.common.AnalysisException e) { + // TVF validation is stricter than a pruned table scan (including unrelated reserved + // names or case collisions). An optional vector optimization must retain scalar SQL. + if (!score) { + return null; + } throw new AnalysisException(e.getMessage(), e); } LogicalTVFRelation relation = new LogicalTVFRelation(scan.getRelationId(), diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java index b052e1a749c0fe..1afc8d2fcfb817 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -333,14 +333,103 @@ private void assertNonNullVectorFullRewrite(Expression expression, boolean ascen @Test void reservedResultColumnKeepsAnOtherwiseValidVectorQueryScalar() { + for (String name : ImmutableList.of("_distance", "_DISTANCE")) { + assertUnprojectedSchemaFallback(ImmutableList.of(Field.nullable(name, + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)))); + } + } + + @Test + void unprojectedReservedPrefixColumnsKeepBothVectorMetricsScalar() { + for (String name : ImmutableList.of(Column.GLOBAL_ROWID_COL, Column.GLOBAL_ROWID_COL + "legacy")) { + assertUnprojectedSchemaFallback(ImmutableList.of(Field.nullable(name, ArrowType.Utf8.INSTANCE))); + } + } + + @Test + void unprojectedCaseInsensitiveDuplicatesKeepBothVectorMetricsScalar() { + assertUnprojectedSchemaFallback(ImmutableList.of(Field.nullable("payload", ArrowType.Utf8.INSTANCE), + Field.nullable("PAYLOAD", ArrowType.Utf8.INSTANCE))); + } + + private void assertUnprojectedSchemaFallback(List additionalFields) { List fields = new ArrayList<>(metadata.getSchema().getFields()); - fields.add(Field.nullable("_distance", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); - LanceTableMetadata collision = LanceTableMetadata.createSnapshotWithIndexes( + fields.addAll(additionalFields); + ImmutableMap.Builder fieldIds = ImmutableMap.builder(); + for (int i = 0; i < fields.size(); i++) { + fieldIds.put(fields.get(i).getName(), i); + } + LanceTableMetadata schema = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + new Schema(fields), Collections.emptyList(), fieldIds.build(), Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(schema); + context.getSessionVariable().feDebug = true; + for (Expression expression : ImmutableList.of(new L2DistanceApproximate(embedding, vector()), + new InnerProductApproximate(embedding, vector()))) { + Alias result = new Alias(expression, "distance"); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(result.toSlot(), + expression instanceof L2DistanceApproximate, false)), 1, 0, + new LogicalProject<>(ImmutableList.of(id, result), scan)); + Plan plan = PlanChecker.from(context, topN).rewrite().getPlan(); + Assertions.assertTrue(plan.collectToList(LogicalTVFRelation.class::isInstance).isEmpty(), + plan.treeString()); + Assertions.assertEquals(1, plan.collectToList(LogicalFileScan.class::isInstance).size()); + Assertions.assertEquals(result.getExprId(), plan.getOutput().get(1).getExprId()); + } + } + + @Test + void fullTextSchemaValidationStillReportsAnExplicitError() { + List fields = new ArrayList<>(metadata.getSchema().getFields()); + fields.add(Field.nullable(Column.GLOBAL_ROWID_COL + "legacy", ArrowType.Utf8.INSTANCE)); + LanceTableMetadata schema = LanceTableMetadata.createSnapshotWithIndexes( new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, new Schema(fields), Collections.emptyList(), - ImmutableMap.of("id", 0, "body", 1, "embedding", 2, "_distance", 3), Collections.emptyList()); - Mockito.when(table.getMetadata(Mockito.any())).thenReturn(collision); - Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, vector()), true, scan) + ImmutableMap.of("id", 0, "body", 1, "embedding", 2, Column.GLOBAL_ROWID_COL + "legacy", 3), + Collections.emptyList()); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(schema); + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral("apple"))), scan); + AnalysisException error = Assertions.assertThrows(AnalysisException.class, + () -> rewrite(new Score(), false, input)); + Assertions.assertTrue(error.getMessage().contains("reserved Doris internal column prefix")); + } + + @Test + void fullTextFullRewritePreservesNullableScoresWithDebugValidation() { + context.getSessionVariable().feDebug = true; + Plan input = new LogicalFilter<>(ImmutableSet.of(new MatchAny(body, new StringLiteral("apple")), + new GreaterThan(new Score(), new FloatLiteral(0.5f))), scan); + Alias result = new Alias(new Score(), "relevance"); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(result.toSlot(), false, false)), + 5, 0, new LogicalProject<>(ImmutableList.of(id, result), input)); + Plan plan = PlanChecker.from(context, topN).rewrite().getPlan(); + search(plan); + Assertions.assertEquals(result.getExprId(), plan.getOutput().get(1).getExprId()); + Assertions.assertTrue(plan.getOutput().get(1).nullable(), plan.treeString()); + Assertions.assertFalse(plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.anyMatch(Score.class::isInstance))), plan.treeString()); + } + + @Test + void secondaryOrderingKeepsBothVectorMetricsScalar() { + for (Expression expression : ImmutableList.of(new L2DistanceApproximate(embedding, vector()), + new InnerProductApproximate(embedding, vector()))) { + Alias result = new Alias(expression, "distance"); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(result.toSlot(), + expression instanceof L2DistanceApproximate, false), new OrderKey(id, true, false)), + 1, 0, new LogicalProject<>(ImmutableList.of(id, result), scan)); + Plan plan = PlanChecker.from(context, topN).rewrite().getPlan(); + Assertions.assertTrue(plan.collectToList(LogicalTVFRelation.class::isInstance).isEmpty(), + plan.treeString()); + } + } + + @Test + void nativeVectorValidationCannotReplaceScalarEvaluationErrors() { + ArrayLiteral wrongDimension = new ArrayLiteral(ImmutableList.of(new FloatLiteral(1))); + Assertions.assertTrue(rewrite(new L2DistanceApproximate(embedding, wrongDimension), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + Assertions.assertTrue(rewrite(new InnerProductApproximate(embedding, wrongDimension), false, scan) .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); } From ef963fb13ab15d7aace56aec666c26d3f1aa49f7 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 10:50:24 +0800 Subject: [PATCH 5/8] [fix](lance) Stop row-id prefilter collection without poisoning cache --- be/src/format_v2/table/lance_reader.cpp | 26 +++++- be/src/format_v2/table/lance_reader.h | 1 + be/test/format_v2/table/lance_reader_test.cpp | 86 ++++++++++++++++++- 3 files changed, 110 insertions(+), 3 deletions(-) diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index 739d61bc188aec..528acdab3dc940 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -38,6 +38,7 @@ #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 { @@ -915,7 +916,17 @@ 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; @@ -945,6 +956,7 @@ Status LanceTableReader::_read_row_id_prefilter(const std::string& filter, RowId 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"); @@ -958,9 +970,12 @@ Status LanceTableReader::_read_row_id_prefilter(const std::string& filter, RowId auto reader = std::move(imported).ValueUnsafe(); std::vector row_ids; while (true) { - RETURN_IF_CANCELLED(_runtime_state); + RETURN_IF_ERROR(_check_row_id_prefilter_stop()); std::shared_ptr 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()); } @@ -997,6 +1012,7 @@ Status LanceTableReader::_apply_row_id_prefilter(LanceScanner* scanner, bool* ap 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; @@ -1011,6 +1027,12 @@ Status LanceTableReader::_apply_row_id_prefilter(LanceScanner* scanner, bool* ap [&] { auto value = std::make_unique(); 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(nullptr); + } return value.release(); }, &hit); @@ -1018,6 +1040,8 @@ Status LanceTableReader::_apply_row_id_prefilter(LanceScanner* scanner, bool* ap } 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); diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 3a548d95f1ac2d..337c768a4ca293 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -87,6 +87,7 @@ class LanceTableReader final : public TableReader { // Empty selects no rows; nullopt keeps the provider's native filter for a broad result. std::optional> 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(); diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 1c5fc552379d0a..0f103217fbd438 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -79,9 +79,11 @@ #include "exprs/vslot_ref.h" #include "format/format_common.h" #include "format_v2/lance/lance_reader_helper.h" +#include "io/io_common.h" #include "runtime/runtime_profile.h" #include "runtime/runtime_state.h" #include "storage/utils.h" +#include "util/debug_points.h" #include "util/defer_op.h" #include "util/timezone_utils.h" #include "util/url_coding.h" @@ -288,13 +290,14 @@ void add_output_columns(Block* block, const Columns& columns) { Status init_reader(LanceTableReader* reader, const Columns& projected_columns, RuntimeState* runtime_state, RuntimeProfile* profile, - TFileScanRangeParams* scan_params, VExprContextSPtrs conjuncts = {}) { + TFileScanRangeParams* scan_params, VExprContextSPtrs conjuncts = {}, + std::shared_ptr io_ctx = nullptr) { return reader->init({ .projected_columns = projected_columns, .conjuncts = std::move(conjuncts), .format = FileFormat::LANCE, .scan_params = scan_params, - .io_ctx = nullptr, + .io_ctx = std::move(io_ctx), .runtime_state = runtime_state, .scanner_profile = profile, }); @@ -768,6 +771,85 @@ TEST(LanceTableReaderFullTextSearchTest, RowIdPrefilterPreservesGlobalScoresAndE } } +TEST(LanceTableReaderFullTextSearchTest, StoppedPrefilterDoesNotPolluteSharedCache) { + const std::filesystem::path uri = "./be/test/format_v2/table/lance/data/fts_indexed.lance"; + LanceFixtureInfo fixture; + ASSERT_TRUE(get_fixture_info(uri, &fixture).ok()); + auto range = make_full_text_search_range(uri, fixture, "body_fts"); + auto params = make_full_text_search_params("lance", 8, 0); + std::string payload; + ASSERT_TRUE(base64_decode( + "Ch4IARIaL2Z1bmN0aW9uc19jb21wYXJpc29uLnlhbWwSERoPCAEaC2d0ZTphbnlfYW55GisKHBoaGgQKAhACIg" + "oaCBIGCgISACIAIgYaBAoCOAMaC2ZpbHRlcl9tYXNrIjsKBnJvd19pZAoFdGl0bGUKBGJvZHkKCGNhdGVnb3J5" + "EhoKBDoCEAIKBGICEAIKBGICEAIKBGICEAIYAjoPEEYqC2RvcmlzLWxhbmNl", + &payload)); + TSearchFilter filter; + filter.__set_format(TSearchFilterFormat::SUBSTRAIT); + filter.__set_payload(payload); + params.lance_scan_params.external_search_request.__set_search_filter(filter); + params.lance_scan_params.external_search_request.__set_prefilter(true); + const Columns columns {projected_column("row_id", TYPE_BIGINT, false), + projected_column("_score", TYPE_FLOAT, true)}; + TQueryOptions options; + options.__set_batch_size(2); + TQueryGlobals globals; + RuntimeState state(globals); + state.set_query_options(options); + // A local scanner stop is not query cancellation and must not affect another reader. + for (int stop_point = 0; stop_point < 3; ++stop_point) { + ShardedKVCache cache(1); + auto io_ctx = std::make_shared(); + io_ctx->should_stop = stop_point == 0; + int reads = 0; + { + const bool previous_debug_points = config::enable_debug_points; + config::enable_debug_points = true; + const std::string point = "LanceTableReader.row_id_prefilter.after_read"; + DebugPoints::instance()->add_with_callback( + point, std::function([&](bool eof) { + ++reads; + if ((stop_point == 1 && reads == 2) || (stop_point == 2 && eof)) { + io_ctx->should_stop = true; + } + })); + Defer cleanup {[&] { + DebugPoints::instance()->remove(point); + config::enable_debug_points = previous_debug_points; + }}; + RuntimeProfile profile("lance_stopped_prefilter"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, ¶ms, {}, io_ctx).ok()); + const auto status = prepare_range(&reader, range, std::nullopt, &cache); + EXPECT_TRUE(status.is()) << status.to_string(); + EXPECT_FALSE(state.is_cancelled()); + EXPECT_EQ(0, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(0, profile.get_counter("LanceRowIdPrefilterFallbacks")->value()); + EXPECT_TRUE(reader.close().ok()); + } + if (stop_point < 2) { + EXPECT_EQ(stop_point == 0 ? 0 : 2, reads); + } else { + EXPECT_GT(reads, 1); + } + for (int iteration = 0; iteration < 2; ++iteration) { + RuntimeProfile profile("lance_active_prefilter"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, ¶ms).ok()); + ASSERT_TRUE(prepare_range(&reader, range, std::nullopt, &cache).ok()); + Block output; + add_output_columns(&output, columns); + const auto rows = read_full_text_search_rows(&reader, &output); + ASSERT_EQ(2, rows.size()); + for (const auto& [id, score] : rows) { + EXPECT_GE(id, 3); + } + EXPECT_EQ(6, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(iteration, profile.get_counter("LanceRowIdPrefilterCacheHits")->value()); + EXPECT_TRUE(reader.close().ok()); + } + } +} + TEST(LanceTableReaderFullTextSearchTest, SupportsMatchAndPhraseQueries) { const std::filesystem::path dataset_uri = "./be/test/format_v2/table/lance/data/fts_indexed.lance"; From f5b14c4eb333a00e55314626740c4ebd493f0411 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 13:13:04 +0800 Subject: [PATCH 6/8] feat(lance): support cosine and Hamming approximate SQL search --- .../array/function_array_distance.cpp | 4 + .../function/array/function_array_distance.h | 49 +++++- .../function/function_array_distance_test.cpp | 64 ++++++++ .../doris/catalog/BuiltinScalarFunctions.java | 8 + .../PushDownSearchTopNIntoLanceScan.java | 57 +++++-- .../scalar/CosineDistanceApproximate.java | 80 ++++++++++ .../scalar/CosineSimilarityApproximate.java | 80 ++++++++++ .../functions/scalar/HammingDistance.java | 100 ++++++++++++ .../scalar/HammingDistanceApproximate.java | 86 ++++++++++ .../visitor/ScalarFunctionVisitor.java | 20 +++ .../LanceSearchExpressionPushDownTest.java | 150 ++++++++++++++++++ ...st_lance_search_expression_pushdown.groovy | 61 +++++++ .../test_array_distance_approximate.groovy | 52 ++++++ 13 files changed, 794 insertions(+), 17 deletions(-) create mode 100644 be/test/exprs/function/function_array_distance_test.cpp create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineDistanceApproximate.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineSimilarityApproximate.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistance.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistanceApproximate.java create mode 100644 regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy diff --git a/be/src/exprs/function/array/function_array_distance.cpp b/be/src/exprs/function/array/function_array_distance.cpp index 89a0dafe1e4aed..ad0daea90e8723 100644 --- a/be/src/exprs/function/array/function_array_distance.cpp +++ b/be/src/exprs/function/array/function_array_distance.cpp @@ -63,6 +63,10 @@ void register_function_array_distance(SimpleFunctionFactory& factory) { factory.register_function>(); factory.register_function>(); factory.register_function>(); + factory.register_function>(); + factory.register_function>(); + factory.register_function>(); + factory.register_function>(); } } // namespace doris diff --git a/be/src/exprs/function/array/function_array_distance.h b/be/src/exprs/function/array/function_array_distance.h index 12969c23e3481f..7840364dfbe12b 100644 --- a/be/src/exprs/function/array/function_array_distance.h +++ b/be/src/exprs/function/array/function_array_distance.h @@ -21,6 +21,8 @@ #include #include +#include + #include "common/exception.h" #include "common/status.h" #include "core/assert_cast.h" @@ -86,7 +88,40 @@ class InnerProductApproximate : public InnerProduct { static constexpr auto name = "inner_product_approximate"; }; -template +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(x[i]) ^ static_cast(y[i])); + } + return static_cast(bits); + } +}; + +class HammingDistanceApproximate : public HammingDistance { +public: + static constexpr auto name = "hamming_distance_approximate"; +}; + +template class FunctionArrayDistance : public IFunction { public: using DataType = PrimitiveTypeTraits::DataType; @@ -94,7 +129,9 @@ class FunctionArrayDistance : public IFunction { static constexpr auto name = DistanceImpl::name; String get_name() const override { return name; } - static FunctionPtr create() { return std::make_shared>(); } + static FunctionPtr create() { + return std::make_shared>(); + } size_t get_number_of_arguments() const override { return 2; } DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { @@ -154,8 +191,8 @@ class FunctionArrayDistance : public IFunction { _validate_no_nulls(col2, "Second argument", get_name()); // Create views — handles Const/Nullable unwrapping automatically - auto view1 = ColumnArrayView::create(col1); - auto view2 = ColumnArrayView::create(col2); + auto view1 = ColumnArrayView::create(col1); + auto view2 = ColumnArrayView::create(col2); auto dst = ColumnType::create(input_rows_count); auto& dst_data = dst->get_data(); @@ -163,8 +200,8 @@ class FunctionArrayDistance : public IFunction { 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(); diff --git a/be/test/exprs/function/function_array_distance_test.cpp b/be/test/exprs/function/function_array_distance_test.cpp new file mode 100644 index 00000000000000..50c73d75275d3d --- /dev/null +++ b/be/test/exprs/function/function_array_distance_test.cpp @@ -0,0 +1,64 @@ +// 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. + +#include "exprs/function/array/function_array_distance.h" + +#include + +#include "core/data_type/data_type_number.h" +#include "function_test_util.h" + +namespace doris { + +TEST(ArrayDistanceFunctionsTest, ApproximateCosineScalarFallbackPreservesZeroSemantics) { + const InputTypeSet types {TYPE_ARRAY, TYPE_FLOAT, TYPE_ARRAY, TYPE_FLOAT}; + const TestArray x {Float32(1), Float32(0)}; + const TestArray y {Float32(0), Float32(1)}; + const TestArray zero {Float32(0), Float32(0)}; + const DataSet distances {{{x, x}, Float32(0)}, {{x, y}, Float32(1)}, {{x, zero}, Float32(2)}}; + const DataSet similarities { + {{x, x}, Float32(1)}, {{x, y}, Float32(0)}, {{x, zero}, Float32(0)}}; + EXPECT_TRUE((check_function("cosine_distance_approximate", types, + distances)) + .ok()); + EXPECT_TRUE((check_function("cosine_similarity_approximate", types, + similarities)) + .ok()); +} + +TEST(ArrayDistanceFunctionsTest, HammingCountsDifferentBitsAndSupportsEmptyArrays) { + const InputTypeSet types {TYPE_ARRAY, TYPE_BIGINT, TYPE_ARRAY, TYPE_BIGINT}; + const TestArray x {Int64(0), Int64(255), Int64(170)}; + const TestArray y {Int64(255), Int64(0), Int64(85)}; + const TestArray empty; + const DataSet rows {{{x, y}, Float32(24)}, {{x, x}, Float32(0)}, {{empty, empty}, Float32(0)}}; + for (const std::string name : {"hamming_distance", "hamming_distance_approximate"}) { + EXPECT_TRUE((check_function(name, types, rows)).ok()); + } +} + +TEST(ArrayDistanceFunctionsTest, HammingRejectsOutOfRangeBeforeByteConversion) { + const int64_t zero[] = {0}; + for (const int64_t value : {-1, 256, 65536}) { + EXPECT_THROW(HammingDistance::distance(&value, zero, 1), Exception); + EXPECT_THROW(HammingDistance::distance(zero, &value, 1), Exception); + } + const int64_t high[] = {255}; + EXPECT_EQ(8, HammingDistance::distance(zero, high, 1)); +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java index 220586eed679e4..aeb683295095c5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java @@ -147,7 +147,9 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Cos; import org.apache.doris.nereids.trees.expressions.functions.scalar.Cosh; import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineDistance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineDistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineSimilarity; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineSimilarityApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Cot; import org.apache.doris.nereids.trees.expressions.functions.scalar.CountEqual; import org.apache.doris.nereids.trees.expressions.functions.scalar.CountSubstring; @@ -237,6 +239,8 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Greatest; import org.apache.doris.nereids.trees.expressions.functions.scalar.Grouping; import org.apache.doris.nereids.trees.expressions.functions.scalar.GroupingId; +import org.apache.doris.nereids.trees.expressions.functions.scalar.HammingDistance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.HammingDistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Hex; import org.apache.doris.nereids.trees.expressions.functions.scalar.HllCardinality; import org.apache.doris.nereids.trees.expressions.functions.scalar.HllEmpty; @@ -731,6 +735,8 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(ConvertTo.class, "convert_to"), scalar(ConvertTz.class, "convert_tz"), scalar(Cos.class, "cos"), + scalar(CosineDistanceApproximate.class, "cosine_distance_approximate"), + scalar(CosineSimilarityApproximate.class, "cosine_similarity_approximate"), scalar(Csc.class, "csc"), scalar(Cosh.class, "cosh"), scalar(Cot.class, "cot"), @@ -818,6 +824,8 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(Greatest.class, "greatest"), scalar(Grouping.class, "grouping"), scalar(GroupingId.class, "grouping_id"), + scalar(HammingDistance.class, "hamming_distance"), + scalar(HammingDistanceApproximate.class, "hamming_distance_approximate"), scalar(Hex.class, "hex"), scalar(HllCardinality.class, "hll_cardinality"), scalar(HllEmpty.class, "hll_empty"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java index 5fd578ab8399c4..d972a72f92f2a1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -43,6 +43,9 @@ import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.Subtract; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineDistanceApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineSimilarityApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.HammingDistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.NonNullable; @@ -59,8 +62,11 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BigIntType; import org.apache.doris.nereids.types.DoubleType; import org.apache.doris.nereids.types.FloatType; +import org.apache.doris.nereids.types.SmallIntType; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.tablefunction.FullTextSearchTableValuedFunction; import org.apache.doris.tablefunction.TableValuedFunctionIf; @@ -117,15 +123,18 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con .map(p -> ((Alias) p).child()).findFirst().orElse(null); boolean l2 = expression instanceof L2DistanceApproximate; boolean dot = expression instanceof InnerProductApproximate; + boolean cosineDistance = expression instanceof CosineDistanceApproximate; + boolean cosineSimilarity = expression instanceof CosineSimilarityApproximate; + boolean hamming = expression instanceof HammingDistanceApproximate; boolean score = expression instanceof Score; - if (!l2 && !dot && !score) { + if (!l2 && !dot && !cosineDistance && !cosineSimilarity && !hamming && !score) { return null; } // The explicit approximate functions select native Lance candidate semantics: NULL // vectors and NaN distances are excluded. Exact functions retain scalar validation. // Each split can prune only the nearest (or highest-scoring) candidates. Reversing // that order or adding tie-breakers would require a different local candidate contract. - if (topN.getOrderKeys().get(0).isAsc() != l2) { + if (topN.getOrderKeys().get(0).isAsc() != (l2 || cosineDistance || hamming)) { if (score) { throw new AnalysisException("Lance BM25 Top-N requires ORDER BY score() DESC"); } @@ -193,26 +202,38 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con filters.removeAll(scoreFilters); } else { Expression column = expression.child(0); - if (column instanceof Cast && column.getDataType().equals(column.child(0).getDataType())) { + // UInt8 maps to SMALLINT in Doris. Widening it to the Hamming function's BIGINT + // arguments preserves every byte; arbitrary element casts must remain scalar. + if (column instanceof Cast && (column.getDataType().equals(column.child(0).getDataType()) + || (hamming && column.getDataType().equals(ArrayType.of(BigIntType.INSTANCE)) + && column.child(0).getDataType().equals(ArrayType.of(SmallIntType.INSTANCE))))) { column = column.child(0); } Expression query = expression.child(1); if (!(column instanceof SlotReference) || !(query instanceof ArrayLiteral)) { return null; } + if ((cosineDistance || cosineSimilarity) && !isUsableCosineQuery((ArrayLiteral) query)) { + return null; + } Field field = metadata.getSchema().findField(physicalName((SlotReference) column)); - // These SQL functions consume ARRAY. Searching a wider physical vector - // directly would bypass Doris's element casts and can change distances/ranking. + // Native search must use the scalar function's physical element type: Float32 for + // floating metrics, UInt8 for Hamming. Other types can change distances/ranking. if (field == null || !(field.getType() instanceof ArrowType.FixedSizeList) - || field.getChildren().size() != 1 - || !(field.getChildren().get(0).getType() instanceof ArrowType.FloatingPoint) - || ((ArrowType.FloatingPoint) field.getChildren().get(0).getType()).getPrecision() - != FloatingPointPrecision.SINGLE) { + || field.getChildren().size() != 1) { + return null; + } + ArrowType element = field.getChildren().get(0).getType(); + if (hamming ? !(element instanceof ArrowType.Int) + || ((ArrowType.Int) element).getBitWidth() != 8 || ((ArrowType.Int) element).getIsSigned() + : !(element instanceof ArrowType.FloatingPoint) + || ((ArrowType.FloatingPoint) element).getPrecision() != FloatingPointPrecision.SINGLE) { return null; } properties.put("column", field.getName()); properties.put("query_vector", query.toSql()); - properties.put("metric", l2 ? "l2" : "dot"); + properties.put("metric", hamming ? "hamming" : cosineDistance || cosineSimilarity ? "cosine" + : l2 ? "l2" : "dot"); } Optional filter = convertFilters(filters, scan, metadata, context); if (!filter.isPresent()) { @@ -267,7 +288,7 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con } Expression value = score ? result : l2 ? new Cast(new Sqrt(new Cast(result, DoubleType.INSTANCE)), FloatType.INSTANCE) - : new Subtract(new FloatLiteral(1.0f), result); + : dot || cosineSimilarity ? new Subtract(new FloatLiteral(1.0f), result) : result; if (!score) { // Preserve the approximate function's NOT NULL contract for existing Top-N and ancestor slots. // The generic TVF result and sqrt are nullable; NonNullable also rejects unexpected runtime NULLs. @@ -289,6 +310,20 @@ private Plan rewrite(LogicalTopN> topN, CascadesContext con return topN.withChildren(project.withProjectsAndChild(projects, child)); } + private boolean isUsableCosineQuery(ArrayLiteral query) { + float squaredNorm = 0; + for (Literal element : query.getValue()) { + if (!(element.getValue() instanceof Number)) { + return false; + } + float value = ((Number) element.getValue()).floatValue(); + squaredNorm += value * value; + } + // Native cosine normalization cannot represent zero/non-finite norms. Keep the existing + // scalar zero-vector results rather than silently turning such a query into no candidates. + return squaredNorm > 0 && Float.isFinite(squaredNorm); + } + private Optional convertFilters(List filters, LogicalFileScan scan, LanceTableMetadata metadata, CascadesContext context) { if (filters.isEmpty()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineDistanceApproximate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineDistanceApproximate.java new file mode 100644 index 00000000000000..17307616e4ec9d --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineDistanceApproximate.java @@ -0,0 +1,80 @@ +// 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.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.FloatType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * cosine_distance_approximate function + */ +public class CosineDistanceApproximate extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNotNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(FloatType.INSTANCE) + .args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE)) + ); + + /** + * Constructor with two vector arguments. + */ + public CosineDistanceApproximate(Expression arg0, Expression arg1) { + super("cosine_distance_approximate", arg0, arg1); + } + + /** constructor for withChildren and reuse signature */ + private CosineDistanceApproximate(ScalarFunctionParams functionParams) { + super(functionParams); + } + + /** + * withChildren. + */ + @Override + public CosineDistanceApproximate withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new CosineDistanceApproximate(getFunctionParams(children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitCosineDistanceApproximate(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public boolean isDeterministic() { + return false; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineSimilarityApproximate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineSimilarityApproximate.java new file mode 100644 index 00000000000000..7240013fd6f9b0 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineSimilarityApproximate.java @@ -0,0 +1,80 @@ +// 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.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.FloatType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * cosine_similarity_approximate function + */ +public class CosineSimilarityApproximate extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNotNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(FloatType.INSTANCE) + .args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE)) + ); + + /** + * Constructor with two vector arguments. + */ + public CosineSimilarityApproximate(Expression arg0, Expression arg1) { + super("cosine_similarity_approximate", arg0, arg1); + } + + /** constructor for withChildren and reuse signature */ + private CosineSimilarityApproximate(ScalarFunctionParams functionParams) { + super(functionParams); + } + + /** + * withChildren. + */ + @Override + public CosineSimilarityApproximate withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new CosineSimilarityApproximate(getFunctionParams(children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitCosineSimilarityApproximate(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public boolean isDeterministic() { + return false; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistance.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistance.java new file mode 100644 index 00000000000000..099827d29a4d97 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistance.java @@ -0,0 +1,100 @@ +// 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.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.types.FloatType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * hamming_distance function + */ +public class HammingDistance extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNotNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(FloatType.INSTANCE) + .args(ArrayType.of(BigIntType.INSTANCE), ArrayType.of(BigIntType.INSTANCE)) + ); + + /** + * Constructor with two vector arguments. + */ + public HammingDistance(Expression arg0, Expression arg1) { + super("hamming_distance", arg0, arg1); + } + + /** constructor for withChildren and reuse signature */ + private HammingDistance(ScalarFunctionParams functionParams) { + super(functionParams); + } + + /** + * withChildren. + */ + @Override + public HammingDistance withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new HammingDistance(getFunctionParams(children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitHammingDistance(this, context); + } + + @Override + public void checkLegalityBeforeTypeCoercion() { + validateByteArguments(getArguments()); + } + + static void validateByteArguments(List arguments) { + // Validate before implicit casts can truncate fractions or wrap integers wider than BIGINT. + for (Expression argument : arguments) { + DataType type = argument.getDataType(); + if (type.isNullType()) { + continue; + } + if (type instanceof ArrayType) { + DataType itemType = ((ArrayType) type).getItemType(); + if (itemType.isNullType() || (itemType.isIntegralType() && !itemType.isLargeIntType())) { + continue; + } + } + throw new AnalysisException("Hamming distance requires integer arrays with elements in [0, 255]"); + } + } + + @Override + public List getSignatures() { + return SIGNATURES; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistanceApproximate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistanceApproximate.java new file mode 100644 index 00000000000000..fdc07feb5bf376 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/HammingDistanceApproximate.java @@ -0,0 +1,86 @@ +// 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.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.FloatType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * hamming_distance_approximate function + */ +public class HammingDistanceApproximate extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNotNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(FloatType.INSTANCE) + .args(ArrayType.of(BigIntType.INSTANCE), ArrayType.of(BigIntType.INSTANCE)) + ); + + /** + * Constructor with two vector arguments. + */ + public HammingDistanceApproximate(Expression arg0, Expression arg1) { + super("hamming_distance_approximate", arg0, arg1); + } + + /** constructor for withChildren and reuse signature */ + private HammingDistanceApproximate(ScalarFunctionParams functionParams) { + super(functionParams); + } + + /** + * withChildren. + */ + @Override + public HammingDistanceApproximate withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new HammingDistanceApproximate(getFunctionParams(children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitHammingDistanceApproximate(this, context); + } + + @Override + public void checkLegalityBeforeTypeCoercion() { + HammingDistance.validateByteArguments(getArguments()); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public boolean isDeterministic() { + return false; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index 77db4cf175d1ed..b282d9bcdc2461 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -157,7 +157,9 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Cos; import org.apache.doris.nereids.trees.expressions.functions.scalar.Cosh; import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineDistance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineDistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineSimilarity; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CosineSimilarityApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Cot; import org.apache.doris.nereids.trees.expressions.functions.scalar.CountEqual; import org.apache.doris.nereids.trees.expressions.functions.scalar.CountSubstring; @@ -250,6 +252,8 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.GetFormat; import org.apache.doris.nereids.trees.expressions.functions.scalar.GetVariantType; import org.apache.doris.nereids.trees.expressions.functions.scalar.Greatest; +import org.apache.doris.nereids.trees.expressions.functions.scalar.HammingDistance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.HammingDistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Hex; import org.apache.doris.nereids.trees.expressions.functions.scalar.HllCardinality; import org.apache.doris.nereids.trees.expressions.functions.scalar.HllEmpty; @@ -1135,6 +1139,14 @@ default R visitCosineDistance(CosineDistance cosineDistance, C context) { return visitScalarFunction(cosineDistance, context); } + default R visitCosineDistanceApproximate(CosineDistanceApproximate function, C context) { + return visitScalarFunction(function, context); + } + + default R visitCosineSimilarityApproximate(CosineSimilarityApproximate function, C context) { + return visitScalarFunction(function, context); + } + default R visitCosineSimilarity(CosineSimilarity cosineSimilarity, C context) { return visitScalarFunction(cosineSimilarity, context); } @@ -1299,6 +1311,14 @@ default R visitMonthsSub(MonthsSub monthsSub, C context) { return visitScalarFunction(monthsSub, context); } + default R visitHammingDistance(HammingDistance function, C context) { + return visitScalarFunction(function, context); + } + + default R visitHammingDistanceApproximate(HammingDistanceApproximate function, C context) { + return visitScalarFunction(function, context); + } + default R visitHoursSub(HoursSub hoursSub, C context) { return visitScalarFunction(hoursSub, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java index 1afc8d2fcfb817..5ab009b104b000 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -18,6 +18,7 @@ package org.apache.doris.nereids.rules.rewrite; import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.FunctionRegistry; import org.apache.doris.catalog.Type; import org.apache.doris.datasource.lance.LanceExternalTable; import org.apache.doris.datasource.lance.metadata.LanceMvccSnapshot; @@ -58,6 +59,8 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BigIntType; import org.apache.doris.nereids.types.DoubleType; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.nereids.util.PlanChecker; @@ -166,6 +169,153 @@ private LogicalTVFRelation search(Plan plan) { return searches.get(0); } + private Expression distanceFunction(String name, Expression column, ArrayLiteral query) { + List args = ImmutableList.of(column, query); + return new FunctionRegistry().findFunctionBuilder(name, args).build(name, args).first; + } + + @Test + void cosineApproximateFunctionsPushBothOrderingAndValueConversions() { + for (String name : ImmutableList.of("cosine_distance_approximate", "cosine_similarity_approximate")) { + boolean ascending = name.equals("cosine_distance_approximate"); + Expression expression = distanceFunction(name, embedding, vector()); + assertNonNullVectorFullRewrite(expression, ascending); + Plan plan = rewrite(expression, ascending, scan); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(plan).getFunction().getCatalogFunction(); + Assertions.assertEquals(TVectorMetric.COSINE, + function.getSearchRequest().getSearchQuery().getVectorSearch().getMetric()); + Assertions.assertEquals(!ascending, plan.anyMatch(p -> ((Plan) p).getExpressions().stream() + .anyMatch(e -> e.anyMatch(Subtract.class::isInstance))), plan.treeString()); + Assertions.assertTrue(rewrite(expression, !ascending, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + + @Test + void zeroCosineQueryRetainsScalarSemantics() { + ArrayLiteral zero = new ArrayLiteral(ImmutableList.of(new FloatLiteral(0), new FloatLiteral(0))); + for (String name : ImmutableList.of("cosine_distance_approximate", "cosine_similarity_approximate")) { + Assertions.assertTrue(rewrite(distanceFunction(name, embedding, zero), + name.equals("cosine_distance_approximate"), scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + + private void useBinaryEmbedding() { + List fields = new ArrayList<>(metadata.getSchema().getFields()); + fields.set(2, new Field("embedding", FieldType.nullable(new ArrowType.FixedSizeList(2)), + ImmutableList.of(Field.nullable("item", new ArrowType.Int(8, false))))); + metadata = LanceTableMetadata.createSnapshotWithIndexes( + new LanceTableAccess("s3://bucket/items.lance", Collections.emptyMap()), 42, + new Schema(fields), Collections.emptyList(), ImmutableMap.of("id", 0, "body", 1, "embedding", 2), + Collections.emptyList()); + Column column = new Column("embedding", org.apache.doris.catalog.ArrayType.create(Type.SMALLINT, true)); + column.setIsAllowNull(true); + Mockito.when(table.getFullSchema(Mockito.any())).thenReturn(ImmutableList.of( + id.getOriginalColumn().get(), body.getOriginalColumn().get(), column)); + Mockito.when(table.getMetadata(Mockito.any())).thenReturn(metadata); + embedding = SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, column, + ImmutableList.of("catalog", "db", "items")); + scan = new LogicalFileScan(new RelationId(1), table, ImmutableList.of("catalog", "db"), + Collections.emptyList(), Optional.empty(), Optional.empty(), Optional.empty(), + Optional.of(ImmutableList.of(id, body, embedding)), Optional.of(new LanceMvccSnapshot(metadata))); + } + + @Test + void hammingApproximateSupportsUnsignedBytesAndSafeWideningCast() { + useBinaryEmbedding(); + ArrayLiteral query = new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(255))); + Expression expression = distanceFunction("hamming_distance_approximate", + new Cast(embedding, ArrayType.of(BigIntType.INSTANCE)), query); + assertNonNullVectorFullRewrite(expression, true); + VectorSearchTableValuedFunction function = (VectorSearchTableValuedFunction) + search(rewrite(expression, true, scan)).getFunction().getCatalogFunction(); + Assertions.assertEquals(TVectorMetric.HAMMING, + function.getSearchRequest().getSearchQuery().getVectorSearch().getMetric()); + Assertions.assertTrue(rewrite(expression, false, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void hammingApproximateCannotReinterpretFloatVectors() { + Expression expression = distanceFunction("hamming_distance_approximate", embedding, + new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(255)))); + Assertions.assertTrue(rewrite(expression, true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void hammingQueryByteValidationRetainsScalarFallback() { + useBinaryEmbedding(); + for (long value : new long[] {-1, 256, 65536}) { + ArrayLiteral query = new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(value))); + Assertions.assertTrue(rewrite(distanceFunction("hamming_distance_approximate", embedding, query), true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + + @Test + void exactHammingFunctionRemainsScalar() { + useBinaryEmbedding(); + Expression expression = distanceFunction("hamming_distance", embedding, + new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(255)))); + Assertions.assertTrue(rewrite(expression, true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void hammingRejectsLossyImplicitArrayConversions() { + for (String name : ImmutableList.of("hamming_distance", "hamming_distance_approximate")) { + Expression function = distanceFunction(name, vector(), vector()); + Assertions.assertThrows(org.apache.doris.nereids.exceptions.AnalysisException.class, + () -> function.checkLegalityBeforeTypeCoercion()); + } + } + + @Test + void hammingAnalyzerWidensUnsignedByteArraysBeforeNativeRewrite() { + useBinaryEmbedding(); + Expression expression = distanceFunction("hamming_distance_approximate", embedding, + new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(255)))); + Alias distance = new Alias(expression, "distance"); + LogicalTopN topN = new LogicalTopN<>(ImmutableList.of(new OrderKey(distance.toSlot(), true, false)), + 5, 0, new LogicalProject<>(ImmutableList.of(id, distance), scan)); + Plan plan = PlanChecker.from(context).analyze(topN).rewrite().getPlan(); + search(plan); + } + + @Test + void hammingCannotBypassNarrowingElementCasts() { + useBinaryEmbedding(); + Expression column = new Cast(embedding, ArrayType.of(org.apache.doris.nereids.types.TinyIntType.INSTANCE)); + Expression expression = distanceFunction("hamming_distance_approximate", column, + new ArrayLiteral(ImmutableList.of(new BigIntLiteral(0), new BigIntLiteral(255)))); + Assertions.assertTrue(rewrite(expression, true, scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + + @Test + void invalidCosineQueryNormKeepsTheScalarPlan() { + for (float value : new float[] {Float.NaN, Float.POSITIVE_INFINITY, Float.MAX_VALUE, Float.MIN_VALUE}) { + ArrayLiteral query = new ArrayLiteral(ImmutableList.of(new FloatLiteral(value), new FloatLiteral(0))); + for (String name : ImmutableList.of("cosine_distance_approximate", "cosine_similarity_approximate")) { + Assertions.assertTrue(rewrite(distanceFunction(name, embedding, query), + name.equals("cosine_distance_approximate"), scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + } + + @Test + void exactCosineFunctionsRemainScalar() { + for (String name : ImmutableList.of("cosine_distance", "cosine_similarity")) { + Assertions.assertTrue(rewrite(distanceFunction(name, embedding, vector()), + name.equals("cosine_distance"), scan) + .collectToList(LogicalTVFRelation.class::isInstance).isEmpty()); + } + } + @Test void l2TopNUsesLanceAndRestoresEuclideanDistance() { Plan plan = rewrite(new L2DistanceApproximate(embedding, vector()), true, scan); diff --git a/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy index d19881e932e150..3e05b376b950a5 100644 --- a/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy +++ b/regression-test/suites/external_table_p0/lance/test_lance_search_expression_pushdown.groovy @@ -70,6 +70,67 @@ suite("test_lance_search_expression_pushdown", "p0,external") { for (int i = 0; i < dotActual.size(); ++i) { assertTrue(Math.abs((dotExpected[i][1] as double) - (dotActual[i][1] as double)) < 0.001) } + def cosineTable = "${catalog}.doris.vs_ivf_flat_f32_cosine" + for (entry in [["cosine_distance_approximate", "ASC", "_distance"], + ["cosine_similarity_approximate", "DESC", "CAST(1 AS FLOAT) - _distance"]]) { + def automatic = """SELECT row_id, ${entry[0]}(embedding, ${vector}) AS value + FROM ${cosineTable} ORDER BY value ${entry[1]} LIMIT 5""" + explain { + sql automatic + contains "externalSearchType=VECTOR" + contains "lanceMetric=cosine" + } + def explicit = """SELECT row_id, ${entry[2]} AS value FROM vector_search( + "table"="${cosineTable}", "column"="embedding", "query_vector"="${vector}", + "top_k"="5", "metric"="cosine") ORDER BY value ${entry[1]}""" + def actual = sql(automatic) + def expected = sql(explicit) + assertEquals(expected.collect { it[0] }, actual.collect { it[0] }) + for (int i = 0; i < actual.size(); ++i) { + assertTrue(Math.abs((expected[i][1] as double) - (actual[i][1] as double)) < 0.0001) + def exact = entry[0].replace("_approximate", "") + def scalar = sql """SELECT ${exact}(embedding, ${vector}) FROM ${cosineTable} + WHERE row_id = ${actual[i][0]}""" + assertTrue(Math.abs((scalar[0][0] as double) - (actual[i][1] as double)) < 0.0001) + } + } + def binaryTable = "${catalog}.doris.vs_ivf_flat_u8" + def binaryVector = "[" + (0..<128).collect { "0" }.join(",") + "]" + def automaticHamming = """SELECT row_id, + hamming_distance_approximate(embedding, ${binaryVector}) AS distance + FROM ${binaryTable} ORDER BY distance LIMIT 5""" + explain { + sql automaticHamming + contains "externalSearchType=VECTOR" + contains "lanceMetric=hamming" + } + def explicitHamming = """SELECT row_id, _distance FROM vector_search( + "table"="${binaryTable}", "column"="embedding", "query_vector"="${binaryVector}", + "top_k"="5", "metric"="hamming") ORDER BY _distance""" + def hammingRows = sql(automaticHamming) + assertEquals(sql(explicitHamming), hammingRows) + for (row in hammingRows) { + def scalar = sql """SELECT hamming_distance(embedding, ${binaryVector}) + FROM ${binaryTable} WHERE row_id = ${row[0]}""" + assertEquals(scalar[0][0], row[1]) + } + for (name in ["cosine_distance", "cosine_similarity"]) { + explain { + sql "SELECT row_id FROM ${cosineTable} ORDER BY ${name}(embedding, ${vector}) LIMIT 5" + notContains "externalSearchType=VECTOR" + } + } + def zeroVector = "[" + (0..<16).collect { "0" }.join(",") + "]" + explain { + sql """SELECT row_id FROM ${cosineTable} + ORDER BY cosine_distance_approximate(embedding, ${zeroVector}) LIMIT 5""" + notContains "externalSearchType=VECTOR" + } + explain { + sql """SELECT row_id FROM ${binaryTable} + ORDER BY hamming_distance(embedding, ${binaryVector}) LIMIT 5""" + notContains "externalSearchType=VECTOR" + } sql "SET fe_debug = ${previousFeDebug}" def textTable = "${catalog}.doris.full_text_search" for (entry in [["MATCH_ANY", "match", "or"], ["MATCH_ALL", "match", "and"], diff --git a/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy new file mode 100644 index 00000000000000..6f8013a29167d1 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy @@ -0,0 +1,52 @@ +// 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_array_distance_approximate") { + for (entry in [["cosine_distance_approximate", "cosine_distance"], + ["cosine_similarity_approximate", "cosine_similarity"]]) { + for (vector in ["[1,0]", "[0,1]", "[0,0]", "[]"]) { + def other = vector == "[]" ? "[]" : "[1,0]" + assertEquals(sql("SELECT ${entry[1]}(${vector}, ${other})"), + sql("SELECT ${entry[0]}(${vector}, ${other})")) + } + } + for (name in ["hamming_distance", "hamming_distance_approximate"]) { + assertEquals(24.0f, sql("SELECT ${name}([0,255,170], [255,0,85])")[0][0] as float) + assertEquals(0.0f, sql("SELECT ${name}([], [])")[0][0] as float) + assertEquals(0.0f, sql("SELECT ${name}([255], [255])")[0][0] as float) + def rows = sql """SELECT ${name}([number], [0]) FROM numbers("number"="4") ORDER BY number""" + assertEquals([0.0f, 1.0f, 1.0f, 2.0f], rows.collect { it[0] as float }) + for (value in [-1, 256, 65536]) { + test { + sql "SELECT ${name}([${value}], [0])" + exception "Hamming vector elements must be integers in [0, 255]" + } + } + test { + sql "SELECT ${name}([1.5], [0])" + exception "Hamming distance requires integer arrays" + } + test { + sql "SELECT ${name}([0], [0,1])" + exception "different input element sizes" + } + test { + sql "SELECT ${name}([NULL], [0])" + exception "null" + } + } +} From d3e3f1b0015a8bfd7c99f88c47d287dde7096101 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 13:34:39 +0800 Subject: [PATCH 7/8] fix(lance): account row-id prefilter buffers to query memory --- be/src/format_v2/table/lance_reader.cpp | 2 +- be/src/format_v2/table/lance_reader.h | 4 +- be/test/format_v2/table/lance_reader_test.cpp | 59 ++++++++++++------- 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index 528acdab3dc940..30ff02df89c0aa 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -968,7 +968,7 @@ Status LanceTableReader::_read_row_id_prefilter(const std::string& filter, RowId imported.status().message()); } auto reader = std::move(imported).ValueUnsafe(); - std::vector row_ids; + DorisVector row_ids; while (true) { RETURN_IF_ERROR(_check_row_id_prefilter_stop()); std::shared_ptr batch; diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 337c768a4ca293..f5383f774462db 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -25,6 +25,7 @@ #include #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" @@ -85,7 +86,8 @@ class LanceTableReader final : public TableReader { struct RowIdPrefilter { Status status; // Empty selects no rows; nullopt keeps the provider's native filter for a broad result. - std::optional> row_ids; + // Charge both collection and cache retention to the query's memory limit. + std::optional> row_ids; }; Status _check_row_id_prefilter_stop() const; Status _read_row_id_prefilter(const std::string& filter, RowIdPrefilter* result); diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 0f103217fbd438..da39eaf9b0a12e 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -80,8 +80,10 @@ #include "format/format_common.h" #include "format_v2/lance/lance_reader_helper.h" #include "io/io_common.h" +#include "runtime/memory/mem_tracker_limiter.h" #include "runtime/runtime_profile.h" #include "runtime/runtime_state.h" +#include "runtime/thread_context.h" #include "storage/utils.h" #include "util/debug_points.h" #include "util/defer_op.h" @@ -2074,27 +2076,44 @@ TEST(LanceTableReaderVectorSearchTest, AppliesSubstraitSearchFilterBeforeTopK) { scan_params.lance_scan_params.external_search_request.__set_search_filter(search_filter); scan_params.lance_scan_params.external_search_request.__set_prefilter(true); - ShardedKVCache cache(1); - for (int iteration = 0; iteration < 2; ++iteration) { - RuntimeProfile profile("lance_row_id_prefilter_cache"); - LanceTableReader reader; - ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); - ASSERT_TRUE( - prepare_range(&reader, - make_lance_range(dataset_uri, fixture.version, fixture.fragment_ids), - std::nullopt, &cache) - .ok()); - Block block; - add_output_columns(&block, columns); - const auto rows = read_vector_search_rows(&reader, &block); - ASSERT_EQ(1U, rows.size()); - EXPECT_EQ(4, rows[0].first); - EXPECT_FLOAT_EQ(8.25F, rows[0].second); - ASSERT_NE(nullptr, profile.get_counter("LanceRowIdPrefilterRows")); - EXPECT_EQ(2, profile.get_counter("LanceRowIdPrefilterRows")->value()); - EXPECT_EQ(iteration, profile.get_counter("LanceRowIdPrefilterCacheHits")->value()); - EXPECT_TRUE(reader.close().ok()); + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "UT-LanceRowIdPrefilter"); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + { + ShardedKVCache cache(1); + int64_t cached_bytes = 0; + for (int iteration = 0; iteration < 2; ++iteration) { + { + RuntimeProfile profile("lance_row_id_prefilter_cache"); + LanceTableReader reader; + ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok()); + ASSERT_TRUE(prepare_range(&reader, + make_lance_range(dataset_uri, fixture.version, + fixture.fragment_ids), + std::nullopt, &cache) + .ok()); + Block block; + add_output_columns(&block, columns); + const auto rows = read_vector_search_rows(&reader, &block); + ASSERT_EQ(1U, rows.size()); + EXPECT_EQ(4, rows[0].first); + EXPECT_FLOAT_EQ(8.25F, rows[0].second); + ASSERT_NE(nullptr, profile.get_counter("LanceRowIdPrefilterRows")); + EXPECT_EQ(2, profile.get_counter("LanceRowIdPrefilterRows")->value()); + EXPECT_EQ(iteration, profile.get_counter("LanceRowIdPrefilterCacheHits")->value()); + EXPECT_TRUE(reader.close().ok()); + } + doris::thread_context()->thread_mem_tracker_mgr->flush_untracked_mem(); + if (iteration == 0) { + cached_bytes = tracker->consumption(); + EXPECT_GE(cached_bytes, 2 * sizeof(uint64_t)); + } else { + EXPECT_EQ(cached_bytes, tracker->consumption()); + } + } } + doris::thread_context()->thread_mem_tracker_mgr->flush_untracked_mem(); + EXPECT_EQ(0, tracker->consumption()); } TEST(LanceTableReaderVectorSearchTest, BroadRowIdPrefilterFallsBackWithoutLosingNearestRow) { From 4ef2fbb7a434c6ee67fa6cf550c3ffc0116f1f45 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 22 Oct 2026 19:11:57 +0800 Subject: [PATCH 8/8] fix(regression): construct Hamming arrays from column expressions --- .../array_functions/test_array_distance_approximate.groovy | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy index 6f8013a29167d1..99f08bd7882976 100644 --- a/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy +++ b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_distance_approximate.groovy @@ -28,7 +28,8 @@ suite("test_array_distance_approximate") { assertEquals(24.0f, sql("SELECT ${name}([0,255,170], [255,0,85])")[0][0] as float) assertEquals(0.0f, sql("SELECT ${name}([], [])")[0][0] as float) assertEquals(0.0f, sql("SELECT ${name}([255], [255])")[0][0] as float) - def rows = sql """SELECT ${name}([number], [0]) FROM numbers("number"="4") ORDER BY number""" + // Bracket array literals require constants; ARRAY() keeps this path backed by a column. + def rows = sql """SELECT ${name}(array(number), [0]) FROM numbers("number"="4") ORDER BY number""" assertEquals([0.0f, 1.0f, 1.0f, 2.0f], rows.collect { it[0] as float }) for (value in [-1, 256, 65536]) { test {