From e88159ce1b824e42e20b1f38da61f1e2b26d2971 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 20 Oct 2026 17:25:22 +0800 Subject: [PATCH] [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 | 202 ++++++++++- .../doris/nereids/jobs/executor/Rewriter.java | 2 + .../apache/doris/nereids/rules/RuleType.java | 1 + .../PushDownSearchTopNIntoLanceScan.java | 318 ++++++++++++++++++ .../functions/table/LanceSearch.java | 43 +++ .../FullTextSearchTableValuedFunction.java | 24 +- ...anceExternalSearchTableValuedFunction.java | 19 ++ .../VectorSearchTableValuedFunction.java | 23 +- .../LanceSearchExpressionPushDownTest.java | 310 +++++++++++++++++ ...st_lance_search_expression_pushdown.groovy | 113 +++++++ thirdparty/patches/lance-c-foyer.patch | 6 +- thirdparty/vars.sh | 10 +- 14 files changed, 1213 insertions(+), 16 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..232c1a712541a5 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,146 @@ 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, 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 9167a0aed681dd..e638216e2421d8 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 @@ -141,6 +141,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 25f46eba8e563c..13c3c095d0accb 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 @@ -346,6 +346,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..47e957b769e521 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownSearchTopNIntoLanceScan.java @@ -0,0 +1,318 @@ +// 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); + } + if (!(input instanceof LogicalFileScan) || !(((LogicalFileScan) input).getTable() + instanceof LanceExternalTable) || project.isDistinct()) { + 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; + } + // 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..b2cbc5a6d89022 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/LanceSearchExpressionPushDownTest.java @@ -0,0 +1,310 @@ +// 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.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 exactDistanceAndReverseOrderRemainScalar() { + 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}" + } +} diff --git a/thirdparty/patches/lance-c-foyer.patch b/thirdparty/patches/lance-c-foyer.patch index 7c7fae2bffc01e..f44a2f8dd0c9eb 100644 --- a/thirdparty/patches/lance-c-foyer.patch +++ b/thirdparty/patches/lance-c-foyer.patch @@ -1,5 +1,5 @@ # Foyer data-cache integration from lance-format/lance-c#73. -# Base: 8e9e09ecc795c91ce7f97513db94c2d91eec799d (lance-c main, including #91 and #94). +# Base: 4fafcd64525b274b113208390bbcca86bbbc093b (lance-c main, including #96 and #97). # PR base: d214faaca89ad7654fb1d85d46cdbce6f3770c95; current PR head: 8a08f171b66a3e01e7ba634bcfc6168a955c7289. # Current PR head includes revert c3857f085d3aa522c79ebc92b28b302e2eb8e865 of cache-origin isolation. # Rebased from https://github.com/lance-format/lance-c/pull/73. @@ -8666,11 +8666,11 @@ index 3bc313ebecbb0de06badd9f8d75f531fc95809c9..298228f7dd68d28452e6436db5c58918 static void test_dataset_schema(const std::string& uri) { TEST(test_dataset_schema); -@@ -1365,6 +1387,7 @@ int main(int argc, char** argv) { +@@ -1402,6 +1424,7 @@ int main(int argc, char** argv) { test_dataset_open(uri); test_shared_session(uri); + test_data_cache_session(uri, write_uri); test_dataset_schema(uri); test_scanner_fluent(uri); - test_distance_range(uri); + test_scanner_row_addr_lists(uri); diff --git a/thirdparty/vars.sh b/thirdparty/vars.sh index f0078fa2a89ecb..55e75715dfae9d 100644 --- a/thirdparty/vars.sh +++ b/thirdparty/vars.sh @@ -580,11 +580,11 @@ PUGIXML_SOURCE=pugixml-1.15 PUGIXML_MD5SUM="3b894c29455eb33a40b165c6e2de5895" # lance-c -# Includes lance-c #91/#93/#94: search filters and synchronous index prewarm. -LANCE_C_DOWNLOAD="https://codeload.github.com/lance-format/lance-c/tar.gz/8e9e09ecc795c91ce7f97513db94c2d91eec799d" -LANCE_C_NAME="lance-c-8e9e09ecc795c91ce7f97513db94c2d91eec799d.tar.gz" -LANCE_C_SOURCE="lance-c-8e9e09ecc795c91ce7f97513db94c2d91eec799d" -LANCE_C_MD5SUM="9e32aa626caf1cea3bfe62d12d60a2c8" +# Includes lance-c #91/#93/#94/#96/#97: search filters, row-ID masks, and index prewarm fixes. +LANCE_C_DOWNLOAD="https://codeload.github.com/lance-format/lance-c/tar.gz/4fafcd64525b274b113208390bbcca86bbbc093b" +LANCE_C_NAME="lance-c-4fafcd64525b274b113208390bbcca86bbbc093b.tar.gz" +LANCE_C_SOURCE="lance-c-4fafcd64525b274b113208390bbcca86bbbc093b" +LANCE_C_MD5SUM="65e924fb62137c5dd4d58af3fb33923a" # paimon-rust PAIMON_RUST_DOWNLOAD="https://github.com/apache/paimon-rust/archive/refs/tags/v0.4.0-rc1.tar.gz"