diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index 3955878e2761f0..f25fbb66bf9983 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -20,6 +20,7 @@ #include #include +#include #include // IWYU pragma: no_include #include @@ -30,6 +31,7 @@ #include // IWYU pragma: keep #include #include +#include #include #include #include @@ -1296,6 +1298,103 @@ DEFINE_Int32(inverted_index_query_cache_shards, "256"); // inverted index match bitmap cache size DEFINE_String(inverted_index_query_cache_limit, "10%"); +namespace { + +constexpr uint64_t COMMON_GRAMS_QUERY_PLAN_ENABLED_MASK = 1; +constexpr uint64_t COMMON_GRAMS_QUERY_PLAN_GENERATION_SHIFT = 1; +std::atomic common_grams_query_plan_state {COMMON_GRAMS_QUERY_PLAN_ENABLED_MASK}; +constexpr uint64_t COMMON_GRAMS_COST_RATIO_MASK = std::numeric_limits::max(); +constexpr uint64_t COMMON_GRAMS_VERIFY_FACTOR_SHIFT = 32; +std::atomic common_grams_cost_model_state {85}; + +bool valid_common_grams_cost_ratio(int32_t value) { + return value >= 0 && value <= 100; +} + +bool valid_common_grams_verify_factor(int32_t value) { + return value >= 0; +} + +void publish_common_grams_query_plan_state(bool enabled) { + uint64_t current = common_grams_query_plan_state.load(std::memory_order_acquire); + while ((current & COMMON_GRAMS_QUERY_PLAN_ENABLED_MASK) != static_cast(enabled)) { + const uint64_t generation = (current >> COMMON_GRAMS_QUERY_PLAN_GENERATION_SHIFT) + 1; + const uint64_t desired = (generation << COMMON_GRAMS_QUERY_PLAN_GENERATION_SHIFT) | + static_cast(enabled); + if (common_grams_query_plan_state.compare_exchange_weak( + current, desired, std::memory_order_acq_rel, std::memory_order_acquire)) { + return; + } + } +} + +void publish_common_grams_cost_ratio(int32_t value) { + uint64_t current = common_grams_cost_model_state.load(std::memory_order_acquire); + if ((current & COMMON_GRAMS_COST_RATIO_MASK) == static_cast(value)) { + return; + } + uint64_t desired = 0; + do { + desired = (current & ~COMMON_GRAMS_COST_RATIO_MASK) | static_cast(value); + } while (!common_grams_cost_model_state.compare_exchange_weak( + current, desired, std::memory_order_acq_rel, std::memory_order_acquire)); +} + +void publish_common_grams_verify_factor(int32_t value) { + uint64_t current = common_grams_cost_model_state.load(std::memory_order_acquire); + if (current >> COMMON_GRAMS_VERIFY_FACTOR_SHIFT == static_cast(value)) { + return; + } + uint64_t desired = 0; + do { + desired = (current & COMMON_GRAMS_COST_RATIO_MASK) | + (static_cast(static_cast(value)) + << COMMON_GRAMS_VERIFY_FACTOR_SHIFT); + } while (!common_grams_cost_model_state.compare_exchange_weak( + current, desired, std::memory_order_acq_rel, std::memory_order_acquire)); +} + +} // namespace + +DEFINE_mBool(enable_common_grams_query_plan, "false"); +DEFINE_ON_UPDATE(enable_common_grams_query_plan, [](bool /*old_value*/, bool new_value) { + publish_common_grams_query_plan_state(new_value); +}); +DEFINE_mBool(enable_common_grams_index_build, "true"); +DEFINE_mInt32(common_grams_plan_cost_ratio_percent, "85"); +DEFINE_Validator(common_grams_plan_cost_ratio_percent, valid_common_grams_cost_ratio); +DEFINE_ON_UPDATE(common_grams_plan_cost_ratio_percent, + [](int32_t /*old_value*/, int32_t new_value) { + publish_common_grams_cost_ratio(new_value); + }); +DEFINE_mInt32(common_grams_position_verify_factor, "0"); +DEFINE_Validator(common_grams_position_verify_factor, valid_common_grams_verify_factor); +DEFINE_ON_UPDATE(common_grams_position_verify_factor, [](int32_t /*old_value*/, int32_t new_value) { + publish_common_grams_verify_factor(new_value); +}); + +namespace { + +void publish_common_grams_query_plan_config() { + publish_common_grams_query_plan_state(enable_common_grams_query_plan); + publish_common_grams_cost_ratio(common_grams_plan_cost_ratio_percent); + publish_common_grams_verify_factor(common_grams_position_verify_factor); +} + +} // namespace + +CommonGramsQueryPlanConfigSnapshot common_grams_query_plan_config_snapshot() { + const uint64_t state = common_grams_query_plan_state.load(std::memory_order_acquire); + const uint64_t cost_model_state = common_grams_cost_model_state.load(std::memory_order_acquire); + return {.enabled = (state & COMMON_GRAMS_QUERY_PLAN_ENABLED_MASK) != 0, + .cache_generation = state >> COMMON_GRAMS_QUERY_PLAN_GENERATION_SHIFT, + .plan_cost_ratio_percent = + static_cast(cost_model_state & COMMON_GRAMS_COST_RATIO_MASK), + .position_verify_factor = + static_cast(cost_model_state >> COMMON_GRAMS_VERIFY_FACTOR_SHIFT), + .cost_model_generation = cost_model_state}; +} + // condition cache limit DEFINE_Int16(condition_cache_limit, "512"); @@ -1309,6 +1408,63 @@ DEFINE_mDouble(inverted_index_ram_buffer_size, "512"); // -1 indicates not working. // Normally we should not change this, it's useful for testing. DEFINE_mInt32(inverted_index_max_buffered_docs, "-1"); +// DIAGNOSTIC (default off): force SNII inverted-index reads onto NO_CACHE +// (precise S3 range GETs) instead of the 1MiB FILE_BLOCK_CACHE. Per-open on the +// SNII reader only -- does NOT touch global enable_file_cache, so cloud mode +// still boots. Measures the block-cache read amplification's true cost. Warm +// queries lose the local cache under this flag, so it is a measurement knob, not +// a production setting. +DEFINE_mBool(inverted_index_read_bypass_file_cache, "false"); +// G16-c: whether plain positions-tier (non-scoring) SNII indexes lay out freq +// regions. Freq bytes serve ONLY BM25 scoring, which the Doris integration +// does not reach yet (scoring_query has no production caller), so the default +// drops them (textbench: -2.2 GB index). Scoring-config indexes always write +// freq regardless. Applies at segment build (write side only); existing +// segments keep whatever layout they were written with (self-describing). +DEFINE_mBool(snii_positions_index_write_freq, "false"); +// G16-h: zstd levels for the SNII dict-block compression and the .prx window +// auto mode. Level 9 (vs the historical 3) shrinks the two largest compressed +// sections -- textbench: index -457 MB (0.918x -> 0.891x V3) -- for an import +// CPU cost inside the run-to-run variance band; zstd decode speed does not +// depend on the level, and warm/cold benches measured no query change. +// Write side only; segments self-describe their compression. +// Default 3 since the all-level-3 evaluation (2026-07-11, 4 corpora): vs +// level 9 the settled index grows only +0.6%..+6.3% (whole table +// +0.3%..+1.9%) while import index CPU drops 17-24% and full-compaction CPU +// 8-24%; settled cold-query latency is unchanged (interleaved A/B). The +// delta+varint-encoded payloads are high-entropy, so level 9's extra search +// buys almost no ratio. Raise only for size-critical deployments. +DEFINE_mInt32(snii_dict_block_zstd_level, "3"); +DEFINE_mInt32(snii_prx_zstd_level, "3"); +// Patch C prx tiering: zstd level for the prx region of DIRECT-LOAD segments +// only (stream/broker load, see IndexColumnWriter::set_direct_load). Inert at +// the defaults (both levels 3); it exists for size-critical deployments that +// RAISE snii_prx_zstd_level (e.g. 9) and still want cheap loads: compaction +// rewrites every segment at snii_prx_zstd_level, so SETTLED data (and the +// cold-query path over it) is unaffected by the load tier -- measured -290s +// (httplogs) / -204s (agentlogs) of import index CPU at 3 vs 9. Same clamp +// [3, 19]. Read at index flush like snii_prx_zstd_level (a mid-load change +// lands on in-flight segments); the direct-load BIT itself is captured once. +DEFINE_mInt32(snii_prx_zstd_level_direct_load, "3"); +// G16-d: target SNII dict block size in bytes; 0 uses the format default +// (64 KiB). Larger blocks compress better under the per-block zstd (the dict +// is the dominant physical section on high-cardinality corpora) at the cost +// of a larger fetch+decompress unit per cold dict-block miss. Write side +// only; the block size is self-described by the on-disk directory. +DEFINE_mInt32(snii_target_dict_block_bytes, "0"); +// Process-wide SNII index-build RAM budget across ALL live segment writers +// (G09); the largest writers are asked to spill early once the sum crosses it. +// 0 disables. Default 8 GiB. +DEFINE_mInt64(snii_index_writer_global_memory_bytes, "0"); +// Minimum reclaimable posting-arena bytes before a G09 forced spill is honored +// (and before a writer is eligible as a spill victim): forced spills reclaim +// ONLY the arena, so smaller triggers cut tiny runs for near-zero relief. +// Default 64 MiB. +DEFINE_mInt64(snii_forced_spill_min_arena_bytes, "67108864"); +// Max spill-run files one SNII writer accumulates before its runs are +// merge-compacted into one (bounds the k-way merge fan-in and its open fds; +// every run is held open for the whole merge). 0 = uncapped. Default 64. +DEFINE_mInt32(snii_spill_max_run_files_per_buffer, "64"); // dict path for chinese analyzer DEFINE_String(inverted_index_dict_path, "${DORIS_HOME}/dict"); DEFINE_Int32(inverted_index_read_buffer_size, "4096"); @@ -2173,6 +2329,10 @@ bool init(const char* conf_file, bool fill_conf_map, bool must_exist, bool set_t SET_FIELD(it.second, std::vector, fill_conf_map, set_to_default); } + // Startup loading does not invoke mutable-config callbacks. Publish after all three fields are + // validated so query threads never seed the atomics by reading mutable config globals. + publish_common_grams_query_plan_config(); + // Emit a warning for every key present in the conf file that does not correspond to a // registered BE config field. Such keys (typos or configs removed in a newer version) // are silently ignored above, so without this warning operators would have no feedback @@ -2239,6 +2399,33 @@ bool init(const char* conf_file, bool fill_conf_map, bool must_exist, bool set_t return Status::OK(); \ } +namespace { + +// UPDATE_FIELD invokes registered validators before assigning the candidate value. Validate the two +// mutable planner coefficients explicitly so their startup and runtime constraints stay identical. +Status validate_common_grams_runtime_config(const std::string& field, const std::string& value) { + bool (*validator)(int32_t) = nullptr; + if (field == "common_grams_plan_cost_ratio_percent") { + validator = valid_common_grams_cost_ratio; + } else if (field == "common_grams_position_verify_factor") { + validator = valid_common_grams_verify_factor; + } else { + return Status::OK(); + } + + int32_t candidate = 0; + if (!convert(value, candidate)) { + return Status::OK(); + } + if (!validator(candidate)) { + return Status::Error("validate {}={} failed", field, + candidate); + } + return Status::OK(); +} + +} // namespace + // write config to be_custom.conf // the caller need to make sure that the given config is valid Status persist_config(const std::string& field, const std::string& value) { @@ -2269,6 +2456,8 @@ Status set_config(const std::string& field, const std::string& value, bool need_ "'{}' is not support to modify", field); } + RETURN_IF_ERROR(validate_common_grams_runtime_config(field, value)); + UPDATE_FIELD(it->second, value, bool, need_persist); UPDATE_FIELD(it->second, value, int16_t, need_persist); UPDATE_FIELD(it->second, value, int32_t, need_persist); diff --git a/be/src/common/config.h b/be/src/common/config.h index cb09dfcb538cff..75ed247cacd20d 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -1336,6 +1336,24 @@ DECLARE_Int32(inverted_index_query_cache_shards); // inverted index match bitmap cache size DECLARE_String(inverted_index_query_cache_limit); +// Process-wide emergency switch for CommonGrams query plans. The generation changes on every +// effective transition so result-cache and single-flight entries from an older state are unreachable. +DECLARE_mBool(enable_common_grams_query_plan); +// Build-only CommonGrams kill switch. Logical index writers snapshot it at construction; changing +// it affects only writers created after the transition and never changes query/cache semantics. +DECLARE_mBool(enable_common_grams_index_build); +// Release-calibrated query-planner coefficients. Both remain mutable for controlled recalibration. +DECLARE_mInt32(common_grams_plan_cost_ratio_percent); +DECLARE_mInt32(common_grams_position_verify_factor); +struct CommonGramsQueryPlanConfigSnapshot { + bool enabled = true; + uint64_t cache_generation = 0; + uint32_t plan_cost_ratio_percent = 85; + uint32_t position_verify_factor = 0; + uint64_t cost_model_generation = 0; +}; +CommonGramsQueryPlanConfigSnapshot common_grams_query_plan_config_snapshot(); + // condition cache limit DECLARE_Int16(condition_cache_limit); @@ -1347,6 +1365,62 @@ DECLARE_Int32(ann_index_result_cache_stale_sweep_time_sec); // inverted index DECLARE_mDouble(inverted_index_ram_buffer_size); DECLARE_mInt32(inverted_index_max_buffered_docs); +// DIAGNOSTIC: force SNII inverted-index reads to bypass the 1MiB FILE_BLOCK_CACHE +// and issue precise S3 range GETs (NO_CACHE) instead. Applies ONLY to the SNII +// index file reader (per-open cache_type), NOT the global enable_file_cache, so +// cloud mode does not FATAL. Default false (keep block cache). Used to measure +// whether the block-cache read amplification hurts cold reads at the cost of +// re-reading S3 on every (warm) query. Not for production: warm loses the local +// cache. Read-side only. +DECLARE_mBool(inverted_index_read_bypass_file_cache); +// G16-c: whether plain positions-tier (non-scoring) SNII indexes lay out freq +// regions. Freq serves ONLY BM25 scoring (no production caller yet), so the +// default (false) drops the layout; scoring-config indexes always keep freq. +// Write-side only; segments are self-describing either way. +DECLARE_mBool(snii_positions_index_write_freq); +// G16-h: zstd levels for SNII dict blocks / prx windows. Default 3 (the +// all-level-3 evaluation showed level 9 buys <=6.3% index size for 17-24% +// import CPU; see the DEFINEs in config.cpp). +DECLARE_mInt32(snii_dict_block_zstd_level); +DECLARE_mInt32(snii_prx_zstd_level); +// Patch C: prx zstd level for DIRECT-LOAD segments only (default 3, cheaper +// import); compaction rewrites at snii_prx_zstd_level so settled segments are +// unaffected. Full contract at the DEFINE in config.cpp. +DECLARE_mInt32(snii_prx_zstd_level_direct_load); +// G16-d: target SNII dict block size in bytes; 0 = format default (64 KiB). +// Bigger blocks -> better per-block zstd on the dict region, larger cold +// fetch+decompress unit per dict-block miss. Write side only. +DECLARE_mInt32(snii_target_dict_block_bytes); +// PROCESS-WIDE budget in bytes for SNII index-build RAM, summed across every +// live SNII segment writer of the BE (all tablets x all concurrent loads; G09). +// The per-writer inverted_index_ram_buffer_size is a reclaimable-buffer spill +// threshold, not a hard cap on persistent vocabulary bytes. A concurrent load +// keeps (tablets x concurrency) writers alive at once, none of which may reach +// that threshold, while their SUM can still be large. When the registered total +// exceeds this process-wide budget, the LARGEST writers are asked to spill their +// posting buffers to disk early (async-safe advisory requests, honored on each +// writer's own thread; output stays byte-identical). 0 disables the global +// limiter (per-writer spilling only). Checked at index-writer creation: a change +// applies to writers created afterwards. +DECLARE_mInt64(snii_index_writer_global_memory_bytes); +// G09 forced-spill floor: minimum reclaimable posting-arena bytes a SNII +// writer must hold before a process-wide forced-spill request is honored, and +// before the global limiter selects it as a spill victim. A forced spill +// reclaims ONLY the posting arena -- the persistent vocab / pair-map +// structures survive it -- so honoring below a real floor degenerates into a +// storm of tiny runs whenever the budget is dominated by persistent bytes +// (each run then costs a file, a sort and a merge-fd for near-zero memory +// relief). The global budget therefore bounds SPILLABLE memory, not +// persistent memory. Default 64 MiB. +DECLARE_mInt64(snii_forced_spill_min_arena_bytes); +// G09 run-file cap: maximum spill-run files one SNII writer may accumulate; +// on the next spill past the cap, the existing runs are merge-compacted into +// a single run first (term stream unchanged). Bounds the final k-way merge's +// fan-in and, decisively, its simultaneously-open file descriptors -- every +// run of a buffer is reopened and held open for the whole merge, so unbounded +// run counts across ~100 concurrent writers can exhaust the BE nofile rlimit +// ("Too many open files" at run reopen). 0 disables the cap. Default 64. +DECLARE_mInt32(snii_spill_max_run_files_per_buffer); // dict path for chinese analyzer DECLARE_String(inverted_index_dict_path); DECLARE_Int32(inverted_index_read_buffer_size); diff --git a/be/src/common/status.h b/be/src/common/status.h index d29b66459a3832..37ff6917b4f567 100644 --- a/be/src/common/status.h +++ b/be/src/common/status.h @@ -299,6 +299,7 @@ namespace ErrorCode { E(INVERTED_INDEX_COMPACTION_ERROR, -6010, false); \ E(INVERTED_INDEX_ANALYZER_ERROR, -6011, false); \ E(INVERTED_INDEX_FILE_CORRUPTED, -6012, false); \ + E(INVERTED_INDEX_SNII_NOT_FOUND, -6013, false); \ E(KEY_NOT_FOUND, -7000, false); \ E(KEY_ALREADY_EXISTS, -7001, false); \ E(ENTRY_NOT_FOUND, -7002, false); \ diff --git a/be/src/exec/operator/olap_scan_operator.cpp b/be/src/exec/operator/olap_scan_operator.cpp index 50d41d3dcc979a..7524a49875aa2b 100644 --- a/be/src/exec/operator/olap_scan_operator.cpp +++ b/be/src/exec/operator/olap_scan_operator.cpp @@ -264,6 +264,10 @@ Status OlapScanLocalState::_init_profile() { ADD_COUNTER_WITH_LEVEL(_segment_profile, "InvertedIndexQueryCacheHit", TUnit::UNIT, 1); _inverted_index_query_cache_miss_counter = ADD_COUNTER_WITH_LEVEL(_segment_profile, "InvertedIndexQueryCacheMiss", TUnit::UNIT, 1); + _inverted_index_query_cache_lookup_counter = ADD_COUNTER_WITH_LEVEL( + _segment_profile, "InvertedIndexQueryCacheLookup", TUnit::UNIT, 1); + _inverted_index_query_cache_insert_counter = ADD_COUNTER_WITH_LEVEL( + _segment_profile, "InvertedIndexQueryCacheInsert", TUnit::UNIT, 1); _inverted_index_query_timer = ADD_TIMER_WITH_LEVEL(_segment_profile, "InvertedIndexQueryTime", 1); _inverted_index_query_null_bitmap_timer = @@ -348,6 +352,8 @@ Status OlapScanLocalState::_init_profile() { _index_filter_profile = std::make_unique("IndexFilter"); _scanner_profile->add_child(_index_filter_profile.get(), true, nullptr); + _snii_prx_profile_counters.initialize(_index_filter_profile.get()); + _snii_phrase_profile_counters.initialize(_index_filter_profile.get()); /* SegmentIterator: - AnnIndexLoadCosts: 102.262us diff --git a/be/src/exec/operator/olap_scan_operator.h b/be/src/exec/operator/olap_scan_operator.h index 1d2d2015039d6a..e72fe1a4d75455 100644 --- a/be/src/exec/operator/olap_scan_operator.h +++ b/be/src/exec/operator/olap_scan_operator.h @@ -28,6 +28,7 @@ #include "exec/operator/operator.h" #include "exec/operator/scan_operator.h" #include "runtime/runtime_profile.h" +#include "storage/index/snii/snii_prx_profile.h" #include "storage/olap_scan_common.h" #include "storage/tablet/tablet_reader.h" @@ -145,6 +146,8 @@ class OlapScanLocalState final : public ScanLocalState { std::unique_ptr _segment_profile; std::unique_ptr _index_filter_profile; + snii::SniiPrxRuntimeProfileCounters _snii_prx_profile_counters; + snii::SniiPhraseRuntimeProfileCounters _snii_phrase_profile_counters; RuntimeProfile::Counter* _tablet_counter = nullptr; RuntimeProfile::Counter* _key_range_counter = nullptr; @@ -239,6 +242,8 @@ class OlapScanLocalState final : public ScanLocalState { RuntimeProfile::Counter* _inverted_index_query_null_bitmap_timer = nullptr; RuntimeProfile::Counter* _inverted_index_query_cache_hit_counter = nullptr; RuntimeProfile::Counter* _inverted_index_query_cache_miss_counter = nullptr; + RuntimeProfile::Counter* _inverted_index_query_cache_lookup_counter = nullptr; + RuntimeProfile::Counter* _inverted_index_query_cache_insert_counter = nullptr; RuntimeProfile::Counter* _inverted_index_query_timer = nullptr; RuntimeProfile::Counter* _inverted_index_query_bitmap_copy_timer = nullptr; RuntimeProfile::Counter* _inverted_index_searcher_open_timer = nullptr; diff --git a/be/src/exec/scan/olap_scanner.cpp b/be/src/exec/scan/olap_scanner.cpp index 834ee012737ecd..acc9e4f96d3fc4 100644 --- a/be/src/exec/scan/olap_scanner.cpp +++ b/be/src/exec/scan/olap_scanner.cpp @@ -56,6 +56,7 @@ #include "storage/binlog.h" #include "storage/id_manager.h" #include "storage/index/inverted/inverted_index_profile.h" +#include "storage/index/inverted/similarity/collection_statistics.h" #include "storage/iterator/block_reader.h" #include "storage/olap_common.h" #include "storage/olap_tuple.h" @@ -154,7 +155,10 @@ static bool has_file_cache_statistics(const io::FileCacheStatistics& stats) { stats.inverted_index_bytes_read_from_remote != 0 || stats.inverted_index_bytes_read_from_peer != 0 || stats.inverted_index_local_io_timer != 0 || stats.inverted_index_remote_io_timer != 0 || - stats.inverted_index_peer_io_timer != 0 || stats.inverted_index_io_timer != 0; + stats.inverted_index_peer_io_timer != 0 || stats.inverted_index_io_timer != 0 || + stats.inverted_index_request_bytes != 0 || stats.inverted_index_read_bytes != 0 || + stats.inverted_index_range_read_count != 0 || + stats.inverted_index_serial_read_rounds != 0; } io::IOContext build_score_runtime_collection_io_context(RuntimeState* state, ReaderType reader_type, @@ -989,6 +993,10 @@ void OlapScanner::_collect_profile_before_close() { stats.inverted_index_query_cache_hit); COUNTER_UPDATE(local_state->_inverted_index_query_cache_miss_counter, stats.inverted_index_query_cache_miss); + COUNTER_UPDATE(local_state->_inverted_index_query_cache_lookup_counter, + stats.inverted_index_query_cache_lookup); + COUNTER_UPDATE(local_state->_inverted_index_query_cache_insert_counter, + stats.inverted_index_query_cache_insert); COUNTER_UPDATE(local_state->_inverted_index_query_timer, stats.inverted_index_query_timer); COUNTER_UPDATE(local_state->_inverted_index_query_null_bitmap_timer, stats.inverted_index_query_null_bitmap_timer); @@ -1011,6 +1019,8 @@ void OlapScanner::_collect_profile_before_close() { COUNTER_UPDATE(local_state->_inverted_index_analyzer_timer, stats.inverted_index_analyzer_timer); COUNTER_UPDATE(local_state->_inverted_index_lookup_timer, stats.inverted_index_lookup_timer); + local_state->_snii_prx_profile_counters.update(stats); + local_state->_snii_phrase_profile_counters.update(stats); COUNTER_UPDATE(local_state->_variant_scan_sparse_column_timer, stats.variant_scan_sparse_column_timer_ns); COUNTER_UPDATE(local_state->_variant_scan_sparse_column_bytes, diff --git a/be/src/exprs/function/function_search.cpp b/be/src/exprs/function/function_search.cpp index e87aff95ee075c..eb9c3e253d2179 100644 --- a/be/src/exprs/function/function_search.cpp +++ b/be/src/exprs/function/function_search.cpp @@ -115,6 +115,50 @@ static std::string extract_segment_prefix( return ""; } +static void collect_referenced_fields(const TSearchClause& clause, + std::unordered_set* fields) { + DORIS_CHECK(fields != nullptr); + if (clause.__isset.field_name && !clause.field_name.empty()) { + fields->insert(clause.field_name); + } + for (const auto& child : clause.children) { + collect_referenced_fields(child, fields); + } +} + +static bool referenced_fields_contain_snii_reader( + const TSearchClause& root, + const std::unordered_map& iterators) { + std::unordered_set referenced_fields; + collect_referenced_fields(root, &referenced_fields); + for (const auto& field_name : referenced_fields) { + auto iterator_it = iterators.find(field_name); + if (iterator_it == iterators.end()) { + continue; + } + auto* inv_iter = dynamic_cast(iterator_it->second); + if (inv_iter == nullptr) { + continue; + } + for (auto type : {InvertedIndexReaderType::FULLTEXT, InvertedIndexReaderType::STRING_TYPE, + InvertedIndexReaderType::BKD}) { + IndexReaderType reader_type = type; + auto reader = inv_iter->get_reader(reader_type); + if (reader == nullptr) { + continue; + } + auto inv_reader = std::dynamic_pointer_cast(reader); + DORIS_CHECK(inv_reader != nullptr); + auto file_reader = inv_reader->get_index_file_reader(); + DORIS_CHECK(file_reader != nullptr); + if (file_reader->get_storage_format() == InvertedIndexStorageFormatPB::SNII) { + return true; + } + } + } + return false; +} + namespace { bool is_nested_group_search_supported() { @@ -216,6 +260,14 @@ InvertedIndexQueryType direct_index_query_type_for_clause(const std::string& cla return InvertedIndexQueryType::UNKNOWN_QUERY; } +std::string normalize_wildcard_pattern(const std::string& value, + const std::map& index_properties) { + const bool has_parser = + inverted_index::InvertedIndexAnalyzer::should_analyzer(index_properties); + const std::string lowercase_setting = get_parser_lowercase_from_properties(index_properties); + return has_parser && lowercase_setting == INVERTED_INDEX_PARSER_TRUE ? to_lower(value) : value; +} + } // namespace Status FunctionSearch::execute_impl(FunctionContext* /*context*/, Block& /*block*/, @@ -274,13 +326,16 @@ Status FunctionSearch::evaluate_inverted_index_with_search_param( OlapReaderStatistics* outer_stats = index_query_context ? index_query_context->stats : nullptr; SCOPED_RAW_TIMER(outer_stats ? &outer_stats->inverted_index_query_timer : &query_timer_dummy); - const bool need_similarity_score = - index_query_context && index_query_context->collection_similarity; - // DSL result cache only stores bitmap/null bitmap. It does not store BM25 scores, // so score() queries must execute scorers to populate CollectionSimilarity. - auto* dsl_cache = (enable_cache && !need_similarity_score) ? InvertedIndexQueryCache::instance() - : nullptr; + const bool enable_scoring = + index_query_context != nullptr && index_query_context->collection_similarity != nullptr; + // Also bypass the DSL cache when any referenced field is served by an SNII reader. + auto* dsl_cache = + enable_cache && !enable_scoring && + !referenced_fields_contain_snii_reader(search_param.root, iterators) + ? InvertedIndexQueryCache::instance() + : nullptr; std::string seg_prefix; std::string dsl_sig; InvertedIndexQueryCache::CacheKey dsl_cache_key; @@ -394,11 +449,9 @@ Status FunctionSearch::evaluate_inverted_index_with_search_param( query_v2::QueryExecutionContext exec_ctx = build_variant_search_query_execution_context(num_rows, resolver, &null_resolver); - bool enable_scoring = false; bool is_asc = false; size_t top_k = 0; if (index_query_context) { - enable_scoring = index_query_context->collection_similarity != nullptr; is_asc = index_query_context->is_asc; top_k = index_query_context->query_limit; } @@ -734,6 +787,40 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, *binding_key = binding.binding_key; } + if (binding.use_snii_native_reader()) { + DORIS_CHECK(binding.inverted_reader != nullptr); + if (clause_type != "WILDCARD") { + return Status::NotSupported( + "SNII native SEARCH supports only WILDCARD clauses; got '{}'", clause_type); + } + + auto data_bitmap = std::make_shared(); + if (value == "*") { + data_bitmap->addRange(0, num_rows); + } else { + std::string pattern = normalize_wildcard_pattern(value, binding.index_properties); + Field query_value = Field::create_field(pattern); + RETURN_IF_ERROR(binding.inverted_reader->query( + context, binding.stored_field_name, query_value, + InvertedIndexQueryType::WILDCARD_QUERY, data_bitmap, nullptr)); + VLOG_DEBUG << "search: SNII WILDCARD clause processed, field=" << field_name + << ", pattern='" << pattern << "' (original='" << value << "')"; + } + + auto null_bitmap = std::make_shared(); + if (binding.inverted_reader->has_null()) { + segment_v2::InvertedIndexQueryCacheHandle null_bitmap_cache_handle; + RETURN_IF_ERROR(binding.inverted_reader->read_null_bitmap( + context, &null_bitmap_cache_handle, nullptr)); + auto cached_null_bitmap = null_bitmap_cache_handle.get_bitmap(); + DORIS_CHECK(cached_null_bitmap != nullptr); + null_bitmap = std::move(cached_null_bitmap); + } + *data_bitmap -= *null_bitmap; + return finish_leaf_query(std::make_shared(std::move(data_bitmap), + std::move(null_bitmap))); + } + if (binding.use_direct_index_reader()) { auto direct_query_type = direct_index_query_type_for_clause(clause_type); if (direct_query_type == InvertedIndexQueryType::UNKNOWN_QUERY) { @@ -800,7 +887,8 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, std::vector term_infos = inverted_index::InvertedIndexAnalyzer::get_analyse_result( - value, binding.index_properties); + value, binding.index_properties, + inverted_index::AnalysisPurpose::kPlainQuery); if (term_infos.empty()) { LOG(WARNING) << "search: No terms found after tokenization for TERM query, field=" << field_name << ", value='" << value @@ -864,7 +952,8 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, std::vector term_infos = inverted_index::InvertedIndexAnalyzer::get_analyse_result( - value, binding.index_properties); + value, binding.index_properties, + inverted_index::AnalysisPurpose::kPlainQuery); if (term_infos.empty()) { LOG(WARNING) << "search: No terms found after tokenization for PHRASE query, field=" << field_name << ", value='" << value @@ -879,8 +968,7 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, const auto& term_info = phrase_term_infos[0]; if (term_info.is_single_term()) { std::wstring term_wstr = StringHelper::to_wstring(term_info.get_single_term()); - return finish_leaf_query( - std::make_shared(context, field_wstr, term_wstr)); + return finish_leaf_query(make_term_query(term_wstr)); } else { auto builder = create_operator_boolean_query_builder(query_v2::OperatorType::OP_OR); @@ -922,7 +1010,8 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, std::vector term_infos = inverted_index::InvertedIndexAnalyzer::get_analyse_result( - value, binding.index_properties); + value, binding.index_properties, + inverted_index::AnalysisPurpose::kPlainQuery); if (term_infos.empty()) { LOG(WARNING) << "search: tokenization yielded no terms for clause '" << clause_type << "', field=" << field_name << ", returning empty BitSetQuery"; @@ -994,8 +1083,7 @@ Status FunctionSearch::build_leaf_query(const TSearchClause& clause, binding.index_properties); std::string lowercase_setting = get_parser_lowercase_from_properties(binding.index_properties); - bool should_lowercase = has_parser && (lowercase_setting == INVERTED_INDEX_PARSER_TRUE); - std::string pattern = should_lowercase ? to_lower(value) : value; + std::string pattern = normalize_wildcard_pattern(value, binding.index_properties); VLOG_DEBUG << "search: WILDCARD clause processed, field=" << field_name << ", pattern='" << pattern << "' (original='" << value << "', has_parser=" << has_parser << ", lower_case=" << lowercase_setting << ")"; diff --git a/be/src/exprs/function/function_tokenize.cpp b/be/src/exprs/function/function_tokenize.cpp index a3d616d635763c..b7aae728cfe504 100644 --- a/be/src/exprs/function/function_tokenize.cpp +++ b/be/src/exprs/function/function_tokenize.cpp @@ -187,7 +187,7 @@ Status FunctionTokenize::execute_impl(FunctionContext* /*context*/, Block& block try { analyzer_holder = doris::segment_v2::inverted_index::InvertedIndexAnalyzer::create_analyzer( - &config); + &config, AnalysisPurpose::kPlainQuery); } catch (CLuceneError& e) { return Status::Error( "inverted index create analyzer failed: {}", e.what()); diff --git a/be/src/exprs/function/match.cpp b/be/src/exprs/function/match.cpp index a280dd035e25b4..b507793785f78c 100644 --- a/be/src/exprs/function/match.cpp +++ b/be/src/exprs/function/match.cpp @@ -87,6 +87,8 @@ Status FunctionMatchBase::evaluate_inverted_index( param.query_type = get_query_type_from_fn_name(); param.num_rows = num_rows; param.roaring = std::make_shared(); + segment_v2::InvertedIndexQueryCacheHandle null_bitmap_cache_handle; + param.null_bitmap_cache_handle = &null_bitmap_cache_handle; param.analyzer_ctx = analyzer_ctx; if (is_string_type(param_type)) { RETURN_IF_ERROR(iter->read_from_index(¶m)); @@ -95,11 +97,11 @@ Status FunctionMatchBase::evaluate_inverted_index( "invalid params type for FunctionMatchBase::evaluate_inverted_index {}", param_type); } - std::shared_ptr null_bitmap = std::make_shared(); - if (iter->has_null()) { - segment_v2::InvertedIndexQueryCacheHandle null_bitmap_cache_handle; - RETURN_IF_ERROR(iter->read_null_bitmap(&null_bitmap_cache_handle)); - null_bitmap = null_bitmap_cache_handle.get_bitmap(); + std::shared_ptr null_bitmap = null_bitmap_cache_handle.get_bitmap(); + if (null_bitmap == nullptr) { + // query_with_null_bitmap leaves the handle empty only when the selected reader proves that + // the index has no null rows. + null_bitmap = std::make_shared(); } segment_v2::InvertedIndexResultBitmap result(param.roaring, null_bitmap); bitmap_result = result; diff --git a/be/src/exprs/function/variant_inverted_index_search.cpp b/be/src/exprs/function/variant_inverted_index_search.cpp index cf3fc0505188c6..1e5d8f244898d4 100644 --- a/be/src/exprs/function/variant_inverted_index_search.cpp +++ b/be/src/exprs/function/variant_inverted_index_search.cpp @@ -36,6 +36,7 @@ #include "storage/index/inverted/inverted_index_compound_reader.h" #include "storage/index/inverted/inverted_index_parser.h" #include "storage/index/inverted/inverted_index_searcher.h" +#include "storage/index/inverted/inverted_index_selector.h" #include "storage/index/inverted/query_v2/bit_set_query/bit_set_scorer.h" #include "storage/index/inverted/query_v2/doc_set.h" #include "storage/index/inverted/query_v2/scorer.h" @@ -138,18 +139,18 @@ Status FieldReaderResolver::resolve(const std::string& field_name, InvertedIndexQueryType effective_query_type = query_type; const auto& column_type = data_it->second.second; - const bool is_text_field = - column_type != nullptr && is_string_type(column_type->get_storage_field_type()); + const bool is_text_field = column_type != nullptr && + is_string_type(get_inverted_index_leaf_field_type(column_type)); auto fb_it = _field_binding_map.find(field_name); std::string analyzer_key; - if (is_text_field && is_variant_sub && fb_it != _field_binding_map.end() && - fb_it->second->__isset.index_properties && !fb_it->second->index_properties.empty()) { + if (is_text_field && effective_query_type != InvertedIndexQueryType::EQUAL_QUERY && + fb_it != _field_binding_map.end() && fb_it->second->__isset.index_properties && + !fb_it->second->index_properties.empty()) { analyzer_key = normalize_analyzer_key( build_analyzer_key_from_properties(fb_it->second->index_properties)); if (inverted_index::InvertedIndexAnalyzer::should_analyzer( fb_it->second->index_properties) && - (effective_query_type == InvertedIndexQueryType::EQUAL_QUERY || - effective_query_type == InvertedIndexQueryType::WILDCARD_QUERY)) { + effective_query_type == InvertedIndexQueryType::WILDCARD_QUERY) { effective_query_type = InvertedIndexQueryType::MATCH_ANY_QUERY; } } @@ -200,12 +201,7 @@ Status FieldReaderResolver::resolve(const std::string& field_name, resolved.inverted_reader = inverted_reader; resolved.binding_key = binding_key; resolved.state = SearchFieldBindingState::BOUND; - if (fb_it != _field_binding_map.end() && fb_it->second->__isset.index_properties && - !fb_it->second->index_properties.empty()) { - resolved.index_properties = fb_it->second->index_properties; - } else { - resolved.index_properties = inverted_reader->get_index_properties(); - } + resolved.index_properties = inverted_reader->get_index_properties(); resolved.analyzer_key = normalize_analyzer_key(build_analyzer_key_from_properties(resolved.index_properties)); @@ -225,6 +221,7 @@ Status FieldReaderResolver::resolve(const std::string& field_name, } if (inverted_reader->type() == InvertedIndexReaderType::BKD) { + resolved.execution_mode = SearchFieldExecutionMode::DIRECT_INDEX; _cache.emplace(binding_key, resolved); if (is_variant_sub) { bool index_file_exists = false; @@ -249,6 +246,29 @@ Status FieldReaderResolver::resolve(const std::string& field_name, return Status::OK(); } + if (index_file_reader->get_storage_format() == InvertedIndexStorageFormatPB::SNII) { + resolved.execution_mode = SearchFieldExecutionMode::SNII_NATIVE; + _cache.emplace(binding_key, resolved); + if (is_variant_sub) { + add_search_binding_diagnostic( + _context, + fmt::format("[VariantSearchBinding] phase=field_resolve " + "result=selected_snii_native logical_field={} stored_field={} " + "query_type={} effective_query_type={} index_id={} suffix={} " + "reader_type={} analyzer_key={} index_file={}", + field_name, stored_field_name, query_type_to_string(query_type), + query_type_to_string(effective_query_type), + inverted_reader->get_index_id(), + inverted_reader->get_index_meta().get_index_suffix(), + reader_type_to_string(inverted_reader->type()), + resolved.analyzer_key, + index_file_reader->get_index_file_path( + &inverted_reader->get_index_meta()))); + } + *binding = resolved; + return Status::OK(); + } + auto index_file_key = index_file_reader->get_index_file_cache_key(&inverted_reader->get_index_meta()); InvertedIndexSearcherCache::CacheKey searcher_cache_key(index_file_key); @@ -291,7 +311,6 @@ Status FieldReaderResolver::resolve(const std::string& field_name, index_file_reader->init(config::inverted_index_read_buffer_size, _context->io_ctx)); auto directory = DORIS_TRY( index_file_reader->open(&inverted_reader->get_index_meta(), _context->io_ctx)); - auto index_searcher_builder = DORIS_TRY( IndexSearcherBuilder::create_index_searcher_builder(inverted_reader->type())); auto searcher_result = @@ -334,6 +353,7 @@ Status FieldReaderResolver::resolve(const std::string& field_name, _searcher_cache_handles.push_back(std::move(searcher_cache_handle)); resolved.lucene_reader = reader_holder; + resolved.execution_mode = SearchFieldExecutionMode::CLUCENE; _binding_readers[binding_key] = reader_holder; _field_readers[resolved.stored_field_wstr] = reader_holder; _readers.emplace_back(reader_holder); diff --git a/be/src/exprs/function/variant_inverted_index_search.h b/be/src/exprs/function/variant_inverted_index_search.h index 973c9c8c826c55..d03dbc486d8e18 100644 --- a/be/src/exprs/function/variant_inverted_index_search.h +++ b/be/src/exprs/function/variant_inverted_index_search.h @@ -68,6 +68,13 @@ enum class SearchFieldBindingState { MISSING_IN_SEGMENT, }; +enum class SearchFieldExecutionMode { + UNBOUND, + CLUCENE, + DIRECT_INDEX, + SNII_NATIVE, +}; + struct FieldReaderBinding { std::string logical_field_name; std::string stored_field_name; @@ -80,13 +87,17 @@ struct FieldReaderBinding { std::string binding_key; std::string analyzer_key; SearchFieldBindingState state = SearchFieldBindingState::MISSING_IN_SEGMENT; + SearchFieldExecutionMode execution_mode = SearchFieldExecutionMode::UNBOUND; bool is_bound() const { return state == SearchFieldBindingState::BOUND || inverted_reader != nullptr || lucene_reader != nullptr; } bool use_direct_index_reader() const { - return is_bound() && inverted_reader != nullptr && lucene_reader == nullptr; + return is_bound() && execution_mode == SearchFieldExecutionMode::DIRECT_INDEX; + } + bool use_snii_native_reader() const { + return is_bound() && execution_mode == SearchFieldExecutionMode::SNII_NATIVE; } }; diff --git a/be/src/exprs/vcompound_pred.h b/be/src/exprs/vcompound_pred.h index 95a32ce40025da..97ec671fbb23b5 100644 --- a/be/src/exprs/vcompound_pred.h +++ b/be/src/exprs/vcompound_pred.h @@ -44,6 +44,15 @@ inline std::string compound_operator_to_string(TExprOpcode::type op) { } } +inline bool inverted_index_status_allows_row_fallback(const Status& status) { + DORIS_CHECK(!status.ok()); + return status.is() || + status.is() || + status.is() || + status.is() || + status.is(); +} + class VCompoundPred : public VectorizedFnCall { ENABLE_FACTORY_CREATOR(VCompoundPred); @@ -212,6 +221,9 @@ class VCompoundPred : public VectorizedFnCall { !st.ok()) { LOG(ERROR) << "expr:" << child->expr_name() << " evaluate_inverted_index error:" << st.to_string(); + if (!inverted_index_status_allows_row_fallback(st)) { + return st; + } all_pass = false; continue; } @@ -241,6 +253,9 @@ class VCompoundPred : public VectorizedFnCall { !st.ok()) { LOG(ERROR) << "expr:" << child->expr_name() << " evaluate_inverted_index error:" << st.to_string(); + if (!inverted_index_status_allows_row_fallback(st)) { + return st; + } all_pass = false; continue; } diff --git a/be/src/exprs/vexpr.h b/be/src/exprs/vexpr.h index 77520377af6f93..485aeb54a0d6c2 100644 --- a/be/src/exprs/vexpr.h +++ b/be/src/exprs/vexpr.h @@ -222,6 +222,10 @@ class VExpr { return empty; } + [[nodiscard]] virtual const InvertedIndexAnalyzerCtx* query_analyzer_ctx() const { + return nullptr; + } + Status _evaluate_inverted_index(VExprContext* context, const FunctionBasePtr& function, uint32_t segment_num_rows); diff --git a/be/src/exprs/vmatch_predicate.cpp b/be/src/exprs/vmatch_predicate.cpp index cf4f87e45d3b7e..fb9f9f23a651be 100644 --- a/be/src/exprs/vmatch_predicate.cpp +++ b/be/src/exprs/vmatch_predicate.cpp @@ -77,7 +77,8 @@ VMatchPredicate::VMatchPredicate(const TExprNode& node) : VExpr(node) { // Always create analyzer based on parser_type for slow path (tables without index). // For index path, FullTextIndexReader will check analyzer_name to decide whether // to use this analyzer or fallback to index's own analyzer. - _analyzer = inverted_index::InvertedIndexAnalyzer::create_analyzer(&config); + _analyzer_provider = inverted_index::InvertedIndexAnalyzer::create_analyzer_provider(&config); + _analyzer = _analyzer_provider->get_analyzer(inverted_index::AnalysisPurpose::kPlainQuery); // Step 3: Create runtime context (only extract runtime-needed info) _analyzer_ctx = std::make_shared(); @@ -85,6 +86,7 @@ VMatchPredicate::VMatchPredicate(const TExprNode& node) : VExpr(node) { _analyzer_ctx->parser_type = config.parser_type; _analyzer_ctx->char_filter_map = std::move(config.char_filter_map); _analyzer_ctx->analyzer = _analyzer; + _analyzer_ctx->analyzer_provider = _analyzer_provider; } VMatchPredicate::~VMatchPredicate() = default; diff --git a/be/src/exprs/vmatch_predicate.h b/be/src/exprs/vmatch_predicate.h index e36b695e3cfb1a..4b6792c44ed88d 100644 --- a/be/src/exprs/vmatch_predicate.h +++ b/be/src/exprs/vmatch_predicate.h @@ -57,6 +57,9 @@ class VMatchPredicate final : public VExpr { const std::string& expr_name() const override; const std::string& function_name() const; [[nodiscard]] const std::string& get_analyzer_key() const override; + [[nodiscard]] const InvertedIndexAnalyzerCtx* query_analyzer_ctx() const override { + return _analyzer_ctx.get(); + } std::string debug_string() const override; @@ -69,6 +72,7 @@ class VMatchPredicate final : public VExpr { // Lifecycle management: holds ownership of the analyzer std::shared_ptr _analyzer; + segment_v2::inverted_index::AnalyzerProviderPtr _analyzer_provider; // Runtime context: holds raw pointer to analyzer and necessary runtime info InvertedIndexAnalyzerCtxSPtr _analyzer_ctx; diff --git a/be/src/exprs/vsearch.cpp b/be/src/exprs/vsearch.cpp index fdfc36e20988dc..74b9c69678e185 100644 --- a/be/src/exprs/vsearch.cpp +++ b/be/src/exprs/vsearch.cpp @@ -250,6 +250,24 @@ Status collect_search_inputs(const VSearchExpr& expr, VExprContext* context, return Status::OK(); } +bool search_status_allows_row_fallback(const Status& status) { + DORIS_CHECK(!status.ok()); + return status.is() || + status.is() || + status.is() || + status.is() || + status.is(); +} + +Status prevent_search_row_fallback(Status status) { + DORIS_CHECK(!status.ok()); + if (!search_status_allows_row_fallback(status)) { + return status; + } + return Status::Error( + "SEARCH cannot fall back to row execution: {}", status.to_string()); +} + } // namespace VSearchExpr::VSearchExpr(const TExprNode& node) : VExpr(node) { @@ -286,7 +304,7 @@ Status VSearchExpr::execute_column_impl(VExprContext* context, const Block* bloc Status VSearchExpr::evaluate_inverted_index(VExprContext* context, uint32_t segment_num_rows) { if (_search_param.original_dsl.empty()) { - return Status::InvalidArgument("search DSL is empty"); + return prevent_search_row_fallback(Status::InvalidArgument("search DSL is empty")); } auto index_context = context->get_index_context(); @@ -296,7 +314,9 @@ Status VSearchExpr::evaluate_inverted_index(VExprContext* context, uint32_t segm } SearchInputBundle bundle; - RETURN_IF_ERROR(collect_search_inputs(*this, context, &bundle)); + if (auto status = collect_search_inputs(*this, context, &bundle); !status.ok()) { + return prevent_search_row_fallback(std::move(status)); + } VLOG_DEBUG << "VSearchExpr: bundle.iterators.size()=" << bundle.iterators.size(); @@ -326,7 +346,7 @@ Status VSearchExpr::evaluate_inverted_index(VExprContext* context, uint32_t segm if (!status.ok()) { LOG(WARNING) << "VSearchExpr: Function evaluation failed: " << status.to_string(); - return status; + return prevent_search_row_fallback(std::move(status)); } index_context->set_index_result_for_expr(this, result_bitmap); diff --git a/be/src/io/cache/block_file_cache_profile.cpp b/be/src/io/cache/block_file_cache_profile.cpp index 7aed33ccaa6e07..40a031f7bcb658 100644 --- a/be/src/io/cache/block_file_cache_profile.cpp +++ b/be/src/io/cache/block_file_cache_profile.cpp @@ -17,6 +17,7 @@ #include "io/cache/block_file_cache_profile.h" +#include #include #include #include @@ -106,12 +107,17 @@ FileCacheStatistics diff_file_cache_statistics(const FileCacheStatistics& curren SUBTRACT_FIELD(inverted_index_bytes_read_from_local); SUBTRACT_FIELD(inverted_index_bytes_read_from_remote); SUBTRACT_FIELD(inverted_index_bytes_read_from_peer); + SUBTRACT_FIELD(inverted_index_remote_physical_read_bytes); + SUBTRACT_FIELD(inverted_index_bytes_write_into_cache); SUBTRACT_FIELD(inverted_index_local_io_timer); SUBTRACT_FIELD(inverted_index_remote_io_timer); SUBTRACT_FIELD(inverted_index_peer_io_timer); SUBTRACT_FIELD(inverted_index_io_timer); SUBTRACT_FIELD(inverted_index_write_cache_io_timer); - SUBTRACT_FIELD(inverted_index_bytes_write_into_cache); + SUBTRACT_FIELD(inverted_index_request_bytes); + SUBTRACT_FIELD(inverted_index_read_bytes); + SUBTRACT_FIELD(inverted_index_range_read_count); + SUBTRACT_FIELD(inverted_index_serial_read_rounds); SUBTRACT_FIELD(segment_footer_index_num_local_io_total); SUBTRACT_FIELD(segment_footer_index_num_remote_io_total); @@ -193,6 +199,10 @@ FileCacheProfileReporter::FileCacheProfileReporter(RuntimeProfile* profile, profile, "InvertedIndexBytesScannedFromRemote", TUnit::BYTES, cache_profile, 1); inverted_index_bytes_scanned_from_peer = ADD_CHILD_COUNTER_WITH_LEVEL( profile, "InvertedIndexBytesScannedFromPeer", TUnit::BYTES, cache_profile, 1); + inverted_index_remote_physical_read_bytes = ADD_CHILD_COUNTER_WITH_LEVEL( + profile, "InvertedIndexRemotePhysicalReadBytes", TUnit::BYTES, cache_profile, 1); + inverted_index_bytes_write_into_cache = ADD_CHILD_COUNTER_WITH_LEVEL( + profile, "InvertedIndexBytesWriteIntoCache", TUnit::BYTES, cache_profile, 1); inverted_index_local_io_timer = ADD_CHILD_TIMER_WITH_LEVEL(profile, "InvertedIndexLocalIOUseTimer", cache_profile, 1); inverted_index_remote_io_timer = @@ -203,8 +213,14 @@ FileCacheProfileReporter::FileCacheProfileReporter(RuntimeProfile* profile, ADD_CHILD_TIMER_WITH_LEVEL(profile, "InvertedIndexIOTimer", cache_profile, 1); inverted_index_write_cache_io_timer = ADD_CHILD_TIMER_WITH_LEVEL( profile, "InvertedIndexWriteCacheIOUseTimer", cache_profile, 1); - inverted_index_bytes_write_into_cache = ADD_CHILD_COUNTER_WITH_LEVEL( - profile, "InvertedIndexBytesWriteIntoCache", TUnit::BYTES, cache_profile, 1); + inverted_index_request_bytes = ADD_CHILD_COUNTER_WITH_LEVEL( + profile, "InvertedIndexRequestBytes", TUnit::BYTES, cache_profile, 1); + inverted_index_read_bytes = ADD_CHILD_COUNTER_WITH_LEVEL(profile, "InvertedIndexReadBytes", + TUnit::BYTES, cache_profile, 1); + inverted_index_range_read_count = ADD_CHILD_COUNTER_WITH_LEVEL( + profile, "InvertedIndexRangeReadCount", TUnit::UNIT, cache_profile, 1); + inverted_index_serial_read_rounds = ADD_CHILD_COUNTER_WITH_LEVEL( + profile, "InvertedIndexSerialReadRounds", TUnit::UNIT, cache_profile, 1); segment_footer_index_num_local_io_total = ADD_CHILD_COUNTER_WITH_LEVEL( profile, "SegmentFooterIndexNumLocalIOTotal", TUnit::UNIT, cache_profile, 1); @@ -290,14 +306,16 @@ void FileCacheProfileReporter::update(const FileCacheStatistics* statistics) con statistics->inverted_index_bytes_read_from_remote); COUNTER_UPDATE(inverted_index_bytes_scanned_from_peer, statistics->inverted_index_bytes_read_from_peer); + COUNTER_UPDATE(inverted_index_remote_physical_read_bytes, + statistics->inverted_index_remote_physical_read_bytes); + COUNTER_UPDATE(inverted_index_bytes_write_into_cache, + statistics->inverted_index_bytes_write_into_cache); COUNTER_UPDATE(inverted_index_local_io_timer, statistics->inverted_index_local_io_timer); COUNTER_UPDATE(inverted_index_remote_io_timer, statistics->inverted_index_remote_io_timer); COUNTER_UPDATE(inverted_index_peer_io_timer, statistics->inverted_index_peer_io_timer); COUNTER_UPDATE(inverted_index_io_timer, statistics->inverted_index_io_timer); COUNTER_UPDATE(inverted_index_write_cache_io_timer, statistics->inverted_index_write_cache_io_timer); - COUNTER_UPDATE(inverted_index_bytes_write_into_cache, - statistics->inverted_index_bytes_write_into_cache); COUNTER_UPDATE(segment_footer_index_num_local_io_total, statistics->segment_footer_index_num_local_io_total); @@ -343,6 +361,11 @@ void FileCacheProfileReporter::update(const FileCacheStatistics* statistics) con } _profile->add_info_string("PeerCacheNodes", peer_nodes); } + COUNTER_UPDATE(inverted_index_request_bytes, statistics->inverted_index_request_bytes); + COUNTER_UPDATE(inverted_index_read_bytes, statistics->inverted_index_read_bytes); + COUNTER_UPDATE(inverted_index_range_read_count, statistics->inverted_index_range_read_count); + COUNTER_UPDATE(inverted_index_serial_read_rounds, + statistics->inverted_index_serial_read_rounds); } } // namespace doris::io diff --git a/be/src/io/cache/block_file_cache_profile.h b/be/src/io/cache/block_file_cache_profile.h index d3fd31033649c8..5e7a9e2a267f1a 100644 --- a/be/src/io/cache/block_file_cache_profile.h +++ b/be/src/io/cache/block_file_cache_profile.h @@ -19,6 +19,7 @@ #include +#include #include #include #include @@ -60,7 +61,6 @@ class FileCacheMetrics { void register_entity(); void update_metrics_callback(); -private: std::mutex _mtx; // use shared_ptr for concurrent std::shared_ptr _statistics; @@ -99,12 +99,17 @@ struct FileCacheProfileReporter { RuntimeProfile::Counter* inverted_index_bytes_scanned_from_cache = nullptr; RuntimeProfile::Counter* inverted_index_bytes_scanned_from_remote = nullptr; RuntimeProfile::Counter* inverted_index_bytes_scanned_from_peer = nullptr; + RuntimeProfile::Counter* inverted_index_remote_physical_read_bytes = nullptr; + RuntimeProfile::Counter* inverted_index_bytes_write_into_cache = nullptr; RuntimeProfile::Counter* inverted_index_local_io_timer = nullptr; RuntimeProfile::Counter* inverted_index_remote_io_timer = nullptr; RuntimeProfile::Counter* inverted_index_peer_io_timer = nullptr; RuntimeProfile::Counter* inverted_index_io_timer = nullptr; RuntimeProfile::Counter* inverted_index_write_cache_io_timer = nullptr; - RuntimeProfile::Counter* inverted_index_bytes_write_into_cache = nullptr; + RuntimeProfile::Counter* inverted_index_request_bytes = nullptr; + RuntimeProfile::Counter* inverted_index_read_bytes = nullptr; + RuntimeProfile::Counter* inverted_index_range_read_count = nullptr; + RuntimeProfile::Counter* inverted_index_serial_read_rounds = nullptr; RuntimeProfile::Counter* segment_footer_index_num_local_io_total = nullptr; RuntimeProfile::Counter* segment_footer_index_num_remote_io_total = nullptr; diff --git a/be/src/io/cache/cached_remote_file_reader.cpp b/be/src/io/cache/cached_remote_file_reader.cpp index 47fe02ee2c3e92..dcd8dde92946b3 100644 --- a/be/src/io/cache/cached_remote_file_reader.cpp +++ b/be/src/io/cache/cached_remote_file_reader.cpp @@ -325,7 +325,11 @@ Status execute_s3_read(size_t empty_start, size_t& size, std::unique_ptr s3_read_counter << 1; SCOPED_RAW_TIMER(&stats.remote_read_timer); stats.from_peer_cache = false; - return remote_file_reader->read_at(empty_start, Slice(buffer.get(), size), &size, io_ctx); + auto st = remote_file_reader->read_at(empty_start, Slice(buffer.get(), size), &size, io_ctx); + if (st.ok()) { + stats.remote_physical_read_bytes += size; + } + return st; } CloudWarmUpManager& get_warm_up_manager() { @@ -343,6 +347,10 @@ struct RaceState { Status peer_status; Status s3_status; std::unique_ptr s3_buf; + // Actual bytes fetched by the winning S3 leg; merged into the caller's + // ReadStatistics only in the winner==1 branch of collect_race_result (the + // losing S3 leg may outlive the caller's stack, so it must not touch stats). + size_t s3_read_size = 0; PeerFetchResult peer_res; std::string peer_winner_cg_id; // compute_group_id of the winning peer candidate std::string peer_winner_host; // host of the winning peer candidate @@ -466,6 +474,7 @@ void launch_s3_race(std::shared_ptr race, size_t empty_start, size_t if (st.ok() && race->winner < 0) { race->winner = 1; race->s3_buf = std::move(s3_buf); + race->s3_read_size = read_size; } race->cv.notify_all(); }; @@ -548,9 +557,11 @@ Status collect_race_result(std::shared_ptr race, size_t span_size, } return Status::OK(); } else if (race->winner == 1) { - // S3 won. + // S3 won: this was a real storage GET, so account it as physical remote IO + // exactly like the non-race download path does. buffer = std::move(race->s3_buf); stats.from_peer_cache = false; + stats.remote_physical_read_bytes += race->s3_read_size; g_peer_race_s3_win << 1; if (io_ctx != nullptr && io_ctx->file_cache_stats != nullptr) { io_ctx->file_cache_stats->num_peer_race_s3_win++; @@ -1038,6 +1049,8 @@ Status CachedRemoteFileReader::_read_remaining_blocks_from_cache( &remote_bytes_read, io_ctx)); indirect_read_bytes += read_size; source_read_breakdown.remote_bytes += remote_bytes_read; + // Self-heal fell back to a real storage GET; count it as physical remote IO. + stats.remote_physical_read_bytes += remote_bytes_read; DCHECK(remote_bytes_read == read_size); } @@ -1112,6 +1125,10 @@ Status CachedRemoteFileReader::_read_remote_only_on_cache_miss( *bytes_read = remote_bytes_read; DCHECK_EQ(*bytes_read, bytes_req); source_read_breakdown.remote_bytes += remote_bytes_read; + // This is a real storage GET, so it must count as physical remote IO just like + // the block-download path; otherwise profiles show scanned-from-remote bytes + // with zero physical reads whenever remote-only-on-miss is active. + stats.remote_physical_read_bytes += remote_bytes_read; g_read_cache_indirect_bytes << remote_bytes_read; g_read_cache_indirect_total_bytes << remote_bytes_read; return Status::OK(); @@ -1413,6 +1430,7 @@ void CachedRemoteFileReader::_update_stats(const ReadStatistics& read_stats, statis->inverted_index_remote_io_timer, statis->inverted_index_peer_io_timer, statis->inverted_index_write_cache_io_timer, statis->inverted_index_bytes_write_into_cache); + statis->inverted_index_remote_physical_read_bytes += read_stats.remote_physical_read_bytes; break; case FileCacheReadType::SEGMENT_FOOTER_INDEX: update_index_stats(statis->segment_footer_index_num_local_io_total, diff --git a/be/src/io/cache/file_cache_common.h b/be/src/io/cache/file_cache_common.h index 8b52af8d161b9a..7357f1a431acc2 100644 --- a/be/src/io/cache/file_cache_common.h +++ b/be/src/io/cache/file_cache_common.h @@ -75,6 +75,7 @@ struct ReadStatistics { int64_t bytes_read_from_local = 0; int64_t bytes_read_from_remote = 0; int64_t bytes_read_from_peer = 0; + int64_t remote_physical_read_bytes = 0; int64_t bytes_write_into_file_cache = 0; int64_t remote_read_timer = 0; int64_t peer_read_timer = 0; diff --git a/be/src/io/io_common.h b/be/src/io/io_common.h index f42e486eb40aea..c5779668c93e8a 100644 --- a/be/src/io/io_common.h +++ b/be/src/io/io_common.h @@ -19,6 +19,9 @@ #include +#include +#include +#include #include #include @@ -80,12 +83,17 @@ struct FileCacheStatistics { int64_t inverted_index_bytes_read_from_local = 0; int64_t inverted_index_bytes_read_from_remote = 0; int64_t inverted_index_bytes_read_from_peer = 0; + int64_t inverted_index_remote_physical_read_bytes = 0; + int64_t inverted_index_bytes_write_into_cache = 0; int64_t inverted_index_local_io_timer = 0; int64_t inverted_index_remote_io_timer = 0; int64_t inverted_index_peer_io_timer = 0; int64_t inverted_index_io_timer = 0; int64_t inverted_index_write_cache_io_timer = 0; - int64_t inverted_index_bytes_write_into_cache = 0; + int64_t inverted_index_request_bytes = 0; + int64_t inverted_index_read_bytes = 0; + int64_t inverted_index_range_read_count = 0; + int64_t inverted_index_serial_read_rounds = 0; int64_t segment_footer_index_num_local_io_total = 0; int64_t segment_footer_index_num_remote_io_total = 0; @@ -141,12 +149,18 @@ struct FileCacheStatistics { inverted_index_bytes_read_from_local += other.inverted_index_bytes_read_from_local; inverted_index_bytes_read_from_remote += other.inverted_index_bytes_read_from_remote; inverted_index_bytes_read_from_peer += other.inverted_index_bytes_read_from_peer; + inverted_index_remote_physical_read_bytes += + other.inverted_index_remote_physical_read_bytes; inverted_index_local_io_timer += other.inverted_index_local_io_timer; inverted_index_remote_io_timer += other.inverted_index_remote_io_timer; inverted_index_peer_io_timer += other.inverted_index_peer_io_timer; inverted_index_io_timer += other.inverted_index_io_timer; inverted_index_write_cache_io_timer += other.inverted_index_write_cache_io_timer; inverted_index_bytes_write_into_cache += other.inverted_index_bytes_write_into_cache; + inverted_index_request_bytes += other.inverted_index_request_bytes; + inverted_index_read_bytes += other.inverted_index_read_bytes; + inverted_index_range_read_count += other.inverted_index_range_read_count; + inverted_index_serial_read_rounds += other.inverted_index_serial_read_rounds; segment_footer_index_num_local_io_total += other.segment_footer_index_num_local_io_total; segment_footer_index_num_remote_io_total += other.segment_footer_index_num_remote_io_total; diff --git a/be/src/runtime/index_policy/index_policy_mgr.cpp b/be/src/runtime/index_policy/index_policy_mgr.cpp index 5ab41a409ba921..e7ec8dfd2b9d0e 100644 --- a/be/src/runtime/index_policy/index_policy_mgr.cpp +++ b/be/src/runtime/index_policy/index_policy_mgr.cpp @@ -17,6 +17,7 @@ #include "runtime/index_policy/index_policy_mgr.h" +#include #include #include #include @@ -24,6 +25,24 @@ #include namespace doris { +namespace { + +class PurposeInsensitiveAnalyzerProvider final + : public segment_v2::inverted_index::AnalyzerProvider { +public: + explicit PurposeInsensitiveAnalyzerProvider(AnalyzerPtr analyzer) + : _analyzer(std::move(analyzer)) {} + + AnalyzerPtr get_analyzer( + segment_v2::inverted_index::AnalysisPurpose /*purpose*/) const override { + return _analyzer; + } + +private: + const AnalyzerPtr _analyzer; +}; + +} // namespace const std::unordered_set IndexPolicyMgr::BUILTIN_NORMALIZERS = {"lowercase"}; @@ -49,11 +68,9 @@ void IndexPolicyMgr::apply_policy_changes(const std::vector& polic LOG(INFO) << "Deleting policy - " << "ID: " << id << ", " << "Name: " << it->second.name; - - // Use normalized name for deletion _name_to_id.erase(normalize_name(it->second.name)); _policys.erase(it); - success_deletes++; + ++success_deletes; } else { LOG(WARNING) << "Delete failed - Policy ID not found: " << id; } @@ -66,8 +83,6 @@ void IndexPolicyMgr::apply_policy_changes(const std::vector& polic << " | New name: " << policy.name; continue; } - - // Use normalized name for case-insensitive lookup std::string normalized_name = normalize_name(policy.name); if (_name_to_id.contains(normalized_name)) { LOG(ERROR) << "Reject update - Duplicate policy name: " << policy.name @@ -77,10 +92,8 @@ void IndexPolicyMgr::apply_policy_changes(const std::vector& polic } _policys.emplace(policy.id, policy); - // Store with normalized key for case-insensitive lookup _name_to_id.emplace(normalized_name, policy.id); - success_updates++; - + ++success_updates; LOG(INFO) << "Successfully applied policy - " << "ID: " << policy.id << ", " << "Name: " << policy.name << ", " @@ -131,7 +144,79 @@ AnalyzerPtr IndexPolicyMgr::get_policy_by_name(const std::string& name) { throw Exception(ErrorCode::INVALID_ARGUMENT, "Policy not found with type: " + name); } -AnalyzerPtr IndexPolicyMgr::build_analyzer_from_policy(const TIndexPolicy& index_policy_analyzer) { +AnalyzerPtr IndexPolicyMgr::get_analyzer_by_name( + const std::string& name, segment_v2::inverted_index::AnalysisPurpose purpose) { + std::shared_lock lock(_mutex); + const std::string normalized_name = normalize_name(name); + auto name_it = _name_to_id.find(normalized_name); + if (name_it == _name_to_id.end()) { + if (is_builtin_normalizer(normalized_name)) { + return build_builtin_normalizer(name); + } + throw Exception(ErrorCode::INVALID_ARGUMENT, "Policy not found with name: " + name); + } + auto policy_it = _policys.find(name_it->second); + if (policy_it == _policys.end()) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "Policy not found with id: " + name); + } + if (policy_it->second.type == TIndexPolicyType::ANALYZER) { + return build_analyzer_provider_from_config( + build_analyzer_config_from_policy(policy_it->second), {}) + ->get_analyzer(purpose); + } + if (policy_it->second.type == TIndexPolicyType::NORMALIZER) { + return build_normalizer_from_policy(policy_it->second); + } + throw Exception(ErrorCode::INVALID_ARGUMENT, "Analyzer policy not found: " + name); +} + +AnalyzerProviderPtr IndexPolicyMgr::get_analyzer_provider_by_name( + const std::string& name, const std::map& outer_char_filter_map) { + std::shared_lock lock(_mutex); + const std::string normalized_name = normalize_name(name); + auto name_it = _name_to_id.find(normalized_name); + if (name_it == _name_to_id.end()) { + if (is_builtin_normalizer(normalized_name)) { + return std::make_shared( + build_builtin_normalizer(name)); + } + throw Exception(ErrorCode::INVALID_ARGUMENT, "Policy not found with name: " + name); + } + auto policy_it = _policys.find(name_it->second); + if (policy_it == _policys.end()) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "Policy not found with id: " + name); + } + if (policy_it->second.type == TIndexPolicyType::ANALYZER) { + return build_analyzer_provider_from_config( + build_analyzer_config_from_policy(policy_it->second), outer_char_filter_map); + } + if (policy_it->second.type == TIndexPolicyType::NORMALIZER) { + return std::make_shared( + build_normalizer_from_policy(policy_it->second)); + } + throw Exception(ErrorCode::INVALID_ARGUMENT, "Analyzer policy not found: " + name); +} + +AnalyzerProviderPtr IndexPolicyMgr::get_analyzer_provider_by_base_fingerprint( + std::string_view base_analyzer_fingerprint, + const std::map& outer_char_filter_map) { + std::shared_lock lock(_mutex); + for (const auto& [_, policy] : _policys) { + if (policy.type != TIndexPolicyType::ANALYZER) { + continue; + } + auto config = build_analyzer_config_from_policy(policy); + if (segment_v2::inverted_index::CustomAnalyzerProvider::calculate_base_analyzer_fingerprint( + config, outer_char_filter_map) != base_analyzer_fingerprint) { + continue; + } + return build_analyzer_provider_from_config(std::move(config), outer_char_filter_map); + } + return nullptr; +} + +segment_v2::inverted_index::CustomAnalyzerConfigPtr +IndexPolicyMgr::build_analyzer_config_from_policy(const TIndexPolicy& index_policy_analyzer) { segment_v2::inverted_index::CustomAnalyzerConfig::Builder builder; auto tokenizer_it = index_policy_analyzer.properties.find(PROP_TOKENIZER); @@ -175,9 +260,23 @@ AnalyzerPtr IndexPolicyMgr::build_analyzer_from_policy(const TIndexPolicy& index builder.add_token_filter_config(name, settings); }); - auto custom_analyzer_config = builder.build(); - return segment_v2::inverted_index::CustomAnalyzer::build_custom_analyzer( - custom_analyzer_config); + return builder.build(); +} + +AnalyzerProviderPtr IndexPolicyMgr::build_analyzer_provider_from_config( + segment_v2::inverted_index::CustomAnalyzerConfigPtr config, + const std::map& outer_char_filter_map) { + // One shape for every policy: the provider sources its CommonGrams word list from the + // BE-local default, so there is no per-policy word set to look up and no "not yet prepared" + // state to represent. + return std::make_shared( + std::move(config), outer_char_filter_map); +} + +AnalyzerPtr IndexPolicyMgr::build_analyzer_from_policy(const TIndexPolicy& index_policy_analyzer) { + return build_analyzer_provider_from_config( + build_analyzer_config_from_policy(index_policy_analyzer), {}) + ->get_analyzer(segment_v2::inverted_index::AnalysisPurpose::kIndex); } AnalyzerPtr IndexPolicyMgr::build_normalizer_from_policy( @@ -224,7 +323,8 @@ void IndexPolicyMgr::process_filter_configs( std::string normalized_filter_name = normalize_name(filter_name); if (_name_to_id.contains(normalized_filter_name)) { // Nested filter policy - const auto& filter_policy = _policys[_name_to_id[normalized_filter_name]]; + const int64_t filter_policy_id = _name_to_id.at(normalized_filter_name); + const auto& filter_policy = _policys.at(filter_policy_id); auto type_it = filter_policy.properties.find(PROP_TYPE); if (type_it == filter_policy.properties.end()) { throw Exception( diff --git a/be/src/runtime/index_policy/index_policy_mgr.h b/be/src/runtime/index_policy/index_policy_mgr.h index edbdee938b86ad..0fb1dc3b3a87c1 100644 --- a/be/src/runtime/index_policy/index_policy_mgr.h +++ b/be/src/runtime/index_policy/index_policy_mgr.h @@ -19,7 +19,13 @@ #include +#include +#include +#include #include +#include +#include +#include #include #include "storage/index/inverted/analyzer/custom_analyzer.h" @@ -29,6 +35,7 @@ namespace doris { using Policys = std::unordered_map; using AnalyzerPtr = std::shared_ptr; +using AnalyzerProviderPtr = segment_v2::inverted_index::AnalyzerProviderPtr; class IndexPolicyMgr { public: @@ -40,8 +47,21 @@ class IndexPolicyMgr { Policys get_index_policys(); AnalyzerPtr get_policy_by_name(const std::string& name); + AnalyzerPtr get_analyzer_by_name(const std::string& name, + segment_v2::inverted_index::AnalysisPurpose purpose); + AnalyzerProviderPtr get_analyzer_provider_by_name( + const std::string& name, + const std::map& outer_char_filter_map = {}); + AnalyzerProviderPtr get_analyzer_provider_by_base_fingerprint( + std::string_view base_analyzer_fingerprint, + const std::map& outer_char_filter_map = {}); private: + segment_v2::inverted_index::CustomAnalyzerConfigPtr build_analyzer_config_from_policy( + const TIndexPolicy& index_policy_analyzer); + AnalyzerProviderPtr build_analyzer_provider_from_config( + segment_v2::inverted_index::CustomAnalyzerConfigPtr config, + const std::map& outer_char_filter_map); AnalyzerPtr build_analyzer_from_policy(const TIndexPolicy& index_policy_analyzer); AnalyzerPtr build_normalizer_from_policy(const TIndexPolicy& index_policy_normalizer); diff --git a/be/src/storage/compaction/collection_statistics.cpp b/be/src/storage/compaction/collection_statistics.cpp deleted file mode 100644 index 9afbe5f0944d43..00000000000000 --- a/be/src/storage/compaction/collection_statistics.cpp +++ /dev/null @@ -1,290 +0,0 @@ -// 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 "storage/compaction/collection_statistics.h" - -#include -#include - -#include "common/exception.h" -#include "exprs/vexpr.h" -#include "exprs/vexpr_context.h" -#include "exprs/vliteral.h" -#include "exprs/vslot_ref.h" -#include "storage/index/index_file_reader.h" -#include "storage/index/index_reader_helper.h" -#include "storage/index/inverted/analyzer/analyzer.h" -#include "storage/index/inverted/util/string_helper.h" -#include "storage/index/inverted/util/term_iterator.h" -#include "storage/rowset/rowset.h" -#include "storage/rowset/rowset_reader.h" -#include "util/uid_util.h" - -namespace doris { - -Status CollectionStatistics::collect(RuntimeState* state, - const std::vector& rs_splits, - const TabletSchemaSPtr& tablet_schema, - const VExprContextSPtrs& common_expr_ctxs_push_down, - io::IOContext* io_ctx) { - std::unordered_map collect_infos; - RETURN_IF_ERROR( - extract_collect_info(state, common_expr_ctxs_push_down, tablet_schema, &collect_infos)); - if (collect_infos.empty()) { - LOG(WARNING) << "Index statistics collection: no collect info extracted."; - return Status::OK(); - } - - for (const auto& rs_split : rs_splits) { - const auto& rs_reader = rs_split.rs_reader; - auto rowset = rs_reader->rowset(); - auto num_segments = rowset->num_segments(); - for (int32_t seg_id = 0; seg_id < num_segments; ++seg_id) { - auto status = - process_segment(rowset, seg_id, tablet_schema.get(), collect_infos, io_ctx); - if (!status.ok()) { - if (status.code() == ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND || - status.code() == ErrorCode::INVERTED_INDEX_BYPASS) { - LOG(ERROR) << "Index statistics collection failed: " << status.to_string(); - } else { - return status; - } - } - } - } - - // Build a single-line log with query_id, tablet_ids, and per-field term statistics - if (VLOG_IS_ON(1)) { - std::set tablet_ids; - for (const auto& rs_split : rs_splits) { - if (rs_split.rs_reader && rs_split.rs_reader->rowset()) { - tablet_ids.insert(rs_split.rs_reader->rowset()->rowset_meta()->tablet_id()); - } - } - - std::ostringstream oss; - oss << "CollectionStatistics: query_id=" << print_id(state->query_id()); - - oss << ", tablet_ids=["; - bool first_tablet = true; - for (int64_t tid : tablet_ids) { - if (!first_tablet) oss << ","; - oss << tid; - first_tablet = false; - } - oss << "]"; - - oss << ", total_num_docs=" << _total_num_docs; - - for (const auto& [ws_field_name, num_tokens] : _total_num_tokens) { - oss << ", {field=" << StringHelper::to_string(ws_field_name) - << ", num_tokens=" << num_tokens << ", terms=["; - - bool first_term = true; - for (const auto& [term, doc_freq] : _term_doc_freqs.at(ws_field_name)) { - if (!first_term) oss << ", "; - oss << "(" << StringHelper::to_string(term) << ":" << doc_freq << ")"; - first_term = false; - } - oss << "]}"; - } - - VLOG(1) << oss.str(); - } - - return Status::OK(); -} - -Status CollectionStatistics::extract_collect_info( - RuntimeState* state, const VExprContextSPtrs& common_expr_ctxs_push_down, - const TabletSchemaSPtr& tablet_schema, CollectInfoMap* collect_infos) { - DCHECK(collect_infos != nullptr); - - std::unordered_map collectors; - collectors[TExprNodeType::MATCH_PRED] = std::make_unique(); - collectors[TExprNodeType::SEARCH_EXPR] = std::make_unique(); - - for (const auto& root_expr_ctx : common_expr_ctxs_push_down) { - const auto& root_expr = root_expr_ctx->root(); - if (root_expr == nullptr) { - continue; - } - - std::stack stack; - stack.emplace(root_expr); - - while (!stack.empty()) { - auto expr = stack.top(); - stack.pop(); - - if (!expr) { - continue; - } - - auto collector_it = collectors.find(expr->node_type()); - if (collector_it != collectors.end()) { - RETURN_IF_ERROR( - collector_it->second->collect(state, tablet_schema, expr, collect_infos)); - } - - const auto& children = expr->children(); - for (const auto& child : children) { - stack.push(child); - } - } - } - - LOG(INFO) << "Extracted collect info for " << collect_infos->size() << " fields"; - - return Status::OK(); -} - -Status CollectionStatistics::process_segment(const RowsetSharedPtr& rowset, int32_t seg_id, - const TabletSchema* tablet_schema, - const CollectInfoMap& collect_infos, - io::IOContext* io_ctx) { - auto seg_path = DORIS_TRY(rowset->segment_path(seg_id)); - auto rowset_meta = rowset->rowset_meta(); - - auto idx_file_reader = std::make_unique( - rowset_meta->fs(), - std::string {InvertedIndexDescriptor::get_index_file_path_prefix(seg_path)}, - tablet_schema->get_inverted_index_storage_format(), - rowset_meta->inverted_index_file_info(seg_id), rowset_meta->tablet_id()); - RETURN_IF_ERROR(idx_file_reader->init(config::inverted_index_read_buffer_size, io_ctx)); - - int32_t total_seg_num_docs = 0; - - for (const auto& [ws_field_name, collect_info] : collect_infos) { - lucene::search::IndexSearcher* index_searcher = nullptr; - lucene::index::IndexReader* index_reader = nullptr; - -#ifdef BE_TEST - auto compound_reader = DORIS_TRY(idx_file_reader->open(collect_info.index_meta, io_ctx)); - auto* reader = lucene::index::IndexReader::open(compound_reader.get()); - auto searcher_ptr = std::make_shared(reader, true); - index_searcher = searcher_ptr.get(); - index_reader = index_searcher->getReader(); -#else - InvertedIndexCacheHandle inverted_index_cache_handle; - auto index_file_key = idx_file_reader->get_index_file_cache_key(collect_info.index_meta); - InvertedIndexSearcherCache::CacheKey searcher_cache_key(index_file_key); - - if (!InvertedIndexSearcherCache::instance()->lookup(searcher_cache_key, - &inverted_index_cache_handle)) { - auto compound_reader = - DORIS_TRY(idx_file_reader->open(collect_info.index_meta, io_ctx)); - auto* reader = lucene::index::IndexReader::open(compound_reader.get()); - size_t reader_size = reader->getTermInfosRAMUsed(); - auto searcher_ptr = std::make_shared(reader, true); - auto* cache_value = new InvertedIndexSearcherCache::CacheValue( - std::move(searcher_ptr), reader_size, UnixMillis()); - InvertedIndexSearcherCache::instance()->insert(searcher_cache_key, cache_value, - &inverted_index_cache_handle); - } - - auto searcher_variant = inverted_index_cache_handle.get_index_searcher(); - auto index_searcher_ptr = std::get(searcher_variant); - index_searcher = index_searcher_ptr.get(); - index_reader = index_searcher->getReader(); -#endif - total_seg_num_docs = std::max(total_seg_num_docs, index_reader->maxDoc()); - - _total_num_tokens[ws_field_name] += - index_reader->sumTotalTermFreq(ws_field_name.c_str()).value_or(0); - - for (const auto& term_info : collect_info.term_infos) { - auto iter = TermIterator::create(io_ctx, false, index_reader, ws_field_name, - term_info.get_single_term()); - _term_doc_freqs[ws_field_name][iter->term()] += iter->doc_freq(); - } - } - - _total_num_docs += total_seg_num_docs; - - return Status::OK(); -} - -uint64_t CollectionStatistics::get_term_doc_freq_by_col(const std::wstring& lucene_col_name, - const std::wstring& term) { - if (!_term_doc_freqs.contains(lucene_col_name)) { - throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, - "Index statistics collection failed: Not such column {}", - StringHelper::to_string(lucene_col_name)); - } - - if (!_term_doc_freqs[lucene_col_name].contains(term)) { - throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, - "Index statistics collection failed: Not such term {}", - StringHelper::to_string(term)); - } - - return _term_doc_freqs[lucene_col_name][term]; -} - -uint64_t CollectionStatistics::get_total_term_cnt_by_col(const std::wstring& lucene_col_name) { - if (!_total_num_tokens.contains(lucene_col_name)) { - throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, - "Index statistics collection failed: Not such column {}", - StringHelper::to_string(lucene_col_name)); - } - - return _total_num_tokens[lucene_col_name]; -} - -uint64_t CollectionStatistics::get_doc_num() const { - if (_total_num_docs == 0) { - throw Exception( - ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, - "Index statistics collection failed: No data available for SimilarityCollector"); - } - - return _total_num_docs; -} - -float CollectionStatistics::get_or_calculate_avg_dl(const std::wstring& lucene_col_name) { - auto iter = _avg_dl_by_col.find(lucene_col_name); - if (iter != _avg_dl_by_col.end()) { - return iter->second; - } - - const uint64_t total_term_cnt = get_total_term_cnt_by_col(lucene_col_name); - const uint64_t total_doc_cnt = get_doc_num(); - float avg_dl = total_doc_cnt > 0 ? float((double)total_term_cnt / (double)total_doc_cnt) : 0.0F; - _avg_dl_by_col[lucene_col_name] = avg_dl; - return avg_dl; -} - -float CollectionStatistics::get_or_calculate_idf(const std::wstring& lucene_col_name, - const std::wstring& term) { - auto iter = _idf_by_col_term.find(lucene_col_name); - if (iter != _idf_by_col_term.end()) { - auto term_iter = iter->second.find(term); - if (term_iter != iter->second.end()) { - return term_iter->second; - } - } - - const uint64_t doc_num = get_doc_num(); - const uint64_t doc_freq = get_term_doc_freq_by_col(lucene_col_name, term); - auto idf = (float)std::log(1 + ((double)doc_num - (double)doc_freq + (double)0.5) / - ((double)doc_freq + (double)0.5)); - _idf_by_col_term[lucene_col_name][term] = idf; - return idf; -} - -} // namespace doris \ No newline at end of file diff --git a/be/src/storage/compaction/compaction.cpp b/be/src/storage/compaction/compaction.cpp index f27f55d309661d..6384dc5d55d080 100644 --- a/be/src/storage/compaction/compaction.cpp +++ b/be/src/storage/compaction/compaction.cpp @@ -22,9 +22,12 @@ #include #include +#include #include #include #include +#include +#include #include #include #include @@ -41,6 +44,7 @@ #include "cloud/cloud_tablet.h" #include "cloud/config.h" #include "cloud/pb_convert.h" +#include "common/check.h" #include "common/config.h" #include "common/metrics/doris_metrics.h" #include "common/status.h" @@ -54,7 +58,6 @@ #include "runtime/memory/mem_tracker_limiter.h" #include "runtime/thread_context.h" #include "storage/compaction/binlog_compaction_policy.h" -#include "storage/compaction/collection_statistics.h" #include "storage/compaction/cumulative_compaction.h" #include "storage/compaction/cumulative_compaction_policy.h" #include "storage/compaction/cumulative_compaction_time_series_policy.h" @@ -65,6 +68,10 @@ #include "storage/index/inverted/inverted_index_compaction.h" #include "storage/index/inverted/inverted_index_desc.h" #include "storage/index/inverted/inverted_index_fs_directory.h" +#include "storage/index/inverted/similarity/collection_statistics.h" +#include "storage/index/snii/compaction/eligibility.h" +#include "storage/index/snii/compaction/snii_index_compaction.h" +#include "storage/index/snii/writer/memory_reporter.h" #include "storage/olap_common.h" #include "storage/olap_define.h" #include "storage/rowset/beta_rowset.h" @@ -84,6 +91,7 @@ #include "storage/txn/txn_manager.h" #include "storage/utils.h" #include "util/pretty_printer.h" +#include "util/stopwatch.hpp" #include "util/time.h" #include "util/trace.h" @@ -119,6 +127,8 @@ bool should_enable_compaction_cache_index_only(bool write_file_cache, ReaderType namespace { +constexpr size_t kSniiCompactionReadAheadBudgetBytes = 64ULL << 20; + bool is_rowset_tidy(std::string& pre_max_key, bool& pre_rs_key_bounds_truncated, const RowsetSharedPtr& rhs) { size_t min_tidy_size = config::ordered_data_compaction_min_segment_size; @@ -822,6 +832,7 @@ Status Compaction::do_inverted_index_compaction() { } OlapStopWatch inverted_watch; + ThreadCpuStopWatch inverted_cpu_watch(true); // translation vec // <> @@ -847,7 +858,8 @@ Status Compaction::do_inverted_index_compaction() { if (dest_segment_num <= 0) { LOG(INFO) << "skip doing index compaction due to no output segments" << ". tablet=" << _tablet->tablet_id() << ", input row number=" << _input_row_num - << ". elapsed time=" << inverted_watch.get_elapse_second() << "s."; + << ". elapsed time=" << inverted_watch.get_elapse_second() + << "s. thread cpu time=" << inverted_cpu_watch.elapsed_time() / 1e9 << "s."; return Status::OK(); } @@ -913,6 +925,7 @@ Status Compaction::do_inverted_index_compaction() { // src index dirs std::vector> index_file_readers(src_segment_num); + std::vector source_rowsets(src_segment_num, nullptr); for (const auto& m : src_seg_to_id_map) { const auto& [rowset_id, seg_id] = m.first; @@ -977,6 +990,7 @@ Status Compaction::do_inverted_index_compaction() { _tablet->tablet_id(), rowset_id.to_string(), seg_id); } index_file_readers[m.second] = std::move(index_file_reader); + source_rowsets[m.second] = rowset; } // dest index files @@ -1006,6 +1020,15 @@ Status Compaction::do_inverted_index_compaction() { << ", destination index size=" << dest_segment_num << "."; Status status = Status::OK(); + std::shared_ptr snii_merge_memory_reporter; + std::unique_ptr validated_snii_rowid_conversion; + if (_cur_tablet_schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + const size_t spill_threshold = + static_cast(config::inverted_index_ram_buffer_size * 1024 * 1024); + snii_merge_memory_reporter = std::make_shared( + nullptr, spill_threshold, snii::writer::MemoryReporter::CapPolicy::kHardLimit); + } for (auto&& column_uniq_id : ctx.columns_to_do_index_compaction) { auto col = _cur_tablet_schema->column_by_uid(column_uniq_id); auto index_metas = _cur_tablet_schema->inverted_indexs(col); @@ -1021,6 +1044,152 @@ Status Compaction::do_inverted_index_compaction() { break; } for (const auto& index_meta : index_metas) { + if (_cur_tablet_schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + std::vector> source_indexes; + std::vector plan_sources; + std::vector eligibility_sources; + source_indexes.reserve(src_segment_num); + plan_sources.reserve(src_segment_num); + eligibility_sources.reserve(src_segment_num); + Status merge_status = Status::OK(); + for (size_t source_ordinal = 0; source_ordinal < src_segment_num; + ++source_ordinal) { + DORIS_CHECK(source_rowsets[source_ordinal] != nullptr); + const auto source_index_metas = + source_rowsets[source_ordinal]->tablet_schema()->inverted_indexs( + column_uniq_id); + const auto source_index_it = std::find_if( + source_index_metas.begin(), source_index_metas.end(), + [&index_meta](const TabletIndex* source_index) { + return source_index->index_id() == index_meta->index_id(); + }); + if (source_index_it == source_index_metas.end()) { + merge_status = Status::Error( + "source SNII index metadata disappeared after eligibility"); + break; + } + const TabletIndex* source_index_meta = *source_index_it; + auto source_index = index_file_readers[source_ordinal]->open_snii_index( + source_index_meta, nullptr, + snii::reader::LogicalIndexOpenMode::kCompaction); + if (!source_index.has_value()) { + merge_status = source_index.error(); + break; + } + source_indexes.push_back(std::move(source_index.value())); + plan_sources.push_back(source_indexes.back().get()); + eligibility_sources.push_back({.reader = std::cref(*source_indexes.back()), + .index_meta = std::cref(*source_index_meta)}); + } + + snii::compaction::SniiCompactionEligibility merge_eligibility; + if (merge_status.ok()) { + merge_status = snii::compaction::validate_snii_compaction_eligibility( + eligibility_sources, *index_meta, &merge_eligibility); + } + + if (merge_status.ok() && validated_snii_rowid_conversion == nullptr) { + std::vector source_segment_doc_counts; + source_segment_doc_counts.reserve(plan_sources.size()); + for (const snii::reader::LogicalIndexReader* source : plan_sources) { + DORIS_CHECK(source != nullptr); + if (source->stats().doc_count > std::numeric_limits::max()) { + merge_status = Status::Error( + "source doc count exceeds the SNII uint32 docid domain"); + break; + } + source_segment_doc_counts.push_back( + static_cast(source->stats().doc_count)); + } + if (merge_status.ok()) { + merge_status = snii::compaction::ValidatedRowIdConversion::create( + &trans_vec, source_segment_doc_counts, dest_segment_num_rows, + &validated_snii_rowid_conversion); + if (merge_status.ok()) { + DBUG_EXECUTE_IF("Compaction::snii_validated_rowid_conversion_created", + DBUG_RUN_CALLBACK()); + } + } + } + + std::unique_ptr merge_plan; + if (merge_status.ok()) { + DORIS_CHECK(validated_snii_rowid_conversion != nullptr); + merge_status = snii::compaction::SniiPlainT2MergePlan::prepare( + std::move(plan_sources), *validated_snii_rowid_conversion, + merge_eligibility, kSniiCompactionReadAheadBudgetBytes, + snii_merge_memory_reporter, &merge_plan); + } + + std::vector destination_sessions( + dest_segment_num, nullptr); + if (merge_status.ok()) { + DORIS_CHECK(snii_merge_memory_reporter != nullptr); + for (size_t destination_ordinal = 0; destination_ordinal < dest_segment_num; + ++destination_ordinal) { + DBUG_EXECUTE_IF("Compaction::before_add_snii_destination_session", + DBUG_RUN_CALLBACK(destination_ordinal, &merge_status)); + if (!merge_status.ok()) { + break; + } + auto* destination_writer = + inverted_index_file_writers[cast_set(destination_ordinal)] + .get(); + if (merge_eligibility.kind == + snii::compaction::SniiStreamedMergeKind::kCommonGramsT3) { + merge_status = destination_writer->add_snii_index_streamed( + index_meta, dest_segment_num_rows[destination_ordinal], + merge_plan->take_destination_null_docids(destination_ordinal), + merge_plan->take_destination_encoded_norms(destination_ordinal), + merge_plan->destination_common_grams_metadata( + destination_ordinal), + merge_plan->destination_common_grams_posting_policy(), + merge_plan->destination_index_config(), + snii_merge_memory_reporter, + &destination_sessions[destination_ordinal]); + } else { + merge_status = destination_writer->add_snii_index_streamed( + index_meta, dest_segment_num_rows[destination_ordinal], + merge_plan->take_destination_null_docids(destination_ordinal), + merge_plan->destination_index_config(), + snii_merge_memory_reporter, + &destination_sessions[destination_ordinal]); + } + if (!merge_status.ok()) { + break; + } + } + } + if (merge_status.ok()) { + DBUG_EXECUTE_IF("Compaction::before_execute_snii_merge", + DBUG_RUN_CALLBACK(&merge_status)); + } + if (!merge_status.ok()) { + for (size_t destination_ordinal = 0; + destination_ordinal < destination_sessions.size(); ++destination_ordinal) { + snii::writer::SniiStreamedIndexSession* session = + destination_sessions[destination_ordinal]; + if (session != nullptr) { + session->abort(merge_status); + DBUG_EXECUTE_IF("Compaction::snii_destination_session_aborted", + DBUG_RUN_CALLBACK(destination_ordinal)); + } + } + } else { + merge_status = merge_plan->execute(destination_sessions); + } + if (!merge_status.ok()) { + if (merge_status.is() || + merge_status.is() || + merge_status.is()) { + error_handler(index_meta->index_id(), column_uniq_id); + } + return merge_status; + } + continue; + } + std::vector dest_index_dirs(dest_segment_num); try { std::vector> src_idx_dirs( @@ -1084,7 +1253,8 @@ Status Compaction::do_inverted_index_compaction() { } LOG(INFO) << "succeed to do index compaction" << ". tablet=" << _tablet->tablet_id() - << ". elapsed time=" << inverted_watch.get_elapse_second() << "s."; + << ". elapsed time=" << inverted_watch.get_elapse_second() + << "s. thread cpu time=" << inverted_cpu_watch.elapsed_time() / 1e9 << "s."; return Status::OK(); } @@ -1222,6 +1392,157 @@ static bool check_rowset_has_inverted_index(const RowsetSharedPtr& src_rs, int32 } void Compaction::construct_index_compaction_columns(RowsetWriterContext& ctx) { + if (_cur_tablet_schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + std::map column_eligibility; + std::map> source_file_readers; + for (const auto& destination_index : _cur_tablet_schema->inverted_indexes()) { + const auto& col_unique_ids = destination_index->col_unique_ids(); + if (col_unique_ids.empty()) { + LOG(WARNING) << "tablet[" << _tablet->tablet_id() << "] index[" + << destination_index->index_id() + << "] has no column unique id, will rebuild its SNII index"; + continue; + } + const int32_t col_unique_id = col_unique_ids[0]; + if (!_cur_tablet_schema->has_column_unique_id(col_unique_id) || + !field_is_slice_type(_cur_tablet_schema->column_by_uid(col_unique_id).type())) { + continue; + } + + bool eligible = true; + size_t source_segment_count = 0; + for (const auto& rowset : _input_rowsets) { + source_segment_count += rowset->num_segments(); + } + if (source_segment_count == 0 || + source_segment_count > kSniiCompactionReadAheadBudgetBytes / + snii::compaction::SniiPlainT2MergePlan:: + kMinReadAheadBudgetPerSource) { + eligible = false; + } + + Status eligibility_status = + eligible ? Status::OK() + : Status::Error( + "source index unavailable or read-ahead budget gate failed"); + std::optional merge_eligibility; + size_t source_ordinal = 0; + + for (const auto& rowset : _input_rowsets) { + if (!eligible) { + break; + } + auto* beta_rowset = static_cast(rowset.get()); + if (beta_rowset->is_skip_index_compaction(col_unique_id)) { + eligible = false; + break; + } + const auto source_index_metas = + rowset->tablet_schema()->inverted_indexs(col_unique_id); + const auto source_index_it = std::find_if( + source_index_metas.begin(), source_index_metas.end(), + [&destination_index](const TabletIndex* source_index) { + return source_index->index_id() == destination_index->index_id(); + }); + if (source_index_it == source_index_metas.end()) { + eligible = false; + break; + } + const TabletIndex* source_index_meta = *source_index_it; + const auto fs = rowset->rowset_meta()->fs(); + if (fs == nullptr) { + eligible = false; + break; + } + + for (uint32_t segment_id = 0; segment_id < rowset->num_segments(); ++segment_id) { + auto segment_path = rowset->segment_path(segment_id); + if (!segment_path.has_value()) { + eligible = false; + break; + } + const std::string index_file_path_prefix = + std::string {InvertedIndexDescriptor::get_index_file_path_prefix( + segment_path.value())}; + auto source_file_reader_it = source_file_readers.find(index_file_path_prefix); + if (source_file_reader_it == source_file_readers.end()) { + auto source_file_reader = std::make_unique( + fs, index_file_path_prefix, InvertedIndexStorageFormatPB::SNII, + rowset->rowset_meta()->inverted_index_file_info(segment_id), + _tablet->tablet_id()); + const Status init_status = + source_file_reader->init(config::inverted_index_read_buffer_size); + if (!init_status.ok()) { + eligible = false; + break; + } + DBUG_EXECUTE_IF("Compaction::snii_eligibility_reader_initialized", + DBUG_RUN_CALLBACK()); + source_file_reader_it = source_file_readers + .emplace(index_file_path_prefix, + std::move(source_file_reader)) + .first; + } + auto source_index = source_file_reader_it->second->open_snii_index( + source_index_meta, nullptr, + snii::reader::LogicalIndexOpenMode::kCompaction); + if (!source_index.has_value()) { + eligible = false; + break; + } + if (!merge_eligibility.has_value()) { + const std::array source = { + snii::compaction::PlainT2CompactionSource { + .reader = std::cref(*source_index.value()), + .index_meta = std::cref(*source_index_meta)}}; + snii::compaction::SniiCompactionEligibility eligibility; + eligibility_status = snii::compaction::validate_snii_compaction_eligibility( + source, *destination_index, &eligibility); + if (eligibility_status.ok()) { + merge_eligibility = std::move(eligibility); + } + } else if (source_index_meta->index_id() != destination_index->index_id() || + source_index_meta->get_index_suffix() != + destination_index->get_index_suffix() || + source_index_meta->properties() != destination_index->properties()) { + eligibility_status = Status::Error( + "source SNII index identity or properties differ from destination"); + } else { + eligibility_status = snii::compaction::validate_snii_source_eligibility( + *source_index.value(), source_ordinal, *merge_eligibility); + } + if (!eligibility_status.ok()) { + eligible = false; + break; + } + ++source_ordinal; + } + } + + if (!eligible && eligibility_status.ok()) { + eligibility_status = Status::Error( + "source SNII index file or metadata is unavailable"); + } + auto [column_it, inserted] = column_eligibility.emplace(col_unique_id, eligible); + if (!inserted) { + column_it->second = column_it->second && eligible; + } + if (!eligible) { + LOG(INFO) << "tablet[" << _tablet->tablet_id() << "] index[" + << destination_index->index_id() + << "] is not eligible for SNII postings compaction; rebuild from raw " + "column. reason=" + << eligibility_status; + } + } + for (const auto& [col_unique_id, eligible] : column_eligibility) { + if (eligible) { + ctx.columns_to_do_index_compaction.insert(col_unique_id); + } + } + return; + } for (const auto& index : _cur_tablet_schema->inverted_indexes()) { auto col_unique_ids = index->col_unique_ids(); // check if column unique ids is empty to avoid crash diff --git a/be/src/storage/index/analyzer_key_matcher.cpp b/be/src/storage/index/analyzer_key_matcher.cpp index ef0bdddf75ddb1..81373fa529327f 100644 --- a/be/src/storage/index/analyzer_key_matcher.cpp +++ b/be/src/storage/index/analyzer_key_matcher.cpp @@ -17,8 +17,6 @@ #include "storage/index/analyzer_key_matcher.h" -#include "storage/index/inverted/inverted_index_iterator.h" - namespace doris::segment_v2 { AnalyzerMatchResult AnalyzerKeyMatcher::match( diff --git a/be/src/storage/index/analyzer_key_matcher.h b/be/src/storage/index/analyzer_key_matcher.h index 833193b259d592..e683c2d4dbb480 100644 --- a/be/src/storage/index/analyzer_key_matcher.h +++ b/be/src/storage/index/analyzer_key_matcher.h @@ -23,11 +23,15 @@ #include #include "storage/index/inverted/inverted_index_parser.h" +#include "storage/index/inverted/inverted_index_reader.h" namespace doris::segment_v2 { -// Forward declaration -struct ReaderEntry; +struct ReaderEntry { + InvertedIndexReaderType type; + std::string analyzer_key; + InvertedIndexReaderPtr reader; +}; // Result of analyzer key matching operation. // Contains candidate readers that match the requested analyzer key. diff --git a/be/src/storage/index/index_file_reader.cpp b/be/src/storage/index/index_file_reader.cpp index 348e1399421e5a..95dbee73397107 100644 --- a/be/src/storage/index/index_file_reader.cpp +++ b/be/src/storage/index/index_file_reader.cpp @@ -20,6 +20,8 @@ #include #include +#include "common/cast_set.h" +#include "common/config.h" #include "storage/index/inverted/inverted_index_compound_reader.h" #include "storage/index/inverted/inverted_index_fs_directory.h" #include "storage/tablet/tablet_schema.h" @@ -31,7 +33,9 @@ Status IndexFileReader::init(int32_t read_buffer_size, const io::IOContext* io_c std::unique_lock lock(_mutex); // Lock for writing if (!_inited) { _read_buffer_size = read_buffer_size; - if (_storage_format >= InvertedIndexStorageFormatPB::V2) { + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + RETURN_IF_ERROR(_init_snii(io_ctx)); + } else if (_storage_format >= InvertedIndexStorageFormatPB::V2) { RETURN_IF_ERROR(_init_from(read_buffer_size, io_ctx)); } _inited = true; @@ -136,7 +140,51 @@ Status IndexFileReader::_init_from(int32_t read_buffer_size, const io::IOContext return Status::OK(); } +Status IndexFileReader::_init_snii(const io::IOContext* io_ctx) { + auto index_file_full_path = InvertedIndexDescriptor::get_index_file_path_v2(_index_path_prefix); + int64_t file_size = -1; + if (_idx_file_info.has_index_size()) { + file_size = _idx_file_info.index_size(); + } + file_size = file_size == 0 ? -1 : file_size; + + io::FileReaderOptions opts; + // DIAGNOSTIC: inverted_index_read_bypass_file_cache forces NO_CACHE (precise S3 + // range GETs) for the SNII index reader only, without touching the global + // enable_file_cache (cloud mode requires the latter). Default path keeps the + // block cache. + opts.cache_type = (config::enable_file_cache && !config::inverted_index_read_bypass_file_cache) + ? io::FileCachePolicy::FILE_BLOCK_CACHE + : io::FileCachePolicy::NO_CACHE; + opts.is_doris_table = true; + opts.file_size = file_size; + opts.tablet_id = _tablet_id; + io::FileReaderSPtr reader; + RETURN_IF_ERROR(_fs->open_file(index_file_full_path, &reader, &opts)); + // With NO_CACHE on a remote filesystem there is no CachedRemoteFileReader to + // account physical remote bytes, so the adapter must count its own reads. + const bool direct_remote_io = opts.cache_type == io::FileCachePolicy::NO_CACHE && + _fs->type() != io::FileSystemType::LOCAL; + _snii_file_reader = std::make_shared( + std::move(reader), /*io_ctx=*/nullptr, direct_remote_io); + _snii_segment_reader = std::make_unique(); + io::IOContext meta_io_ctx; + if (io_ctx != nullptr) { + meta_io_ctx = *io_ctx; + } + meta_io_ctx.is_inverted_index = true; + meta_io_ctx.is_index_data = true; + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(&meta_io_ctx); + RETURN_IF_ERROR(doris::snii::reader::SniiSegmentReader::open(_snii_file_reader.get(), + _snii_segment_reader.get())); + return Status::OK(); +} + Result IndexFileReader::get_all_directories() { + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + return ResultError(Status::Error( + "SNII format does not expose CLucene directories")); + } InvertedIndexDirectoryMap res; std::shared_lock lock(_mutex); // Lock for reading for (auto& [index, _] : _indices_entries) { @@ -155,6 +203,11 @@ Result> IndexFileReader:: int64_t index_id, const std::string& index_suffix, const io::IOContext* io_ctx) const { std::unique_ptr compound_reader; + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + return ResultError(Status::Error( + "SNII format does not open CLucene compound readers")); + } + if (_storage_format == InvertedIndexStorageFormatPB::V1) { auto index_file_path = InvertedIndexDescriptor::get_index_file_path_v1( _index_path_prefix, index_id, index_suffix); @@ -229,6 +282,35 @@ Result> IndexFileReader:: return compound_reader; } +Result> IndexFileReader::open_snii_index( + const TabletIndex* index_meta, const io::IOContext* io_ctx, + doris::snii::reader::LogicalIndexOpenMode open_mode) const { + DCHECK(_storage_format == InvertedIndexStorageFormatPB::SNII); + std::shared_lock lock(_mutex); + if (_snii_segment_reader == nullptr) { + return ResultError(Status::Error( + "SNII index file {} is not opened", + InvertedIndexDescriptor::get_index_file_path_v2(_index_path_prefix))); + } + io::IOContext meta_io_ctx; + if (io_ctx != nullptr) { + meta_io_ctx = *io_ctx; + } + meta_io_ctx.is_inverted_index = true; + meta_io_ctx.is_index_data = true; + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(&meta_io_ctx); + + auto logical_reader = std::make_unique(); + auto status = _snii_segment_reader->open_index(cast_set(index_meta->index_id()), + index_meta->get_index_suffix(), + logical_reader.get(), open_mode); + auto doris_status = status; + if (!doris_status.ok()) { + return ResultError(doris_status); + } + return logical_reader; +} + Result> IndexFileReader::open( const TabletIndex* index_meta, const io::IOContext* io_ctx) const { auto index_id = index_meta->index_id(); @@ -254,6 +336,19 @@ Status IndexFileReader::index_file_exist(const TabletIndex* index_meta, bool* re auto index_file_path = InvertedIndexDescriptor::get_index_file_path_v1( _index_path_prefix, index_meta->index_id(), index_meta->get_index_suffix()); return _fs->exists(index_file_path, res); + } else if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + auto index_file_path = InvertedIndexDescriptor::get_index_file_path_v2(_index_path_prefix); + RETURN_IF_ERROR(_fs->exists(index_file_path, res)); + if (!*res) { + return Status::OK(); + } + std::shared_lock lock(_mutex); + if (_snii_segment_reader == nullptr) { + *res = false; + return Status::OK(); + } + return _snii_segment_reader->index_exists(cast_set(index_meta->index_id()), + index_meta->get_index_suffix(), res); } else { std::shared_lock lock(_mutex); // Lock for reading if (_stream == nullptr) { @@ -279,6 +374,25 @@ Status IndexFileReader::has_null(const TabletIndex* index_meta, bool* res) const *res = true; return Status::OK(); } + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + std::shared_lock lock(_mutex); + if (_snii_segment_reader == nullptr) { + return Status::Error( + "SNII index file {} is not opened", + InvertedIndexDescriptor::get_index_file_path_v2(_index_path_prefix)); + } + io::IOContext meta_io_ctx; + meta_io_ctx.is_inverted_index = true; + meta_io_ctx.is_index_data = true; + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(&meta_io_ctx); + + doris::snii::format::SectionRefs section_refs; + RETURN_IF_ERROR(_snii_segment_reader->section_refs_for_index( + cast_set(index_meta->index_id()), index_meta->get_index_suffix(), + §ion_refs)); + *res = section_refs.null_bitmap.length > 0; + return Status::OK(); + } std::shared_lock lock(_mutex); // Lock for reading if (_stream == nullptr) { return Status::Error( diff --git a/be/src/storage/index/index_file_reader.h b/be/src/storage/index/index_file_reader.h index fb4ec2b9a62fe3..f09aa9f8476503 100644 --- a/be/src/storage/index/index_file_reader.h +++ b/be/src/storage/index/index_file_reader.h @@ -35,6 +35,9 @@ #include "io/fs/file_system.h" #include "storage/index/index_file_writer.h" #include "storage/index/inverted/inverted_index_desc.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/snii_doris_adapter.h" namespace doris { class TabletIndex; @@ -60,7 +63,7 @@ class IndexFileReader { : _fs(std::move(fs)), _index_path_prefix(std::move(index_path_prefix)), _storage_format(storage_format), - _idx_file_info(idx_file_info), + _idx_file_info(std::move(idx_file_info)), _tablet_id(tablet_id) {} virtual ~IndexFileReader() = default; @@ -68,6 +71,10 @@ class IndexFileReader { const io::IOContext* io_ctx = nullptr); MOCK_FUNCTION Result> open( const TabletIndex* index_meta, const io::IOContext* io_ctx = nullptr) const; + Result> open_snii_index( + const TabletIndex* index_meta, const io::IOContext* io_ctx = nullptr, + doris::snii::reader::LogicalIndexOpenMode open_mode = + doris::snii::reader::LogicalIndexOpenMode::kQuery) const; void debug_file_entries(); std::string get_index_file_cache_key(const TabletIndex* index_meta) const; std::string get_index_file_path(const TabletIndex* index_meta) const; @@ -75,12 +82,19 @@ class IndexFileReader { Status has_null(const TabletIndex* index_meta, bool* res) const; Result get_all_directories(); // open file v2, init _stream - int64_t get_inverted_file_size() const { return _stream == nullptr ? 0 : _stream->length(); } + int64_t get_inverted_file_size() const { + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + return _snii_file_reader == nullptr ? 0 : _snii_file_reader->size(); + } + return _stream == nullptr ? 0 : _stream->length(); + } const std::string& get_index_path_prefix() const { return _index_path_prefix; } + InvertedIndexStorageFormatPB get_storage_format() const { return _storage_format; } friend IndexFileWriter; protected: Status _init_from(int32_t read_buffer_size, const io::IOContext* io_ctx); + Status _init_snii(const io::IOContext* io_ctx); Result> _open( int64_t index_id, const std::string& index_suffix, const io::IOContext* io_ctx = nullptr) const; @@ -88,6 +102,8 @@ class IndexFileReader { private: IndicesEntriesMap _indices_entries; std::unique_ptr _stream = nullptr; + std::shared_ptr _snii_file_reader; + std::unique_ptr _snii_segment_reader; const io::FileSystemSPtr _fs; std::string _index_path_prefix; int32_t _read_buffer_size = -1; diff --git a/be/src/storage/index/index_file_writer.cpp b/be/src/storage/index/index_file_writer.cpp index afd09c84620bb5..3cc6b4bd495096 100644 --- a/be/src/storage/index/index_file_writer.cpp +++ b/be/src/storage/index/index_file_writer.cpp @@ -19,12 +19,14 @@ #include +#include #include #include +#include "common/cast_set.h" +#include "common/config.h" #include "common/status.h" #include "io/fs/packed_file_writer.h" -#include "io/fs/s3_file_writer.h" #include "io/fs/stream_sink_file_writer.h" #include "storage/index/ann/ann_index_files.h" #include "storage/index/index_file_reader.h" @@ -34,10 +36,53 @@ #include "storage/index/inverted/inverted_index_desc.h" #include "storage/index/inverted/inverted_index_fs_directory.h" #include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/snii/snii_doris_adapter.h" #include "storage/tablet/tablet_schema.h" namespace doris::segment_v2 { +// Resolves whether one segment index lays out freq regions (G16-c). Freq +// serves ONLY BM25 scoring: a scoring config always keeps it; a plain +// positions config keeps it only when the escape-hatch config asks for the +// full T2 layout. NOT in the anonymous namespace on purpose -- the UT covers +// this production policy line directly (a flipped operator or inverted flag +// here would otherwise stay green: no BE test drives add_snii_index). +bool snii_effective_write_freq(doris::snii::format::IndexConfig index_config) { + return doris::snii::format::has_scoring(index_config) || + config::snii_positions_index_write_freq; +} + +// Shared write-parameter resolution for one SNII index flush; `input->config` +// must already be set. BOTH the build path (add_snii_index) and the T2.2 +// compaction-merge streamed session resolve through this single helper, so the +// merge fast path can never drift from the rebuild contract (the T2 semantic +// golden invariant depends on parameter parity). NOT in the anonymous +// namespace on purpose -- the UT pins the resolved values directly. +void snii_resolve_index_write_params(bool is_direct_load, + doris::snii::writer::SniiIndexInput* input) { + // G16-c: freq regions serve only BM25 scoring; a plain positions index + // drops them unless the escape hatch asks for the full T2 layout. + input->write_freq = snii_effective_write_freq(input->config); + // G16-h: zstd levels. dict blocks accept zstd's full sane range; the prx + // level floor is 3 because the writer passes -level into the prx builders + // and -1 is the historic "auto at default level 3" sentinel -- a + // configured level 1 would silently resolve to 3 anyway (levels 1-2 buy + // nothing over 3 on these payloads). + input->dict_block_zstd_level = std::clamp(config::snii_dict_block_zstd_level, 1, 19); + // Patch C prx tiering: a direct load compresses prx at the cheaper load + // level; compaction / schema change / ADD INDEX keep snii_prx_zstd_level + // and compaction rewrites every segment with it, so settled segments (and + // the cold-query path over them) are byte-for-byte unaffected. + input->prx_zstd_level = std::clamp( + is_direct_load ? config::snii_prx_zstd_level_direct_load : config::snii_prx_zstd_level, + 3, 19); + // G16-d: dict block size experiment knob; <= 0 keeps the format default. + if (config::snii_target_dict_block_bytes > 0) { + input->target_dict_block_bytes = + static_cast(config::snii_target_dict_block_bytes); + } +} + IndexFileWriter::IndexFileWriter(io::FileSystemSPtr fs, std::string index_path_prefix, std::string rowset_id, int64_t seg_id, InvertedIndexStorageFormatPB storage_format, @@ -56,7 +101,7 @@ IndexFileWriter::IndexFileWriter(io::FileSystemSPtr fs, std::string index_path_p _tmp_dir = tmp_file_dir.native(); if (_storage_format == InvertedIndexStorageFormatPB::V1) { _index_storage_format = std::make_unique(this); - } else { + } else if (_storage_format != InvertedIndexStorageFormatPB::SNII) { _index_storage_format = std::make_unique(this); } } @@ -84,6 +129,10 @@ Status IndexFileWriter::_insert_directory_into_map(int64_t index_id, } Result> IndexFileWriter::open(const TabletIndex* index_meta) { + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + return ResultError(Status::Error( + "SNII format does not open CLucene directories")); + } auto local_fs_index_path = InvertedIndexDescriptor::get_temporary_index_path( _tmp_dir, _rowset_id, _seg_id, index_meta->index_id(), index_meta->get_index_suffix()); auto dir = std::shared_ptr(DorisFSDirectoryFactory::getDirectory( @@ -97,6 +146,121 @@ Result> IndexFileWriter::open(const TabletInde return dir; } +Status IndexFileWriter::add_snii_index(const TabletIndex* index_meta, uint32_t doc_count, + std::vector null_docids, + doris::snii::writer::SpimiTermBuffer* const term_buffer, + doris::snii::format::IndexConfig index_config, + SniiAddIndexOptions options, + doris::snii::writer::MemoryReporter* const mem_reporter) { + DCHECK(_storage_format == InvertedIndexStorageFormatPB::SNII); + DCHECK(index_meta != nullptr); + DCHECK(term_buffer != nullptr); + if (_idx_v2_writer == nullptr) { + return Status::Error( + "SNII index file writer is null for {}", _index_path_prefix); + } + if (_snii_file_writer == nullptr) { + _snii_file_writer = std::make_unique(_idx_v2_writer.get()); + _snii_compound_writer = + std::make_unique(_snii_file_writer.get()); + } + + doris::snii::writer::SniiIndexInput input; + input.index_id = cast_set(index_meta->index_id()); + input.index_suffix = index_meta->get_index_suffix(); + input.config = index_config; + input.doc_count = doc_count; + input.null_docids = std::move(null_docids); + input.encoded_norms = std::move(options.encoded_norms); + input.common_grams_metadata = std::move(options.common_grams_metadata); + input.common_grams_posting_policy = options.common_grams_posting_policy; + input.term_source = term_buffer; + input.mem_reporter = mem_reporter; + snii_resolve_index_write_params(options.is_direct_load, &input); + RETURN_IF_ERROR(_snii_compound_writer->add_logical_index(input)); + ++_snii_index_count; + return Status::OK(); +} + +Status IndexFileWriter::add_snii_index_streamed( + const TabletIndex* index_meta, uint32_t doc_count, + doris::snii::writer::TrackedNullDocids null_docids, + doris::snii::format::IndexConfig index_config, + std::shared_ptr mem_reporter, + doris::snii::writer::SniiStreamedIndexSession** session) { + return add_snii_index_streamed(index_meta, doc_count, std::move(null_docids), + doris::snii::writer::TrackedEncodedNorms(std::vector()), + std::nullopt, + doris::snii::format::CommonGramsPostingPolicy::kNone, + index_config, std::move(mem_reporter), session); +} + +Status IndexFileWriter::add_snii_index_streamed( + const TabletIndex* index_meta, uint32_t doc_count, + doris::snii::writer::TrackedNullDocids null_docids, + doris::snii::writer::TrackedEncodedNorms encoded_norms, + std::optional common_grams_metadata, + doris::snii::format::CommonGramsPostingPolicy common_grams_posting_policy, + doris::snii::format::IndexConfig index_config, + std::shared_ptr mem_reporter, + doris::snii::writer::SniiStreamedIndexSession** session) { + DCHECK(_storage_format == InvertedIndexStorageFormatPB::SNII); + DCHECK(index_meta != nullptr); + if (session == nullptr) { + return Status::Error( + "SNII streamed session out parameter is null for {}", _index_path_prefix); + } + *session = nullptr; + if (_idx_v2_writer == nullptr) { + return Status::Error( + "SNII index file writer is null for {}", _index_path_prefix); + } + const bool has_scoring = doris::snii::format::has_scoring(index_config); + const bool valid_scoring_shape = + has_scoring ? common_grams_metadata.has_value() && encoded_norms.size() == doc_count + : !common_grams_metadata.has_value() && encoded_norms.empty(); + if (!valid_scoring_shape) { + return Status::InternalError( + "SNII streamed merge scoring shape disagrees with eligibility for {}", + _index_path_prefix); + } + if (_snii_file_writer == nullptr) { + _snii_file_writer = std::make_unique(_idx_v2_writer.get()); + _snii_compound_writer = + std::make_unique(_snii_file_writer.get()); + } + + doris::snii::writer::SniiIndexInput input; + input.index_id = cast_set(index_meta->index_id()); + input.index_suffix = index_meta->get_index_suffix(); + input.config = index_config; + input.doc_count = doc_count; + input.mem_reporter = mem_reporter.get(); + input.common_grams_metadata = std::move(common_grams_metadata); + input.common_grams_posting_policy = common_grams_posting_policy; + // Merge output is always the settled-segment shape: COMPACTION prx level. + snii_resolve_index_write_params(/*is_direct_load=*/false, &input); + if (mem_reporter != nullptr) { + constexpr uint64_t kMaxStreamedDictResidentBytes = 64ULL << 20; + DORIS_CHECK_GE(mem_reporter->cap_bytes(), 8); + input.dict_resident_cap_bytes = + std::min(kMaxStreamedDictResidentBytes, mem_reporter->cap_bytes() / 8); + } + RETURN_IF_ERROR(_snii_compound_writer->begin_streamed_index( + std::move(input), std::move(null_docids), std::move(encoded_norms), session)); + if (mem_reporter != nullptr) { + _snii_memory_reporters.push_back(std::move(mem_reporter)); + } + ++_snii_index_count; + return Status::OK(); +} + +void IndexFileWriter::retain_snii_memory_reporter( + std::unique_ptr mem_reporter) { + DCHECK(mem_reporter != nullptr); + _snii_memory_reporters.emplace_back(std::move(mem_reporter)); +} + Status IndexFileWriter::delete_index(const TabletIndex* index_meta) { DBUG_EXECUTE_IF("IndexFileWriter::delete_index_index_meta_nullptr", { index_meta = nullptr; }); if (!index_meta) { @@ -123,6 +287,9 @@ Status IndexFileWriter::delete_index(const TabletIndex* index_meta) { } Status IndexFileWriter::add_into_searcher_cache() { + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + return Status::OK(); + } auto index_file_reader = std::make_unique( _fs, _index_path_prefix, _storage_format, InvertedIndexFileInfo(), _tablet_id); auto st = index_file_reader->init(); @@ -196,13 +363,22 @@ Result> IndexFileWriter::_construct_index_ Status IndexFileWriter::begin_close() { DCHECK(!_closed) << debug_string(); _closed = true; - if (_indices_dirs.empty()) { - // An empty file must still be created even if there are no indexes to write - if (dynamic_cast(_idx_v2_writer.get()) != nullptr || - dynamic_cast(_idx_v2_writer.get()) != nullptr || - dynamic_cast(_idx_v2_writer.get()) != nullptr) { - return _idx_v2_writer->close(true); + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + if (_snii_compound_writer == nullptr) { + if (_idx_v2_writer == nullptr) { + return Status::OK(); + } + _snii_file_writer = + std::make_unique(_idx_v2_writer.get()); + _snii_compound_writer = std::make_unique( + _snii_file_writer.get()); } + RETURN_IF_ERROR(_snii_compound_writer->finish()); + _total_file_size = _idx_v2_writer->bytes_appended(); + _file_info.set_index_size(_total_file_size); + return _idx_v2_writer->close(true); + } + if (_indices_dirs.empty() && _storage_format == InvertedIndexStorageFormatPB::V1) { return Status::OK(); } DBUG_EXECUTE_IF("inverted_index_storage_format_must_be_v2", { @@ -238,15 +414,15 @@ Status IndexFileWriter::begin_close() { Status IndexFileWriter::finish_close() { DCHECK(_closed) << debug_string(); - if (_indices_dirs.empty()) { - // An empty file must still be created even if there are no indexes to write - if (dynamic_cast(_idx_v2_writer.get()) != nullptr || - dynamic_cast(_idx_v2_writer.get()) != nullptr || - dynamic_cast(_idx_v2_writer.get()) != nullptr) { - return _idx_v2_writer->close(false); + if (_storage_format == InvertedIndexStorageFormatPB::SNII) { + if (_idx_v2_writer != nullptr && _idx_v2_writer->state() != io::FileWriter::State::CLOSED) { + RETURN_IF_ERROR(_idx_v2_writer->close(false)); } return Status::OK(); } + if (_indices_dirs.empty() && _storage_format == InvertedIndexStorageFormatPB::V1) { + return Status::OK(); + } if (_idx_v2_writer != nullptr && _idx_v2_writer->state() != io::FileWriter::State::CLOSED) { RETURN_IF_ERROR(_idx_v2_writer->close(false)); } diff --git a/be/src/storage/index/index_file_writer.h b/be/src/storage/index/index_file_writer.h index a303de8b68c156..0c2e8a1734917b 100644 --- a/be/src/storage/index/index_file_writer.h +++ b/be/src/storage/index/index_file_writer.h @@ -22,23 +22,38 @@ #include #include +#include #include #include +#include #include "common/be_mock_util.h" #include "io/fs/file_system.h" #include "io/fs/file_writer.h" #include "io/fs/local_file_system.h" #include "storage/index/index_storage_format.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" #include "storage/index/inverted/inverted_index_common.h" #include "storage/index/inverted/inverted_index_compound_reader.h" #include "storage/index/inverted/inverted_index_searcher.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" + +namespace doris::snii::writer { +class MemoryReporter; +class SpimiTermBuffer; +class SniiCompoundWriter; +} // namespace doris::snii::writer namespace doris { class TabletIndex; namespace segment_v2 { class DorisFSDirectory; +namespace snii_doris { +class DorisSniiFileWriter; +} // namespace snii_doris using InvertedIndexDirectoryMap = std::map, std::shared_ptr>; @@ -55,6 +70,55 @@ class IndexFileWriter { virtual ~IndexFileWriter() = default; MOCK_FUNCTION Result> open(const TabletIndex* index_meta); + // Write-path facts for one SNII index flush. + struct SniiAddIndexOptions { + // This flush serves a stream/broker load (DataWriteType::TYPE_DIRECT): + // the prx region compresses at snii_prx_zstd_level_direct_load; + // compaction / schema change / ADD INDEX keep snii_prx_zstd_level. + bool is_direct_load = false; + // Present only for a CommonGrams writer that has a complete immutable + // capability identity. These are semantic BM25 inputs; physical TTF is + // still derived from every emitted unigram and gram posting. + std::vector encoded_norms; + std::optional common_grams_metadata; + snii::format::CommonGramsPostingPolicy common_grams_posting_policy = + snii::format::CommonGramsPostingPolicy::kNone; + }; + Status add_snii_index(const TabletIndex* index_meta, uint32_t doc_count, + std::vector null_docids, + doris::snii::writer::SpimiTermBuffer* const term_buffer, + doris::snii::format::IndexConfig index_config, + SniiAddIndexOptions options, + doris::snii::writer::MemoryReporter* const mem_reporter); + // T2.2 compaction index merge fast path: begins a STREAMED SNII index + // session on this compound. Unlike add_snii_index (which drains a SPIMI + // term buffer), the caller pushes pre-merged, lexicographically sorted + // terms through *session and seals the index with (*session)->finish(). + // Write parameters resolve through the SAME helper as add_snii_index + // (write_freq / zstd levels / dict block size), always at the COMPACTION + // prx tier (a merge is never a direct load). CommonGrams T3 callers transfer + // a precharged destination norm vector and a validated static metadata seed; + // the streamed session late-binds semantic token_count before finish. Only ONE + // session may be active per compound at a time, and begin_close() with an + // unfinished session fails instead of sealing a half-fed container. The + // handle is owned by this writer and valid until it is destroyed. + Status add_snii_index_streamed( + const TabletIndex* index_meta, uint32_t doc_count, + doris::snii::writer::TrackedNullDocids null_docids, + doris::snii::format::IndexConfig index_config, + std::shared_ptr mem_reporter, + doris::snii::writer::SniiStreamedIndexSession** session); + Status add_snii_index_streamed( + const TabletIndex* index_meta, uint32_t doc_count, + doris::snii::writer::TrackedNullDocids null_docids, + doris::snii::writer::TrackedEncodedNorms encoded_norms, + std::optional common_grams_metadata, + doris::snii::format::CommonGramsPostingPolicy common_grams_posting_policy, + doris::snii::format::IndexConfig index_config, + std::shared_ptr mem_reporter, + doris::snii::writer::SniiStreamedIndexSession** session); + void retain_snii_memory_reporter( + std::unique_ptr mem_reporter); Status delete_index(const TabletIndex* index_meta); Status initialize(InvertedIndexDirectoryMap& indices_dirs); Status add_into_searcher_cache(); @@ -113,6 +177,10 @@ class IndexFileWriter { IndexStorageFormatPtr _index_storage_format; int64_t _tablet_id = -1; + std::unique_ptr _snii_file_writer; + std::vector> _snii_memory_reporters; + std::unique_ptr _snii_compound_writer; + size_t _snii_index_count = 0; friend class IndexStorageFormatV1; friend class IndexStorageFormatV2; diff --git a/be/src/storage/index/index_query_context.h b/be/src/storage/index/index_query_context.h index 92fb698f4d3767..14106db4a5b7b7 100644 --- a/be/src/storage/index/index_query_context.h +++ b/be/src/storage/index/index_query_context.h @@ -18,7 +18,7 @@ #pragma once #include "storage/compaction/collection_similarity.h" -#include "storage/compaction/collection_statistics.h" +#include "storage/index/inverted/similarity/collection_statistics.h" namespace doris::segment_v2 { @@ -32,6 +32,24 @@ struct IndexQueryContext { size_t query_limit = 0; bool is_asc = false; + + // G02 count-only fast-path handshake. Set by SegmentIterator ONLY while it + // evaluates the single pushed-down MATCH predicate of a COUNT_ON_INDEX scan + // whose row space is provably unfiltered (no deletes, no other conjuncts, + // full row bitmap, no row-id consumers -- see count_on_index_fastpath.h), + // and reset immediately after. When set, an index reader MAY answer the + // query with a bitmap whose CARDINALITY equals the match count without the + // row ids being real (SNII returns [0, df) straight from dict-entry df, + // skipping the posting decode). Readers must never cache such a bitmap + // under a key a row-accurate query could hit. + bool count_on_index_fastpath = false; + // G03 reply direction of the same handshake. Set by a reader iff it DID + // answer with such a fabricated count bitmap (never on a query-cache hit, + // a single-flight shared result, or any row-accurate decode). Read and + // reset by SegmentIterator right after the index apply; a true value is + // the precondition for the count-emission shortcut that materializes the + // remaining count as default rows without iterating the row bitmap. + bool count_on_index_fastpath_hit = false; }; using IndexQueryContextPtr = std::shared_ptr; diff --git a/be/src/storage/index/index_writer.cpp b/be/src/storage/index/index_writer.cpp index 2325d280471337..6fb23c3c107e51 100644 --- a/be/src/storage/index/index_writer.cpp +++ b/be/src/storage/index/index_writer.cpp @@ -18,6 +18,7 @@ #include "common/exception.h" #include "storage/index/ann/ann_index_writer.h" #include "storage/index/inverted/inverted_index_writer.h" +#include "storage/index/snii/snii_index_writer.h" #include "storage/tablet/tablet_schema.h" #include "storage/types.h" @@ -80,6 +81,22 @@ Status IndexColumnWriter::create(const TabletColumn* column, } } + if (storage_format == InvertedIndexStorageFormatPB::SNII) { + if (!is_string_type(type)) { + return Status::Error( + "SNII inverted index storage format does not support BKD index type {}", + type); + } + *res = std::make_unique(index_file_writer, index_meta, + single_field); + auto st = (*res)->init(); + if (!st.ok()) { + (*res)->close_on_error(); + return st; + } + return Status::OK(); + } + DBUG_EXECUTE_IF("InvertedIndexColumnWriter::create_unsupported_type_for_inverted_index", { type = FieldType::OLAP_FIELD_TYPE_JSONB; }) switch (type) { diff --git a/be/src/storage/index/index_writer.h b/be/src/storage/index/index_writer.h index a0760f99000fcb..def7a8fc4c72cc 100644 --- a/be/src/storage/index/index_writer.h +++ b/be/src/storage/index/index_writer.h @@ -59,6 +59,15 @@ class IndexColumnWriter { virtual Status add_nulls(uint32_t count) = 0; virtual Status add_array_nulls(const uint8_t* null_map, size_t num_rows) = 0; + // Write-path hint from the segment writer: this writer serves a direct load + // (stream/broker load, DataWriteType::TYPE_DIRECT) as opposed to compaction + // / schema change / index build. The column writer forwards it + // unconditionally right after create() succeeds (i.e. after init(), before + // any value is added); creation paths outside the column writer (e.g. ADD + // INDEX in index_builder.cpp) never call it and keep the non-direct default. + // Default no-op: SNII uses it to select the direct-load PRX zstd level. + virtual void set_direct_load(bool /*is_direct_load*/) {} + virtual Status finish() = 0; virtual int64_t size() const = 0; diff --git a/be/src/storage/index/inverted/abstract_analysis_factory.h b/be/src/storage/index/inverted/abstract_analysis_factory.h index f9afc42f19db93..3452f93fef7cdf 100644 --- a/be/src/storage/index/inverted/abstract_analysis_factory.h +++ b/be/src/storage/index/inverted/abstract_analysis_factory.h @@ -21,6 +21,19 @@ namespace doris::segment_v2::inverted_index { +enum class AnalysisPurpose { + kIndex, + kSniiTransientIndex, + kPlainQuery, + kExactPhraseQuery, + kPhrasePrefixQuery, +}; + +enum class PositionCapability { + kUnknown, + kAlwaysUnitIncrement, +}; + class AbstractAnalysisFactory { public: virtual ~AbstractAnalysisFactory() = default; diff --git a/be/src/storage/index/inverted/analysis_factory_mgr.cpp b/be/src/storage/index/inverted/analysis_factory_mgr.cpp index a208a275bda640..d41960fad206a9 100644 --- a/be/src/storage/index/inverted/analysis_factory_mgr.cpp +++ b/be/src/storage/index/inverted/analysis_factory_mgr.cpp @@ -21,6 +21,7 @@ #include "storage/index/inverted/char_filter/empty_char_filter_factory.h" #include "storage/index/inverted/char_filter/icu_normalizer_char_filter_factory.h" #include "storage/index/inverted/token_filter/ascii_folding_filter_factory.h" +#include "storage/index/inverted/token_filter/common_grams_filter_factory.h" #include "storage/index/inverted/token_filter/empty_token_filter_factory.h" #include "storage/index/inverted/token_filter/icu_normalizer_filter_factory.h" #include "storage/index/inverted/token_filter/lower_case_filter_factory.h" @@ -82,6 +83,8 @@ void AnalysisFactoryMgr::initialise() { "pinyin", []() { return std::make_shared(); }); registerFactory( "icu_normalizer", []() { return std::make_shared(); }); + registerFactory( + "common_grams", []() { return std::make_shared(); }); }); } diff --git a/be/src/storage/index/inverted/analyzer/analyzer.cpp b/be/src/storage/index/inverted/analyzer/analyzer.cpp index 6d16613bcd9b00..c9735ee032b4f7 100644 --- a/be/src/storage/index/inverted/analyzer/analyzer.cpp +++ b/be/src/storage/index/inverted/analyzer/analyzer.cpp @@ -40,8 +40,27 @@ #include "storage/index/inverted/analyzer/icu/icu_analyzer.h" #include "storage/index/inverted/analyzer/ik/IKAnalyzer.h" #include "storage/index/inverted/char_filter/char_replace_char_filter_factory.h" +#include "storage/index/inverted/token_filter/common_grams_filter.h" namespace doris::segment_v2::inverted_index { +namespace { + +class BuiltinAnalyzerProvider final : public AnalyzerProvider { +public: + explicit BuiltinAnalyzerProvider(InvertedIndexAnalyzerConfig config) + : _analyzer(InvertedIndexAnalyzer::create_builtin_analyzer( + config.analyzer_name.empty() + ? config.parser_type + : get_inverted_index_parser_type_from_string(config.analyzer_name), + config.parser_mode, config.lower_case, config.stop_words)) {} + + AnalyzerPtr get_analyzer(AnalysisPurpose) const override { return _analyzer; } + +private: + const AnalyzerPtr _analyzer; +}; + +} // namespace ReaderPtr InvertedIndexAnalyzer::create_reader(const CharFilterMap& char_filter_map) { ReaderPtr reader = std::make_shared>(); @@ -130,32 +149,43 @@ AnalyzerPtr InvertedIndexAnalyzer::create_builtin_analyzer(InvertedIndexParserTy } AnalyzerPtr InvertedIndexAnalyzer::create_analyzer(const InvertedIndexAnalyzerConfig* config) { - DCHECK(config != nullptr); - const std::string& analyzer_name = config->analyzer_name; - - // Handle empty analyzer name - use builtin analyzer based on parser_type. - // This is the common case when user does not specify USING ANALYZER. - if (analyzer_name.empty()) { - return create_builtin_analyzer(config->parser_type, config->parser_mode, config->lower_case, - config->stop_words); - } + return create_analyzer(config, AnalysisPurpose::kPlainQuery); +} - // Check if it's a builtin analyzer name (english, chinese, standard, etc.) - if (is_builtin_analyzer(analyzer_name)) { - InvertedIndexParserType parser_type = - get_inverted_index_parser_type_from_string(analyzer_name); +AnalyzerPtr InvertedIndexAnalyzer::create_analyzer(const InvertedIndexAnalyzerConfig* config, + AnalysisPurpose purpose) { + DCHECK(config != nullptr); + if (config->analyzer_name.empty() || is_builtin_analyzer(config->analyzer_name)) { + const InvertedIndexParserType parser_type = + config->analyzer_name.empty() + ? config->parser_type + : get_inverted_index_parser_type_from_string(config->analyzer_name); return create_builtin_analyzer(parser_type, config->parser_mode, config->lower_case, config->stop_words); } - // Custom analyzer - look up in policy manager auto* index_policy_mgr = doris::ExecEnv::GetInstance()->index_policy_mgr(); - if (!index_policy_mgr) { + if (index_policy_mgr == nullptr) { throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, "Index policy manager is not initialized"); } + return index_policy_mgr->get_analyzer_by_name(config->analyzer_name, purpose); +} + +AnalyzerProviderPtr InvertedIndexAnalyzer::create_analyzer_provider( + const InvertedIndexAnalyzerConfig* config) { + DCHECK(config != nullptr); + if (config->analyzer_name.empty() || is_builtin_analyzer(config->analyzer_name)) { + return std::make_shared(*config); + } - return index_policy_mgr->get_policy_by_name(analyzer_name); + auto* index_policy_mgr = doris::ExecEnv::GetInstance()->index_policy_mgr(); + if (index_policy_mgr == nullptr) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "Index policy manager is not initialized"); + } + return index_policy_mgr->get_analyzer_provider_by_name(config->analyzer_name, + config->char_filter_map); } std::vector InvertedIndexAnalyzer::get_analyse_result( @@ -172,6 +202,8 @@ std::vector InvertedIndexAnalyzer::get_analyse_result( t.term = std::string(token.termBuffer(), token.termLength()); position += token.getPositionIncrement(); t.position = position; + t.key_kind = is_common_gram_token_type(token.type()) ? TermKeyKind::kCommonGram + : TermKeyKind::kPlain; analyse_result.emplace_back(std::move(t)); } } @@ -185,6 +217,12 @@ std::vector InvertedIndexAnalyzer::get_analyse_result( std::vector InvertedIndexAnalyzer::get_analyse_result( const std::string& search_str, const std::map& properties) { + return get_analyse_result(search_str, properties, AnalysisPurpose::kPlainQuery); +} + +std::vector InvertedIndexAnalyzer::get_analyse_result( + const std::string& search_str, const std::map& properties, + AnalysisPurpose purpose) { if (!should_analyzer(properties)) { // Keyword index: all strings (including empty) are valid tokens for exact match. // Empty string is a valid value in keyword index and should be matchable. @@ -200,12 +238,26 @@ std::vector InvertedIndexAnalyzer::get_analyse_result( config.lower_case = get_parser_lowercase_from_properties(properties); config.stop_words = get_parser_stopwords_from_properties(properties); config.char_filter_map = get_parser_char_filter_map_from_properties(properties); - auto analyzer = create_analyzer(&config); + auto analyzer = create_analyzer(&config, purpose); auto reader = create_reader(config.char_filter_map); reader->init(search_str.data(), static_cast(search_str.size()), true); return get_analyse_result(reader, analyzer.get()); } +AnalysisPurpose select_analysis_purpose(InvertedIndexQueryType query_type, int32_t slop, + bool is_similarity) { + if (is_similarity) { + return AnalysisPurpose::kPlainQuery; + } + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && slop == 0) { + return AnalysisPurpose::kExactPhraseQuery; + } + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY) { + return AnalysisPurpose::kPhrasePrefixQuery; + } + return AnalysisPurpose::kPlainQuery; +} + bool InvertedIndexAnalyzer::should_analyzer(const std::map& properties) { auto parser_type = get_inverted_index_parser_type_from_string( get_parser_string_from_properties(properties)); diff --git a/be/src/storage/index/inverted/analyzer/analyzer.h b/be/src/storage/index/inverted/analyzer/analyzer.h index 98588a251bf047..4378b2e6c24b94 100644 --- a/be/src/storage/index/inverted/analyzer/analyzer.h +++ b/be/src/storage/index/inverted/analyzer/analyzer.h @@ -20,6 +20,7 @@ #include #include +#include "storage/index/inverted/analyzer/analyzer_provider.h" #include "storage/index/inverted/inverted_index_parser.h" #include "storage/index/inverted/inverted_index_query_type.h" #include "storage/index/inverted/query/query.h" @@ -49,14 +50,23 @@ class InvertedIndexAnalyzer { const std::string& lower_case, const std::string& stop_words); static AnalyzerPtr create_analyzer(const InvertedIndexAnalyzerConfig* config); + static AnalyzerPtr create_analyzer(const InvertedIndexAnalyzerConfig* config, + AnalysisPurpose purpose); + static AnalyzerProviderPtr create_analyzer_provider(const InvertedIndexAnalyzerConfig* config); static std::vector get_analyse_result(ReaderPtr reader, lucene::analysis::Analyzer* analyzer); static std::vector get_analyse_result( const std::string& search_str, const std::map& properties); + static std::vector get_analyse_result( + const std::string& search_str, const std::map& properties, + AnalysisPurpose purpose); static bool should_analyzer(const std::map& properties); }; +AnalysisPurpose select_analysis_purpose(InvertedIndexQueryType query_type, int32_t slop, + bool is_similarity); + } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/src/storage/index/inverted/analyzer/analyzer_provider.h b/be/src/storage/index/inverted/analyzer/analyzer_provider.h new file mode 100644 index 00000000000000..60fe8250b00a85 --- /dev/null +++ b/be/src/storage/index/inverted/analyzer/analyzer_provider.h @@ -0,0 +1,46 @@ +// 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. + +#pragma once + +#include +#include + +#include "storage/index/inverted/abstract_analysis_factory.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" + +namespace lucene::analysis { +class Analyzer; +} + +namespace doris::segment_v2::inverted_index { + +class CommonWordSet; + +class AnalyzerProvider { +public: + virtual ~AnalyzerProvider() = default; + virtual std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const = 0; + virtual std::string_view base_analyzer_fingerprint() const { return {}; } + virtual bool uses_common_grams() const { return false; } + virtual const CommonGramsQueryIdentity* common_grams_identity() const { return nullptr; } + virtual const CommonWordSet* common_grams_word_set() const { return nullptr; } +}; +using AnalyzerProviderPtr = std::shared_ptr; + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/analyzer/custom_analyzer.cpp b/be/src/storage/index/inverted/analyzer/custom_analyzer.cpp index 3970d2e3cac74d..469f5ee2d6921a 100644 --- a/be/src/storage/index/inverted/analyzer/custom_analyzer.cpp +++ b/be/src/storage/index/inverted/analyzer/custom_analyzer.cpp @@ -17,12 +17,116 @@ #include "storage/index/inverted/analyzer/custom_analyzer.h" +#include +#include + #include "common/status.h" #include "runtime/exec_env.h" #include "storage/index/inverted/analysis_factory_mgr.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/inverted/token_filter/common_grams_filter_factory.h" #include "storage/index/inverted/token_stream.h" +#include "util/sha.h" namespace doris::segment_v2::inverted_index { +namespace { + +bool config_uses_common_grams(const ImmutableCustomAnalyzerConfigPtr& config) { + DORIS_CHECK(config != nullptr); + const auto& filter_configs = config->get_token_filter_configs(); + return std::any_of(filter_configs.begin(), filter_configs.end(), + [](const auto& entry) { return entry->get_name() == "common_grams"; }); +} + +void append_canonical_value(std::string_view value, std::string* output) { + output->append(std::to_string(value.size())); + output->push_back(':'); + output->append(value); +} + +void append_component(std::string_view role, const ComponentConfigPtr& component, + std::string* output) { + append_canonical_value(role, output); + append_canonical_value(component->get_name(), output); + const auto entries = component->get_params().sorted_entries(); + append_canonical_value(std::to_string(entries.size()), output); + for (const auto& [key, value] : entries) { + append_canonical_value(key, output); + append_canonical_value(value, output); + } +} + +std::string sha256(std::string_view value) { + SHA256Digest digest; + digest.reset(value.data(), value.size()); + return std::string(digest.digest()); +} + +std::string calculate_base_analyzer_fingerprint_impl( + const ImmutableCustomAnalyzerConfigPtr& config, + const std::map& outer_char_filter_map) { + DORIS_CHECK(config != nullptr); + std::string base; + append_canonical_value("doris-common-grams-base-analyzer:v1", &base); + append_canonical_value("outer_char_filter", &base); + append_canonical_value(std::to_string(outer_char_filter_map.size()), &base); + for (const auto& [key, value] : outer_char_filter_map) { + append_canonical_value(key, &base); + append_canonical_value(value, &base); + } + append_component("tokenizer", config->get_tokenizer_config(), &base); + const auto char_filters = config->get_char_filter_configs(); + append_canonical_value(std::to_string(char_filters.size()), &base); + for (const auto& char_filter : char_filters) { + append_component("char_filter", char_filter, &base); + } + const auto token_filters = config->get_token_filter_configs(); + const auto base_token_filter_count = std::count_if( + token_filters.begin(), token_filters.end(), + [](const auto& token_filter) { return token_filter->get_name() != "common_grams"; }); + append_canonical_value(std::to_string(base_token_filter_count), &base); + for (const auto& token_filter : token_filters) { + if (token_filter->get_name() != "common_grams") { + append_component("token_filter", token_filter, &base); + } + } + return sha256(base); +} + +CommonGramsQueryIdentity build_common_grams_identity(std::string dictionary_identity, + std::string base_analyzer_fingerprint) { + std::string common_grams; + append_canonical_value("doris-common-grams:v1", &common_grams); + append_canonical_value(std::to_string(COMMON_GRAMS_SEMANTICS_VERSION_V1), &common_grams); + append_canonical_value(std::to_string(COMMON_GRAMS_KEY_VERSION_V1), &common_grams); + append_canonical_value(WORDSET_FORMAT_V1, &common_grams); + append_canonical_value(dictionary_identity, &common_grams); + return {.common_grams_dictionary_identity = std::move(dictionary_identity), + .base_analyzer_fingerprint = std::move(base_analyzer_fingerprint), + .common_grams_fingerprint = sha256(common_grams)}; +} + +std::array, 5> build_purpose_analyzers( + const ImmutableCustomAnalyzerConfigPtr& config, + const std::shared_ptr& common_words) { + DORIS_CHECK(config != nullptr); + const bool has_common_grams = config_uses_common_grams(config); + if (!has_common_grams) { + auto analyzer = CustomAnalyzer::build_custom_analyzer(config); + return {analyzer, analyzer, analyzer, analyzer, analyzer}; + } + return {CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kIndex, common_words), + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kSniiTransientIndex, + common_words), + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPlainQuery, + common_words), + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kExactPhraseQuery, + common_words), + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPhrasePrefixQuery, + common_words)}; +} + +} // namespace CustomAnalyzer::CustomAnalyzer(Builder* builder) { _tokenizer = builder->_tokenizer; @@ -74,7 +178,8 @@ TokenStreamComponentsPtr CustomAnalyzer::create_components() { return std::make_shared(tk, ts); } -CustomAnalyzerPtr CustomAnalyzer::build_custom_analyzer(const CustomAnalyzerConfigPtr& config) { +CustomAnalyzerPtr CustomAnalyzer::build_custom_analyzer( + const ImmutableCustomAnalyzerConfigPtr& config) { if (config == nullptr) { throw Exception(ErrorCode::ILLEGAL_STATE, "Null configuration detected."); } @@ -90,6 +195,134 @@ CustomAnalyzerPtr CustomAnalyzer::build_custom_analyzer(const CustomAnalyzerConf return builder.build(); } +CustomAnalyzerPtr CustomAnalyzer::build_custom_analyzer( + const ImmutableCustomAnalyzerConfigPtr& config, AnalysisPurpose purpose) { + return build_custom_analyzer(config, purpose, CommonWordSet::default_word_set()); +} + +CustomAnalyzerPtr CustomAnalyzer::build_custom_analyzer( + const ImmutableCustomAnalyzerConfigPtr& config, AnalysisPurpose purpose, + const std::shared_ptr& common_words) { + if (config == nullptr) { + throw Exception(ErrorCode::ILLEGAL_STATE, "Null configuration detected."); + } + + CustomAnalyzer::Builder builder; + for (const auto& filter_config : config->get_char_filter_configs()) { + builder.add_char_filter(filter_config->get_name(), filter_config->get_params()); + } + builder.with_tokenizer(config->get_tokenizer_config()->get_name(), + config->get_tokenizer_config()->get_params()); + + const auto filter_configs = config->get_token_filter_configs(); + const size_t common_grams_count = + std::count_if(filter_configs.begin(), filter_configs.end(), + [](const auto& entry) { return entry->get_name() == "common_grams"; }); + if (common_grams_count == 0) { + for (const auto& filter_config : filter_configs) { + builder.add_token_filter(filter_config->get_name(), filter_config->get_params()); + } + return builder.build(); + } + if (common_grams_count != 1 || filter_configs.back()->get_name() != "common_grams") { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "common_grams must appear exactly once as the terminal token filter"); + } + if (builder._tokenizer->position_capability() != PositionCapability::kAlwaysUnitIncrement) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "CommonGrams tokenizer does not guarantee unit position increments"); + } + + for (size_t i = 0; i + 1 < filter_configs.size(); ++i) { + auto factory = AnalysisFactoryMgr::instance().create( + filter_configs[i]->get_name(), filter_configs[i]->get_params()); + if (factory->position_capability() != PositionCapability::kAlwaysUnitIncrement) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "CommonGrams token filter '{}' does not guarantee unit position " + "increments", + filter_configs[i]->get_name()); + } + builder._token_filters.push_back(std::move(factory)); + } + + auto common_grams = AnalysisFactoryMgr::instance().create( + filter_configs.back()->get_name(), filter_configs.back()->get_params()); + auto common_grams_factory = std::dynamic_pointer_cast(common_grams); + DORIS_CHECK(common_grams_factory != nullptr); + common_grams_factory->set_common_words(common_words); + switch (purpose) { + case AnalysisPurpose::kIndex: + common_grams_factory->set_output_mode(CommonGramsOutputMode::kEscapedV1Index); + builder._token_filters.push_back(std::move(common_grams)); + break; + case AnalysisPurpose::kSniiTransientIndex: + common_grams_factory->set_output_mode(CommonGramsOutputMode::kEscapedV1SpimiIndex); + builder._token_filters.push_back(std::move(common_grams)); + break; + case AnalysisPurpose::kPlainQuery: { + auto factory = std::make_shared(); + factory->initialize({}); + builder._token_filters.push_back(std::move(factory)); + break; + } + case AnalysisPurpose::kExactPhraseQuery: { + builder._token_filters.push_back(std::move(common_grams)); + auto factory = std::make_shared(common_words); + factory->initialize({}); + builder._token_filters.push_back(std::move(factory)); + break; + } + case AnalysisPurpose::kPhrasePrefixQuery: { + builder._token_filters.push_back(std::move(common_grams)); + auto factory = std::make_shared(common_words); + factory->initialize({}); + builder._token_filters.push_back(std::move(factory)); + break; + } + } + return builder.build(); +} + +CustomAnalyzerProvider::CustomAnalyzerProvider( + ImmutableCustomAnalyzerConfigPtr config, + std::map outer_char_filter_map) + : _config(std::move(config)), + _base_analyzer_fingerprint( + calculate_base_analyzer_fingerprint(_config, outer_char_filter_map)), + _uses_common_grams(config_uses_common_grams(_config)) { + _common_words = CommonWordSet::default_word_set(); + _analyzers = build_purpose_analyzers(_config, _common_words); + if (_uses_common_grams) { + // Content-derived, so a BE reading a segment grammed against a different word list sees a + // mismatched identity and falls back to the plain plan instead of trusting its grams. + _common_grams_identity = + build_common_grams_identity(_common_words->identity(), _base_analyzer_fingerprint); + } +} + +std::string CustomAnalyzerProvider::calculate_base_analyzer_fingerprint( + const ImmutableCustomAnalyzerConfigPtr& config, + const std::map& outer_char_filter_map) { + return calculate_base_analyzer_fingerprint_impl(config, outer_char_filter_map); +} + +std::shared_ptr CustomAnalyzerProvider::get_analyzer( + AnalysisPurpose purpose) const { + switch (purpose) { + case AnalysisPurpose::kIndex: + return _analyzers[0]; + case AnalysisPurpose::kSniiTransientIndex: + return _analyzers[1]; + case AnalysisPurpose::kPlainQuery: + return _analyzers[2]; + case AnalysisPurpose::kExactPhraseQuery: + return _analyzers[3]; + case AnalysisPurpose::kPhrasePrefixQuery: + return _analyzers[4]; + } + __builtin_unreachable(); +} + void CustomAnalyzer::Builder::with_tokenizer(const std::string& name, const Settings& params) { _tokenizer = AnalysisFactoryMgr::instance().create(name, params); } diff --git a/be/src/storage/index/inverted/analyzer/custom_analyzer.h b/be/src/storage/index/inverted/analyzer/custom_analyzer.h index 68565355496d68..9f89a442912222 100644 --- a/be/src/storage/index/inverted/analyzer/custom_analyzer.h +++ b/be/src/storage/index/inverted/analyzer/custom_analyzer.h @@ -17,6 +17,11 @@ #pragma once +#include +#include +#include + +#include "storage/index/inverted/analyzer/analyzer_provider.h" #include "storage/index/inverted/analyzer/custom_analyzer_config.h" #include "storage/index/inverted/char_filter/char_filter_factory.h" #include "storage/index/inverted/setting.h" @@ -25,6 +30,7 @@ namespace doris::segment_v2::inverted_index { +class CommonWordSet; class CustomAnalyzer; using CustomAnalyzerPtr = std::shared_ptr; @@ -60,7 +66,12 @@ class CustomAnalyzer : public Analyzer { TokenStream* tokenStream(const TCHAR* fieldName, const ReaderPtr& reader) override; TokenStream* reusableTokenStream(const TCHAR* fieldName, const ReaderPtr& reader) override; - static CustomAnalyzerPtr build_custom_analyzer(const CustomAnalyzerConfigPtr& config); + static CustomAnalyzerPtr build_custom_analyzer(const ImmutableCustomAnalyzerConfigPtr& config); + static CustomAnalyzerPtr build_custom_analyzer(const ImmutableCustomAnalyzerConfigPtr& config, + AnalysisPurpose purpose); + static CustomAnalyzerPtr build_custom_analyzer( + const ImmutableCustomAnalyzerConfigPtr& config, AnalysisPurpose purpose, + const std::shared_ptr& common_words); private: ReaderPtr init_reader(ReaderPtr reader); @@ -73,4 +84,39 @@ class CustomAnalyzer : public Analyzer { TokenStreamComponentsPtr _reuse_token_stream; }; +class CustomAnalyzerProvider final : public AnalyzerProvider { +public: + // The CommonGrams word list is not a parameter: it is the BE-local + // CommonWordSet::default_word_set(), and the dictionary identity stamped into segments comes + // from that set's content. An index policy cannot choose either one. + explicit CustomAnalyzerProvider(ImmutableCustomAnalyzerConfigPtr config, + std::map outer_char_filter_map = {}); + + std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const override; + std::string_view base_analyzer_fingerprint() const override { + return _base_analyzer_fingerprint; + } + bool uses_common_grams() const override { return _uses_common_grams; } + const CommonGramsQueryIdentity* common_grams_identity() const override { + return _common_grams_identity ? &*_common_grams_identity : nullptr; + } + const CommonWordSet* common_grams_word_set() const override { + return _uses_common_grams ? _common_words.get() : nullptr; + } + const std::shared_ptr& common_words() const { return _common_words; } + + static std::string calculate_base_analyzer_fingerprint( + const ImmutableCustomAnalyzerConfigPtr& config, + const std::map& outer_char_filter_map = {}); + +private: + ImmutableCustomAnalyzerConfigPtr _config; + const std::string _base_analyzer_fingerprint; + std::shared_ptr _common_words; + bool _uses_common_grams = false; + std::optional _common_grams_identity; + std::array, 5> _analyzers; +}; + } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/src/storage/index/inverted/analyzer/custom_analyzer_config.cpp b/be/src/storage/index/inverted/analyzer/custom_analyzer_config.cpp index b946c13e17d207..0a4c6a5da3c09e 100644 --- a/be/src/storage/index/inverted/analyzer/custom_analyzer_config.cpp +++ b/be/src/storage/index/inverted/analyzer/custom_analyzer_config.cpp @@ -27,15 +27,15 @@ CustomAnalyzerConfig::CustomAnalyzerConfig(Builder* builder) { _token_filters = builder->_token_filters; } -ComponentConfigPtr CustomAnalyzerConfig::get_tokenizer_config() { +ComponentConfigPtr CustomAnalyzerConfig::get_tokenizer_config() const { return _tokenizer_config; } -std::vector CustomAnalyzerConfig::get_char_filter_configs() { +std::vector CustomAnalyzerConfig::get_char_filter_configs() const { return _char_filters; } -std::vector CustomAnalyzerConfig::get_token_filter_configs() { +std::vector CustomAnalyzerConfig::get_token_filter_configs() const { return _token_filters; } diff --git a/be/src/storage/index/inverted/analyzer/custom_analyzer_config.h b/be/src/storage/index/inverted/analyzer/custom_analyzer_config.h index 619d9ae78c9dac..7f6d55e8d2d11a 100644 --- a/be/src/storage/index/inverted/analyzer/custom_analyzer_config.h +++ b/be/src/storage/index/inverted/analyzer/custom_analyzer_config.h @@ -29,6 +29,7 @@ using ComponentConfigPtr = std::shared_ptr; class CustomAnalyzerConfig; using CustomAnalyzerConfigPtr = std::shared_ptr; +using ImmutableCustomAnalyzerConfigPtr = std::shared_ptr; class CustomAnalyzerConfig { public: @@ -53,9 +54,9 @@ class CustomAnalyzerConfig { CustomAnalyzerConfig(Builder* builder); ~CustomAnalyzerConfig() = default; - ComponentConfigPtr get_tokenizer_config(); - std::vector get_char_filter_configs(); - std::vector get_token_filter_configs(); + ComponentConfigPtr get_tokenizer_config() const; + std::vector get_char_filter_configs() const; + std::vector get_token_filter_configs() const; private: ComponentConfigPtr _tokenizer_config; diff --git a/be/src/storage/index/inverted/analyzer/segment_analyzer_context.cpp b/be/src/storage/index/inverted/analyzer/segment_analyzer_context.cpp new file mode 100644 index 00000000000000..38b4a320fb17ed --- /dev/null +++ b/be/src/storage/index/inverted/analyzer/segment_analyzer_context.cpp @@ -0,0 +1,110 @@ +// 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 "storage/index/inverted/analyzer/segment_analyzer_context.h" + +#include + +#include "common/exception.h" +#include "runtime/index_policy/index_policy_mgr.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +Result> analyzer_bypass(std::string_view reason) { + return ResultError(Status::Error( + "segment analyzer unavailable: {}", reason)); +} + +} // namespace + +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, + const std::optional& segment_metadata, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr) { + return maybe_rebuild_segment_analyzer_context(request_context, + segment_metadata ? &*segment_metadata : nullptr, + physical_index_properties, index_policy_mgr); +} + +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, + const CommonGramsSegmentMetadata* segment_metadata, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr) { + if (segment_metadata == nullptr) { + return std::optional {}; + } + if (segment_metadata->base_analyzer_fingerprint.empty()) { + return analyzer_bypass("typed metadata has no base analyzer fingerprint"); + } + return maybe_rebuild_segment_analyzer_context(request_context, + segment_metadata->base_analyzer_fingerprint, + physical_index_properties, index_policy_mgr); +} + +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, std::string_view segment_base_fingerprint, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr) { + DORIS_CHECK(!segment_base_fingerprint.empty()); + if (request_context == nullptr || request_context->analyzer_provider == nullptr || + !request_context->should_tokenize()) { + return analyzer_bypass("query context cannot reconstruct the physical token stream"); + } + + std::string_view request_base_fingerprint = + request_context->analyzer_provider->base_analyzer_fingerprint(); + if (request_base_fingerprint.empty()) { + const auto* identity = request_context->get_common_grams_identity(); + if (identity != nullptr) { + request_base_fingerprint = identity->base_analyzer_fingerprint; + } + } + if (request_base_fingerprint == segment_base_fingerprint) { + return std::optional {}; + } + if (index_policy_mgr == nullptr) { + return analyzer_bypass("index policy manager is not initialized"); + } + + const CharFilterMap physical_char_filter_map = + get_parser_char_filter_map_from_properties(physical_index_properties); + AnalyzerProviderPtr provider; + try { + provider = index_policy_mgr->get_analyzer_provider_by_base_fingerprint( + segment_base_fingerprint, physical_char_filter_map); + } catch (const CLuceneError& error) { + return analyzer_bypass(error.what()); + } catch (const Exception& error) { + return analyzer_bypass(error.what()); + } + if (provider == nullptr) { + return analyzer_bypass("no installed policy matches the segment base fingerprint"); + } + DORIS_CHECK(provider->base_analyzer_fingerprint() == segment_base_fingerprint); + + InvertedIndexAnalyzerCtx effective_context = *request_context; + effective_context.char_filter_map = physical_char_filter_map; + effective_context.analyzer.reset(); + effective_context.analyzer_provider = std::move(provider); + effective_context.common_grams_identity.reset(); + return std::optional(std::move(effective_context)); +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/analyzer/segment_analyzer_context.h b/be/src/storage/index/inverted/analyzer/segment_analyzer_context.h new file mode 100644 index 00000000000000..06be7eeb9e4811 --- /dev/null +++ b/be/src/storage/index/inverted/analyzer/segment_analyzer_context.h @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/inverted_index_parser.h" + +namespace doris { + +class IndexPolicyMgr; + +namespace segment_v2::inverted_index { + +// A missing metadata record is a legacy segment and keeps the request analyzer. A present record +// must identify its base analyzer; otherwise the caller must bypass the inverted index. +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, + const std::optional& segment_metadata, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr); + +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, + const CommonGramsSegmentMetadata* segment_metadata, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr); + +// Returns nullopt when the request analyzer already matches the persisted segment analyzer. +// A returned context owns a fresh provider and must remain local to one query execution. +Result> maybe_rebuild_segment_analyzer_context( + const InvertedIndexAnalyzerCtx* request_context, std::string_view segment_base_fingerprint, + const std::map& physical_index_properties, + IndexPolicyMgr* index_policy_mgr); + +} // namespace segment_v2::inverted_index +} // namespace doris diff --git a/be/src/storage/index/inverted/common/single_flight.h b/be/src/storage/index/inverted/common/single_flight.h new file mode 100644 index 00000000000000..3d360afc4a2b26 --- /dev/null +++ b/be/src/storage/index/inverted/common/single_flight.h @@ -0,0 +1,109 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" + +namespace doris::segment_v2::inverted_index { + +// Collapses concurrent operations with the same key into one execution. No work +// or follower wake-up runs while the selected shard mutex is held. +template +class SingleFlight { +public: + using ResultFuture = std::shared_future; + + std::optional join_or_lead(const std::string& key) { + auto& shard = _shard_for(key); + std::lock_guard guard(shard.mutex); + if (auto it = shard.inflight.find(key); it != shard.inflight.end()) { + return it->second->future; + } + auto flight = std::make_shared(); + flight->future = flight->promise.get_future().share(); + shard.inflight.emplace(key, std::move(flight)); + return std::nullopt; + } + + void publish(const std::string& key, Result result) { + auto& shard = _shard_for(key); + std::shared_ptr flight; + { + std::lock_guard guard(shard.mutex); + auto it = shard.inflight.find(key); + if (it == shard.inflight.end() || it->second->publishing) { + return; + } + flight = it->second; + flight->publishing = true; + } + flight->promise.set_value(std::move(result)); + { + std::lock_guard guard(shard.mutex); + auto it = shard.inflight.find(key); + DORIS_CHECK(it != shard.inflight.end()); + DORIS_CHECK(it->second == flight); + shard.inflight.erase(it); + } + } + + size_t inflight_size() const { + std::array, kShardCount> guards; + for (size_t i = 0; i < kShardCount; ++i) { + guards[i] = std::unique_lock(_shards[i].mutex); + } + + size_t size = 0; + for (const auto& shard : _shards) { + size += shard.inflight.size(); + } + return size; + } + +private: + struct Flight { + std::promise promise; + ResultFuture future; + bool publishing = false; + }; + + struct Shard { + mutable std::mutex mutex; + std::unordered_map> inflight; + }; + + static constexpr size_t kShardCount = 64; + + Shard& _shard_for(const std::string& key) { + return _shards[std::hash {}(key) % kShardCount]; + } + + std::array _shards; +}; + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_grams_key_codec.cpp b/be/src/storage/index/inverted/common_grams/common_grams_key_codec.cpp new file mode 100644 index 00000000000000..98f0e3ac918643 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_grams_key_codec.cpp @@ -0,0 +1,279 @@ +// 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 "storage/index/inverted/common_grams/common_grams_key_codec.h" + +#include +#include + +#include "util/utf8_check.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +constexpr size_t COMMON_GRAM_LENGTH_HEX_BYTES = 8; +constexpr size_t COMMON_GRAM_SEPARATOR_BYTES = 1; +constexpr size_t COMMON_GRAM_FIXED_BYTES = + CG_V1_MARKER.size() + COMMON_GRAM_LENGTH_HEX_BYTES + COMMON_GRAM_SEPARATOR_BYTES; +constexpr std::string_view HEX_DIGITS = "0123456789abcdef"; +constexpr std::string_view LEGACY_PHRASE_BIGRAM_MARKER = + "\x1f" + "SNII_PHRASE_BIGRAM" + "\x1f"; + +ResultError analyzer_error(std::string_view message) { + return ResultError(Status::Error("{}", message)); +} + +bool prefixes_overlap(std::string_view left, std::string_view right) { + return left.starts_with(right) || right.starts_with(left); +} + +} // namespace + +Status validate_common_grams_logical_term(std::string_view term, std::string_view component) { + if (term.find('\0') != std::string_view::npos) { + return Status::Error( + "CommonGrams {} contains NUL", component); + } + if (!validate_utf8(term.data(), term.size())) { + return Status::Error( + "CommonGrams {} is not valid UTF-8", component); + } + if (term.size() > COMMON_GRAM_MAX_ENCODED_BYTES) { + return Status::Error( + "CommonGrams {} has {} UTF-8 bytes, exceeding {}", component, term.size(), + COMMON_GRAM_MAX_ENCODED_BYTES); + } + return Status::OK(); +} + +Result encode_plain_term(std::string_view term, PlainTermKeyVersion version) { + std::string encoded; + auto encoded_result = try_encode_plain_term(term, version, &encoded); + if (!encoded_result.has_value()) { + return ResultError(std::move(encoded_result.error())); + } + if (!encoded_result.value()) { + return analyzer_error("escaped plain term would exceed the 16383-byte key limit"); + } + return encoded; +} + +Result try_encode_plain_term(std::string_view term, PlainTermKeyVersion version, + std::string* output) { + DORIS_CHECK(output != nullptr); + output->clear(); + auto status = validate_common_grams_logical_term(term, "plain term"); + if (!status.ok()) { + return ResultError(std::move(status)); + } + + switch (version) { + case PlainTermKeyVersion::kLegacyRaw: + case PlainTermKeyVersion::kRawNoInternal: + output->assign(term); + return true; + case PlainTermKeyVersion::kEscapedV1: + if (term.empty() || (term.front() != PLAIN_ESCAPE_PREFIX && term.front() != '\x1f')) { + output->assign(term); + return true; + } + return try_encode_escaped_plain_term_prevalidated(term, *output); + } + return analyzer_error("unknown plain-term key version"); +} + +Result> try_encode_plain_term_view(std::string_view term, + PlainTermKeyVersion version, + std::string* scratch) { + DORIS_CHECK(scratch != nullptr); + scratch->clear(); + auto status = validate_common_grams_logical_term(term, "plain term"); + if (!status.ok()) { + return ResultError(std::move(status)); + } + + switch (version) { + case PlainTermKeyVersion::kLegacyRaw: + case PlainTermKeyVersion::kRawNoInternal: + return std::optional(term); + case PlainTermKeyVersion::kEscapedV1: + if (term.empty() || (term.front() != PLAIN_ESCAPE_PREFIX && term.front() != '\x1f')) { + return std::optional(term); + } + if (!try_encode_escaped_plain_term_prevalidated(term, *scratch)) { + return std::optional(); + } + return std::optional(*scratch); + } + return analyzer_error("unknown plain-term key version"); +} + +bool try_encode_escaped_plain_term_prevalidated(std::string_view logical_term, + std::string& output) { + DCHECK(!logical_term.empty()); + DCHECK(logical_term.front() == PLAIN_ESCAPE_PREFIX || logical_term.front() == '\x1f'); + DCHECK(validate_common_grams_logical_term(logical_term, "plain term").ok()); + output.clear(); + if (logical_term.size() == COMMON_GRAM_MAX_ENCODED_BYTES) { + return false; + } + output.reserve(logical_term.size() + 1); + output.push_back(PLAIN_ESCAPE_PREFIX); + output.push_back(logical_term.front() == PLAIN_ESCAPE_PREFIX ? 'E' : 'G'); + output.append(logical_term.substr(1)); + return true; +} + +Result decode_plain_term_view(std::string_view term, PlainTermKeyVersion version, + std::string* scratch) { + DORIS_CHECK(scratch != nullptr); + scratch->clear(); + if (term.size() > COMMON_GRAM_MAX_ENCODED_BYTES) { + return analyzer_error("encoded plain term exceeds the 16383-byte key limit"); + } + + switch (version) { + case PlainTermKeyVersion::kLegacyRaw: + case PlainTermKeyVersion::kRawNoInternal: { + auto status = validate_common_grams_logical_term(term, "plain term"); + if (!status.ok()) { + return ResultError(std::move(status)); + } + return term; + } + case PlainTermKeyVersion::kEscapedV1: { + if (term.empty() || term.front() != PLAIN_ESCAPE_PREFIX) { + if (!term.empty() && term.front() == '\x1f') { + return analyzer_error("escaped plain term enters the internal namespace"); + } + auto status = validate_common_grams_logical_term(term, "plain term"); + if (!status.ok()) { + return ResultError(std::move(status)); + } + return term; + } + if (term.size() < 2 || (term[1] != 'E' && term[1] != 'G')) { + return analyzer_error("invalid plain_term_escape:v1 key"); + } + scratch->reserve(term.size() - 1); + scratch->push_back(term[1] == 'E' ? PLAIN_ESCAPE_PREFIX : '\x1f'); + scratch->append(term.substr(2)); + auto status = validate_common_grams_logical_term(*scratch, "plain term"); + if (!status.ok()) { + return ResultError(std::move(status)); + } + return std::string_view(*scratch); + } + } + return analyzer_error("unknown plain-term key version"); +} + +Result decode_plain_term(std::string_view term, PlainTermKeyVersion version) { + std::string scratch; + auto decoded = decode_plain_term_view(term, version, &scratch); + if (!decoded.has_value()) { + return ResultError(std::move(decoded.error())); + } + return std::string(*decoded); +} + +bool is_internal_term_key(std::string_view physical_term) { + return physical_term.starts_with(INTERNAL_TERM_NAMESPACE_BEGIN); +} + +bool legacy_raw_exact_requires_bypass(std::string_view logical_term) { + return logical_term.starts_with(CG_V1_MARKER) || + logical_term.starts_with(LEGACY_PHRASE_BIGRAM_MARKER); +} + +bool legacy_raw_prefix_requires_bypass(std::string_view logical_prefix) { + return prefixes_overlap(logical_prefix, CG_V1_MARKER) || + prefixes_overlap(logical_prefix, LEGACY_PHRASE_BIGRAM_MARKER); +} + +bool is_common_gram_encodable(std::string_view left, std::string_view right) { + if (!validate_common_grams_logical_term(left, "left term").ok() || + !validate_common_grams_logical_term(right, "right term").ok()) { + return false; + } + return is_common_gram_encodable_prevalidated(left, right); +} + +bool is_common_gram_encodable_prevalidated(std::string_view left, std::string_view right) { + DCHECK(validate_common_grams_logical_term(left, "left term").ok()); + DCHECK(validate_common_grams_logical_term(right, "right term").ok()); + return common_gram_component_sizes_encodable(left.size(), right.size()); +} + +bool common_gram_component_sizes_encodable(size_t left_size, size_t right_size) { + return left_size <= COMMON_GRAM_MAX_ENCODED_BYTES - COMMON_GRAM_FIXED_BYTES && + right_size <= COMMON_GRAM_MAX_ENCODED_BYTES - COMMON_GRAM_FIXED_BYTES - left_size; +} + +Result try_encode_common_gram(std::string_view left, std::string_view right, + std::string* output) { + DORIS_CHECK(output != nullptr); + output->clear(); + auto left_status = validate_common_grams_logical_term(left, "left term"); + if (!left_status.ok()) { + return ResultError(std::move(left_status)); + } + auto right_status = validate_common_grams_logical_term(right, "right term"); + if (!right_status.ok()) { + return ResultError(std::move(right_status)); + } + return try_encode_common_gram_prevalidated(left, right, *output); +} + +bool try_encode_common_gram_prevalidated(std::string_view left, std::string_view right, + std::string& output) { + DCHECK(validate_common_grams_logical_term(left, "left term").ok()); + DCHECK(validate_common_grams_logical_term(right, "right term").ok()); + output.clear(); + if (COMMON_GRAM_FIXED_BYTES + left.size() + right.size() > COMMON_GRAM_MAX_ENCODED_BYTES) { + return false; + } + + output.reserve(COMMON_GRAM_FIXED_BYTES + left.size() + right.size()); + output.append(CG_V1_MARKER); + char length_and_separator[COMMON_GRAM_LENGTH_HEX_BYTES + COMMON_GRAM_SEPARATOR_BYTES]; + for (size_t i = 0; i < COMMON_GRAM_LENGTH_HEX_BYTES; ++i) { + const size_t shift = (COMMON_GRAM_LENGTH_HEX_BYTES - i - 1) * 4; + length_and_separator[i] = HEX_DIGITS[(left.size() >> shift) & 0xf]; + } + length_and_separator[COMMON_GRAM_LENGTH_HEX_BYTES] = ':'; + output.append(length_and_separator, sizeof(length_and_separator)); + output.append(left); + output.append(right); + return true; +} + +Result encode_common_gram(std::string_view left, std::string_view right) { + std::string encoded; + auto encoded_result = try_encode_common_gram(left, right, &encoded); + if (!encoded_result.has_value()) { + return ResultError(std::move(encoded_result.error())); + } + if (!encoded_result.value()) { + return analyzer_error("encoded common gram would exceed the 16383-byte key limit"); + } + return encoded; +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_grams_key_codec.h b/be/src/storage/index/inverted/common_grams/common_grams_key_codec.h new file mode 100644 index 00000000000000..07869877a55d01 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_grams_key_codec.h @@ -0,0 +1,72 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris::segment_v2::inverted_index { + +inline constexpr size_t COMMON_GRAM_MAX_ENCODED_BYTES = 16383; +inline constexpr char PLAIN_ESCAPE_PREFIX = '\x1e'; +inline constexpr std::string_view INTERNAL_TERM_NAMESPACE_BEGIN {"\x1f", 1}; +inline constexpr std::string_view INTERNAL_TERM_NAMESPACE_END {"\x20", 1}; +inline constexpr std::string_view CG_V1_MARKER = + "\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f"; +inline constexpr std::string_view CG_V1_MARKER_END = + "\x1f" + "DORIS_COMMON_GRAM_V1" + "\x20"; + +enum class PlainTermKeyVersion : uint8_t { + kLegacyRaw = 0, + kEscapedV1 = 1, + kRawNoInternal = 2, +}; + +Result encode_plain_term(std::string_view term, PlainTermKeyVersion version); +Result try_encode_plain_term(std::string_view term, PlainTermKeyVersion version, + std::string* output); +Result> try_encode_plain_term_view(std::string_view term, + PlainTermKeyVersion version, + std::string* scratch); +bool try_encode_escaped_plain_term_prevalidated(std::string_view logical_term, std::string& output); +Result decode_plain_term_view(std::string_view term, PlainTermKeyVersion version, + std::string* scratch); +Result decode_plain_term(std::string_view term, PlainTermKeyVersion version); +bool is_internal_term_key(std::string_view physical_term); +bool legacy_raw_exact_requires_bypass(std::string_view logical_term); +bool legacy_raw_prefix_requires_bypass(std::string_view logical_prefix); +Status validate_common_grams_logical_term(std::string_view term, std::string_view component); +Result try_encode_common_gram(std::string_view left, std::string_view right, + std::string* output); +bool try_encode_common_gram_prevalidated(std::string_view left, std::string_view right, + std::string& output); +bool common_gram_component_sizes_encodable(size_t left_size, size_t right_size); +bool is_common_gram_encodable_prevalidated(std::string_view left, std::string_view right); +Result encode_common_gram(std::string_view left, std::string_view right); +bool is_common_gram_encodable(std::string_view left, std::string_view right); + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_grams_query_cost.h b/be/src/storage/index/inverted/common_grams/common_grams_query_cost.h new file mode 100644 index 00000000000000..2738c52150a705 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_grams_query_cost.h @@ -0,0 +1,53 @@ +// 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. + +#pragma once + +#include +#include + +namespace doris::segment_v2::inverted_index { + +struct CommonGramsPlanRawCost { + uint64_t posting_bytes_or_df_sum = 0; + uint64_t estimated_candidate_df = 0; + uint32_t clause_count = 0; +}; + +struct CommonGramsPlanCostModel { + uint32_t position_verify_factor = 0; + uint32_t common_grams_cost_ratio_percent = 85; + uint64_t generation = 0; +}; + +inline uint64_t estimate_common_grams_plan_cost(const CommonGramsPlanRawCost& input, + uint32_t position_verify_factor) { + const unsigned __int128 estimate = + static_cast(input.posting_bytes_or_df_sum) + + static_cast(input.estimated_candidate_df) * position_verify_factor * + input.clause_count; + return estimate > std::numeric_limits::max() ? std::numeric_limits::max() + : static_cast(estimate); +} + +inline bool common_grams_plan_cost_wins(uint64_t plain_cost, uint64_t common_grams_cost, + uint32_t common_grams_cost_ratio_percent) { + return static_cast(common_grams_cost) * 100 <= + static_cast(plain_cost) * common_grams_cost_ratio_percent; +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.cpp b/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.cpp new file mode 100644 index 00000000000000..c8dcddecf70608 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.cpp @@ -0,0 +1,158 @@ +// 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 "storage/index/inverted/common_grams/common_grams_segment_metadata.h" + +namespace doris::segment_v2::inverted_index { + +CommonGramsSegmentMetadata make_common_grams_segment_metadata( + const CommonGramsQueryIdentity& identity) { + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = CommonGramsCoverage::kComplete; + metadata.common_grams_semantics_version = COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = identity.common_grams_dictionary_identity; + metadata.base_analyzer_fingerprint = identity.base_analyzer_fingerprint; + metadata.common_grams_fingerprint = identity.common_grams_fingerprint; + metadata.scoring_coverage = ScoringCoverage::kComplete; + metadata.scoring_stats_version = COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + return metadata; +} + +bool common_grams_identity_matches(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity) { + return metadata.common_grams_dictionary_identity == identity.common_grams_dictionary_identity && + metadata.base_analyzer_fingerprint == identity.base_analyzer_fingerprint && + metadata.common_grams_fingerprint == identity.common_grams_fingerprint; +} + +Status validate_common_grams_segment_metadata(const CommonGramsSegmentMetadata& metadata) { + switch (metadata.plain_term_key_version) { + case PlainTermKeyVersion::kLegacyRaw: + case PlainTermKeyVersion::kEscapedV1: + case PlainTermKeyVersion::kRawNoInternal: + break; + default: + return Status::Error( + "common_grams_metadata: invalid plain-term key version"); + } + + switch (metadata.common_grams_coverage) { + case CommonGramsCoverage::kNone: + case CommonGramsCoverage::kComplete: + case CommonGramsCoverage::kMixed: + break; + default: + return Status::Error( + "common_grams_metadata: invalid coverage"); + } + + switch (metadata.scoring_coverage) { + case ScoringCoverage::kNone: + case ScoringCoverage::kComplete: + break; + default: + return Status::Error( + "common_grams_metadata: invalid scoring coverage"); + } + + if (metadata.plain_term_key_version == PlainTermKeyVersion::kRawNoInternal && + metadata.common_grams_coverage != CommonGramsCoverage::kNone) { + return Status::Error( + "common_grams_metadata: raw-no-internal segment has gram coverage"); + } + + if (metadata.common_grams_coverage == CommonGramsCoverage::kComplete && + (metadata.plain_term_key_version != PlainTermKeyVersion::kEscapedV1 || + metadata.common_grams_semantics_version == 0 || metadata.common_grams_key_version == 0 || + metadata.common_grams_dictionary_identity.empty() || + metadata.base_analyzer_fingerprint.empty() || metadata.common_grams_fingerprint.empty())) { + return Status::Error( + "common_grams_metadata: incomplete complete-coverage identity"); + } + + if (metadata.scoring_coverage == ScoringCoverage::kComplete && + (metadata.scoring_stats_version == 0 || metadata.norm_semantics_version == 0 || + metadata.base_analyzer_fingerprint.empty())) { + return Status::Error( + "common_grams_metadata: incomplete scoring identity"); + } + return Status::OK(); +} + +Status validate_snii_scoring_metadata(const CommonGramsSegmentMetadata* metadata, + uint64_t physical_doc_count, + uint64_t physical_sum_total_term_freq, bool has_scoring_tier, + bool has_positions, bool has_norms) { + if (metadata == nullptr) { + return Status::Error( + "SNII semantic scoring metadata is missing"); + } + RETURN_IF_ERROR(validate_common_grams_segment_metadata(*metadata)); + if (metadata->scoring_coverage != ScoringCoverage::kComplete || + metadata->scoring_stats_version != COMMON_GRAMS_SCORING_STATS_VERSION_V1 || + metadata->norm_semantics_version != COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1) { + return Status::Error( + "SNII scoring metadata uses unsupported semantics"); + } + if (!has_scoring_tier || !has_positions || !has_norms) { + return Status::Error( + "SNII complete scoring metadata requires the scoring tier, positions, and norms"); + } + if (metadata->scoring_doc_count != physical_doc_count) { + return Status::Error( + "SNII semantic scoring document count {} differs from physical document count {}", + metadata->scoring_doc_count, physical_doc_count); + } + if (metadata->scoring_token_count > physical_sum_total_term_freq) { + return Status::Error( + "SNII semantic scoring token count {} exceeds physical term frequency {}", + metadata->scoring_token_count, physical_sum_total_term_freq); + } + if (metadata->plain_term_key_version == PlainTermKeyVersion::kRawNoInternal && + metadata->common_grams_coverage == CommonGramsCoverage::kNone && + metadata->scoring_token_count != physical_sum_total_term_freq) { + return Status::Error( + "SNII semantic plain token count {} differs from physical term frequency {}", + metadata->scoring_token_count, physical_sum_total_term_freq); + } + if (physical_sum_total_term_freq != 0 && metadata->scoring_token_count == 0) { + return Status::Error( + "SNII non-empty physical postings have zero semantic scoring tokens"); + } + return Status::OK(); +} + +bool is_common_grams_query_compatible(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity) { + return is_common_grams_query_compatible(metadata, identity, CommonGramsCoverage::kComplete); +} + +bool is_common_grams_query_compatible(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity, + CommonGramsCoverage required_coverage) { + return metadata.plain_term_key_version == PlainTermKeyVersion::kEscapedV1 && + metadata.common_grams_coverage == required_coverage && + required_coverage != CommonGramsCoverage::kNone && + metadata.common_grams_semantics_version == COMMON_GRAMS_SEMANTICS_VERSION_V1 && + metadata.common_grams_key_version == COMMON_GRAMS_KEY_VERSION_V1 && + common_grams_identity_matches(metadata, identity); +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.h b/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.h new file mode 100644 index 00000000000000..b5864bd340fc87 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_grams_segment_metadata.h @@ -0,0 +1,87 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" + +namespace doris::segment_v2::inverted_index { + +inline constexpr uint32_t COMMON_GRAMS_SEGMENT_METADATA_VERSION = 1; +inline constexpr uint32_t COMMON_GRAMS_SEMANTICS_VERSION_V1 = 1; +inline constexpr uint32_t COMMON_GRAMS_KEY_VERSION_V1 = 1; +inline constexpr uint32_t COMMON_GRAMS_SCORING_STATS_VERSION_V1 = 1; +inline constexpr uint32_t COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1 = 1; + +enum class CommonGramsCoverage : uint8_t { + kNone = 0, + kComplete = 1, + kMixed = 2, +}; + +enum class ScoringCoverage : uint8_t { + kNone = 0, + kComplete = 1, +}; + +// SNII metadata. A missing record is a legacy segment; a present record must +// validate before any capability is used. +struct CommonGramsSegmentMetadata { + PlainTermKeyVersion plain_term_key_version = PlainTermKeyVersion::kLegacyRaw; + CommonGramsCoverage common_grams_coverage = CommonGramsCoverage::kNone; + uint32_t common_grams_semantics_version = 0; + uint32_t common_grams_key_version = 0; + std::string common_grams_dictionary_identity; + std::string base_analyzer_fingerprint; + std::string common_grams_fingerprint; + ScoringCoverage scoring_coverage = ScoringCoverage::kNone; + uint32_t scoring_stats_version = 0; + uint32_t norm_semantics_version = 0; + uint64_t scoring_doc_count = 0; + uint64_t scoring_token_count = 0; + + bool operator==(const CommonGramsSegmentMetadata&) const = default; +}; + +struct CommonGramsQueryIdentity { + std::string common_grams_dictionary_identity; + std::string base_analyzer_fingerprint; + std::string common_grams_fingerprint; + + bool operator==(const CommonGramsQueryIdentity&) const = default; +}; + +CommonGramsSegmentMetadata make_common_grams_segment_metadata( + const CommonGramsQueryIdentity& identity); +bool common_grams_identity_matches(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity); +Status validate_common_grams_segment_metadata(const CommonGramsSegmentMetadata& metadata); +Status validate_snii_scoring_metadata(const CommonGramsSegmentMetadata* metadata, + uint64_t physical_doc_count, + uint64_t physical_sum_total_term_freq, bool has_scoring_tier, + bool has_positions, bool has_norms); +bool is_common_grams_query_compatible(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity); +bool is_common_grams_query_compatible(const CommonGramsSegmentMetadata& metadata, + const CommonGramsQueryIdentity& identity, + CommonGramsCoverage required_coverage); + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_word_set.cpp b/be/src/storage/index/inverted/common_grams/common_word_set.cpp new file mode 100644 index 00000000000000..4e74cac590778d --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_word_set.cpp @@ -0,0 +1,197 @@ +// 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 "storage/index/inverted/common_grams/common_word_set.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/logging.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "util/md5.h" +#include "util/utf8_check.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +#ifdef BE_TEST +std::atomic g_common_word_membership_lookups {0}; +std::atomic g_common_word_hash_lookups {0}; +#endif + +ResultError wordset_error(std::string_view message) { + return ResultError(Status::Error("{}", message)); +} + +} // namespace + +CommonWordSet::CommonWordSet(WordContainer words, std::string identity) + : _words(std::move(words)), _identity(std::move(identity)) { + DORIS_CHECK(!_identity.empty()); + for (const std::string& word : _words) { + DORIS_CHECK(!word.empty()); + DORIS_CHECK_LE(word.size(), std::numeric_limits::max()); + const auto bytes = static_cast(word.size()); + _min_word_bytes = std::min(_min_word_bytes, bytes); + _max_word_bytes = std::max(_max_word_bytes, bytes); + const uint8_t first = static_cast(word.front()); + _first_byte_mask[first >> 6] |= uint64_t {1} << (first & 63); + } +} + +const CommonWordSet& CommonWordSet::builtin_english_stop_words_v1() { + static const CommonWordSet words( + WordContainer { + "a", "an", "and", "are", "as", "at", "be", "but", "by", + "for", "if", "in", "into", "is", "it", "no", "not", "of", + "on", "or", "such", "that", "the", "their", "then", "there", "these", + "they", "this", "to", "was", "will", "with", + }, + std::string(BUILTIN_COMMON_WORDS_RESOURCE)); + return words; +} + +Result CommonWordSet::parse_words(std::string_view content) { + // Digest first: the parse loop below advances `content`, so hashing it afterwards would hash + // an empty view. Content-derived so a segment records exactly which list grammed it, and + // digesting the raw bytes (not the parsed set) means a comment-only edit also yields a new + // identity -- the safe direction, since that re-plans instead of risking a stale match. + Md5Digest digest; + digest.update(content.data(), content.size()); + digest.digest(); + const std::string identity = "wordset:md5:" + digest.hex(); + + if (content.find('\0') != std::string_view::npos) { + return wordset_error("CommonGrams word list contains NUL"); + } + if (!validate_utf8(content.data(), content.size())) { + return wordset_error("CommonGrams word list is not valid UTF-8"); + } + + WordContainer words; + while (!content.empty()) { + const size_t newline = content.find('\n'); + std::string_view term = content.substr(0, newline); + if (newline != std::string_view::npos && !term.empty() && term.back() == '\r') { + term.remove_suffix(1); + } + if (!term.empty() && term.front() != '#') { + if (term.size() > COMMON_GRAM_MAX_ENCODED_BYTES) { + return wordset_error("wordset:v1 term exceeds the 16383-byte token limit"); + } + words.emplace(term); + } + if (newline == std::string_view::npos) { + break; + } + content.remove_prefix(newline + 1); + } + return CommonWordSet(std::move(words), identity); +} + +std::string CommonWordSet::default_word_set_path() { + return config::inverted_index_dict_path + "/common_grams/default_words.txt"; +} + +std::shared_ptr CommonWordSet::default_word_set() { + // Lazy singleton: the word list is immutable for the process, so every analyzer shares one + // copy and the fallback warning below is emitted at most once no matter how many analyzers + // are built. A later edit to inverted_index_dict_path is therefore ignored, which is intended + // -- swapping the list under live segments would change what their grams mean. + static const std::shared_ptr instance = [] { + auto builtin = [] { + return std::shared_ptr(&builtin_english_stop_words_v1(), + [](const CommonWordSet*) {}); + }; + const std::string path = default_word_set_path(); + std::ifstream input(path, std::ios::binary); + if (!input.is_open()) { + // Expected on a stock install: shipping no file means "use the built-in list", so this + // is INFO. A file that exists but cannot be read or parsed is an operator mistake and + // warns below. + LOG(INFO) << "No CommonGrams word list at " << path + << ", using the built-in English stop words"; + return builtin(); + } + std::ostringstream buffer; + buffer << input.rdbuf(); + if (input.bad()) { + LOG(WARNING) << "CommonGrams word list at " << path + << " could not be read, falling back to the built-in English stop words"; + return builtin(); + } + const std::string content = buffer.str(); + auto parsed = parse_words(content); + if (!parsed.has_value()) { + LOG(WARNING) << "CommonGrams word list at " << path << " is invalid: " << parsed.error() + << ", falling back to the built-in English stop words"; + return builtin(); + } + LOG(INFO) << "Loaded " << parsed.value().size() << " CommonGrams words from " << path; + return std::shared_ptr( + std::make_shared(std::move(parsed.value()))); + }(); + return instance; +} + +bool CommonWordSet::contains(std::string_view term) const { +#ifdef BE_TEST + g_common_word_membership_lookups.fetch_add(1, std::memory_order_relaxed); +#endif + if (term.size() < _min_word_bytes || term.size() > _max_word_bytes) { + return false; + } + const auto first = static_cast(term.front()); + if ((_first_byte_mask[first >> 6] & (uint64_t {1} << (first & 63))) == 0) { + return false; + } +#ifdef BE_TEST + g_common_word_hash_lookups.fetch_add(1, std::memory_order_relaxed); +#endif + return _words.contains(term); +} + +#ifdef BE_TEST +namespace common_grams_testing { + +uint64_t common_word_membership_lookup_count() { + return g_common_word_membership_lookups.load(std::memory_order_relaxed); +} + +void reset_common_word_membership_lookup_count() { + g_common_word_membership_lookups.store(0, std::memory_order_relaxed); +} + +uint64_t common_word_hash_lookup_count() { + return g_common_word_hash_lookups.load(std::memory_order_relaxed); +} + +void reset_common_word_hash_lookup_count() { + g_common_word_hash_lookups.store(0, std::memory_order_relaxed); +} + +} // namespace common_grams_testing +#endif + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/common_grams/common_word_set.h b/be/src/storage/index/inverted/common_grams/common_word_set.h new file mode 100644 index 00000000000000..b6aac21f48b501 --- /dev/null +++ b/be/src/storage/index/inverted/common_grams/common_word_set.h @@ -0,0 +1,94 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris::segment_v2::inverted_index { + +inline constexpr std::string_view BUILTIN_COMMON_WORDS_RESOURCE = "builtin:lucene_english_stop:v1"; +inline constexpr std::string_view WORDSET_FORMAT_V1 = "wordset:v1"; + +class CommonWordSet { +public: + static const CommonWordSet& builtin_english_stop_words_v1(); + + // Parses a newline-separated word list. Blank lines and '#' comments are skipped. + static Result parse_words(std::string_view content); + + // The BE-local word list every CommonGrams analyzer grams against, read once from + // /common_grams/default_words.txt -- the same layout the icu, ik and + // pinyin dictionaries use. Falls back to builtin_english_stop_words_v1() when the file is + // absent or unparseable. The set is deliberately not selectable per index policy: every replica + // of a tablet must gram identically, and a policy-supplied list would have to be distributed + // and acknowledged before any index could use it. + static std::shared_ptr default_word_set(); + + // Exposed so tests can assert the layout without duplicating the literal. + static std::string default_word_set_path(); + + // Stamped into a segment's CommonGrams metadata and compared against the querying analyzer's + // identity, so a segment is only read with gram expectations that match how it was written. + // Because the word list is now a BE-local file, this MUST derive from the content: two BEs + // pointed at different files, or one BE after the file is edited, would otherwise stamp the + // same identity onto incompatible segments and silently mis-plan phrase queries. + const std::string& identity() const { return _identity; } + + bool contains(std::string_view term) const; + size_t size() const { return _words.size(); } + +private: + struct TransparentStringHash { + using is_transparent = void; + + size_t operator()(std::string_view value) const { + return std::hash {}(value); + } + }; + + using WordContainer = phmap::flat_hash_set>; + + CommonWordSet(WordContainer words, std::string identity); + + WordContainer _words; + // Never empty. See identity() for why this is content-derived rather than a fixed constant. + std::string _identity; + std::array _first_byte_mask {}; + uint16_t _min_word_bytes = std::numeric_limits::max(); + uint16_t _max_word_bytes = 0; +}; + +#ifdef BE_TEST +namespace common_grams_testing { +uint64_t common_word_membership_lookup_count(); +void reset_common_word_membership_lookup_count(); +} // namespace common_grams_testing +#endif + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/inverted_index_cache.cpp b/be/src/storage/index/inverted/inverted_index_cache.cpp index e759a0163ac7d3..402637c3798c6f 100644 --- a/be/src/storage/index/inverted/inverted_index_cache.cpp +++ b/be/src/storage/index/inverted/inverted_index_cache.cpp @@ -27,10 +27,50 @@ #include "runtime/exec_env.h" #include "runtime/thread_context.h" +#include "util/coding.h" #include "util/defer_op.h" namespace doris::segment_v2 { +namespace { + +void append_length_prefixed(std::string_view value, std::string* output) { + put_fixed64_le(output, value.size()); + output->append(value); +} + +} // namespace + +std::string InvertedIndexRawQuerySemantic::encode() const { + std::string output; + output.reserve(sizeof(cache_semantics_version) + sizeof(uint64_t) + raw_query_bytes.size() + + sizeof(query_type) + sizeof(slop) + sizeof(ordered) + sizeof(max_expansions) + + sizeof(common_grams_cache_generation) + sizeof(uint64_t)); + put_fixed32_le(&output, cache_semantics_version); + append_length_prefixed(raw_query_bytes, &output); + put_fixed32_le(&output, static_cast(query_type)); + put_fixed32_le(&output, static_cast(slop)); + output.push_back(static_cast(ordered)); + put_fixed32_le(&output, static_cast(max_expansions)); + put_fixed64_le(&output, common_grams_cache_generation); + return output; +} + +std::string InvertedIndexQueryCache::CacheKey::encode() const { + if (query_type_to_string(query_type).empty()) { + return {}; + } + std::string output; + const std::string index_path_string = index_path.string(); + output.reserve(3 * sizeof(uint64_t) + index_path_string.size() + column_name.size() + + sizeof(query_type) + value.size()); + append_length_prefixed(index_path_string, &output); + append_length_prefixed(column_name, &output); + put_fixed32_le(&output, static_cast(query_type)); + append_length_prefixed(value, &output); + return output; +} + InvertedIndexSearcherCache* InvertedIndexSearcherCache::create_global_instance( size_t capacity, uint32_t num_shards) { return new InvertedIndexSearcherCache(capacity, num_shards); diff --git a/be/src/storage/index/inverted/inverted_index_cache.h b/be/src/storage/index/inverted/inverted_index_cache.h index 0e33fe747b3523..476a44fae7b26e 100644 --- a/be/src/storage/index/inverted/inverted_index_cache.h +++ b/be/src/storage/index/inverted/inverted_index_cache.h @@ -26,6 +26,7 @@ #include #include #include +#include #include "common/config.h" #include "common/status.h" @@ -35,6 +36,7 @@ #include "runtime/memory/lru_cache_policy.h" #include "runtime/memory/mem_tracker.h" #include "storage/index/inverted/inverted_index_searcher.h" +#include "storage/index/snii/reader/logical_index_reader.h" #include "util/lru_cache.h" #include "util/slice.h" #include "util/time.h" @@ -42,6 +44,7 @@ namespace doris { namespace segment_v2 { class InvertedIndexCacheHandle; +class IndexFileReader; class InvertedIndexSearcherCache { public: @@ -56,6 +59,8 @@ class InvertedIndexSearcherCache { class CacheValue : public LRUCacheValueBase { public: IndexSearcherPtr index_searcher; + std::shared_ptr snii_index_file_reader; + std::unique_ptr snii_logical_reader; size_t size = 0; int64_t last_visit_time; @@ -65,6 +70,14 @@ class InvertedIndexSearcherCache { size = mem_size; last_visit_time = visit_time; } + explicit CacheValue(std::unique_ptr logical_reader, + size_t mem_size, int64_t visit_time, + std::shared_ptr index_file_reader) + : snii_index_file_reader(std::move(index_file_reader)), + snii_logical_reader(std::move(logical_reader)) { + size = mem_size; + last_visit_time = visit_time; + } }; // Create global instance of this class. // "capacity" is the capacity of lru cache. @@ -166,6 +179,11 @@ class InvertedIndexCacheHandle { return ((InvertedIndexSearcherCache::CacheValue*)_cache->value(_handle))->index_searcher; } + doris::snii::reader::LogicalIndexReader* get_snii_logical_reader() { + return ((InvertedIndexSearcherCache::CacheValue*)_cache->value(_handle)) + ->snii_logical_reader.get(); + } + InvertedIndexSearcherCache::CacheValue* get_index_cache_value() { return ((InvertedIndexSearcherCache::CacheValue*)_cache->value(_handle)); } @@ -180,6 +198,23 @@ class InvertedIndexCacheHandle { class InvertedIndexQueryCacheHandle; +inline constexpr uint32_t INVERTED_INDEX_QUERY_CACHE_SEMANTICS_VERSION = 1; + +// Stable identity shared by result-cache and row-accurate single-flight. It intentionally contains +// no analyzer output or internal plan kind: those are segment-local implementation details below +// the cache lookup. +struct InvertedIndexRawQuerySemantic { + std::string_view raw_query_bytes; + InvertedIndexQueryType query_type; + int32_t slop = 0; + bool ordered = false; + int32_t max_expansions = 0; + uint32_t cache_semantics_version = INVERTED_INDEX_QUERY_CACHE_SEMANTICS_VERSION; + uint64_t common_grams_cache_generation = 0; + + std::string encode() const; +}; + class InvertedIndexQueryCache : public LRUCachePolicy { public: using LRUCachePolicy::insert; @@ -191,21 +226,8 @@ class InvertedIndexQueryCache : public LRUCachePolicy { InvertedIndexQueryType query_type; // query type std::string value; // query value - // Encode to a flat binary which can be used as LRUCache's key - std::string encode() const { - std::string key_buf(index_path.string()); - key_buf.append("/"); - key_buf.append(column_name); - key_buf.append("/"); - auto query_type_str = query_type_to_string(query_type); - if (query_type_str.empty()) { - return ""; - } - key_buf.append(query_type_str); - key_buf.append("/"); - key_buf.append(value); - return key_buf; - } + // Encode to an unambiguous flat binary which can be used as LRUCache's key. + std::string encode() const; }; class CacheValue : public LRUCacheValueBase { diff --git a/be/src/storage/index/inverted/inverted_index_fs_directory.cpp b/be/src/storage/index/inverted/inverted_index_fs_directory.cpp index 30f168e8b14e02..e332e3cde8069e 100644 --- a/be/src/storage/index/inverted/inverted_index_fs_directory.cpp +++ b/be/src/storage/index/inverted/inverted_index_fs_directory.cpp @@ -97,11 +97,22 @@ bool DorisFSDirectory::FSIndexInput::open(const io::FileSystemSPtr& fs, const ch auto h = std::make_shared(path); io::FileReaderOptions reader_options; - reader_options.cache_type = config::enable_file_cache ? io::FileCachePolicy::FILE_BLOCK_CACHE - : io::FileCachePolicy::NO_CACHE; + // DIAGNOSTIC: inverted_index_read_bypass_file_cache forces NO_CACHE (precise + // S3 range GETs) for inverted-index reads (both CLucene here and SNII), so a + // fair direct-vs-direct engine comparison can bypass the 1MiB block cache + // without disabling the global enable_file_cache (cloud mode requires it). + reader_options.cache_type = + (config::enable_file_cache && !config::inverted_index_read_bypass_file_cache) + ? io::FileCachePolicy::FILE_BLOCK_CACHE + : io::FileCachePolicy::NO_CACHE; reader_options.is_doris_table = true; reader_options.file_size = file_size; reader_options.tablet_id = tablet_id; + // NO_CACHE on a remote fs means every range GET bypasses CachedRemoteFileReader, + // the only layer that normally accounts physical remote bytes; readInternal then + // counts them itself. Local NO_CACHE reads must stay excluded. + h->_direct_remote_io = reader_options.cache_type == io::FileCachePolicy::NO_CACHE && + fs->type() != io::FileSystemType::LOCAL; Status st = fs->open_file(path, &h->_reader, &reader_options); DBUG_EXECUTE_IF("inverted file read error: index file not found", { st = Status::Error("index file not found"); }) @@ -179,20 +190,19 @@ void DorisFSDirectory::FSIndexInput::close() { } void DorisFSDirectory::FSIndexInput::setIoContext(const void* io_ctx) { + // Copy the caller's full IOContext (expiration_time for TTL cache classification, + // is_warmup, is_disposable, read_file_cache, bypass_peer_read, the miss policy and + // the remote-scan cache-write limiter all propagate); a field whitelist here silently + // drops newly added flags. Only the per-stream identity bits are re-stamped below. + const bool is_index_data = _io_ctx.is_index_data; if (io_ctx) { const auto& ctx = static_cast(io_ctx); - _io_ctx.reader_type = ctx->reader_type; - _io_ctx.query_id = ctx->query_id; - _io_ctx.file_cache_stats = ctx->file_cache_stats; - _io_ctx.file_cache_miss_policy = ctx->file_cache_miss_policy; - _io_ctx.remote_scan_cache_write_limiter = ctx->remote_scan_cache_write_limiter; + _io_ctx = *ctx; } else { - _io_ctx.reader_type = ReaderType::UNKNOWN; - _io_ctx.query_id = nullptr; - _io_ctx.file_cache_stats = nullptr; - _io_ctx.file_cache_miss_policy = io::FileCacheMissPolicy::READ_THROUGH_AND_WRITE_BACK; - _io_ctx.remote_scan_cache_write_limiter = nullptr; + _io_ctx = io::IOContext {}; } + _io_ctx.is_index_data = is_index_data; + _io_ctx.is_inverted_index = true; } const void* DorisFSDirectory::FSIndexInput::getIoContext() { @@ -251,6 +261,15 @@ void DorisFSDirectory::FSIndexInput::readInternal(uint8_t* b, const int32_t len) if (_io_ctx.file_cache_stats != nullptr) { _io_ctx.file_cache_stats->inverted_index_io_timer += inverted_index_io_timer; + _io_ctx.file_cache_stats->inverted_index_request_bytes += len; + _io_ctx.file_cache_stats->inverted_index_read_bytes += len; + ++_io_ctx.file_cache_stats->inverted_index_range_read_count; + ++_io_ctx.file_cache_stats->inverted_index_serial_read_rounds; + if (_handle->_direct_remote_io) { + // Cache-bypassed remote read: no CachedRemoteFileReader below to count + // the GET, so account the physical remote bytes here. + _io_ctx.file_cache_stats->inverted_index_remote_physical_read_bytes += len; + } } } diff --git a/be/src/storage/index/inverted/inverted_index_fs_directory.h b/be/src/storage/index/inverted/inverted_index_fs_directory.h index 79854df88d235d..64719880e7caa9 100644 --- a/be/src/storage/index/inverted/inverted_index_fs_directory.h +++ b/be/src/storage/index/inverted/inverted_index_fs_directory.h @@ -170,6 +170,11 @@ class DorisFSDirectory::FSIndexInput : public lucene::store::BufferedIndexInput std::mutex _shared_lock; //std::mutex* _shared_lock = nullptr; char path[4096]; + // True when _reader serves ranges straight from remote storage with no + // CachedRemoteFileReader in between (NO_CACHE on a non-local fs): only + // that layer would normally account physical remote bytes, so + // readInternal then counts them itself. + bool _direct_remote_io = false; SharedHandle(const char* path); ~SharedHandle() override; }; diff --git a/be/src/storage/index/inverted/inverted_index_iterator.cpp b/be/src/storage/index/inverted/inverted_index_iterator.cpp index 6e93d70d025964..6fbaa5bc944ec3 100644 --- a/be/src/storage/index/inverted/inverted_index_iterator.cpp +++ b/be/src/storage/index/inverted/inverted_index_iterator.cpp @@ -46,12 +46,13 @@ void InvertedIndexIterator::add_reader(InvertedIndexReaderType type, VLOG_DEBUG << "InvertedIndexIterator add_reader: type=" << static_cast(type) << ", analyzer_key=" << analyzer_key; - const size_t entry_index = _reader_entries.size(); - _reader_entries.push_back( - ReaderEntry {.type = type, .analyzer_key = std::move(analyzer_key), .reader = reader}); - - // Update index for O(1) lookup - _key_to_entries[_reader_entries.back().analyzer_key].push_back(entry_index); + auto status = add_inverted_index_selection_candidate( + InvertedIndexSelectionCandidate {.index_id = cast_set(reader->get_index_id()), + .reader_type = type, + .analyzer_key = std::move(analyzer_key)}, + &_selection_candidates, &_key_to_entries); + DORIS_CHECK(status.ok()) << status; + _readers.push_back(reader); } Status InvertedIndexIterator::read_from_index(const IndexParam& param) { @@ -101,6 +102,11 @@ Status InvertedIndexIterator::read_from_index(const IndexParam& param) { // Note: analyzer_ctx is now passed via i_param->analyzer_ctx auto execute_query = [&]() { + if (i_param->null_bitmap_cache_handle != nullptr) { + return reader->query_with_null_bitmap( + _context, i_param->column_name, i_param->query_value, i_param->query_type, + i_param->roaring, i_param->null_bitmap_cache_handle, i_param->analyzer_ctx); + } return reader->query(_context, i_param->column_name, i_param->query_value, i_param->query_type, i_param->roaring, i_param->analyzer_ctx); }; @@ -122,14 +128,12 @@ Status InvertedIndexIterator::read_from_index(const IndexParam& param) { } Status InvertedIndexIterator::read_null_bitmap(InvertedIndexQueryCacheHandle* cache_handle) { - // For null bitmap, use any available reader (empty = auto-select) - auto reader = DORIS_TRY(select_best_reader("")); + auto reader = DORIS_TRY(select_any_reader()); return reader->read_null_bitmap(_context, cache_handle, nullptr); } Result InvertedIndexIterator::has_null() { - // For has_null check, use any available reader (empty = auto-select) - auto reader = DORIS_TRY(select_best_reader("")); + auto reader = DORIS_TRY(select_any_reader()); return reader->has_null(); } @@ -149,177 +153,48 @@ Status InvertedIndexIterator::try_read_from_inverted_index(const InvertedIndexRe return Status::OK(); } -// When multiple candidates of the preferred type exist, pick the one with -// the smallest index_id so that the choice is deterministic regardless of -// the order indexes appear in the rowset schema. Different segments may -// have different index orderings (e.g. after sequential BUILD INDEX -// operations), and relying on iteration order would cause inconsistent -// query results across segments. -static const ReaderEntry* pick_preferred(const std::vector& candidates, - InvertedIndexReaderType preferred_type) { - const ReaderEntry* best = nullptr; - for (const auto* entry : candidates) { - if (entry->type == preferred_type) { - if (best == nullptr || entry->reader->get_index_id() < best->reader->get_index_id()) { - best = entry; - } - } - } - return best; -} - -static const ReaderEntry* pick_smallest_index_id( - const std::vector& candidates) { - const ReaderEntry* best = candidates.front(); - for (const auto* entry : candidates) { - if (entry->reader->get_index_id() < best->reader->get_index_id()) { - best = entry; - } - } - return best; -} - -Result InvertedIndexIterator::select_for_text( - const AnalyzerMatchResult& match, InvertedIndexQueryType query_type, - const std::string& analyzer_key) { - // Bypass: explicit analyzer specified but not found - if (match.empty() && AnalyzerKeyMatcher::is_explicit(analyzer_key)) { - return ResultError(Status::Error( - "No inverted index reader found for analyzer '{}'. " - "The index for this analyzer may not be built yet.", - analyzer_key)); - } - - if (match.empty()) { - return ResultError(Status::Error( - "No available inverted index readers for text column.")); - } - - // MATCH queries prefer FULLTEXT - if (is_match_query(query_type)) { - if (auto* best = pick_preferred(match.candidates, InvertedIndexReaderType::FULLTEXT)) { - return best->reader; - } - } - - // EQUAL queries prefer STRING_TYPE for exact match - if (is_equal_query(query_type)) { - if (auto* best = pick_preferred(match.candidates, InvertedIndexReaderType::STRING_TYPE)) { - return best->reader; - } - } - - // Default: smallest index_id for deterministic selection - return pick_smallest_index_id(match.candidates)->reader; -} - -Result InvertedIndexIterator::select_for_numeric( - const AnalyzerMatchResult& match, InvertedIndexQueryType query_type) { - if (match.empty()) { - return ResultError(Status::Error( - "No available inverted index readers for numeric column.")); - } - - // RANGE queries prefer BKD - if (is_range_query(query_type)) { - if (const auto* best = pick_preferred(match.candidates, InvertedIndexReaderType::BKD)) { - return best->reader; - } - } - - // Fallback priority: BKD > STRING_TYPE > smallest index_id - if (const auto* best = pick_preferred(match.candidates, InvertedIndexReaderType::BKD)) { - return best->reader; - } - if (const auto* best = pick_preferred(match.candidates, InvertedIndexReaderType::STRING_TYPE)) { - return best->reader; - } - - // Last resort: smallest index_id for deterministic selection - return pick_smallest_index_id(match.candidates)->reader; -} - Result InvertedIndexIterator::select_best_reader( const DataTypePtr& column_type, InvertedIndexQueryType query_type, const std::string& analyzer_key) { - if (_reader_entries.empty()) { - return ResultError(Status::Error( - "No available inverted index readers. Check if index is properly initialized.")); - } - - // Normalize once at entry point const std::string normalized_key = ensure_normalized_key(analyzer_key); + const auto field_type = get_inverted_index_leaf_field_type(column_type); + auto selection = select_best_inverted_index_candidate(_selection_candidates, _key_to_entries, + field_type, query_type, normalized_key); + if (!selection.has_value()) { + return ResultError(std::move(selection.error())); + } + const size_t selected = *selection; + DORIS_CHECK(selected < _readers.size()); + return _readers[selected]; +} - // Single reader optimization - if (_reader_entries.size() == 1) { - const auto& entry = _reader_entries.front(); - if (AnalyzerKeyMatcher::is_explicit(normalized_key) && - entry.analyzer_key != normalized_key) { - return ResultError(Status::Error( - "No inverted index reader found for analyzer '{}'. " - "Available analyzer: '{}'.", - normalized_key, entry.analyzer_key)); - } - return entry.reader; - } - - // Match analyzer key using AnalyzerKeyMatcher - auto match = AnalyzerKeyMatcher::match(normalized_key, _reader_entries, _key_to_entries); - - // Dispatch by column type - const auto field_type = column_type->get_storage_field_type(); - - if (is_string_type(field_type)) { - return select_for_text(match, query_type, normalized_key); - } - - if (field_is_numeric_type(field_type)) { - return select_for_numeric(match, query_type); - } - - // Default: return deterministic candidate or error - if (match.empty()) { - return ResultError(Status::Error( - "No available inverted index readers for column type.")); +Result InvertedIndexIterator::select_any_reader() { + auto selection = select_best_inverted_index_candidate( + _selection_candidates, _key_to_entries, FieldType::OLAP_FIELD_TYPE_UNKNOWN, + InvertedIndexQueryType::UNKNOWN_QUERY, ""); + if (!selection.has_value()) { + return ResultError(std::move(selection.error())); } - return pick_smallest_index_id(match.candidates)->reader; + const size_t selected = *selection; + DORIS_CHECK(selected < _readers.size()); + return _readers[selected]; } Result InvertedIndexIterator::select_best_reader( const std::string& analyzer_key) { - if (_reader_entries.empty()) { - return ResultError(Status::Error( - "No available inverted index readers. Check if index is properly initialized.")); + if (analyzer_key.empty()) { + return select_any_reader(); } - const std::string normalized_key = ensure_normalized_key(analyzer_key); - - // Single reader optimization - if (_reader_entries.size() == 1) { - const auto& entry = _reader_entries.front(); - if (AnalyzerKeyMatcher::is_explicit(normalized_key) && - entry.analyzer_key != normalized_key) { - return ResultError(Status::Error( - "No inverted index reader found for analyzer '{}'. " - "Available analyzer: '{}'.", - normalized_key, entry.analyzer_key)); - } - return entry.reader; - } - - // Match and return deterministic candidate - auto match = AnalyzerKeyMatcher::match(normalized_key, _reader_entries, _key_to_entries); - - if (match.empty()) { - if (AnalyzerKeyMatcher::is_explicit(normalized_key)) { - return ResultError(Status::Error( - "No inverted index reader found for analyzer '{}'.", normalized_key)); - } - return ResultError(Status::Error( - "No available inverted index readers.")); - } - - return pick_smallest_index_id(match.candidates)->reader; + auto selection = select_best_inverted_index_candidate( + _selection_candidates, _key_to_entries, FieldType::OLAP_FIELD_TYPE_UNKNOWN, + InvertedIndexQueryType::UNKNOWN_QUERY, normalized_key); + if (!selection.has_value()) { + return ResultError(std::move(selection.error())); + } + const size_t selected = *selection; + DORIS_CHECK(selected < _readers.size()); + return _readers[selected]; } IndexReaderPtr InvertedIndexIterator::get_reader(IndexReaderType type) const { @@ -327,9 +202,10 @@ IndexReaderPtr InvertedIndexIterator::get_reader(IndexReaderType type) const { if (inverted_type == nullptr) { return nullptr; } - for (const auto& entry : _reader_entries) { - if (entry.type == *inverted_type) { - return entry.reader; + for (size_t i = 0; i < _selection_candidates.size(); ++i) { + if (_selection_candidates[i].reader_type == *inverted_type) { + DORIS_CHECK(i < _readers.size()); + return _readers[i]; } } return nullptr; diff --git a/be/src/storage/index/inverted/inverted_index_iterator.h b/be/src/storage/index/inverted/inverted_index_iterator.h index afc4a663670633..5f7c2250cc0836 100644 --- a/be/src/storage/index/inverted/inverted_index_iterator.h +++ b/be/src/storage/index/inverted/inverted_index_iterator.h @@ -20,10 +20,10 @@ #include #include "core/field.h" -#include "storage/index/analyzer_key_matcher.h" #include "storage/index/index_iterator.h" #include "storage/index/inverted/inverted_index_parser.h" #include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/inverted/inverted_index_selector.h" namespace doris::segment_v2 { @@ -35,19 +35,13 @@ struct InvertedIndexParam { uint32_t num_rows; std::shared_ptr roaring; bool skip_try = false; + // Non-null only when the caller consumes both the query result and this reader's null bitmap. + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle = nullptr; // Pointer to analyzer context (can be nullptr if not needed) // Used by FullTextIndexReader for tokenization const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr; }; -// Entry representing an inverted index reader with its type and analyzer key. -// Used by InvertedIndexIterator and AnalyzerKeyMatcher for reader selection. -struct ReaderEntry { - InvertedIndexReaderType type; - std::string analyzer_key; - InvertedIndexReaderPtr reader; -}; - class InvertedIndexIterator : public IndexIterator { public: InvertedIndexIterator(); @@ -67,6 +61,10 @@ class InvertedIndexIterator : public IndexIterator { [[nodiscard]] Result select_best_reader( const DataTypePtr& column_type, InvertedIndexQueryType query_type, const std::string& analyzer_key); + + [[nodiscard]] Result select_any_reader(); + + // Temporary compatibility for variant fields whose runtime binding has no type. [[nodiscard]] Result select_best_reader( const std::string& analyzer_key); @@ -81,27 +79,16 @@ class InvertedIndexIterator : public IndexIterator { // Empty input stays empty (means "user did not specify"). static std::string ensure_normalized_key(const std::string& analyzer_key); - // Select best reader for text (string) columns. - // Handles FULLTEXT vs STRING_TYPE priority based on query type. - // Returns BYPASS error if explicit analyzer not found. - [[nodiscard]] Result select_for_text(const AnalyzerMatchResult& match, - InvertedIndexQueryType query_type, - const std::string& analyzer_key); - - // Select best reader for numeric columns. - // Handles BKD priority for range queries. - [[nodiscard]] Result select_for_numeric( - const AnalyzerMatchResult& match, InvertedIndexQueryType query_type); - - // THREAD SAFETY: _reader_entries and _key_to_entries are populated during initialization + // THREAD SAFETY: reader metadata and _key_to_entries are populated during initialization // phase (via add_reader) and only read during query phase (via read_from_index/select_best_reader). // These two phases are guaranteed not to overlap, so no synchronization is needed. // Do NOT call add_reader() after any read_from_index() call on the same iterator. - std::vector _reader_entries; + std::vector _selection_candidates; + std::vector _readers; - // Index for O(1) lookup by analyzer_key. Maps normalized key to indices in _reader_entries. + // Index for O(1) lookup by analyzer_key. Maps normalized key to candidate indices. // Built incrementally in add_reader(). - std::unordered_map> _key_to_entries; + InvertedIndexSelectionKeyIndex _key_to_entries; }; } // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/storage/index/inverted/inverted_index_parser.cpp b/be/src/storage/index/inverted/inverted_index_parser.cpp index 47819cc62f6397..1bd0a41454de80 100644 --- a/be/src/storage/index/inverted/inverted_index_parser.cpp +++ b/be/src/storage/index/inverted/inverted_index_parser.cpp @@ -190,16 +190,11 @@ std::string normalize_analyzer_key(std::string_view analyzer) { std::string build_analyzer_key_from_properties( const std::map& properties) { - // Build analyzer key from index properties for reader registration. - // This determines how the index is stored/identified. - - // 1. Check for custom analyzer name - auto custom_it = properties.find(INVERTED_INDEX_ANALYZER_NAME_KEY); - if (custom_it != properties.end() && !custom_it->second.empty()) { - return to_lower(custom_it->second); + auto analyzer = properties.find(INVERTED_INDEX_ANALYZER_NAME_KEY); + if (analyzer != properties.end() && !analyzer->second.empty()) { + return normalize_analyzer_key(analyzer->second); } - // 2. Fall back to parser type std::string parser; auto parser_it = properties.find(INVERTED_INDEX_PARSER_KEY); if (parser_it != properties.end()) { @@ -211,11 +206,10 @@ std::string build_analyzer_key_from_properties( } } - // 3. Return normalized parser or "" for no explicit configuration if (parser.empty()) { - return ""; // No explicit parser - empty key means "no configuration" + return ""; } - return to_lower(parser); + return normalize_analyzer_key(parser); } // ============================================================================ diff --git a/be/src/storage/index/inverted/inverted_index_parser.h b/be/src/storage/index/inverted/inverted_index_parser.h index d2d3df47abd0a3..c2ab8aa6727927 100644 --- a/be/src/storage/index/inverted/inverted_index_parser.h +++ b/be/src/storage/index/inverted/inverted_index_parser.h @@ -19,9 +19,12 @@ #include #include +#include #include #include +#include "storage/index/inverted/analyzer/analyzer_provider.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" #include "util/debug_points.h" namespace lucene { @@ -119,6 +122,36 @@ struct InvertedIndexAnalyzerCtx { // Used for creating reader and tokenization CharFilterMap char_filter_map; std::shared_ptr analyzer; + segment_v2::inverted_index::AnalyzerProviderPtr analyzer_provider; + std::optional common_grams_identity; + + std::shared_ptr get_analyzer( + segment_v2::inverted_index::AnalysisPurpose purpose) const { + if (analyzer_provider != nullptr) { + return analyzer_provider->get_analyzer(purpose); + } + return analyzer; + } + + const segment_v2::inverted_index::CommonGramsQueryIdentity* get_common_grams_identity() const { + if (common_grams_identity.has_value()) { + return &*common_grams_identity; + } + return analyzer_provider == nullptr ? nullptr : analyzer_provider->common_grams_identity(); + } + + bool has_complete_common_grams_identity() const { + const auto* identity = get_common_grams_identity(); + return identity != nullptr && !identity->common_grams_dictionary_identity.empty() && + !identity->base_analyzer_fingerprint.empty() && + !identity->common_grams_fingerprint.empty(); + } + + // Raw-query cache and single-flight keys intentionally exclude analyzer output. A tokenizing + // provider therefore needs a complete immutable identity before those results may be shared. + bool can_share_raw_query_semantics() const { + return !should_tokenize() || has_complete_common_grams_identity(); + } // Returns true if tokenization should be performed. // Decision is based on parser_type (from index properties): @@ -170,7 +203,7 @@ std::string get_parser_dict_compression_from_properties( std::string get_analyzer_name_from_properties(const std::map& properties); // Build a normalized analyzer key from index properties. -// Checks custom_analyzer first, then falls back to parser type. +// Precedence is analyzer, then parser type. Normalizer is not an index-selection key. std::string build_analyzer_key_from_properties( const std::map& properties); diff --git a/be/src/storage/index/inverted/inverted_index_reader.cpp b/be/src/storage/index/inverted/inverted_index_reader.cpp index 0e83238064297e..23d568aa8dcd8f 100644 --- a/be/src/storage/index/inverted/inverted_index_reader.cpp +++ b/be/src/storage/index/inverted/inverted_index_reader.cpp @@ -154,6 +154,20 @@ std::string InvertedIndexReader::get_index_file_path() { return _index_file_reader->get_index_file_path(&_index_meta); } +Status InvertedIndexReader::query_with_null_bitmap( + const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx) { + DORIS_CHECK(null_bitmap_cache_handle != nullptr); + RETURN_IF_ERROR(query(context, column_name, query_value, query_type, bit_map, analyzer_ctx)); + if (!has_null()) { + return Status::OK(); + } + return read_null_bitmap(context, null_bitmap_cache_handle); +} + Status InvertedIndexReader::read_null_bitmap(const IndexQueryContextPtr& context, InvertedIndexQueryCacheHandle* cache_handle, lucene::store::Directory* dir) { @@ -219,15 +233,16 @@ bool InvertedIndexReader::handle_query_cache(const IndexQueryContextPtr& context InvertedIndexQueryCache* cache, const InvertedIndexQueryCache::CacheKey& cache_key, InvertedIndexQueryCacheHandle* cache_handler, - std::shared_ptr& bit_map) { + std::shared_ptr& bit_map, + bool enabled) { const auto& query_options = context->runtime_state->query_options(); - - bool cache_hit = false; - if (query_options.enable_inverted_index_query_cache) { - SCOPED_RAW_TIMER(&context->stats->inverted_index_lookup_timer); - cache_hit = cache->lookup(cache_key, cache_handler); + if (!enabled || !query_options.enable_inverted_index_query_cache) { + return false; } + context->stats->inverted_index_query_cache_lookup++; + SCOPED_RAW_TIMER(&context->stats->inverted_index_lookup_timer); + const bool cache_hit = cache->lookup(cache_key, cache_handler); if (cache_hit) { DBUG_EXECUTE_IF("InvertedIndexReader.handle_query_cache_hit", { return Status::Error("handle query cache hit"); @@ -245,6 +260,19 @@ bool InvertedIndexReader::handle_query_cache(const IndexQueryContextPtr& context return false; } +void InvertedIndexReader::insert_query_cache(const IndexQueryContextPtr& context, + InvertedIndexQueryCache* cache, + const InvertedIndexQueryCache::CacheKey& cache_key, + std::shared_ptr bit_map, + InvertedIndexQueryCacheHandle* cache_handler, + bool enabled) { + if (!enabled || !context->runtime_state->query_options().enable_inverted_index_query_cache) { + return; + } + cache->insert(cache_key, std::move(bit_map), cache_handler); + context->stats->inverted_index_query_cache_insert++; +} + Status InvertedIndexReader::handle_searcher_cache( const IndexQueryContextPtr& context, InvertedIndexCacheHandle* inverted_index_cache_handle) { diff --git a/be/src/storage/index/inverted/inverted_index_reader.h b/be/src/storage/index/inverted/inverted_index_reader.h index 0e2f6a120d41e3..6960e135e572b5 100644 --- a/be/src/storage/index/inverted/inverted_index_reader.h +++ b/be/src/storage/index/inverted/inverted_index_reader.h @@ -226,13 +226,19 @@ class InvertedIndexReader : public IndexReader { const Field& query_value, InvertedIndexQueryType query_type, std::shared_ptr& bit_map, const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr) = 0; + virtual Status query_with_null_bitmap(const IndexQueryContextPtr& context, + const std::string& column_name, const Field& query_value, + InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr); virtual Status try_query(const IndexQueryContextPtr& context, const std::string& column_name, const Field& query_value, InvertedIndexQueryType query_type, size_t* count) = 0; - Status read_null_bitmap(const IndexQueryContextPtr& context, - InvertedIndexQueryCacheHandle* cache_handle, - lucene::store::Directory* dir = nullptr); + virtual Status read_null_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryCacheHandle* cache_handle, + lucene::store::Directory* dir = nullptr); virtual InvertedIndexReaderType type() = 0; @@ -249,7 +255,11 @@ class InvertedIndexReader : public IndexReader { bool handle_query_cache(const IndexQueryContextPtr& context, InvertedIndexQueryCache* cache, const InvertedIndexQueryCache::CacheKey& cache_key, InvertedIndexQueryCacheHandle* cache_handler, - std::shared_ptr& bit_map); + std::shared_ptr& bit_map, bool enabled = true); + void insert_query_cache(const IndexQueryContextPtr& context, InvertedIndexQueryCache* cache, + const InvertedIndexQueryCache::CacheKey& cache_key, + std::shared_ptr bit_map, + InvertedIndexQueryCacheHandle* cache_handler, bool enabled = true); virtual Status handle_searcher_cache(const IndexQueryContextPtr& context, InvertedIndexCacheHandle* inverted_index_cache_handle); @@ -335,7 +345,6 @@ class InvertedIndexVisitor : public lucene::util::bkd::bkd_reader::intersect_vis std::string query_min; std::string query_max; -public: InvertedIndexVisitor(const void* io_ctx, lucene::util::bkd::bkd_reader* r, roaring::Roaring* hits, bool only_count = false); ~InvertedIndexVisitor() override = default; diff --git a/be/src/storage/index/inverted/inverted_index_selector.cpp b/be/src/storage/index/inverted/inverted_index_selector.cpp new file mode 100644 index 00000000000000..dac0ec9b41b1cb --- /dev/null +++ b/be/src/storage/index/inverted/inverted_index_selector.cpp @@ -0,0 +1,145 @@ +// 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 "storage/index/inverted/inverted_index_selector.h" + +#include + +#include "common/logging.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/utils.h" + +namespace doris::segment_v2 { + +Status add_inverted_index_selection_candidate( + InvertedIndexSelectionCandidate candidate, + std::vector* candidates, + InvertedIndexSelectionKeyIndex* key_index) { + DORIS_CHECK(candidates != nullptr); + DORIS_CHECK(key_index != nullptr); + for (const auto& existing : *candidates) { + if (existing.index_id == candidate.index_id) { + return Status::Error( + "Duplicate inverted index id {} in one field", candidate.index_id); + } + } + + const size_t candidate_index = candidates->size(); + candidates->push_back(std::move(candidate)); + (*key_index)[candidates->back().analyzer_key].push_back(candidate_index); + return Status::OK(); +} + +Result select_best_inverted_index_candidate( + const std::vector& candidates, + const InvertedIndexSelectionKeyIndex& key_index, FieldType field_type, + InvertedIndexQueryType query_type, std::string_view normalized_analyzer_key) { + if (candidates.empty()) { + return ResultError(Status::Error( + "No available inverted index candidates")); + } + + const std::vector* exact_candidates = nullptr; + if (!normalized_analyzer_key.empty()) { + const auto exact = key_index.find(std::string(normalized_analyzer_key)); + if (exact == key_index.end() || exact->second.empty()) { + return ResultError(Status::Error( + "No inverted index found for analyzer '{}'", normalized_analyzer_key)); + } + exact_candidates = &exact->second; + } + + const size_t candidate_count = + exact_candidates == nullptr ? candidates.size() : exact_candidates->size(); + auto candidate_at = + [&](size_t ordinal) -> std::pair { + const size_t index = exact_candidates == nullptr ? ordinal : (*exact_candidates)[ordinal]; + DORIS_CHECK(index < candidates.size()); + return {index, candidates[index]}; + }; + auto pick = + [&](std::optional preferred_type) -> std::optional { + std::optional best; + for (size_t ordinal = 0; ordinal < candidate_count; ++ordinal) { + const auto [index, candidate] = candidate_at(ordinal); + if (preferred_type.has_value() && candidate.reader_type != *preferred_type) { + continue; + } + if (!best.has_value() || candidate.index_id < candidates[*best].index_id) { + best = index; + } + } + return best; + }; + + if (is_string_type(field_type)) { + if (is_match_query(query_type)) { + if (auto best = pick(InvertedIndexReaderType::FULLTEXT); best.has_value()) { + return *best; + } + } + if (is_equal_query(query_type)) { + if (auto best = pick(InvertedIndexReaderType::STRING_TYPE); best.has_value()) { + return *best; + } + } + } else if (field_is_numeric_type(field_type)) { + if (is_range_query(query_type)) { + if (auto best = pick(InvertedIndexReaderType::BKD); best.has_value()) { + return *best; + } + } + if (auto best = pick(InvertedIndexReaderType::BKD); best.has_value()) { + return *best; + } + if (auto best = pick(InvertedIndexReaderType::STRING_TYPE); best.has_value()) { + return *best; + } + } + + auto best = pick(std::nullopt); + DORIS_CHECK(best.has_value()); + return *best; +} + +FieldType get_inverted_index_leaf_field_type(const DataTypePtr& column_type) { + DORIS_CHECK(column_type != nullptr); + DataTypePtr leaf_type = remove_nullable(column_type); + while (leaf_type->get_storage_field_type() == FieldType::OLAP_FIELD_TYPE_ARRAY) { + const auto* array_type = dynamic_cast(leaf_type.get()); + DORIS_CHECK(array_type != nullptr); + leaf_type = remove_nullable(array_type->get_nested_type()); + } + return leaf_type->get_storage_field_type(); +} + +InvertedIndexReaderType infer_inverted_index_reader_type( + FieldType field_type, const std::map& properties) { + if (is_string_type(field_type)) { + return inverted_index::InvertedIndexAnalyzer::should_analyzer(properties) + ? InvertedIndexReaderType::FULLTEXT + : InvertedIndexReaderType::STRING_TYPE; + } + if (field_is_numeric_type(field_type)) { + return InvertedIndexReaderType::BKD; + } + return InvertedIndexReaderType::UNKNOWN; +} + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/inverted/inverted_index_selector.h b/be/src/storage/index/inverted/inverted_index_selector.h new file mode 100644 index 00000000000000..c295a9db3bc33e --- /dev/null +++ b/be/src/storage/index/inverted/inverted_index_selector.h @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "core/data_type/data_type.h" +#include "storage/index/inverted/inverted_index_query_type.h" +#include "storage/olap_common.h" + +namespace doris::segment_v2 { + +struct InvertedIndexSelectionCandidate { + int64_t index_id; + InvertedIndexReaderType reader_type; + std::string analyzer_key; +}; + +using InvertedIndexSelectionKeyIndex = std::unordered_map>; + +Status add_inverted_index_selection_candidate( + InvertedIndexSelectionCandidate candidate, + std::vector* candidates, + InvertedIndexSelectionKeyIndex* key_index); + +[[nodiscard]] Result select_best_inverted_index_candidate( + const std::vector& candidates, + const InvertedIndexSelectionKeyIndex& key_index, FieldType field_type, + InvertedIndexQueryType query_type, std::string_view normalized_analyzer_key); + +FieldType get_inverted_index_leaf_field_type(const DataTypePtr& column_type); + +InvertedIndexReaderType infer_inverted_index_reader_type( + FieldType field_type, const std::map& properties); + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/inverted/query/conjunction_query.cpp b/be/src/storage/index/inverted/query/conjunction_query.cpp index 8fc3ba0b16148d..eb25605d938803 100644 --- a/be/src/storage/index/inverted/query/conjunction_query.cpp +++ b/be/src/storage/index/inverted/query/conjunction_query.cpp @@ -17,8 +17,8 @@ #include "storage/index/inverted/query/conjunction_query.h" -#include "storage/compaction/collection_statistics.h" #include "storage/index/inverted/query/query_helper.h" +#include "storage/index/inverted/similarity/collection_statistics.h" #include "storage/index/inverted/util/mock_iterator.h" #include "storage/index/inverted/util/string_helper.h" diff --git a/be/src/storage/index/inverted/query/query_info.h b/be/src/storage/index/inverted/query/query_info.h index 829faaca4d7f24..685d3eee951b84 100644 --- a/be/src/storage/index/inverted/query/query_info.h +++ b/be/src/storage/index/inverted/query/query_info.h @@ -17,17 +17,25 @@ #pragma once +#include +#include #include #include namespace doris::segment_v2 { +enum class TermKeyKind : uint8_t { + kPlain = 0, + kCommonGram = 1, +}; + class TermInfo { public: using Term = std::variant>; Term term; int32_t position = 0; + TermKeyKind key_kind = TermKeyKind::kPlain; bool is_single_term() const { return std::holds_alternative(term); } bool is_multi_terms() const { return std::holds_alternative>(term); } @@ -49,6 +57,15 @@ class InvertedIndexQueryInfo { // for test bool use_mock_iter = false; + bool has_common_gram() const { + for (const auto& term_info : term_infos) { + if (term_info.key_kind == TermKeyKind::kCommonGram) { + return true; + } + } + return false; + } + std::string generate_tokens_key() const { std::string key; for (const auto& token : term_infos) { diff --git a/be/src/storage/index/inverted/setting.h b/be/src/storage/index/inverted/setting.h index 51782ab0b2de5d..8e2a7ce1a72dc4 100644 --- a/be/src/storage/index/inverted/setting.h +++ b/be/src/storage/index/inverted/setting.h @@ -19,12 +19,14 @@ #include +#include #include #include #include #include #include #include +#include #include "common/exception.h" @@ -168,6 +170,12 @@ class Settings { return result; } + std::vector> sorted_entries() const { + std::vector> entries(_args.begin(), _args.end()); + std::sort(entries.begin(), entries.end()); + return entries; + } + private: std::unordered_map _args; }; diff --git a/be/src/storage/index/inverted/similarity/collection_statistics.cpp b/be/src/storage/index/inverted/similarity/collection_statistics.cpp new file mode 100644 index 00000000000000..87f42e31f8e75e --- /dev/null +++ b/be/src/storage/index/inverted/similarity/collection_statistics.cpp @@ -0,0 +1,529 @@ +// 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 "storage/index/inverted/similarity/collection_statistics.h" + +#include +#include +#include + +#include "common/exception.h" +#include "exprs/vexpr.h" +#include "exprs/vexpr_context.h" +#include "exprs/vliteral.h" +#include "exprs/vslot_ref.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/index_reader_helper.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/util/string_helper.h" +#include "storage/index/inverted/util/term_iterator.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/reader/dict_block_cache.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/rowset/rowset.h" +#include "storage/rowset/rowset_reader.h" +#include "util/uid_util.h" + +namespace doris { +namespace { + +struct SniiScoringSegmentStats { + uint64_t doc_count = 0; + uint64_t token_count = 0; + segment_v2::inverted_index::PlainTermKeyVersion plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw; + std::string base_analyzer_fingerprint; +}; + +Result resolve_snii_scoring_segment( + const std::optional& metadata, + uint64_t index_doc_count, uint64_t physical_sum_total_term_freq, bool has_scoring_tier, + bool has_positions, bool has_semantic_norms) { + using namespace segment_v2::inverted_index; + const auto* metadata_ptr = metadata ? &*metadata : nullptr; + auto validation_status = validate_snii_scoring_metadata( + metadata_ptr, index_doc_count, physical_sum_total_term_freq, has_scoring_tier, + has_positions, has_semantic_norms); + if (!validation_status.ok()) { + return ResultError(std::move(validation_status)); + } + DORIS_CHECK(metadata_ptr != nullptr); + + return SniiScoringSegmentStats { + .doc_count = metadata_ptr->scoring_doc_count, + .token_count = metadata_ptr->scoring_token_count, + .plain_term_key_version = metadata_ptr->plain_term_key_version, + .base_analyzer_fingerprint = metadata_ptr->base_analyzer_fingerprint}; +} + +void add_term_doc_frequency( + std::unordered_map>* + logical_frequencies, + const std::wstring& field, const std::wstring& logical_term, uint64_t doc_frequency) { + DORIS_CHECK(logical_frequencies != nullptr); + (*logical_frequencies)[field][logical_term] += doc_frequency; +} + +} // namespace + +Status CollectionStatistics::collect(RuntimeState* state, + const std::vector& rs_splits, + const TabletSchemaSPtr& tablet_schema, + const VExprContextSPtrs& common_expr_ctxs_push_down, + io::IOContext* io_ctx) { + std::vector rowsets; + rowsets.reserve(rs_splits.size()); + for (const auto& rs_split : rs_splits) { + DORIS_CHECK(rs_split.rs_reader != nullptr); + auto rowset = rs_split.rs_reader->rowset(); + DORIS_CHECK(rowset != nullptr); + rowsets.emplace_back(std::move(rowset)); + } + return collect_full_collection(state, rowsets, tablet_schema, common_expr_ctxs_push_down, + io_ctx); +} + +Status CollectionStatistics::collect_full_collection( + RuntimeState* state, const std::vector& rowsets, + const TabletSchemaSPtr& tablet_schema, const VExprContextSPtrs& common_expr_ctxs_push_down, + io::IOContext* io_ctx) { + clear(); + std::unordered_map collect_infos; + RETURN_IF_ERROR( + extract_collect_info(state, common_expr_ctxs_push_down, tablet_schema, &collect_infos)); + if (collect_infos.empty()) { + LOG(WARNING) << "Index statistics collection: no collect info extracted."; + return Status::OK(); + } + + for (const auto& rowset : rowsets) { + DORIS_CHECK(rowset != nullptr); + const auto num_segments = rowset->num_segments(); + for (int64_t seg_id = 0; seg_id < num_segments; ++seg_id) { + auto status = + process_segment(rowset, seg_id, tablet_schema.get(), collect_infos, io_ctx); + if (!status.ok()) { + if (tablet_schema->get_inverted_index_storage_format() != + InvertedIndexStorageFormatPB::SNII && + (status.code() == ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND || + status.code() == ErrorCode::INVERTED_INDEX_BYPASS)) { + LOG(ERROR) << "Index statistics collection failed: " << status.to_string(); + continue; + } + clear(); + return status; + } + } + } + + // Build a single-line log with query_id, tablet_ids, and per-field term statistics + if (VLOG_IS_ON(1)) { + std::set tablet_ids; + for (const auto& rowset : rowsets) { + DORIS_CHECK(rowset != nullptr); + tablet_ids.insert(rowset->rowset_meta()->tablet_id()); + } + + std::ostringstream oss; + oss << "CollectionStatistics: query_id=" << print_id(state->query_id()); + + oss << ", tablet_ids=["; + bool first_tablet = true; + for (int64_t tid : tablet_ids) { + if (!first_tablet) oss << ","; + oss << tid; + first_tablet = false; + } + oss << "]"; + + oss << ", total_num_docs=" << _total_num_docs; + + for (const auto& [ws_field_name, num_tokens] : _total_num_tokens) { + oss << ", {field=" << StringHelper::to_string(ws_field_name) + << ", num_tokens=" << num_tokens << ", terms=["; + + auto field_term_doc_freqs = _term_doc_freqs.find(ws_field_name); + if (field_term_doc_freqs != _term_doc_freqs.end()) { + bool first_term = true; + for (const auto& [term, doc_freq] : field_term_doc_freqs->second) { + if (!first_term) oss << ", "; + oss << "(" << StringHelper::to_string(term) << ":" << doc_freq << ")"; + first_term = false; + } + } + oss << "]}"; + } + + VLOG(1) << oss.str(); + } + + return Status::OK(); +} + +Status CollectionStatistics::extract_collect_info( + RuntimeState* state, const VExprContextSPtrs& common_expr_ctxs_push_down, + const TabletSchemaSPtr& tablet_schema, CollectInfoMap* collect_infos) { + DCHECK(collect_infos != nullptr); + + std::unordered_map collectors; + collectors[TExprNodeType::MATCH_PRED] = std::make_unique(); + collectors[TExprNodeType::SEARCH_EXPR] = std::make_unique(); + + for (const auto& root_expr_ctx : common_expr_ctxs_push_down) { + const auto& root_expr = root_expr_ctx->root(); + if (root_expr == nullptr) { + continue; + } + + std::stack stack; + stack.emplace(root_expr); + + while (!stack.empty()) { + auto expr = stack.top(); + stack.pop(); + + if (!expr) { + continue; + } + + auto collector_it = collectors.find(expr->node_type()); + if (collector_it != collectors.end()) { + RETURN_IF_ERROR( + collector_it->second->collect(state, tablet_schema, expr, collect_infos)); + } + + const auto& children = expr->children(); + for (auto child = children.rbegin(); child != children.rend(); ++child) { + stack.push(*child); + } + } + } + + LOG(INFO) << "Extracted collect info for " << collect_infos->size() << " fields"; + + return Status::OK(); +} + +Status CollectionStatistics::process_segment(const RowsetSharedPtr& rowset, int64_t seg_id, + const TabletSchema* tablet_schema, + const CollectInfoMap& collect_infos, + io::IOContext* io_ctx) { + auto seg_path = DORIS_TRY(rowset->segment_path(seg_id)); + auto rowset_meta = rowset->rowset_meta(); + + auto idx_file_reader = std::make_unique( + rowset_meta->fs(), + std::string {InvertedIndexDescriptor::get_index_file_path_prefix(seg_path)}, + tablet_schema->get_inverted_index_storage_format(), + rowset_meta->inverted_index_file_info(static_cast(seg_id)), + rowset_meta->tablet_id()); + const bool is_snii = + idx_file_reader->get_storage_format() == InvertedIndexStorageFormatPB::SNII; + auto init_status = idx_file_reader->init(config::inverted_index_read_buffer_size, io_ctx); + if (!init_status.ok()) { + if (is_snii && (init_status.code() == ErrorCode::NOT_FOUND || + init_status.code() == ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND || + init_status.code() == ErrorCode::INVERTED_INDEX_BYPASS)) { + return Status::Error( + "SNII scoring requires every collection segment: {}", init_status.msg()); + } + return init_status; + } + + if (is_snii) { + segment_v2::snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(io_ctx); + SniiScoringSegmentAccumulator segment_accumulator; + for (const auto& [ws_field_name, collect_info] : collect_infos) { + auto logical_reader_result = + idx_file_reader->open_snii_index(collect_info.index_meta, io_ctx); + if (!logical_reader_result.has_value()) { + auto status = std::move(logical_reader_result.error()); + if (status.code() == ErrorCode::INVERTED_INDEX_SNII_NOT_FOUND || + status.code() == ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND || + status.code() == ErrorCode::INVERTED_INDEX_BYPASS) { + return Status::Error( + "SNII scoring requires every logical index in the collection: {}", + status.msg()); + } + return status; + } + auto logical_reader = std::move(logical_reader_result.value()); + const auto* common_grams_metadata = logical_reader->common_grams_metadata(); + segment_v2::inverted_index::PlainTermKeyVersion key_version; + RETURN_IF_ERROR(admit_snii_scoring_segment( + ws_field_name, + common_grams_metadata == nullptr + ? std::nullopt + : std::optional( + *common_grams_metadata), + collect_info.expected_base_analyzer_fingerprint, + logical_reader->stats().doc_count, logical_reader->stats().sum_total_term_freq, + logical_reader->tier() == ::doris::snii::format::IndexTier::kT3, + logical_reader->has_positions(), + logical_reader->section_refs().norms.length != 0, &key_version, + &segment_accumulator)); + DORIS_CHECK(common_grams_metadata != nullptr); + + ::doris::snii::reader::DictBlockCache dict_block_cache; + for (const auto& logical_term_bytes : collect_info.unique_terms) { + std::string physical_term; + const bool term_present = + DORIS_TRY(segment_v2::inverted_index::try_encode_plain_term( + logical_term_bytes, key_version, &physical_term)); + const auto logical_term = + segment_v2::inverted_index::StringHelper::to_wstring(logical_term_bytes); + if (!term_present) { + add_term_doc_frequency(&segment_accumulator.term_doc_freqs, ws_field_name, + logical_term, 0); + continue; + } + + bool found = false; + ::doris::snii::format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(logical_reader->lookup(physical_term, &found, &entry, &frq_base, + &prx_base, &dict_block_cache)); + if (found && entry.df > common_grams_metadata->scoring_doc_count) { + return Status::Error( + "SNII term document frequency {} exceeds scoring document count {}", + entry.df, common_grams_metadata->scoring_doc_count); + } + add_term_doc_frequency(&segment_accumulator.term_doc_freqs, ws_field_name, + logical_term, found ? entry.df : 0); + } + } + + commit_snii_scoring_segment(std::move(segment_accumulator)); + return Status::OK(); + } + + int32_t total_segment_docs = 0; + + for (const auto& [ws_field_name, collect_info] : collect_infos) { + lucene::search::IndexSearcher* index_searcher = nullptr; + lucene::index::IndexReader* index_reader = nullptr; + +#ifdef BE_TEST + auto compound_reader = DORIS_TRY(idx_file_reader->open(collect_info.index_meta, io_ctx)); + auto* reader = lucene::index::IndexReader::open(compound_reader.get()); + auto owned_index_searcher = std::make_shared(reader, true); + index_searcher = owned_index_searcher.get(); + index_reader = index_searcher->getReader(); +#else + InvertedIndexCacheHandle inverted_index_cache_handle; + auto index_file_key = idx_file_reader->get_index_file_cache_key(collect_info.index_meta); + InvertedIndexSearcherCache::CacheKey searcher_cache_key(index_file_key); + + if (!InvertedIndexSearcherCache::instance()->lookup(searcher_cache_key, + &inverted_index_cache_handle)) { + auto compound_reader = + DORIS_TRY(idx_file_reader->open(collect_info.index_meta, io_ctx)); + auto* reader = lucene::index::IndexReader::open(compound_reader.get()); + size_t reader_size = reader->getTermInfosRAMUsed(); + auto searcher_ptr = std::make_shared(reader, true); + auto* cache_value = new InvertedIndexSearcherCache::CacheValue( + std::move(searcher_ptr), reader_size, UnixMillis()); + InvertedIndexSearcherCache::instance()->insert(searcher_cache_key, cache_value, + &inverted_index_cache_handle); + } + + auto searcher_variant = inverted_index_cache_handle.get_index_searcher(); + auto index_searcher_ptr = std::get(searcher_variant); + index_searcher = index_searcher_ptr.get(); + index_reader = index_searcher->getReader(); +#endif + total_segment_docs = std::max(total_segment_docs, index_reader->maxDoc()); + _total_num_tokens[ws_field_name] += + index_reader->sumTotalTermFreq(ws_field_name.c_str()).value_or(0); + + for (const auto& logical_term_bytes : collect_info.unique_terms) { + const auto logical_term = + segment_v2::inverted_index::StringHelper::to_wstring(logical_term_bytes); + auto iter = TermIterator::create(io_ctx, false, index_reader, ws_field_name, + logical_term_bytes); + add_term_doc_frequency(&_term_doc_freqs, ws_field_name, logical_term, iter->doc_freq()); + } + } + + _total_num_docs += static_cast(total_segment_docs); + _avg_dl_by_col.clear(); + _idf_by_col_term.clear(); + + return Status::OK(); +} + +Status CollectionStatistics::admit_snii_scoring_segment( + const std::wstring& field_name, + const std::optional& metadata, + std::string_view expected_base_analyzer_fingerprint, uint64_t index_doc_count, + uint64_t physical_sum_total_term_freq, bool has_scoring_tier, bool has_positions, + bool has_semantic_norms, + segment_v2::inverted_index::PlainTermKeyVersion* plain_term_key_version, + SniiScoringSegmentAccumulator* segment_accumulator) { + DORIS_CHECK(plain_term_key_version != nullptr); + DORIS_CHECK(segment_accumulator != nullptr); + auto segment_stats = + resolve_snii_scoring_segment(metadata, index_doc_count, physical_sum_total_term_freq, + has_scoring_tier, has_positions, has_semantic_norms); + if (!segment_stats.has_value()) { + clear(); + return segment_stats.error(); + } + + const auto& base_analyzer_fingerprint = segment_stats->base_analyzer_fingerprint; + if (base_analyzer_fingerprint != expected_base_analyzer_fingerprint) { + clear(); + return Status::Error( + "SNII scoring segment base analyzer does not match the request analyzer for field " + "{}", + StringHelper::to_string(field_name)); + } + auto staged_fingerprint = segment_accumulator->base_analyzer_fingerprints.find(field_name); + if (staged_fingerprint != segment_accumulator->base_analyzer_fingerprints.end() && + staged_fingerprint->second != base_analyzer_fingerprint) { + clear(); + return Status::Error( + "SNII scoring cannot combine segments with different base analyzers for field {}", + StringHelper::to_string(field_name)); + } + auto collected_fingerprint = _snii_base_analyzer_fingerprints.find(field_name); + if (collected_fingerprint != _snii_base_analyzer_fingerprints.end() && + collected_fingerprint->second != base_analyzer_fingerprint) { + clear(); + return Status::Error( + "SNII scoring cannot combine segments with different base analyzers for field {}", + StringHelper::to_string(field_name)); + } + segment_accumulator->base_analyzer_fingerprints.insert_or_assign(field_name, + base_analyzer_fingerprint); + + if (!segment_accumulator->token_counts.empty() && + segment_accumulator->doc_count != segment_stats->doc_count) { + clear(); + return Status::Error( + "SNII scoring fields in one segment have different document counts: {} and {}", + segment_accumulator->doc_count, segment_stats->doc_count); + } + segment_accumulator->doc_count = segment_stats->doc_count; + segment_accumulator->token_counts[field_name] += segment_stats->token_count; + *plain_term_key_version = segment_stats->plain_term_key_version; + return Status::OK(); +} + +void CollectionStatistics::commit_snii_scoring_segment( + SniiScoringSegmentAccumulator&& segment_accumulator) { + for (const auto& [field_name, base_analyzer_fingerprint] : + segment_accumulator.base_analyzer_fingerprints) { + auto [collected_fingerprint, inserted] = + _snii_base_analyzer_fingerprints.try_emplace(field_name, base_analyzer_fingerprint); + DORIS_CHECK(inserted || collected_fingerprint->second == base_analyzer_fingerprint); + } + for (const auto& [field_name, token_count] : segment_accumulator.token_counts) { + _total_num_tokens[field_name] += token_count; + } + for (const auto& [field_name, term_doc_freqs] : segment_accumulator.term_doc_freqs) { + for (const auto& [term, doc_freq] : term_doc_freqs) { + _term_doc_freqs[field_name][term] += doc_freq; + } + } + _total_num_docs += segment_accumulator.doc_count; + _avg_dl_by_col.clear(); + _idf_by_col_term.clear(); +} + +void CollectionStatistics::clear() { + _total_num_docs = 0; + _total_num_tokens.clear(); + _term_doc_freqs.clear(); + _snii_base_analyzer_fingerprints.clear(); + _avg_dl_by_col.clear(); + _idf_by_col_term.clear(); +} + +uint64_t CollectionStatistics::get_term_doc_freq_by_col(const std::wstring& lucene_col_name, + const std::wstring& term) { + const auto field = _term_doc_freqs.find(lucene_col_name); + if (field == _term_doc_freqs.end()) { + throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, + "Index statistics collection failed: Not such column {}", + StringHelper::to_string(lucene_col_name)); + } + + const auto term_frequency = field->second.find(term); + if (term_frequency == field->second.end()) { + throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, + "Index statistics collection failed: Not such term {}", + StringHelper::to_string(term)); + } + + return term_frequency->second; +} + +uint64_t CollectionStatistics::get_total_term_cnt_by_col(const std::wstring& lucene_col_name) { + const auto token_count = _total_num_tokens.find(lucene_col_name); + if (token_count == _total_num_tokens.end()) { + throw Exception(ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, + "Index statistics collection failed: Not such column {}", + StringHelper::to_string(lucene_col_name)); + } + + return token_count->second; +} + +uint64_t CollectionStatistics::get_doc_num() const { + if (_total_num_docs == 0) { + throw Exception( + ErrorCode::INVERTED_INDEX_CLUCENE_ERROR, + "Index statistics collection failed: No data available for SimilarityCollector"); + } + + return _total_num_docs; +} + +float CollectionStatistics::get_or_calculate_avg_dl(const std::wstring& lucene_col_name) { + auto iter = _avg_dl_by_col.find(lucene_col_name); + if (iter != _avg_dl_by_col.end()) { + return iter->second; + } + + const uint64_t total_term_cnt = get_total_term_cnt_by_col(lucene_col_name); + const uint64_t total_doc_cnt = get_doc_num(); + float avg_dl = total_doc_cnt > 0 ? float((double)total_term_cnt / (double)total_doc_cnt) : 0.0F; + _avg_dl_by_col[lucene_col_name] = avg_dl; + return avg_dl; +} + +float CollectionStatistics::get_or_calculate_idf(const std::wstring& lucene_col_name, + const std::wstring& term) { + auto iter = _idf_by_col_term.find(lucene_col_name); + if (iter != _idf_by_col_term.end()) { + auto term_iter = iter->second.find(term); + if (term_iter != iter->second.end()) { + return term_iter->second; + } + } + + const uint64_t doc_num = get_doc_num(); + const uint64_t doc_freq = get_term_doc_freq_by_col(lucene_col_name, term); + auto idf = (float)std::log(1 + ((double)doc_num - (double)doc_freq + (double)0.5) / + ((double)doc_freq + (double)0.5)); + _idf_by_col_term[lucene_col_name][term] = idf; + return idf; +} + +} // namespace doris diff --git a/be/src/storage/compaction/collection_statistics.h b/be/src/storage/index/inverted/similarity/collection_statistics.h similarity index 65% rename from be/src/storage/compaction/collection_statistics.h rename to be/src/storage/index/inverted/similarity/collection_statistics.h index b93d5a424ae1ab..151bb9945d13a7 100644 --- a/be/src/storage/compaction/collection_statistics.h +++ b/be/src/storage/index/inverted/similarity/collection_statistics.h @@ -16,16 +16,24 @@ // under the License. #pragma once +#include #include +#include +#include +#include #include +#include #include +#include +#include #include "common/be_mock_util.h" #include "exprs/vexpr_fwd.h" #include "runtime/runtime_state.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" #include "storage/index/inverted/query/query_info.h" +#include "storage/index/inverted/similarity/predicate_collector.h" #include "storage/olap_common.h" -#include "storage/predicate_collector.h" namespace doris { @@ -52,19 +60,40 @@ class CollectionStatistics { Status collect(RuntimeState* state, const std::vector& rs_splits, const TabletSchemaSPtr& tablet_schema, const VExprContextSPtrs& common_expr_ctxs_push_down, io::IOContext* io_ctx); + Status collect_full_collection(RuntimeState* state, const std::vector& rowsets, + const TabletSchemaSPtr& tablet_schema, + const VExprContextSPtrs& common_expr_ctxs_push_down, + io::IOContext* io_ctx); MOCK_FUNCTION float get_or_calculate_idf(const std::wstring& lucene_col_name, const std::wstring& term); MOCK_FUNCTION float get_or_calculate_avg_dl(const std::wstring& lucene_col_name); private: + struct SniiScoringSegmentAccumulator { + uint64_t doc_count = 0; + std::unordered_map token_counts; + std::unordered_map> term_doc_freqs; + std::unordered_map base_analyzer_fingerprints; + }; + Status extract_collect_info(RuntimeState* state, const VExprContextSPtrs& common_expr_ctxs_push_down, const TabletSchemaSPtr& tablet_schema, CollectInfoMap* collect_infos); - Status process_segment(const RowsetSharedPtr& rowset, int32_t seg_id, + Status process_segment(const RowsetSharedPtr& rowset, int64_t seg_id, const TabletSchema* tablet_schema, const CollectInfoMap& collect_infos, io::IOContext* io_ctx); + Status admit_snii_scoring_segment( + const std::wstring& field_name, + const std::optional& metadata, + std::string_view expected_base_analyzer_fingerprint, uint64_t index_doc_count, + uint64_t physical_sum_total_term_freq, bool has_scoring_tier, bool has_positions, + bool has_semantic_norms, + segment_v2::inverted_index::PlainTermKeyVersion* plain_term_key_version, + SniiScoringSegmentAccumulator* segment_accumulator); + void commit_snii_scoring_segment(SniiScoringSegmentAccumulator&& segment_accumulator); + void clear(); uint64_t get_term_doc_freq_by_col(const std::wstring& lucene_col_name, const std::wstring& term); @@ -74,6 +103,7 @@ class CollectionStatistics { uint64_t _total_num_docs = 0; std::unordered_map _total_num_tokens; std::unordered_map> _term_doc_freqs; + std::unordered_map _snii_base_analyzer_fingerprints; std::unordered_map _avg_dl_by_col; std::unordered_map> _idf_by_col_term; diff --git a/be/src/storage/index/inverted/similarity/predicate_collector.cpp b/be/src/storage/index/inverted/similarity/predicate_collector.cpp new file mode 100644 index 00000000000000..223e676b7a9254 --- /dev/null +++ b/be/src/storage/index/inverted/similarity/predicate_collector.cpp @@ -0,0 +1,594 @@ +// 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 "storage/index/inverted/similarity/predicate_collector.h" + +#include + +#include + +#include "exec/common/variant_util.h" +#include "exprs/vexpr.h" +#include "exprs/vexpr_context.h" +#include "exprs/vliteral.h" +#include "exprs/vsearch.h" +#include "exprs/vslot_ref.h" +#include "gen_cpp/Exprs_types.h" +#include "storage/index/index_reader_helper.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/inverted_index_selector.h" +#include "storage/index/inverted/util/string_helper.h" +#include "storage/tablet/tablet_schema.h" +#include "storage/utils.h" + +namespace doris { + +using namespace segment_v2; + +namespace { + +InvertedIndexAnalyzerCtx analyzer_context_from_properties( + const std::map& properties) { + InvertedIndexAnalyzerConfig config; + config.analyzer_name = get_analyzer_name_from_properties(properties); + config.parser_type = get_inverted_index_parser_type_from_string( + get_parser_string_from_properties(properties)); + config.parser_mode = get_parser_mode_string_from_properties(properties); + config.lower_case = get_parser_lowercase_from_properties(properties); + config.stop_words = get_parser_stopwords_from_properties(properties); + config.char_filter_map = get_parser_char_filter_map_from_properties(properties); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.analyzer_name = config.analyzer_name; + analyzer_ctx.parser_type = config.parser_type; + analyzer_ctx.char_filter_map = config.char_filter_map; + analyzer_ctx.analyzer_provider = + inverted_index::InvertedIndexAnalyzer::create_analyzer_provider(&config); + return analyzer_ctx; +} + +std::vector analyze_plain_query(const std::string& value, + const InvertedIndexAnalyzerCtx& analyzer_ctx) { + DORIS_CHECK(analyzer_ctx.analyzer_provider != nullptr); + auto analyzer = analyzer_ctx.analyzer_provider->get_analyzer( + inverted_index::AnalysisPurpose::kPlainQuery); + auto reader = + inverted_index::InvertedIndexAnalyzer::create_reader(analyzer_ctx.char_filter_map); + reader->init(value.data(), static_cast(value.size()), true); + return inverted_index::InvertedIndexAnalyzer::get_analyse_result(reader, analyzer.get()); +} + +Status append_scoring_leaf(CollectInfo* collect_info, const std::vector& term_infos, + std::string_view base_analyzer_fingerprint) { + DORIS_CHECK(collect_info != nullptr); + if (!collect_info->logical_scoring_leaves.empty() && + collect_info->expected_base_analyzer_fingerprint != base_analyzer_fingerprint) { + return Status::Error( + "Scoring predicates for one field use different base analyzers"); + } + if (collect_info->logical_scoring_leaves.empty()) { + collect_info->expected_base_analyzer_fingerprint = base_analyzer_fingerprint; + } + + LogicalScoringLeaf leaf; + leaf.clauses.reserve(term_infos.size()); + for (const auto& term_info : term_infos) { + DORIS_CHECK(term_info.is_single_term()); + DORIS_CHECK(term_info.key_kind == TermKeyKind::kPlain); + const auto& term = term_info.get_single_term(); + auto [slot, inserted] = collect_info->unique_term_slots.try_emplace( + term, static_cast(collect_info->unique_terms.size())); + if (inserted) { + collect_info->unique_terms.push_back(term); + } + leaf.clauses.emplace_back( + LogicalScoringClause {.df_slot = slot->second, .position = term_info.position}); + } + collect_info->logical_scoring_leaves.emplace_back(std::move(leaf)); + return Status::OK(); +} + +InvertedIndexQueryType match_query_type(TExprOpcode::type opcode) { + switch (opcode) { + case TExprOpcode::MATCH_ANY: + return InvertedIndexQueryType::MATCH_ANY_QUERY; + case TExprOpcode::MATCH_ALL: + return InvertedIndexQueryType::MATCH_ALL_QUERY; + case TExprOpcode::MATCH_PHRASE: + return InvertedIndexQueryType::MATCH_PHRASE_QUERY; + case TExprOpcode::MATCH_PHRASE_PREFIX: + return InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY; + case TExprOpcode::MATCH_REGEXP: + return InvertedIndexQueryType::MATCH_REGEXP_QUERY; + case TExprOpcode::MATCH_PHRASE_EDGE: + return InvertedIndexQueryType::MATCH_PHRASE_EDGE_QUERY; + default: + return InvertedIndexQueryType::UNKNOWN_QUERY; + } +} + +InvertedIndexQueryType search_query_type(std::string_view clause_type) { + if (clause_type == "EXACT") { + return InvertedIndexQueryType::EQUAL_QUERY; + } + if (clause_type == "PHRASE") { + return InvertedIndexQueryType::MATCH_PHRASE_QUERY; + } + if (clause_type == "ALL") { + return InvertedIndexQueryType::MATCH_ALL_QUERY; + } + return InvertedIndexQueryType::MATCH_ANY_QUERY; +} + +Result select_index_meta(const std::vector& index_metas, + FieldType field_type, + InvertedIndexQueryType query_type, + std::string_view analyzer_key) { + std::vector candidates; + candidates.reserve(index_metas.size()); + InvertedIndexSelectionKeyIndex key_index; + for (const auto* index_meta : index_metas) { + auto status = add_inverted_index_selection_candidate( + InvertedIndexSelectionCandidate {.index_id = index_meta->index_id(), + .reader_type = infer_inverted_index_reader_type( + field_type, index_meta->properties()), + .analyzer_key = build_analyzer_key_from_properties( + index_meta->properties())}, + &candidates, &key_index); + if (!status.ok()) { + return ResultError(std::move(status)); + } + } + + auto selection = select_best_inverted_index_candidate( + candidates, key_index, field_type, query_type, normalize_analyzer_key(analyzer_key)); + if (!selection.has_value()) { + return ResultError(std::move(selection.error())); + } + const size_t selected = *selection; + DORIS_CHECK(selected < index_metas.size()); + return index_metas[selected]; +} + +Status validate_same_physical_index(const CollectInfo& collect_info, + const TabletIndex& selected_index) { + DORIS_CHECK(collect_info.index_meta != nullptr); + if (collect_info.index_meta->index_id() != selected_index.index_id() || + collect_info.index_meta->get_index_suffix() != selected_index.get_index_suffix()) { + return Status::Error( + "Scoring predicates for one field select different inverted indexes: {} and {}", + collect_info.index_meta->index_id(), selected_index.index_id()); + } + return Status::OK(); +} + +struct ScoringIndexCandidates { + FieldType field_type = FieldType::OLAP_FIELD_TYPE_UNKNOWN; + std::string index_suffix_path; + std::vector index_metas; + std::vector> owned_index_metas; +}; + +FieldType scoring_leaf_type(const TabletColumn& column) { + const TabletColumn* leaf = &column; + while (leaf->is_array_type()) { + DORIS_CHECK_EQ(leaf->get_subtype_count(), 1); + leaf = &leaf->get_sub_column(0); + } + return leaf->type(); +} + +ScoringIndexCandidates resolve_text_scoring_index_candidates(const TabletSchemaSPtr& tablet_schema, + const TabletColumn& column) { + ScoringIndexCandidates candidates {.field_type = scoring_leaf_type(column), + .index_suffix_path = column.suffix_path(), + .index_metas = tablet_schema->inverted_indexs(column), + .owned_index_metas = {}}; + + // The collector has tablet-schema context but no segment-side variant + // inference. Resolve only shapes that are deterministic from schema: + // typed/materialized paths, field-pattern templates, and a plain parent + // index inherited by the dynamic VARIANT placeholder. + if (!candidates.index_metas.empty() || !column.is_extracted_column()) { + return candidates; + } + + TabletSchema::SubColumnInfo sub_column_info; + const std::string relative_path = column.path_info_ptr()->copy_pop_front().get_path(); + if (variant_util::generate_sub_column_info(*tablet_schema, column.parent_unique_id(), + relative_path, &sub_column_info) && + !sub_column_info.indexes.empty()) { + candidates.field_type = scoring_leaf_type(sub_column_info.column); + candidates.index_suffix_path = sub_column_info.column.suffix_path(); + for (auto& index : sub_column_info.indexes) { + candidates.index_metas.push_back(index.get()); + candidates.owned_index_metas.emplace_back(std::move(index)); + } + return candidates; + } + + if (!column.is_variant_type()) { + return candidates; + } + + // MATCH and score-bearing SEARCH clauses have text semantics. When a dynamic + // VARIANT path has no materialized type, those semantics provide the missing + // leaf-type proof for selecting its plain parent full-text index. Typed paths + // returned above keep their schema type, so numeric BKD leaves remain rejected. + candidates.field_type = FieldType::OLAP_FIELD_TYPE_STRING; + const auto parent_indexes = tablet_schema->inverted_indexs(column.parent_unique_id()); + for (const auto* index : parent_indexes) { + if (!index->field_pattern().empty()) { + continue; + } + auto owned_index = std::make_shared(*index); + owned_index->set_escaped_escaped_index_suffix_path(column.path_info_ptr()->get_path()); + candidates.index_metas.push_back(owned_index.get()); + candidates.owned_index_metas.emplace_back(std::move(owned_index)); + } + return candidates; +} + +Status validate_scoring_leaf_type(const ScoringIndexCandidates& candidates, + std::string_view field_name) { + if (candidates.field_type == FieldType::OLAP_FIELD_TYPE_VARIANT) { + return Status::Error( + "Index statistics collection failed: Cannot prove scoring leaf type for field={}", + field_name); + } + return Status::OK(); +} + +void preserve_selected_index_metadata(const ScoringIndexCandidates& candidates, + const TabletIndex* selected_index, + CollectInfo* collect_info) { + DORIS_CHECK(selected_index != nullptr); + DORIS_CHECK(collect_info != nullptr); + collect_info->index_meta = selected_index; + for (const auto& owned_index : candidates.owned_index_metas) { + if (owned_index.get() == selected_index) { + collect_info->owned_index_meta = owned_index; + return; + } + } +} + +Result resolve_search_scoring_index_candidates( + const TabletSchemaSPtr& tablet_schema, const std::string& field_name, + const TSearchFieldBinding* field_binding) { + const int32_t column_index = tablet_schema->field_index(field_name); + if (column_index >= 0) { + return resolve_text_scoring_index_candidates(tablet_schema, + tablet_schema->column(column_index)); + } + + if (field_binding == nullptr || !field_binding->__isset.is_variant_subcolumn || + !field_binding->is_variant_subcolumn || !field_binding->__isset.parent_field_name || + field_binding->parent_field_name.empty() || !field_binding->__isset.subcolumn_path || + field_binding->subcolumn_path.empty()) { + return ResultError(Status::Error( + "Index statistics collection failed: Cannot resolve search field={}", field_name)); + } + + const int32_t parent_column_index = + tablet_schema->field_index(field_binding->parent_field_name); + if (parent_column_index < 0) { + return ResultError(Status::Error( + "Index statistics collection failed: Cannot resolve parent={} for search field={}", + field_binding->parent_field_name, field_name)); + } + const auto& parent_column = tablet_schema->column(parent_column_index); + if (!parent_column.is_variant_type()) { + return ResultError(Status::Error( + "Index statistics collection failed: Search field={} parent={} is not VARIANT", + field_name, field_binding->parent_field_name)); + } + + TabletColumn dynamic_column; + dynamic_column.set_unique_id(-1); + dynamic_column.set_name(field_name); + dynamic_column.set_type(FieldType::OLAP_FIELD_TYPE_VARIANT); + dynamic_column.set_parent_unique_id(parent_column.unique_id()); + dynamic_column.set_path_info( + PathInData(field_binding->parent_field_name + "." + field_binding->subcolumn_path)); + return resolve_text_scoring_index_candidates(tablet_schema, dynamic_column); +} + +} // namespace + +VSlotRef* PredicateCollector::find_slot_ref(const VExprSPtr& expr) const { + if (!expr) { + return nullptr; + } + + auto cur = VExpr::expr_without_cast(expr); + if (cur->node_type() == TExprNodeType::SLOT_REF) { + return static_cast(cur.get()); + } + + for (const auto& ch : cur->children()) { + if (auto* s = find_slot_ref(ch)) { + return s; + } + } + + return nullptr; +} + +std::string PredicateCollector::build_field_name(int32_t col_unique_id, + const std::string& suffix_path) const { + std::string field_name = std::to_string(col_unique_id); + if (!suffix_path.empty()) { + field_name += "." + suffix_path; + } + return field_name; +} + +Status MatchPredicateCollector::collect(RuntimeState* state, const TabletSchemaSPtr& tablet_schema, + const VExprSPtr& expr, CollectInfoMap* collect_infos) { + DCHECK(collect_infos != nullptr); + + auto* left_slot_ref = find_slot_ref(expr->children()[0]); + if (left_slot_ref == nullptr) { + return Status::Error( + "Index statistics collection failed: Cannot find slot reference in match predicate " + "left expression"); + } + + auto* right_literal = static_cast(expr->children()[1].get()); + DCHECK(right_literal != nullptr); + + const auto* sd = state->desc_tbl().get_slot_descriptor(left_slot_ref->slot_id()); + if (sd == nullptr) { + return Status::Error( + "Index statistics collection failed: Cannot find slot descriptor for slot_id={}", + left_slot_ref->slot_id()); + } + + int32_t col_idx = tablet_schema->field_index(left_slot_ref->column_name()); + if (col_idx == -1) { + return Status::Error( + "Index statistics collection failed: Cannot find column index for column={}", + left_slot_ref->column_name()); + } + + const auto& column = tablet_schema->column(col_idx); + auto candidates = resolve_text_scoring_index_candidates(tablet_schema, column); + RETURN_IF_ERROR(validate_scoring_leaf_type(candidates, left_slot_ref->column_name())); + +#ifndef BE_TEST + if (candidates.index_metas.empty()) { + return Status::Error( + "Index statistics collection failed: Score query is not supported without inverted " + "index for column={}", + left_slot_ref->column_name()); + } +#else + if (candidates.index_metas.empty()) { + return Status::OK(); + } +#endif + + const auto* analyzer_ctx = expr->query_analyzer_ctx(); + DORIS_CHECK(analyzer_ctx != nullptr); + DORIS_CHECK(analyzer_ctx->analyzer_provider != nullptr); + const auto query_type = match_query_type(expr->op()); + DORIS_CHECK(query_type != InvertedIndexQueryType::UNKNOWN_QUERY); + const auto* index_meta = + DORIS_TRY(select_index_meta(candidates.index_metas, candidates.field_type, query_type, + analyzer_ctx->analyzer_name)); + if (!InvertedIndexAnalyzer::should_analyzer(index_meta->properties()) || + !IndexReaderHelper::is_need_similarity_score(expr->op(), index_meta)) { + return Status::OK(); + } + + auto options = DataTypeSerDe::get_default_format_options(); + options.timezone = &state->timezone_obj(); + auto term_infos = analyze_plain_query(right_literal->value(options), *analyzer_ctx); + if (expr->op() == TExprOpcode::MATCH_PHRASE_PREFIX && !term_infos.empty()) { + term_infos.pop_back(); + } + const auto base_analyzer_fingerprint = + analyzer_ctx->analyzer_provider->base_analyzer_fingerprint(); + + std::string field_name = + build_field_name(index_meta->col_unique_ids()[0], candidates.index_suffix_path); + std::wstring ws_field_name = StringHelper::to_wstring(field_name); + + auto iter = collect_infos->find(ws_field_name); + if (iter == collect_infos->end()) { + CollectInfo collect_info; + RETURN_IF_ERROR(append_scoring_leaf(&collect_info, term_infos, base_analyzer_fingerprint)); + preserve_selected_index_metadata(candidates, index_meta, &collect_info); + (*collect_infos)[ws_field_name] = std::move(collect_info); + } else { + RETURN_IF_ERROR(validate_same_physical_index(iter->second, *index_meta)); + RETURN_IF_ERROR(append_scoring_leaf(&iter->second, term_infos, base_analyzer_fingerprint)); + } + + return Status::OK(); +} + +Status SearchPredicateCollector::collect(RuntimeState* state, const TabletSchemaSPtr& tablet_schema, + const VExprSPtr& expr, CollectInfoMap* collect_infos) { + DCHECK(collect_infos != nullptr); + + auto* search_expr = dynamic_cast(expr.get()); + if (search_expr == nullptr) { + return Status::InternalError("SearchPredicateCollector: expr is not VSearchExpr type"); + } + + const TSearchParam& search_param = search_expr->get_search_param(); + FieldBindingMap field_bindings; + field_bindings.reserve(search_param.field_bindings.size()); + for (const auto& field_binding : search_param.field_bindings) { + field_bindings[field_binding.field_name] = &field_binding; + } + + RETURN_IF_ERROR(collect_from_clause(search_param.root, state, tablet_schema, field_bindings, + collect_infos)); + + return Status::OK(); +} + +Status SearchPredicateCollector::collect_from_clause(const TSearchClause& clause, + RuntimeState* state, + const TabletSchemaSPtr& tablet_schema, + const FieldBindingMap& field_bindings, + CollectInfoMap* collect_infos) { + const std::string& clause_type = clause.clause_type; + if (clause_type == "NESTED") { + return Status::Error( + "Scoring nested search clauses is not supported"); + } + ClauseTypeCategory category = get_clause_type_category(clause_type); + + if (category == ClauseTypeCategory::COMPOUND) { + if (clause.__isset.children) { + for (const auto& child_clause : clause.children) { + RETURN_IF_ERROR(collect_from_clause(child_clause, state, tablet_schema, + field_bindings, collect_infos)); + } + } + return Status::OK(); + } + + return collect_from_leaf(clause, state, tablet_schema, field_bindings, collect_infos); +} + +Status SearchPredicateCollector::collect_from_leaf(const TSearchClause& clause, RuntimeState* state, + const TabletSchemaSPtr& tablet_schema, + const FieldBindingMap& field_bindings, + CollectInfoMap* collect_infos) { + if (!clause.__isset.field_name || !clause.__isset.value) { + return Status::InvalidArgument("Search clause missing field_name or value"); + } + + const std::string& field_name = clause.field_name; + const std::string& value = clause.value; + const std::string& clause_type = clause.clause_type; + + if (!is_score_query_type(clause_type)) { + return Status::OK(); + } + + const auto field_binding_iter = field_bindings.find(field_name); + const auto* field_binding = + field_binding_iter == field_bindings.end() ? nullptr : field_binding_iter->second; + auto candidates = DORIS_TRY( + resolve_search_scoring_index_candidates(tablet_schema, field_name, field_binding)); + RETURN_IF_ERROR(validate_scoring_leaf_type(candidates, field_name)); + if (candidates.index_metas.empty()) { + return Status::Error( + "Index statistics collection failed: Score query is not supported without " + "inverted index for search field={}", + field_name); + } + + ClauseTypeCategory category = get_clause_type_category(clause_type); + auto query_type = search_query_type(clause_type); + std::string analyzer_key; + if (query_type != InvertedIndexQueryType::EQUAL_QUERY && field_binding != nullptr && + field_binding->__isset.index_properties && !field_binding->index_properties.empty() && + is_string_type(candidates.field_type)) { + analyzer_key = build_analyzer_key_from_properties(field_binding->index_properties); + } + + auto selected_index = select_index_meta(candidates.index_metas, candidates.field_type, + query_type, analyzer_key); + if (!selected_index.has_value()) { + return Status::Error( + "Index statistics collection failed: Cannot select scoring index for search " + "field={}: {}", + field_name, selected_index.error().to_string()); + } + const auto* index_meta = *selected_index; + if (infer_inverted_index_reader_type(candidates.field_type, index_meta->properties()) == + InvertedIndexReaderType::BKD) { + return Status::Error( + "Index statistics collection failed: BM25 scoring does not support numeric BKD " + "search field={}", + field_name); + } + + const auto& analysis_properties = index_meta->properties(); + + std::vector term_infos; + std::string_view base_analyzer_fingerprint; + std::optional analyzer_ctx; + if (InvertedIndexAnalyzer::should_analyzer(analysis_properties)) { + analyzer_ctx.emplace(analyzer_context_from_properties(analysis_properties)); + base_analyzer_fingerprint = analyzer_ctx->analyzer_provider->base_analyzer_fingerprint(); + } + + if (clause_type == "MATCH") { + term_infos.emplace_back(value); + } else if (category == ClauseTypeCategory::TOKENIZED) { + if (analyzer_ctx.has_value()) { + term_infos = analyze_plain_query(value, *analyzer_ctx); + } else { + term_infos.emplace_back(value); + } + } else if (category == ClauseTypeCategory::NON_TOKENIZED) { + if (clause_type == "TERM" && analyzer_ctx.has_value()) { + term_infos = analyze_plain_query(value, *analyzer_ctx); + } else { + term_infos.emplace_back(value); + } + } + + std::string lucene_field_name = + build_field_name(index_meta->col_unique_ids()[0], candidates.index_suffix_path); + std::wstring ws_field_name = StringHelper::to_wstring(lucene_field_name); + + auto iter = collect_infos->find(ws_field_name); + if (iter == collect_infos->end()) { + CollectInfo collect_info; + RETURN_IF_ERROR(append_scoring_leaf(&collect_info, term_infos, base_analyzer_fingerprint)); + preserve_selected_index_metadata(candidates, index_meta, &collect_info); + (*collect_infos)[ws_field_name] = std::move(collect_info); + } else { + RETURN_IF_ERROR(validate_same_physical_index(iter->second, *index_meta)); + RETURN_IF_ERROR(append_scoring_leaf(&iter->second, term_infos, base_analyzer_fingerprint)); + } + + return Status::OK(); +} + +bool SearchPredicateCollector::is_score_query_type(const std::string& clause_type) const { + return clause_type == "TERM" || clause_type == "EXACT" || clause_type == "PHRASE" || + clause_type == "MATCH" || clause_type == "ANY" || clause_type == "ALL"; +} + +SearchPredicateCollector::ClauseTypeCategory SearchPredicateCollector::get_clause_type_category( + const std::string& clause_type) const { + if (clause_type == "AND" || clause_type == "OR" || clause_type == "NOT" || + clause_type == "OCCUR_BOOLEAN") { + return ClauseTypeCategory::COMPOUND; + } else if (clause_type == "TERM" || clause_type == "EXACT") { + return ClauseTypeCategory::NON_TOKENIZED; + } else if (clause_type == "PHRASE" || clause_type == "MATCH" || clause_type == "ANY" || + clause_type == "ALL") { + return ClauseTypeCategory::TOKENIZED; + } else { + LOG(WARNING) << "Unknown clause type '" << clause_type + << "', defaulting to NON_TOKENIZED category"; + return ClauseTypeCategory::NON_TOKENIZED; + } +} + +} // namespace doris diff --git a/be/src/storage/predicate_collector.h b/be/src/storage/index/inverted/similarity/predicate_collector.h similarity index 80% rename from be/src/storage/predicate_collector.h rename to be/src/storage/index/inverted/similarity/predicate_collector.h index c96e0af9c45ed5..c43865a63a52d3 100644 --- a/be/src/storage/predicate_collector.h +++ b/be/src/storage/index/inverted/similarity/predicate_collector.h @@ -17,10 +17,11 @@ #pragma once +#include #include -#include #include #include +#include #include "common/status.h" #include "exprs/vexpr_fwd.h" @@ -35,14 +36,20 @@ class TabletIndex; class TabletSchema; using TabletSchemaSPtr = std::shared_ptr; -struct TermInfoComparer { - bool operator()(const segment_v2::TermInfo& lhs, const segment_v2::TermInfo& rhs) const { - return lhs.term < rhs.term; - } +struct LogicalScoringClause { + uint32_t df_slot = 0; + int32_t position = 0; +}; + +struct LogicalScoringLeaf { + std::vector clauses; }; struct CollectInfo { - std::set term_infos; + std::vector unique_terms; + std::unordered_map unique_term_slots; + std::vector logical_scoring_leaves; + std::string expected_base_analyzer_fingerprint; std::shared_ptr owned_index_meta; const TabletIndex* index_meta = nullptr; }; @@ -73,16 +80,19 @@ class SearchPredicateCollector : public PredicateCollector { private: enum class ClauseTypeCategory { NON_TOKENIZED, TOKENIZED, COMPOUND }; + using FieldBindingMap = std::unordered_map; Status collect_from_clause(const TSearchClause& clause, RuntimeState* state, const TabletSchemaSPtr& tablet_schema, + const FieldBindingMap& field_bindings, CollectInfoMap* collect_infos); Status collect_from_leaf(const TSearchClause& clause, RuntimeState* state, - const TabletSchemaSPtr& tablet_schema, CollectInfoMap* collect_infos); + const TabletSchemaSPtr& tablet_schema, + const FieldBindingMap& field_bindings, CollectInfoMap* collect_infos); bool is_score_query_type(const std::string& clause_type) const; ClauseTypeCategory get_clause_type_category(const std::string& clause_type) const; }; using PredicateCollectorPtr = std::unique_ptr; -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/be/src/storage/index/inverted/token_filter/ascii_folding_filter_factory.h b/be/src/storage/index/inverted/token_filter/ascii_folding_filter_factory.h index 9710f0f3cadede..ee3321d96a27c7 100644 --- a/be/src/storage/index/inverted/token_filter/ascii_folding_filter_factory.h +++ b/be/src/storage/index/inverted/token_filter/ascii_folding_filter_factory.h @@ -35,6 +35,11 @@ class ASCIIFoldingFilterFactory : public TokenFilterFactory { return std::make_shared(in, _preserve_original); } + PositionCapability position_capability() const override { + return _preserve_original ? PositionCapability::kUnknown + : PositionCapability::kAlwaysUnitIncrement; + } + private: bool _preserve_original = false; }; diff --git a/be/src/storage/index/inverted/token_filter/common_grams_filter.cpp b/be/src/storage/index/inverted/token_filter/common_grams_filter.cpp new file mode 100644 index 00000000000000..2781ea58842975 --- /dev/null +++ b/be/src/storage/index/inverted/token_filter/common_grams_filter.cpp @@ -0,0 +1,384 @@ +// 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 "storage/index/inverted/token_filter/common_grams_filter.h" + +#include +#include + +#include "common/exception.h" +#include "common/logging.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +void validate_input_token_shape(const Token& token, std::string_view term) { + if (token.getPositionIncrement() != 1) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "CommonGrams requires position increment 1, got {}", + token.getPositionIncrement()); + } + if (term.empty()) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "CommonGrams requires non-empty input tokens"); + } +} + +void validate_input_token(const Token& token, std::string_view term) { + validate_input_token_shape(token, term); + auto status = validate_common_grams_logical_term(term, "input token"); + if (!status.ok()) { + throw Exception(status); + } +} + +void set_output_token(Token* token, std::string_view term, int32_t position_increment, + int32_t start_offset, int32_t end_offset, const TCHAR* type) { + token->clear(); + token->setTextNoCopy(term.data(), static_cast(term.size())); + token->setPositionIncrement(position_increment); + token->setStartOffset(start_offset); + token->setEndOffset(end_offset); + token->setType(type); +} + +bool is_common_word(const CommonGramsBufferedToken& token, const CommonWordSet& common_words, + std::optional* cached_membership) { + if (!cached_membership->has_value()) { + cached_membership->emplace(common_words.contains(token.term)); + } + return cached_membership->value(); +} + +} // namespace + +bool common_grams_query_may_use_gram(std::span terms, CommonGramsQueryMode mode, + const CommonWordSet& common_words) { + if (terms.size() < 2) { + return false; + } + const size_t relevant_term_count = + mode == CommonGramsQueryMode::kPhrasePrefix ? terms.size() - 1 : terms.size(); + for (size_t i = 0; i < relevant_term_count; ++i) { + if (common_words.contains(terms[i])) { + return true; + } + } + return false; +} + +CommonGramsFilter::CommonGramsFilter(TokenStreamPtr in, + std::shared_ptr common_words, + CommonGramsOutputMode output_mode) + : DorisTokenFilter(std::move(in)), + _common_words(std::move(common_words)), + _output_mode(output_mode) { + DORIS_CHECK(_common_words != nullptr); +} + +bool CommonGramsFilter::read_input(CommonGramsBufferedToken* token, + std::optional* is_common) { + if (_in->next(&_input_token) == nullptr) { + return false; + } + std::string_view term(_input_token.termBuffer(), _input_token.termLength()); + validate_input_token(_input_token, term); + token->term.assign(term); + token->start_offset = _input_token.startOffset(); + token->end_offset = _input_token.endOffset(); + token->type = _input_token.type(); + is_common->reset(); + return true; +} + +Token* CommonGramsFilter::emit_unigram(Token* token, const CommonGramsBufferedToken& buffered) { + const std::string_view output_term = encode_plain_term(buffered.term); + set_output_token(token, output_term, 1, buffered.start_offset, buffered.end_offset, + buffered.type); + return token; +} + +std::string_view CommonGramsFilter::encode_plain_term(std::string_view term) { + if (_output_mode != CommonGramsOutputMode::kLogical && !term.empty() && + (term.front() == PLAIN_ESCAPE_PREFIX || term.front() == '\x1f')) { + if (!try_encode_escaped_plain_term_prevalidated(term, _encoded_plain_term)) { + throw Exception(Status::Error( + "CommonGrams escaped plain term would exceed the 16383-byte key limit; " + "set enable_common_grams_index_build=false and retry the import in a new " + "transaction")); + } + return _encoded_plain_term; + } + return term; +} + +std::string_view CommonGramsFilter::encode_snii_plain_term(std::string_view term) { + DORIS_CHECK(_output_mode == CommonGramsOutputMode::kEscapedV1SpimiIndex); + DORIS_CHECK(!term.empty()); + if (term.front() != PLAIN_ESCAPE_PREFIX && term.front() != '\x1f') { + return term; + } + if (term.size() == COMMON_GRAM_MAX_ENCODED_BYTES) { + throw Exception(Status::Error( + "CommonGrams escaped plain term would exceed the 16383-byte key limit; " + "set enable_common_grams_index_build=false and retry the import in a new " + "transaction")); + } + DORIS_CHECK_LT(term.size(), COMMON_GRAM_MAX_ENCODED_BYTES); + _encoded_plain_term.clear(); + _encoded_plain_term.reserve(term.size() + 1); + _encoded_plain_term.push_back(PLAIN_ESCAPE_PREFIX); + _encoded_plain_term.push_back(term.front() == PLAIN_ESCAPE_PREFIX ? 'E' : 'G'); + _encoded_plain_term.append(term.substr(1)); + return _encoded_plain_term; +} + +Token* CommonGramsFilter::emit_gram(Token* token, int32_t start_offset, int32_t end_offset) { + set_output_token(token, _gram, 0, start_offset, end_offset, COMMON_GRAM_TOKEN_TYPE); + return token; +} + +Token* CommonGramsFilter::next(Token* token) { + if (_emit_current) { + _emit_current = false; + return emit_unigram(token, _current); + } + + if (!_has_current) { + if (!read_input(&_current, &_current_is_common)) { + return nullptr; + } + _has_current = true; + return emit_unigram(token, _current); + } + + if (!read_input(&_lookahead, &_lookahead_is_common)) { + _has_current = false; + return nullptr; + } + + const bool uses_gram = is_common_word(_current, *_common_words, &_current_is_common) || + is_common_word(_lookahead, *_common_words, &_lookahead_is_common); + if (uses_gram) { + const bool encoded = + try_encode_common_gram_prevalidated(_current.term, _lookahead.term, _gram); + if (encoded) { + const int32_t start_offset = _current.start_offset; + const int32_t end_offset = _lookahead.end_offset; + std::swap(_current, _lookahead); + std::swap(_current_is_common, _lookahead_is_common); + _emit_current = true; + return emit_gram(token, start_offset, end_offset); + } + } + + std::swap(_current, _lookahead); + std::swap(_current_is_common, _lookahead_is_common); + return emit_unigram(token, _current); +} + +bool CommonGramsFilter::next_snii_index_event(SniiCommonGramsIndexEvent* event) { + DORIS_CHECK(event != nullptr); + DORIS_CHECK(_output_mode == CommonGramsOutputMode::kEscapedV1SpimiIndex); + if (_in->next(&_input_token) == nullptr) { + return false; + } + + const std::string_view logical_term(_input_token.termBuffer(), + _input_token.termLength()); + validate_input_token_shape(_input_token, logical_term); + const bool requires_prevalidation = + logical_term.size() > COMMON_GRAM_MAX_ENCODED_BYTES || + (logical_term.size() == COMMON_GRAM_MAX_ENCODED_BYTES && + (logical_term.front() == PLAIN_ESCAPE_PREFIX || logical_term.front() == '\x1f')); + if (requires_prevalidation) { + auto status = validate_common_grams_logical_term(logical_term, "input token"); + if (!status.ok()) { + throw Exception(status); + } + } + event->logical_term = logical_term; + event->plain_term = encode_snii_plain_term(logical_term); + return true; +} + +void CommonGramsFilter::reset() { + DorisTokenFilter::reset(); + _current.term.clear(); + _current.start_offset = 0; + _current.end_offset = 0; + _current.type = Token::getDefaultType(); + _lookahead.term.clear(); + _lookahead.start_offset = 0; + _lookahead.end_offset = 0; + _lookahead.type = Token::getDefaultType(); + _current_is_common.reset(); + _lookahead_is_common.reset(); + _has_current = false; + _emit_current = false; + _gram.clear(); + _encoded_plain_term.clear(); +} + +Token* CommonGramsPositionFilter::next(Token* token) { + if (_in->next(token) == nullptr) { + return nullptr; + } + const std::string_view term(token->termBuffer(), token->termLength()); + validate_input_token(*token, term); + return token; +} + +CommonGramsQueryFilterBase::CommonGramsQueryFilterBase( + TokenStreamPtr in, std::shared_ptr common_words, + CommonGramsQueryMode mode) + : DorisTokenFilter(std::move(in)), _common_words(std::move(common_words)), _mode(mode) { + DORIS_CHECK(_common_words != nullptr); +} + +bool CommonGramsQueryFilterBase::pair_uses_gram( + const std::vector& unigrams, std::optional* left_is_common, + size_t pair_index) const { + const bool is_prefix_boundary = + _mode == CommonGramsQueryMode::kPhrasePrefix && pair_index + 2 == unigrams.size(); + if (is_prefix_boundary) { + return is_common_word(unigrams[pair_index], *_common_words, left_is_common); + } + + std::optional right_is_common; + const bool uses_gram = + is_common_word(unigrams[pair_index], *_common_words, left_is_common) || + is_common_word(unigrams[pair_index + 1], *_common_words, &right_is_common); + *left_is_common = right_is_common; + return uses_gram; +} + +void CommonGramsQueryFilterBase::append_plain_output(std::vector* output, + const CommonGramsBufferedToken& token) { + output->push_back(token); +} + +bool CommonGramsQueryFilterBase::append_gram_output( + std::vector* output, + const std::optional& indexed_gram, + const CommonGramsBufferedToken& left, const CommonGramsBufferedToken& right) { + if (!indexed_gram.has_value()) { + DCHECK(!is_common_gram_encodable(left.term, right.term)); + return false; + } + CommonGramsBufferedToken query_gram = *indexed_gram; + query_gram.type = COMMON_GRAM_TOKEN_TYPE; + output->push_back(std::move(query_gram)); + return true; +} + +void CommonGramsQueryFilterBase::prepare_output() { + std::vector unigrams; + std::vector> indexed_grams; + Token token; + while (_in->next(&token) != nullptr) { + const std::string_view term(token.termBuffer(), token.termLength()); + if (token.getPositionIncrement() == 0) { + if (!is_common_gram_token_type(token.type()) || unigrams.empty() || + indexed_grams.back().has_value()) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "Invalid indexed CommonGrams token sequence"); + } + auto status = validate_common_grams_logical_term(term, "indexed gram"); + if (!status.ok()) { + throw Exception(status); + } + indexed_grams.back() = CommonGramsBufferedToken {.term = std::string(term), + .start_offset = token.startOffset(), + .end_offset = token.endOffset(), + .type = token.type()}; + continue; + } + validate_input_token(token, term); + unigrams.push_back({.term = std::string(term), + .start_offset = token.startOffset(), + .end_offset = token.endOffset(), + .type = token.type()}); + indexed_grams.emplace_back(); + } + + std::vector output; + if (unigrams.size() < 2) { + _output = std::move(unigrams); + _prepared = true; + return; + } + + bool last_pair_used_gram = false; + std::optional left_is_common; + for (size_t i = 0; i + 1 < unigrams.size(); ++i) { + last_pair_used_gram = pair_uses_gram(unigrams, &left_is_common, i); + if (last_pair_used_gram) { + if (!append_gram_output(&output, indexed_grams[i], unigrams[i], unigrams[i + 1])) { + output.clear(); + for (const auto& unigram : unigrams) { + append_plain_output(&output, unigram); + } + _output = std::move(output); + _prepared = true; + return; + } + } else { + append_plain_output(&output, unigrams[i]); + } + } + if (!last_pair_used_gram) { + append_plain_output(&output, unigrams.back()); + } + _output = std::move(output); + _prepared = true; +} + +Token* CommonGramsQueryFilterBase::emit(Token* token, const CommonGramsBufferedToken& buffered) { + set_output_token(token, buffered.term, 1, buffered.start_offset, buffered.end_offset, + buffered.type); + return token; +} + +Token* CommonGramsQueryFilterBase::next(Token* token) { + if (_failure != nullptr) { + std::rethrow_exception(_failure); + } + if (!_prepared) { + try { + prepare_output(); + } catch (...) { + _failure = std::current_exception(); + throw; + } + } + if (_next_output == _output.size()) { + return nullptr; + } + return emit(token, _output[_next_output++]); +} + +void CommonGramsQueryFilterBase::reset() { + DorisTokenFilter::reset(); + _output.clear(); + _next_output = 0; + _prepared = false; + _failure = nullptr; +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/token_filter/common_grams_filter.h b/be/src/storage/index/inverted/token_filter/common_grams_filter.h new file mode 100644 index 00000000000000..5fae9ce79890ac --- /dev/null +++ b/be/src/storage/index/inverted/token_filter/common_grams_filter.h @@ -0,0 +1,152 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/inverted/token_filter/token_filter.h" + +namespace doris::segment_v2::inverted_index { + +inline constexpr const TCHAR* COMMON_GRAM_TOKEN_TYPE = L"common_gram"; + +inline bool is_common_gram_token_type(const TCHAR* type) { + return std::wstring_view(type) == COMMON_GRAM_TOKEN_TYPE; +} + +struct CommonGramsBufferedToken { + std::string term; + int32_t start_offset = 0; + int32_t end_offset = 0; + const TCHAR* type = Token::getDefaultType(); +}; + +enum class CommonGramsOutputMode { + kLogical, + kEscapedV1Index, + kEscapedV1SpimiIndex, +}; + +struct SniiCommonGramsIndexEvent { + // Analyzer output before the EscapedV1 namespace transform. The view has the + // same lifetime as plain_term and is used for validate-on-first-intern. + std::string_view logical_term; + // EscapedV1 physical plain key. The view remains valid until the next event + // call or reset; the SNII writer interns it synchronously. + std::string_view plain_term; +}; + +class CommonGramsFilter final : public DorisTokenFilter { +public: + CommonGramsFilter(TokenStreamPtr in, std::shared_ptr common_words, + CommonGramsOutputMode output_mode = CommonGramsOutputMode::kLogical); + + Token* next(Token* token) override; + bool next_snii_index_event(SniiCommonGramsIndexEvent* event); + void reset() override; + const CommonWordSet& common_words() const { return *_common_words; } + +private: + bool read_input(CommonGramsBufferedToken* token, std::optional* is_common); + std::string_view encode_plain_term(std::string_view term); + std::string_view encode_snii_plain_term(std::string_view term); + Token* emit_unigram(Token* token, const CommonGramsBufferedToken& buffered); + Token* emit_gram(Token* token, int32_t start_offset, int32_t end_offset); + + std::shared_ptr _common_words; + CommonGramsOutputMode _output_mode; + Token _input_token; + CommonGramsBufferedToken _current; + CommonGramsBufferedToken _lookahead; + std::optional _current_is_common; + std::optional _lookahead_is_common; + bool _has_current = false; + bool _emit_current = false; + std::string _gram; + std::string _encoded_plain_term; +}; + +class CommonGramsPositionFilter final : public DorisTokenFilter { +public: + explicit CommonGramsPositionFilter(TokenStreamPtr in) : DorisTokenFilter(std::move(in)) {} + + Token* next(Token* token) override; +}; + +enum class CommonGramsQueryMode { + kExact, + kPhrasePrefix, +}; + +// A false result proves that the purpose-specific query filter cannot select a +// gram. A true result stays conservative because key encoding may still force +// the filter to replay the complete plain stream. +bool common_grams_query_may_use_gram(std::span terms, CommonGramsQueryMode mode, + const CommonWordSet& common_words); + +class CommonGramsQueryFilterBase : public DorisTokenFilter { +public: + CommonGramsQueryFilterBase(TokenStreamPtr in, std::shared_ptr common_words, + CommonGramsQueryMode mode); + + Token* next(Token* token) override; + void reset() override; + +private: + void prepare_output(); + bool pair_uses_gram(const std::vector& unigrams, + std::optional* left_is_common, size_t pair_index) const; + static void append_plain_output(std::vector* output, + const CommonGramsBufferedToken& token); + static bool append_gram_output(std::vector* output, + const std::optional& indexed_gram, + const CommonGramsBufferedToken& left, + const CommonGramsBufferedToken& right); + static Token* emit(Token* token, const CommonGramsBufferedToken& buffered); + + std::shared_ptr _common_words; + CommonGramsQueryMode _mode; + std::vector _output; + size_t _next_output = 0; + bool _prepared = false; + std::exception_ptr _failure; +}; + +class CommonGramsQueryFilter final : public CommonGramsQueryFilterBase { +public: + CommonGramsQueryFilter(TokenStreamPtr in, std::shared_ptr common_words) + : CommonGramsQueryFilterBase(std::move(in), std::move(common_words), + CommonGramsQueryMode::kExact) {} +}; + +class CommonGramsPhrasePrefixFilter final : public CommonGramsQueryFilterBase { +public: + CommonGramsPhrasePrefixFilter(TokenStreamPtr in, + std::shared_ptr common_words) + : CommonGramsQueryFilterBase(std::move(in), std::move(common_words), + CommonGramsQueryMode::kPhrasePrefix) {} +}; + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/token_filter/common_grams_filter_factory.h b/be/src/storage/index/inverted/token_filter/common_grams_filter_factory.h new file mode 100644 index 00000000000000..7803c7f00b6004 --- /dev/null +++ b/be/src/storage/index/inverted/token_filter/common_grams_filter_factory.h @@ -0,0 +1,107 @@ +// 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. + +#pragma once + +#include + +#include "common/exception.h" +#include "storage/index/inverted/token_filter/common_grams_filter.h" +#include "storage/index/inverted/token_filter/token_filter_factory.h" + +namespace doris::segment_v2::inverted_index { + +class CommonGramsFilterFactory final : public TokenFilterFactory { +public: + explicit CommonGramsFilterFactory(std::shared_ptr common_words = nullptr) + : _common_words(std::move(common_words)) {} + + void initialize(const Settings& settings) override { + if (!settings.empty()) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "common_grams does not accept settings"); + } + if (_common_words == nullptr) { + _common_words = CommonWordSet::default_word_set(); + } + } + + TokenFilterPtr create(const TokenStreamPtr& in) override { + return std::make_shared(in, _common_words, _output_mode); + } + + void set_common_words(std::shared_ptr common_words) { + _common_words = std::move(common_words); + } + + void set_output_mode(CommonGramsOutputMode output_mode) { _output_mode = output_mode; } + + const std::shared_ptr& common_words() const { return _common_words; } + +private: + std::shared_ptr _common_words; + CommonGramsOutputMode _output_mode = CommonGramsOutputMode::kLogical; +}; + +class CommonGramsPositionFilterFactory final : public TokenFilterFactory { +public: + void initialize(const Settings&) override {} + + TokenFilterPtr create(const TokenStreamPtr& in) override { + return std::make_shared(in); + } + + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } +}; + +class CommonGramsQueryFilterFactory final : public TokenFilterFactory { +public: + explicit CommonGramsQueryFilterFactory(std::shared_ptr common_words) + : _common_words(std::move(common_words)) {} + + void initialize(const Settings&) override {} + + TokenFilterPtr create(const TokenStreamPtr& in) override { + return std::make_shared(in, _common_words); + } + + const std::shared_ptr& common_words() const { return _common_words; } + +private: + std::shared_ptr _common_words; +}; + +class CommonGramsPhrasePrefixFilterFactory final : public TokenFilterFactory { +public: + explicit CommonGramsPhrasePrefixFilterFactory(std::shared_ptr common_words) + : _common_words(std::move(common_words)) {} + + void initialize(const Settings&) override {} + + TokenFilterPtr create(const TokenStreamPtr& in) override { + return std::make_shared(in, _common_words); + } + + const std::shared_ptr& common_words() const { return _common_words; } + +private: + std::shared_ptr _common_words; +}; + +} // namespace doris::segment_v2::inverted_index diff --git a/be/src/storage/index/inverted/token_filter/empty_token_filter_factory.h b/be/src/storage/index/inverted/token_filter/empty_token_filter_factory.h index 2855546e78bdb1..d1c42966f14a11 100644 --- a/be/src/storage/index/inverted/token_filter/empty_token_filter_factory.h +++ b/be/src/storage/index/inverted/token_filter/empty_token_filter_factory.h @@ -46,6 +46,10 @@ class EmptyTokenFilterFactory : public TokenFilterFactory { token_filter->initialize(); return token_filter; } + + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } }; } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/src/storage/index/inverted/token_filter/icu_normalizer_filter_factory.h b/be/src/storage/index/inverted/token_filter/icu_normalizer_filter_factory.h index c0cedb8c483cbd..27dc897e818a34 100644 --- a/be/src/storage/index/inverted/token_filter/icu_normalizer_filter_factory.h +++ b/be/src/storage/index/inverted/token_filter/icu_normalizer_filter_factory.h @@ -82,6 +82,10 @@ class ICUNormalizerFilterFactory : public TokenFilterFactory { return std::make_shared(in, _normalizer); } + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } + private: static const icu::Normalizer2* get_normalizer(const std::string& name, UErrorCode& status) { std::string lower_name = to_lower_copy(trim_copy(name)); diff --git a/be/src/storage/index/inverted/token_filter/lower_case_filter.h b/be/src/storage/index/inverted/token_filter/lower_case_filter.h index dd55d48060d03a..ff737209a1ea09 100644 --- a/be/src/storage/index/inverted/token_filter/lower_case_filter.h +++ b/be/src/storage/index/inverted/token_filter/lower_case_filter.h @@ -19,10 +19,40 @@ #include +#ifdef BE_TEST +#include +#endif + +#include "common/cast_set.h" +#include "common/exception.h" #include "storage/index/inverted/token_filter/token_filter.h" +#include "util/utf8_check.h" namespace doris::segment_v2::inverted_index { +#ifdef BE_TEST +namespace lower_case_testing { + +inline std::atomic& unicode_path_counter() { + static std::atomic counter {0}; + return counter; +} + +inline uint64_t unicode_path_count() { + return unicode_path_counter().load(std::memory_order_relaxed); +} + +inline void reset_unicode_path_count() { + unicode_path_counter().store(0, std::memory_order_relaxed); +} + +inline void note_unicode_path() { + unicode_path_counter().fetch_add(1, std::memory_order_relaxed); +} + +} // namespace lower_case_testing +#endif + /** * @brief A token filter that converts Unicode text to lowercase using ICU library. * @@ -40,7 +70,7 @@ class LowerCaseFilter : public DorisTokenFilter { UErrorCode status = U_ZERO_ERROR; auto* ucsm = ucasemap_open("", 0, &status); if (U_FAILURE(status)) { - throw Exception(ErrorCode::RUNTIME_ERROR, + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, "Failed to open UCaseMap. ICU Error: " + std::to_string(status) + " - " + u_errorName(status)); } @@ -52,21 +82,53 @@ class LowerCaseFilter : public DorisTokenFilter { return nullptr; } std::string_view term(t->termBuffer(), t->termLength()); - - size_t max_len = term.size() * 2; - if (_lower_term.size() < max_len) { - _lower_term.resize(max_len); + bool has_ascii_upper = false; + bool all_ascii = true; + for (const char value : term) { + const auto byte = static_cast(value); + if (byte >= 0x80) { + all_ascii = false; + break; + } + has_ascii_upper |= byte >= 'A' && byte <= 'Z'; + } + if (all_ascii) { + if (!has_ascii_upper) { + return t; + } + _lower_term.resize(term.size()); + for (size_t i = 0; i < term.size(); ++i) { + const auto byte = static_cast(term[i]); + _lower_term[i] = + static_cast(byte >= 'A' && byte <= 'Z' ? byte + ('a' - 'A') : byte); + } + set_text(t, _lower_term); + return t; + } +#ifdef BE_TEST + lower_case_testing::note_unicode_path(); +#endif + if (!validate_utf8(term.data(), term.size())) { + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "Failed to lowercase token: invalid UTF-8"); } + _lower_term.resize(term.size()); UErrorCode status = U_ZERO_ERROR; - int32_t result_len = ucasemap_utf8ToLower(_ucsm.get(), _lower_term.data(), max_len, - term.data(), term.size(), &status); + int32_t result_len = ucasemap_utf8ToLower( + _ucsm.get(), _lower_term.data(), cast_set(_lower_term.size()), term.data(), + cast_set(term.size()), &status); + if (status == U_BUFFER_OVERFLOW_ERROR) { + _lower_term.resize(cast_set(result_len)); + status = U_ZERO_ERROR; + result_len = ucasemap_utf8ToLower(_ucsm.get(), _lower_term.data(), + cast_set(_lower_term.size()), term.data(), + cast_set(term.size()), &status); + } if (U_FAILURE(status)) { - LOG(WARNING) << "Failed to convert to lowercase. " - << "Term: '" << term << "', " - << "ICU Error: " << status << " - " << u_errorName(status) - << ", Buffer size: " << max_len << std::endl; - return nullptr; + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "Failed to convert token to lowercase. ICU Error: {} - {}", + static_cast(status), u_errorName(status)); } set_text(t, std::string_view(_lower_term.data(), result_len)); diff --git a/be/src/storage/index/inverted/token_filter/lower_case_filter_factory.h b/be/src/storage/index/inverted/token_filter/lower_case_filter_factory.h index 03467c5114a6b4..b24ad4b0822944 100644 --- a/be/src/storage/index/inverted/token_filter/lower_case_filter_factory.h +++ b/be/src/storage/index/inverted/token_filter/lower_case_filter_factory.h @@ -34,6 +34,10 @@ class LowerCaseFilterFactory : public TokenFilterFactory { token_filter->initialize(); return token_filter; } + + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } }; } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/src/storage/index/inverted/token_filter/token_filter_factory.h b/be/src/storage/index/inverted/token_filter/token_filter_factory.h index ebbd8836715e8d..5e1bd93236977a 100644 --- a/be/src/storage/index/inverted/token_filter/token_filter_factory.h +++ b/be/src/storage/index/inverted/token_filter/token_filter_factory.h @@ -28,6 +28,7 @@ class TokenFilterFactory : public AbstractAnalysisFactory { ~TokenFilterFactory() override = default; virtual TokenFilterPtr create(const TokenStreamPtr& in) = 0; + virtual PositionCapability position_capability() const { return PositionCapability::kUnknown; } }; using TokenFilterFactoryPtr = std::shared_ptr; diff --git a/be/src/storage/index/inverted/tokenizer/char/char_group_tokenizer_factory.h b/be/src/storage/index/inverted/tokenizer/char/char_group_tokenizer_factory.h index 030c0c2512f850..68273ad1366783 100644 --- a/be/src/storage/index/inverted/tokenizer/char/char_group_tokenizer_factory.h +++ b/be/src/storage/index/inverted/tokenizer/char/char_group_tokenizer_factory.h @@ -30,6 +30,10 @@ class CharGroupTokenizerFactory : public TokenizerFactory { TokenizerPtr create() override; + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } + private: static UChar32 parse_escaped_char(const icu::UnicodeString& unicode_str); diff --git a/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.cpp b/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.cpp index d7b774ef0966a4..54ac1d72c9de82 100644 --- a/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.cpp +++ b/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.cpp @@ -17,10 +17,32 @@ #include "storage/index/inverted/tokenizer/char/char_tokenizer.h" +#ifdef BE_TEST +#include +#endif + #include "common/exception.h" namespace doris::segment_v2::inverted_index { +#ifdef BE_TEST +namespace { +std::atomic g_non_ascii_decode_count {0}; +} // namespace + +namespace char_tokenizer_testing { + +uint64_t non_ascii_decode_count() { + return g_non_ascii_decode_count.load(std::memory_order_relaxed); +} + +void reset_non_ascii_decode_count() { + g_non_ascii_decode_count.store(0, std::memory_order_relaxed); +} + +} // namespace char_tokenizer_testing +#endif + void CharTokenizer::initialize(int32_t max_token_len) { if (max_token_len > MAX_TOKEN_LENGTH_LIMIT || max_token_len <= 0) { throw Exception(ErrorCode::INVALID_ARGUMENT, @@ -29,6 +51,33 @@ void CharTokenizer::initialize(int32_t max_token_len) { " passed: " + std::to_string(max_token_len)); } _max_token_len = max_token_len; + for (size_t value = 0; value < _ascii_char_classes.size(); ++value) { + const auto c = static_cast(value); + _ascii_char_classes[value] = + is_cjk_char(c) + ? AsciiCharClass::kCjk + : (is_token_char(c) ? AsciiCharClass::kToken : AsciiCharClass::kDelimiter); + } +} + +CharTokenizer::AsciiCharClass CharTokenizer::read_next_char_class() { + const auto first_byte = static_cast(_char_buffer[_buffer_index]); + if (first_byte < _ascii_char_classes.size()) { + ++_buffer_index; + return _ascii_char_classes[first_byte]; + } +#ifdef BE_TEST + g_non_ascii_decode_count.fetch_add(1, std::memory_order_relaxed); +#endif + UChar32 c = U_UNASSIGNED; + U8_NEXT(_char_buffer, _buffer_index, _data_len, c); + if (c < 0) { + return AsciiCharClass::kInvalid; + } + if (is_cjk_char(c)) { + return AsciiCharClass::kCjk; + } + return is_token_char(c) ? AsciiCharClass::kToken : AsciiCharClass::kDelimiter; } Token* CharTokenizer::next(Token* token) { @@ -46,14 +95,13 @@ Token* CharTokenizer::next(Token* token) { break; } - UChar32 c = U_UNASSIGNED; const int32_t prev_i = _buffer_index; - U8_NEXT(_char_buffer, _buffer_index, _data_len, c); - if (c < 0) { + const AsciiCharClass char_class = read_next_char_class(); + if (char_class == AsciiCharClass::kInvalid) { continue; } - if (is_cjk_char(c)) { + if (char_class == AsciiCharClass::kCjk) { if (start == -1) { start = prev_i; end = _buffer_index - 1; @@ -61,7 +109,7 @@ Token* CharTokenizer::next(Token* token) { _buffer_index = prev_i; } break; - } else if (is_token_char(c)) { + } else if (char_class == AsciiCharClass::kToken) { if (start == -1) { start = prev_i; } diff --git a/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.h b/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.h index 66a9979a9fc97b..bae2f6c55d734f 100644 --- a/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.h +++ b/be/src/storage/index/inverted/tokenizer/char/char_tokenizer.h @@ -17,6 +17,9 @@ #pragma once +#include +#include + #include "storage/index/inverted/tokenizer/tokenizer.h" namespace doris::segment_v2::inverted_index { @@ -36,8 +39,13 @@ class CharTokenizer : public DorisTokenizer { static constexpr int32_t DEFAULT_MAX_WORD_LEN = 255; private: + enum class AsciiCharClass : uint8_t { kDelimiter, kToken, kCjk, kInvalid }; + static constexpr int32_t MAX_TOKEN_LENGTH_LIMIT = 16383; + AsciiCharClass read_next_char_class(); + + std::array _ascii_char_classes {}; int32_t _max_token_len = 0; int32_t _buffer_index = 0; diff --git a/be/src/storage/index/inverted/tokenizer/empty/empty_tokenizer_factory.h b/be/src/storage/index/inverted/tokenizer/empty/empty_tokenizer_factory.h index 27229ba99f1184..b540dfa45a3f61 100644 --- a/be/src/storage/index/inverted/tokenizer/empty/empty_tokenizer_factory.h +++ b/be/src/storage/index/inverted/tokenizer/empty/empty_tokenizer_factory.h @@ -67,6 +67,10 @@ class EmptyTokenizerFactory : public TokenizerFactory { tokenizer->initialize(); return tokenizer; } + + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } }; } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/src/storage/index/inverted/tokenizer/ngram/edge_ngram_tokenizer_factory.h b/be/src/storage/index/inverted/tokenizer/ngram/edge_ngram_tokenizer_factory.h index 8d63b5cd230582..a0b88153a7a11f 100644 --- a/be/src/storage/index/inverted/tokenizer/ngram/edge_ngram_tokenizer_factory.h +++ b/be/src/storage/index/inverted/tokenizer/ngram/edge_ngram_tokenizer_factory.h @@ -56,6 +56,10 @@ class EdgeNGramTokenizerFactory : public TokenizerFactory { } } + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } + private: int32_t _min_gram = 0; int32_t _max_gram = 0; diff --git a/be/src/storage/index/inverted/tokenizer/ngram/ngram_tokenizer_factory.h b/be/src/storage/index/inverted/tokenizer/ngram/ngram_tokenizer_factory.h index d064749d9d51a5..2ee428e32ff77f 100644 --- a/be/src/storage/index/inverted/tokenizer/ngram/ngram_tokenizer_factory.h +++ b/be/src/storage/index/inverted/tokenizer/ngram/ngram_tokenizer_factory.h @@ -50,6 +50,10 @@ class NGramTokenizerFactory : public TokenizerFactory { } } + PositionCapability position_capability() const override { + return PositionCapability::kAlwaysUnitIncrement; + } + static void initialize_matchers(); static CharMatcherPtr parse_token_chars(const Settings& settings); diff --git a/be/src/storage/index/inverted/tokenizer/tokenizer_factory.h b/be/src/storage/index/inverted/tokenizer/tokenizer_factory.h index b788990ecf09be..654726a9b976bb 100644 --- a/be/src/storage/index/inverted/tokenizer/tokenizer_factory.h +++ b/be/src/storage/index/inverted/tokenizer/tokenizer_factory.h @@ -28,6 +28,7 @@ class TokenizerFactory : public AbstractAnalysisFactory { ~TokenizerFactory() override = default; virtual TokenizerPtr create() = 0; + virtual PositionCapability position_capability() const { return PositionCapability::kUnknown; } }; using TokenizerFactoryPtr = std::shared_ptr; diff --git a/be/src/storage/index/snii/common/slice.h b/be/src/storage/index/snii/common/slice.h new file mode 100644 index 00000000000000..e5b80932944df3 --- /dev/null +++ b/be/src/storage/index/snii/common/slice.h @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include + +namespace doris::snii { + +// Read-only byte view (does not own memory). Lifetime is managed by the underlying buffer. +class Slice { +public: + Slice() = default; + Slice(const uint8_t* d, size_t n) : data_(d), size_(n) {} + explicit Slice(const std::vector& v) : data_(v.data()), size_(v.size()) {} + explicit Slice(std::string_view sv) + : data_(reinterpret_cast(sv.data())), size_(sv.size()) {} + + const uint8_t* data() const { return data_; } + size_t size() const { return size_; } + bool empty() const { return size_ == 0; } + + uint8_t operator[](size_t i) const { + assert(i < size_); + return data_[i]; + } + + Slice subslice(size_t off, size_t n) const { + assert(off + n <= size_); + return Slice(data_ + off, n); + } + +private: + const uint8_t* data_ = nullptr; + size_t size_ = 0; +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/common/uninitialized_buffer.h b/be/src/storage/index/snii/common/uninitialized_buffer.h new file mode 100644 index 00000000000000..1b1855db566d0c --- /dev/null +++ b/be/src/storage/index/snii/common/uninitialized_buffer.h @@ -0,0 +1,82 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +// Single auditable seam for "resize a decode buffer that is about to be fully +// overwritten". Collapses the scattered `assign(n, 0)` / `resize(n)` calls in the +// PFOR/zstd decode paths into one intent-revealing entry point and removes the +// redundant zero-fill on warm-reused buffers. +namespace doris::snii { + +// Resize a vector of trivially-copyable T to `n` without any zero-fill beyond what +// std::vector itself mandates. +// +// CONTRACT: the caller MUST fully overwrite [0, n) before reading any element. The +// PFOR (`pfor_decode`/`bitunpack`, including the w==0 memset and exception patch +// paths) and zstd (`ZSTD_decompress` + length check) decoders all satisfy this. +// +// For a warm-reused buffer whose current size() >= n this performs NO +// value-initialization (resize shrinks in place). For a cold/grown buffer +// std::vector still value-initializes the new tail -- that is unavoidable for +// std::vector and is intentional here. Do NOT use on any buffer that may be read +// before being fully overwritten. +template +inline void resize_uninitialized(std::vector& v, std::size_t n) { + static_assert(std::is_trivially_copyable_v, + "resize_uninitialized requires a trivially-copyable element type"); + v.resize(n); +} + +// Allocator that default-initializes (instead of value-initializes) on the no-arg +// construct path: for trivial T this skips zeroing entirely. Use only for buffers +// that are fully overwritten before any read. +template +struct default_init_allocator : std::allocator { + template + struct rebind { + using other = default_init_allocator; + }; + using std::allocator::allocator; + + template + void construct(U* p, Args&&... args) { + if constexpr (sizeof...(Args) == 0) { + ::new (static_cast(p)) U; // default-init: no zeroing for trivial U + } else { + ::new (static_cast(p)) U(std::forward(args)...); + } + } +}; + +// A vector whose grow path default-initializes, so even a cold grow avoids the +// zero-fill that std::vector would perform for trivial T. +template +using uninitialized_vector = std::vector>; + +template +inline void resize_uninitialized(uninitialized_vector& v, std::size_t n) { + v.resize(n); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/compaction/eligibility.cpp b/be/src/storage/index/snii/compaction/eligibility.cpp new file mode 100644 index 00000000000000..92b8676e979ad7 --- /dev/null +++ b/be/src/storage/index/snii/compaction/eligibility.cpp @@ -0,0 +1,386 @@ +// 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 "storage/index/snii/compaction/eligibility.h" + +#include +#include +#include +#include + +#include "CLucene.h" +#include "common/check.h" +#include "common/config.h" +#include "common/exception.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/tablet/tablet_schema.h" + +namespace doris::snii::compaction { + +namespace inverted_index = segment_v2::inverted_index; + +namespace { + +Status reject(std::string_view reason) { + return Status::Error( + "SNII streamed compaction is not eligible: {}", reason); +} + +InvertedIndexAnalyzerConfig analyzer_config_from_properties( + const std::map& properties) { + InvertedIndexAnalyzerConfig config; + config.analyzer_name = get_analyzer_name_from_properties(properties); + config.parser_type = get_inverted_index_parser_type_from_string( + get_parser_string_from_properties(properties)); + config.parser_mode = get_parser_mode_string_from_properties(properties); + config.char_filter_map = get_parser_char_filter_map_from_properties(properties); + config.lower_case = get_parser_lowercase_from_properties(properties); + config.stop_words = get_parser_stopwords_from_properties(properties); + return config; +} + +Status validate_source_shape(const reader::LogicalIndexReader& source, size_t source_ordinal) { + if (source.tier() != format::IndexTier::kT2) { + return reject(fmt::format("source {} is not T2", source_ordinal)); + } + if (!source.has_positions()) { + return reject(fmt::format("source {} has no positions", source_ordinal)); + } + const auto& norms = source.section_refs().norms; + if (norms.offset != 0 || norms.length != 0) { + return reject(fmt::format("source {} carries scoring norms", source_ordinal)); + } + if (source.common_grams_metadata() != nullptr) { + return reject( + fmt::format("source {} carries CommonGrams/scoring metadata", source_ordinal)); + } + + const auto& stats = source.stats(); + if (stats.indexed_doc_count > stats.doc_count || + stats.null_count != stats.doc_count - stats.indexed_doc_count) { + return reject(fmt::format("source {} document statistics violate indexed + null = doc", + source_ordinal)); + } + return Status::OK(); +} + +inverted_index::CommonGramsSegmentMetadata static_common_grams_metadata( + const inverted_index::CommonGramsSegmentMetadata& metadata) { + auto seed = metadata; + seed.scoring_doc_count = 0; + seed.scoring_token_count = 0; + return seed; +} + +Status validate_common_grams_source_shape(const reader::LogicalIndexReader& source, + size_t source_ordinal) { + if (source.tier() != format::IndexTier::kT3 || !source.has_positions()) { + return reject(fmt::format("source {} is not CommonGrams T3", source_ordinal)); + } + if (source.section_refs().norms.length == 0) { + return reject(fmt::format("source {} has no scoring norms", source_ordinal)); + } + const auto* metadata = source.common_grams_metadata(); + if (metadata == nullptr) { + return reject(fmt::format("source {} has no CommonGrams metadata", source_ordinal)); + } + const Status metadata_status = + inverted_index::validate_common_grams_segment_metadata(*metadata); + const bool complete_shape = + metadata->common_grams_coverage == inverted_index::CommonGramsCoverage::kComplete && + source.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kNone; + const bool hybrid_shape = + metadata->common_grams_coverage == inverted_index::CommonGramsCoverage::kMixed && + source.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1; + if (!metadata_status.ok() || (!complete_shape && !hybrid_shape) || + metadata->plain_term_key_version != inverted_index::PlainTermKeyVersion::kEscapedV1 || + metadata->common_grams_semantics_version != + inverted_index::COMMON_GRAMS_SEMANTICS_VERSION_V1 || + metadata->common_grams_key_version != inverted_index::COMMON_GRAMS_KEY_VERSION_V1 || + metadata->scoring_coverage != inverted_index::ScoringCoverage::kComplete || + metadata->scoring_stats_version != inverted_index::COMMON_GRAMS_SCORING_STATS_VERSION_V1 || + metadata->norm_semantics_version != + inverted_index::COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1) { + return reject(fmt::format("source {} has incomplete or unsupported CommonGrams metadata", + source_ordinal)); + } + const Status scoring_status = inverted_index::validate_snii_scoring_metadata( + metadata, source.stats().doc_count, source.stats().sum_total_term_freq, + /*has_scoring_tier=*/true, /*has_positions=*/true, /*has_norms=*/true); + if (!scoring_status.ok()) { + return reject(fmt::format("source {} has incomplete scoring metadata: {}", source_ordinal, + scoring_status.to_string())); + } + const auto& stats = source.stats(); + if (stats.indexed_doc_count > stats.doc_count || + stats.null_count != stats.doc_count - stats.indexed_doc_count) { + return reject(fmt::format("source {} document statistics violate indexed + null = doc", + source_ordinal)); + } + return Status::OK(); +} + +Status reject_legacy_bigram(const reader::LogicalIndexReader& source, size_t source_ordinal) { + bool found_legacy_bigram = false; + RETURN_IF_ERROR(source.visit_prefix_terms( + format::kPhraseBigramTermMarker, + [&found_legacy_bigram](reader::LogicalIndexReader::PrefixHit&&, bool* stop) { + found_legacy_bigram = true; + *stop = true; + return Status::OK(); + })); + if (found_legacy_bigram) { + return reject( + fmt::format("source {} contains a legacy phrase-bigram term", source_ordinal)); + } + return Status::OK(); +} + +Status resolve_destination_analyzer(const TabletIndex& destination_index, + const AnalyzerProviderFactory& analyzer_provider_factory, + inverted_index::AnalyzerProviderPtr* analyzer_provider) { + analyzer_provider->reset(); + const auto& properties = destination_index.properties(); + if (!inverted_index::InvertedIndexAnalyzer::should_analyzer(properties)) { + return Status::OK(); + } + const InvertedIndexAnalyzerConfig analyzer_config = analyzer_config_from_properties(properties); + try { + *analyzer_provider = + analyzer_provider_factory + ? analyzer_provider_factory(analyzer_config) + : inverted_index::InvertedIndexAnalyzer::create_analyzer_provider( + &analyzer_config); + } catch (const CLuceneError& error) { + return reject(fmt::format("destination analyzer resolution failed: {}", error.what())); + } catch (const Exception& error) { + return reject(fmt::format("destination analyzer resolution failed: {}", error.what())); + } + DORIS_CHECK(*analyzer_provider != nullptr); + return Status::OK(); +} + +Status validate_destination_policy(const TabletIndex& destination_index, + const AnalyzerProviderFactory& analyzer_provider_factory) { + const auto& properties = destination_index.properties(); + if (get_parser_phrase_support_string_from_properties(properties) != + INVERTED_INDEX_PARSER_PHRASE_SUPPORT_YES) { + return reject("destination does not request phrase positions"); + } + if (!inverted_index::InvertedIndexAnalyzer::should_analyzer(properties)) { + return Status::OK(); + } + + inverted_index::AnalyzerProviderPtr analyzer_provider; + RETURN_IF_ERROR(resolve_destination_analyzer(destination_index, analyzer_provider_factory, + &analyzer_provider)); + DORIS_CHECK(analyzer_provider != nullptr); + if (config::enable_common_grams_index_build && analyzer_provider->uses_common_grams()) { + return reject("destination analyzer policy would build CommonGrams"); + } + return Status::OK(); +} + +} // namespace + +Status validate_plain_t2_source(const reader::LogicalIndexReader& source, size_t source_ordinal) { + return validate_source_shape(source, source_ordinal); +} + +Status validate_plain_t2_source_eligibility(const reader::LogicalIndexReader& source, + size_t source_ordinal) { + RETURN_IF_ERROR(validate_source_shape(source, source_ordinal)); + return reject_legacy_bigram(source, source_ordinal); +} + +Status validate_snii_source_eligibility(const reader::LogicalIndexReader& source, + size_t source_ordinal, + const SniiCompactionEligibility& eligibility) { + if (eligibility.kind == SniiStreamedMergeKind::kPlainT2) { + return validate_plain_t2_source_eligibility(source, source_ordinal); + } + RETURN_IF_ERROR(validate_common_grams_source_shape(source, source_ordinal)); + RETURN_IF_ERROR(reject_legacy_bigram(source, source_ordinal)); + if (!eligibility.common_grams_metadata_seed.has_value()) { + return reject("CommonGrams eligibility has no metadata identity seed"); + } + if (source.common_grams_posting_policy() != eligibility.common_grams_posting_policy) { + return reject(fmt::format("source {} CommonGrams posting policy differs from eligibility", + source_ordinal)); + } + if (static_common_grams_metadata(*source.common_grams_metadata()) != + *eligibility.common_grams_metadata_seed) { + return reject(fmt::format("source {} CommonGrams static identity differs from eligibility", + source_ordinal)); + } + return Status::OK(); +} + +Status validate_plain_t2_compaction_eligibility( + std::span sources, const TabletIndex& destination_index, + const AnalyzerProviderFactory& analyzer_provider_factory) { + if (sources.empty()) { + return reject("no source logical indexes"); + } + if (!destination_index.is_inverted_index()) { + return reject("destination is not an inverted index"); + } + + const TabletIndex& first_index_meta = sources.front().index_meta.get(); + if (!first_index_meta.is_inverted_index()) { + return reject("source 0 is not an inverted index"); + } + const auto& source_properties = first_index_meta.properties(); + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + const TabletIndex& index_meta = sources[source_ordinal].index_meta.get(); + if (!index_meta.is_inverted_index()) { + return reject(fmt::format("source {} is not an inverted index", source_ordinal)); + } + if (index_meta.properties() != source_properties) { + return reject(fmt::format("source {} properties differ from source 0", source_ordinal)); + } + if (index_meta.index_id() != destination_index.index_id()) { + return reject( + fmt::format("source {} index id differs from destination", source_ordinal)); + } + if (index_meta.get_index_suffix() != destination_index.get_index_suffix()) { + return reject( + fmt::format("source {} index suffix differs from destination", source_ordinal)); + } + } + if (destination_index.properties() != source_properties) { + return reject("destination properties differ from source properties"); + } + + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + RETURN_IF_ERROR(validate_plain_t2_source_eligibility(sources[source_ordinal].reader.get(), + source_ordinal)); + } + return validate_destination_policy(destination_index, analyzer_provider_factory); +} + +Status validate_snii_compaction_eligibility( + std::span sources, const TabletIndex& destination_index, + SniiCompactionEligibility* out, const AnalyzerProviderFactory& analyzer_provider_factory) { + if (out == nullptr) { + return Status::InvalidArgument("SNII compaction eligibility has null output"); + } + *out = SniiCompactionEligibility {}; + if (sources.empty()) { + return reject("no source logical indexes"); + } + if (!destination_index.is_inverted_index()) { + return reject("destination is not an inverted index"); + } + if (get_parser_phrase_support_string_from_properties(destination_index.properties()) != + INVERTED_INDEX_PARSER_PHRASE_SUPPORT_YES) { + return reject("destination does not request phrase positions"); + } + + const TabletIndex& first_index_meta = sources.front().index_meta.get(); + if (!first_index_meta.is_inverted_index()) { + return reject("source 0 is not an inverted index"); + } + const auto& source_properties = first_index_meta.properties(); + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + const TabletIndex& index_meta = sources[source_ordinal].index_meta.get(); + if (!index_meta.is_inverted_index()) { + return reject(fmt::format("source {} is not an inverted index", source_ordinal)); + } + if (index_meta.properties() != source_properties) { + return reject(fmt::format("source {} properties differ from source 0", source_ordinal)); + } + if (index_meta.index_id() != destination_index.index_id()) { + return reject( + fmt::format("source {} index id differs from destination", source_ordinal)); + } + if (index_meta.get_index_suffix() != destination_index.get_index_suffix()) { + return reject( + fmt::format("source {} index suffix differs from destination", source_ordinal)); + } + } + if (destination_index.properties() != source_properties) { + return reject("destination properties differ from source properties"); + } + + const format::IndexTier source_tier = sources.front().reader.get().tier(); + if (source_tier == format::IndexTier::kT2) { + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + if (sources[source_ordinal].reader.get().tier() != format::IndexTier::kT2) { + return reject("source streamed-merge shapes are not homogeneous"); + } + RETURN_IF_ERROR(validate_plain_t2_source_eligibility( + sources[source_ordinal].reader.get(), source_ordinal)); + } + RETURN_IF_ERROR(validate_destination_policy(destination_index, analyzer_provider_factory)); + out->kind = SniiStreamedMergeKind::kPlainT2; + return Status::OK(); + } + if (source_tier != format::IndexTier::kT3) { + return reject("source streamed-merge shape is neither plain T2 nor CommonGrams T3"); + } + + std::optional source_seed; + std::optional source_policy; + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + const auto& source = sources[source_ordinal].reader.get(); + if (source.tier() != format::IndexTier::kT3) { + return reject("source streamed-merge shapes are not homogeneous"); + } + RETURN_IF_ERROR(validate_common_grams_source_shape(source, source_ordinal)); + RETURN_IF_ERROR(reject_legacy_bigram(source, source_ordinal)); + const auto seed = static_common_grams_metadata(*source.common_grams_metadata()); + if (!source_seed.has_value()) { + source_seed = seed; + source_policy = source.common_grams_posting_policy(); + } else if (source.common_grams_posting_policy() != *source_policy) { + return reject("source CommonGrams posting policies are not homogeneous"); + } else if (*source_seed != seed) { + return reject(fmt::format("source {} CommonGrams static identity differs from source 0", + source_ordinal)); + } + } + + if (!config::enable_common_grams_index_build) { + return reject("destination CommonGrams index build is disabled"); + } + inverted_index::AnalyzerProviderPtr analyzer_provider; + RETURN_IF_ERROR(resolve_destination_analyzer(destination_index, analyzer_provider_factory, + &analyzer_provider)); + if (analyzer_provider == nullptr || !analyzer_provider->uses_common_grams()) { + return reject("destination analyzer does not build CommonGrams"); + } + const auto* destination_identity = analyzer_provider->common_grams_identity(); + if (destination_identity == nullptr) { + return reject("destination analyzer has no complete CommonGrams identity"); + } + if (!source_seed.has_value() || + !inverted_index::common_grams_identity_matches(*source_seed, *destination_identity)) { + return reject("destination CommonGrams identity differs from source identity"); + } + out->kind = SniiStreamedMergeKind::kCommonGramsT3; + out->common_grams_metadata_seed = std::move(source_seed); + DORIS_CHECK(source_policy.has_value()); + out->common_grams_posting_policy = *source_policy; + return Status::OK(); +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/eligibility.h b/be/src/storage/index/snii/compaction/eligibility.h new file mode 100644 index 00000000000000..bbbfae59204d22 --- /dev/null +++ b/be/src/storage/index/snii/compaction/eligibility.h @@ -0,0 +1,114 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/analyzer/analyzer_provider.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/inverted_index_parser.h" +#include "storage/index/snii/format/core_metadata.h" + +namespace doris { + +class TabletIndex; + +namespace snii::reader { +class LogicalIndexReader; +} + +namespace snii::compaction { + +// One already-opened SNII source logical index and the TabletIndex metadata +// that produced it. References make null inputs unrepresentable; both objects +// must outlive validate_plain_t2_compaction_eligibility(). +struct PlainT2CompactionSource { + std::reference_wrapper reader; + std::reference_wrapper index_meta; +}; + +// Injectable only at analyzer-provider construction. Production callers omit +// it and use InvertedIndexAnalyzer::create_analyzer_provider; focused tests can +// supply an immutable provider without mutating the process IndexPolicyMgr. +using AnalyzerProviderFactory = std::function; + +enum class SniiStreamedMergeKind : uint8_t { + kPlainT2, + kCommonGramsT3, +}; + +// Validated destination shape for one streamed merge. CommonGrams metadata is a +// static identity seed: destination doc_count is bound when its session starts, +// and semantic token_count is bound after the single postings pass. +struct SniiCompactionEligibility { + SniiStreamedMergeKind kind = SniiStreamedMergeKind::kPlainT2; + std::optional + common_grams_metadata_seed; + format::CommonGramsPostingPolicy common_grams_posting_policy = + format::CommonGramsPostingPolicy::kNone; +}; + +// O(1) physical/semantic validation for one source. The merge planner may call +// this while preparing opened sources; the aggregate validator below reuses the +// exact same predicate. +Status validate_plain_t2_source(const reader::LogicalIndexReader& source, size_t source_ordinal); + +// Gate-only validation. In addition to the O(1) physical shape checks above, +// performs one bounded DICT prefix seek for the legacy hidden-bigram namespace +// so the owner can select raw rebuild before creating streamed output. +Status validate_plain_t2_source_eligibility(const reader::LogicalIndexReader& source, + size_t source_ordinal); + +// Revalidates one opened source against an already selected streamed shape. +// CommonGrams checks include exact current versions and static identity equality +// with the eligibility seed. +Status validate_snii_source_eligibility(const reader::LogicalIndexReader& source, + size_t source_ordinal, + const SniiCompactionEligibility& eligibility); + +// O(source-count), O(1)-metadata validation for the SNII postings-merge fast +// path. OK means every source is a plain positions-only T2 index, all source +// properties are byte-for-byte equal to the current destination properties, +// source/destination index ids and escaped suffixes match, and the current +// destination analyzer policy will not build CommonGrams. +// Rejections use INVERTED_INDEX_NOT_SUPPORTED with a concrete reason so the +// caller can select raw-column rebuild before creating output. +// +// Each source also receives one bounded seek for the legacy hidden-bigram +// marker namespace. This avoids discovering an unsupported legacy segment only +// after the raw-column compaction path has already been skipped. +Status validate_plain_t2_compaction_eligibility( + std::span sources, const TabletIndex& destination_index, + const AnalyzerProviderFactory& analyzer_provider_factory = {}); + +// Accepts either the existing plain positions-only T2 shape or homogeneous, +// complete CommonGrams T3 sources. Mixed shapes and any CommonGrams identity or +// destination-build-policy mismatch return INVERTED_INDEX_NOT_SUPPORTED so the +// caller can rebuild from raw columns before creating streamed output. +Status validate_snii_compaction_eligibility( + std::span sources, const TabletIndex& destination_index, + SniiCompactionEligibility* out, + const AnalyzerProviderFactory& analyzer_provider_factory = {}); + +} // namespace snii::compaction +} // namespace doris diff --git a/be/src/storage/index/snii/compaction/indexed_winner_tree.h b/be/src/storage/index/snii/compaction/indexed_winner_tree.h new file mode 100644 index 00000000000000..e16205a484832c --- /dev/null +++ b/be/src/storage/index/snii/compaction/indexed_winner_tree.h @@ -0,0 +1,103 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/check.h" + +namespace doris::snii::compaction { + +// Complete binary winner tree over dense source ordinals. Updating one source +// touches only its leaf-to-root path. runner_up() inspects the sibling winners +// on the current winner's path without changing the tree. +template +class IndexedWinnerTree { +public: + static constexpr size_t kNoSource = std::numeric_limits::max(); + + explicit IndexedWinnerTree(Before before) : before_(std::move(before)), nodes_(2, kNoSource) {} + + template + void build(size_t source_count, IsLive&& is_live) { + source_count_ = source_count; + leaf_base_ = 1; + while (leaf_base_ < source_count_) { + DORIS_CHECK_LE(leaf_base_, std::numeric_limits::max() / 2); + leaf_base_ *= 2; + } + nodes_.assign(leaf_base_ * 2, kNoSource); + for (size_t source = 0; source < source_count_; ++source) { + if (is_live(source)) { + nodes_[leaf_base_ + source] = source; + } + } + for (size_t node = leaf_base_; node-- > 1;) { + nodes_[node] = select(nodes_[node * 2], nodes_[node * 2 + 1]); + } + } + + bool empty() const noexcept { return nodes_[1] == kNoSource; } + + size_t winner() const { + DCHECK(!empty()); + return nodes_[1]; + } + + size_t runner_up() const { + const size_t current_winner = winner(); + size_t candidate = kNoSource; + size_t node = leaf_base_ + current_winner; + while (node > 1) { + candidate = select(candidate, nodes_[node ^ 1]); + node /= 2; + } + return candidate; + } + + void update(size_t source, bool live) { + DCHECK_LT(source, source_count_); + size_t node = leaf_base_ + source; + nodes_[node] = live ? source : kNoSource; + while (node > 1) { + node /= 2; + nodes_[node] = select(nodes_[node * 2], nodes_[node * 2 + 1]); + } + } + +private: + size_t select(size_t candidate, size_t challenger) const { + if (candidate == kNoSource) { + return challenger; + } + if (challenger == kNoSource) { + return candidate; + } + return before_(challenger, candidate) ? challenger : candidate; + } + + Before before_; + size_t source_count_ = 0; + size_t leaf_base_ = 1; + boost::container::small_vector nodes_; +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/posting_cursor.cpp b/be/src/storage/index/snii/compaction/posting_cursor.cpp new file mode 100644 index 00000000000000..fbb25efc7701a6 --- /dev/null +++ b/be/src/storage/index/snii/compaction/posting_cursor.cpp @@ -0,0 +1,1025 @@ +// 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 "storage/index/snii/compaction/posting_cursor.h" + +#include + +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/reader/windowed_posting.h" + +namespace doris::snii::compaction { + +namespace { + +Status checked_add(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (rhs > std::numeric_limits::max() - lhs) { + return Status::Error(message); + } + *out = lhs + rhs; + return Status::OK(); +} + +Status posting_corruption(const char* message, uint32_t source_ordinal) { + return Status::Error( + "posting_cursor: {} (src_ord={})", message, source_ordinal); +} + +} // namespace + +Status validate_posting_region(const format::RegionRef& region, uint64_t file_size) { + if (region.offset > file_size || region.length > file_size - region.offset) { + return Status::Error( + "posting_cursor: posting region outside source file"); + } + return Status::OK(); +} + +bool posting_entry_has_positions(const format::DictEntry& entry) { + return entry.kind == format::DictEntryKind::kInline ? !entry.prx_bytes.empty() + : entry.prx_len != 0; +} + +SniiPostingReadContext::TermLease::~TermLease() { + if (context_ != nullptr) { + context_->release_term(); + } +} + +size_t SniiPostingReadContext::DecoderWorkspace::capacity_bytes() const { + return docs_scratch.capacity() + prx_scratch.capacity() + decompressed.capacity() + + sizeof(uint32_t) * (docids.capacity() + positions_flat.capacity() + + position_offsets.capacity() + frequencies.capacity()) + + sizeof(DestinationPostingRun) * destination_runs.capacity(); +} + +void SniiPostingReadContext::DecoderWorkspace::init_memory_reporter( + writer::MemoryReporter* memory_reporter) { + if (memory_reporter == nullptr) return; + docs_scratch_reservation = memory_reporter->make_reservation(); + prx_scratch_reservation = memory_reporter->make_reservation(); + docids_reservation = memory_reporter->make_reservation(); + positions_reservation = memory_reporter->make_reservation(); + position_offsets_reservation = memory_reporter->make_reservation(); + destination_runs_reservation = memory_reporter->make_reservation(); + frequencies_reservation = memory_reporter->make_reservation(); + decompressed_reservation = memory_reporter->make_reservation(); + reservations_enabled = true; +} + +Status SniiPostingReadContext::DecoderWorkspace::reserve_remapped(size_t document_count, + size_t run_count, + bool retain_frequencies) { + const bool grow_runs = destination_runs.capacity() < run_count; + const bool grow_frequencies = retain_frequencies && frequencies.capacity() < document_count; + if (!reservations_enabled) { + destination_runs.reserve(run_count); + if (retain_frequencies) { + frequencies.reserve(document_count); + } + return Status::OK(); + } + if (!grow_runs && !grow_frequencies) { + DCHECK_EQ(destination_runs_reservation.bytes(), + destination_runs.capacity() * sizeof(DestinationPostingRun)); + DCHECK_EQ(frequencies_reservation.bytes(), frequencies.capacity() * sizeof(uint32_t)); + return Status::OK(); + } + if (run_count > std::numeric_limits::max() / sizeof(DestinationPostingRun) || + document_count > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "posting cursor: remapped workspace size overflows size_t"); + } + writer::MemoryReporter::Reservation replacement_runs; + writer::MemoryReporter::Reservation replacement_frequencies; + if (grow_runs) { + RETURN_IF_ERROR(destination_runs_reservation.prepare_replacement( + run_count * sizeof(DestinationPostingRun), &replacement_runs)); + } + if (grow_frequencies) { + RETURN_IF_ERROR(frequencies_reservation.prepare_replacement( + document_count * sizeof(uint32_t), &replacement_frequencies)); + } + { + std::vector new_runs; + std::vector new_frequencies; + if (grow_runs) { + new_runs.reserve(run_count); + DCHECK_EQ(new_runs.capacity(), run_count); + destination_runs.swap(new_runs); + } + if (grow_frequencies) { + new_frequencies.reserve(document_count); + DCHECK_EQ(new_frequencies.capacity(), document_count); + frequencies.swap(new_frequencies); + } + } + if (grow_runs) { + destination_runs_reservation = std::move(replacement_runs); + } + if (grow_frequencies) { + frequencies_reservation = std::move(replacement_frequencies); + } + return Status::OK(); +} + +Status SniiPostingReadContext::DecoderWorkspace::reserve_docids(size_t count) { + if (!reservations_enabled || docids.capacity() >= count) { + if (reservations_enabled) { + DCHECK_EQ(docids_reservation.bytes(), docids.capacity() * sizeof(uint32_t)); + } + return Status::OK(); + } + const size_t target_bytes = count * sizeof(uint32_t); + writer::MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(docids_reservation.prepare_replacement(target_bytes, &replacement)); + docids.reserve(count); + DCHECK_EQ(docids.capacity(), count); + docids_reservation = std::move(replacement); + return Status::OK(); +} + +Status SniiPostingReadContext::DecoderWorkspace::reserve_csr(std::vector* pos_flat, + size_t position_count, + std::vector* pos_off, + size_t offset_count) { + DCHECK(reservations_enabled); + DCHECK_EQ(pos_flat, &positions_flat); + DCHECK_EQ(pos_off, &position_offsets); + if (position_count > std::numeric_limits::max() / sizeof(uint32_t) || + offset_count > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "posting cursor: position workspace size overflows size_t"); + } + + const bool grow_positions = positions_flat.capacity() < position_count; + const bool grow_offsets = position_offsets.capacity() < offset_count; + if (!grow_positions && !grow_offsets) { + DCHECK_EQ(positions_reservation.bytes(), positions_flat.capacity() * sizeof(uint32_t)); + DCHECK_EQ(position_offsets_reservation.bytes(), + position_offsets.capacity() * sizeof(uint32_t)); + return Status::OK(); + } + + writer::MemoryReporter::Reservation replacement_positions; + writer::MemoryReporter::Reservation replacement_offsets; + if (grow_positions) { + RETURN_IF_ERROR(positions_reservation.prepare_replacement(position_count * sizeof(uint32_t), + &replacement_positions)); + } + if (grow_offsets) { + RETURN_IF_ERROR(position_offsets_reservation.prepare_replacement( + offset_count * sizeof(uint32_t), &replacement_offsets)); + } + + { + std::vector new_positions; + std::vector new_offsets; + if (grow_positions) { + new_positions.reserve(position_count); + DCHECK_EQ(new_positions.capacity(), position_count); + positions_flat.swap(new_positions); + } + if (grow_offsets) { + new_offsets.reserve(offset_count); + DCHECK_EQ(new_offsets.capacity(), offset_count); + position_offsets.swap(new_offsets); + } + } + if (grow_positions) { + positions_reservation = std::move(replacement_positions); + } + if (grow_offsets) { + position_offsets_reservation = std::move(replacement_offsets); + } + return Status::OK(); +} + +Status SniiPostingReadContext::DecoderWorkspace::reserve_decompression( + size_t bytes, std::vector** buffer) { + DCHECK(reservations_enabled); + DCHECK(buffer != nullptr); + if (decompressed.capacity() < bytes) { + writer::MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(decompressed_reservation.prepare_replacement(bytes, &replacement)); + { + std::vector new_decompressed; + new_decompressed.reserve(bytes); + DCHECK_EQ(new_decompressed.capacity(), bytes); + decompressed.swap(new_decompressed); + } + decompressed_reservation = std::move(replacement); + } else { + DCHECK_EQ(decompressed_reservation.bytes(), decompressed.capacity()); + } + *buffer = &decompressed; + return Status::OK(); +} + +void SniiPostingReadContext::DecoderWorkspace::release_large_buffers( + size_t retained_capacity_limit_bytes) { + prelude = format::FrqPreludeReader(); + if (capacity_bytes() <= retained_capacity_limit_bytes) { + return; + } + std::vector().swap(docs_scratch); + std::vector().swap(prx_scratch); + std::vector().swap(decompressed); + std::vector().swap(docids); + std::vector().swap(positions_flat); + std::vector().swap(position_offsets); + std::vector().swap(destination_runs); + std::vector().swap(frequencies); + docs_scratch_reservation.reset(); + prx_scratch_reservation.reset(); + docids_reservation.reset(); + positions_reservation.reset(); + position_offsets_reservation.reset(); + destination_runs_reservation.reset(); + frequencies_reservation.reset(); + decompressed_reservation.reset(); + DCHECK_LE(capacity_bytes(), retained_capacity_limit_bytes); +} + +Status SniiPostingReadContext::poison(Status status) { + DCHECK(!status.ok()); + if (failed_.ok()) { + failed_ = std::move(status); + } + return failed_; +} + +Status SniiPostingReadContext::init() { + if (initialized_) { + return Status::Error( + "posting_read_context: init called twice"); + } + if (index_ == nullptr || index_->reader() == nullptr) { + return Status::Error( + "posting_read_context: null source index"); + } + if (total_read_ahead_budget_bytes_ < 2 || + total_read_ahead_budget_bytes_ > kMaxReadAheadBudgetBytes) { + return Status::Error( + "posting_read_context: total read-ahead budget outside [2, {}]", + kMaxReadAheadBudgetBytes); + } + + posting_region_ = index_->section_refs().posting_region; + RETURN_IF_ERROR(validate_posting_region(posting_region_, index_->reader()->size())); + decoder_workspace_.init_memory_reporter(memory_reporter_); + posting_cache_ = std::make_unique( + index_->reader(), posting_region_.offset, posting_region_.length, + total_read_ahead_budget_bytes_, memory_reporter_); + RETURN_IF_ERROR(posting_cache_->init()); + initialized_ = true; + return Status::OK(); +} + +Status SniiPostingReadContext::validate_next_range(const format::RegionRef& range, + bool has_previous, uint64_t previous_end, + const char* stream, uint64_t* end) const { + DCHECK(end != nullptr); + if (range.length == 0 || range.offset < posting_region_.offset) { + return Status::Error( + "posting_read_context: invalid {} term range", stream); + } + const uint64_t relative_offset = range.offset - posting_region_.offset; + if (relative_offset > posting_region_.length || + range.length > posting_region_.length - relative_offset) { + return Status::Error( + "posting_read_context: {} term range outside posting region", stream); + } + if (has_previous && range.offset < previous_end) { + return Status::Error( + "posting_read_context: {} term ranges are not monotone", stream); + } + *end = range.offset + range.length; + return Status::OK(); +} + +Status SniiPostingReadContext::acquire_term(bool has_docs_range, bool has_prx_range, + const format::RegionRef& docs_range, + const format::RegionRef& prx_range, + std::unique_ptr* lease) { + DCHECK(lease != nullptr); + DCHECK(*lease == nullptr); + if (!initialized_) { + return Status::Error( + "posting_read_context: acquire before init"); + } + if (!failed_.ok()) { + return failed_; + } + if (term_active_) { + return Status::Error( + "posting_read_context: concurrent term cursor"); + } + + uint64_t docs_end = 0; + uint64_t prx_end = 0; + if (has_docs_range) { + Status status = + validate_next_range(docs_range, has_docs_range_, last_docs_end_, "docs", &docs_end); + if (!status.ok()) { + return poison(status); + } + last_docs_end_ = docs_end; + has_docs_range_ = true; + } + if (has_prx_range) { + Status status = + validate_next_range(prx_range, has_prx_range_, last_prx_end_, "prx", &prx_end); + if (!status.ok()) { + return poison(status); + } + last_prx_end_ = prx_end; + has_prx_range_ = true; + } + + term_active_ = true; + lease->reset(new TermLease(this)); + return Status::OK(); +} + +void SniiPostingReadContext::release_term() { + DCHECK(term_active_); + term_active_ = false; + decoder_workspace_.release_large_buffers(retained_decoder_workspace_limit_bytes()); +} + +Status SniiPostingReadContext::poison_active_term(Status status, TermLease* lease) { + DCHECK(lease != nullptr); + DCHECK_EQ(lease->context_, this); + DCHECK(term_active_); + Status first = poison(std::move(status)); + lease->context_ = nullptr; + release_term(); + return first; +} + +uint64_t SniiPostingReadContext::docs_read_calls() const { + DCHECK(initialized_); + return posting_cache_->read_calls(PostingStream::kDocs); +} + +uint64_t SniiPostingReadContext::prx_read_calls() const { + DCHECK(initialized_); + return posting_cache_->read_calls(PostingStream::kPrx); +} + +uint64_t SniiPostingReadContext::docs_buffer_hits() const { + DCHECK(initialized_); + return posting_cache_->buffer_hits(PostingStream::kDocs); +} + +uint64_t SniiPostingReadContext::prx_buffer_hits() const { + DCHECK(initialized_); + return posting_cache_->buffer_hits(PostingStream::kPrx); +} + +uint64_t SniiPostingReadContext::physical_read_ranges() const { + DCHECK(initialized_); + return posting_cache_->physical_read_ranges(); +} + +uint64_t SniiPostingReadContext::physical_read_bytes() const { + DCHECK(initialized_); + return posting_cache_->physical_read_bytes(); +} + +size_t SniiPostingReadContext::resident_read_ahead_capacity_bytes() const { + DCHECK(initialized_); + return posting_cache_->resident_capacity_bytes(); +} + +size_t SniiPostingReadContext::decoder_workspace_capacity_bytes() const { + DCHECK(initialized_); + return decoder_workspace_.capacity_bytes(); +} + +Status SniiPostingCursor::poison(Status status) { + DCHECK(!status.ok()); + if (failed_.ok()) { + if (term_lease_ != nullptr) { + failed_ = read_context_->poison_active_term(std::move(status), term_lease_.get()); + term_lease_.reset(); + } else { + failed_ = std::move(status); + } + } + return failed_; +} + +Status SniiPostingCursor::validate_entry_geometry() { + if (entry_.df == 0) { + return posting_corruption("zero-df dictionary entry", source_ordinal_); + } + if (entry_.kind == format::DictEntryKind::kInline) { + if (entry_.enc != format::DictEntryEnc::kSlim) { + return posting_corruption("inline entry is not slim", source_ordinal_); + } + if (entry_.inline_dd_disk_len != entry_.dd_meta.disk_len || + entry_.inline_dd_disk_len > entry_.frq_bytes.size()) { + return posting_corruption("inline dd geometry mismatch", source_ordinal_); + } + const uint64_t freq_len = entry_.frq_bytes.size() - entry_.inline_dd_disk_len; + if (entry_.freq_meta.disk_len != freq_len) { + return posting_corruption("inline freq geometry mismatch", source_ordinal_); + } + shape_ = Shape::kFlat; + return Status::OK(); + } + + if (entry_.kind != format::DictEntryKind::kPodRef) { + return posting_corruption("unknown dictionary entry kind", source_ordinal_); + } + if (entry_.enc == format::DictEntryEnc::kWindowed) { + if (entry_.prelude_len == 0 || entry_.prelude_len > entry_.frq_docs_len || + entry_.frq_docs_len > entry_.frq_len) { + return posting_corruption("invalid windowed frq geometry", source_ordinal_); + } + shape_ = Shape::kWindowed; + return Status::OK(); + } + if (entry_.enc != format::DictEntryEnc::kSlim) { + return posting_corruption("unknown dictionary entry encoding", source_ordinal_); + } + if (entry_.prelude_len != 0 || entry_.frq_docs_len != entry_.dd_meta.disk_len || + entry_.frq_docs_len > entry_.frq_len) { + return posting_corruption("invalid slim frq geometry", source_ordinal_); + } + if (entry_.freq_meta.disk_len != entry_.frq_len - entry_.frq_docs_len) { + return posting_corruption("slim freq geometry mismatch", source_ordinal_); + } + shape_ = Shape::kFlat; + return Status::OK(); +} + +Status SniiPostingCursor::prepare_flat_ranges() { + if (entry_.kind == format::DictEntryKind::kInline) { + flat_dd_len_ = entry_.inline_dd_disk_len; + flat_prx_len_ = entry_.prx_bytes.size(); + return Status::OK(); + } + + uint64_t frq_len = 0; + RETURN_IF_ERROR(index_->resolve_frq_window(entry_, frq_base_, &flat_dd_abs_, &frq_len)); + if (frq_len != entry_.frq_len) { + return posting_corruption("resolved slim frq length mismatch", source_ordinal_); + } + flat_dd_len_ = entry_.frq_docs_len; + if (term_has_positions_) { + RETURN_IF_ERROR( + index_->resolve_prx_window(entry_, prx_base_, &flat_prx_abs_, &flat_prx_len_)); + if (flat_prx_len_ != entry_.prx_len) { + return posting_corruption("resolved slim prx length mismatch", source_ordinal_); + } + } + return Status::OK(); +} + +Status SniiPostingCursor::prepare_windowed_ranges() { + RETURN_IF_ERROR(index_->resolve_frq_window(entry_, frq_base_, &flat_dd_abs_, &flat_dd_len_)); + if (flat_dd_len_ != entry_.frq_len - entry_.prelude_len || flat_dd_abs_ < entry_.prelude_len) { + return posting_corruption("resolved windowed frq geometry mismatch", source_ordinal_); + } + if (term_has_positions_) { + RETURN_IF_ERROR( + index_->resolve_prx_window(entry_, prx_base_, &flat_prx_abs_, &flat_prx_len_)); + if (flat_prx_len_ != entry_.prx_len) { + return posting_corruption("resolved windowed prx length mismatch", source_ordinal_); + } + } + return Status::OK(); +} + +Status SniiPostingCursor::prepare_windowed() { + DCHECK(workspace_ != nullptr); + Slice prelude_bytes; + RETURN_IF_ERROR(read_context_->posting_cache_->resolve( + PostingStream::kDocs, flat_dd_abs_ - entry_.prelude_len, entry_.prelude_len, + &workspace_->docs_scratch, &prelude_bytes, + read_context_->memory_reporter_ == nullptr ? nullptr + : &workspace_->docs_scratch_reservation)); + RETURN_IF_ERROR(format::FrqPreludeReader::open(prelude_bytes, &workspace_->prelude)); + if (workspace_->prelude.has_prx() != term_has_positions_) { + return posting_corruption("windowed prelude position shape differs from entry", + source_ordinal_); + } + + uint64_t docs_prefix_len = 0; + RETURN_IF_ERROR(checked_add(entry_.prelude_len, workspace_->prelude.dd_block_len(), + "posting_cursor: windowed docs prefix overflow", &docs_prefix_len)); + if (docs_prefix_len != entry_.frq_docs_len) { + return posting_corruption("windowed docs prefix mismatch", source_ordinal_); + } + uint64_t encoded_frq_len = 0; + RETURN_IF_ERROR(checked_add(docs_prefix_len, workspace_->prelude.freq_block_len(), + "posting_cursor: windowed frq length overflow", &encoded_frq_len)); + if (encoded_frq_len != entry_.frq_len) { + return posting_corruption("windowed frq blocks do not tile entry", source_ordinal_); + } + + uint64_t dd_bytes = 0; + uint64_t freq_bytes = 0; + uint64_t prx_bytes = 0; + uint64_t docs = 0; + uint32_t previous_last_docid = 0; + bool has_previous_window = false; + for (uint32_t window = 0; window < workspace_->prelude.n_windows(); ++window) { + format::WindowMeta meta; + RETURN_IF_ERROR(workspace_->prelude.window(window, &meta)); + if (meta.doc_count == 0 || meta.dd_off != dd_bytes || meta.prx_off != prx_bytes || + (workspace_->prelude.has_freq() && meta.freq_off != freq_bytes)) { + return posting_corruption("non-contiguous window metadata", source_ordinal_); + } + if ((!has_previous_window && meta.win_base != 0) || + (has_previous_window && + (meta.win_base != previous_last_docid || meta.last_docid <= previous_last_docid))) { + return posting_corruption("invalid window docid chain", source_ordinal_); + } + if (meta.last_docid >= index_->stats().doc_count) { + return posting_corruption("window last docid outside index", source_ordinal_); + } + RETURN_IF_ERROR(checked_add(dd_bytes, meta.dd_disk_len, + "posting_cursor: window dd bytes overflow", &dd_bytes)); + RETURN_IF_ERROR(checked_add(freq_bytes, meta.freq_disk_len, + "posting_cursor: window freq bytes overflow", &freq_bytes)); + RETURN_IF_ERROR(checked_add(prx_bytes, meta.prx_len, + "posting_cursor: window prx bytes overflow", &prx_bytes)); + RETURN_IF_ERROR(checked_add(docs, meta.doc_count, + "posting_cursor: window doc count overflow", &docs)); + previous_last_docid = meta.last_docid; + has_previous_window = true; + } + if (dd_bytes != workspace_->prelude.dd_block_len() || + freq_bytes != workspace_->prelude.freq_block_len() || prx_bytes != entry_.prx_len || + docs != entry_.df) { + return posting_corruption("window directory totals mismatch", source_ordinal_); + } + return Status::OK(); +} + +Status SniiPostingCursor::init() { + if (initialized_) { + return Status::Error( + "posting_cursor: init called twice"); + } + if (read_context_ == nullptr || index_ == nullptr || rowid_conversion_ == nullptr) { + return Status::Error( + "posting_cursor: null read context or rowid conversion"); + } + if (!read_context_->initialized()) { + return Status::Error( + "posting_cursor: source read context not initialized"); + } + if (!read_context_->failed_status().ok()) { + return read_context_->failed_status(); + } + if (index_->tier() == format::IndexTier::kT1) { + return Status::Error( + "posting_cursor: positions index is required"); + } + if (!index_->has_positions()) { + return posting_corruption("positions tier lacks positions capability", source_ordinal_); + } + if (source_ordinal_ >= rowid_conversion_->source_segment_count()) { + return Status::Error( + "posting_cursor: source ordinal outside rowid conversion"); + } + if (index_->stats().doc_count != + rowid_conversion_->source_segment_doc_counts()[source_ordinal_]) { + return posting_corruption("rowid conversion size differs from source doc count", + source_ordinal_); + } + source_mapping_ = rowid_conversion_->source_mapping(source_ordinal_); + source_has_deletions_ = rowid_conversion_->source_has_deletions(source_ordinal_); + + RETURN_IF_ERROR(validate_entry_geometry()); + const bool has_pod_ranges = entry_.kind == format::DictEntryKind::kPodRef; + format::RegionRef docs_range; + format::RegionRef prx_range; + if (shape_ == Shape::kWindowed) { + RETURN_IF_ERROR(prepare_windowed_ranges()); + docs_range.offset = flat_dd_abs_ - entry_.prelude_len; + docs_range.length = entry_.frq_docs_len; + prx_range.offset = flat_prx_abs_; + prx_range.length = flat_prx_len_; + } else { + RETURN_IF_ERROR(prepare_flat_ranges()); + if (has_pod_ranges) { + docs_range.offset = flat_dd_abs_; + docs_range.length = flat_dd_len_; + prx_range.offset = flat_prx_abs_; + prx_range.length = flat_prx_len_; + } + } + RETURN_IF_ERROR(read_context_->acquire_term(has_pod_ranges, + has_pod_ranges && term_has_positions_, docs_range, + prx_range, &term_lease_)); + workspace_ = &read_context_->decoder_workspace_; + workspace_->destination_runs.clear(); + workspace_->frequencies.clear(); + next_destination_run_ = 0; + if (shape_ == Shape::kWindowed) { + const Status status = prepare_windowed(); + if (!status.ok()) { + return poison(status); + } + } + initialized_ = true; + return Status::OK(); +} + +Status SniiPostingCursor::decode_prx(Slice bytes, format::PrxDecodedShape* shape) { + DCHECK(workspace_ != nullptr); + DCHECK(shape != nullptr); + ByteSource source(bytes); + format::PrxDecodeContext decode_context { + .shape = shape, + .allocation_gate = workspace_->reservations_enabled ? workspace_ : nullptr}; + RETURN_IF_ERROR(format::read_prx_window_csr(&source, &workspace_->positions_flat, + &workspace_->position_offsets, &decode_context)); + if (!source.eof()) { + return posting_corruption("trailing bytes after prx frame", source_ordinal_); + } + return Status::OK(); +} + +Status SniiPostingCursor::decode_dd(Slice bytes, const format::FrqRegionMeta& meta, + uint64_t win_base, uint32_t expected_doc_count) { + DCHECK(workspace_ != nullptr); + if (!workspace_->reservations_enabled) { + return format::decode_dd_region(bytes, meta, win_base, expected_doc_count, + &workspace_->docids); + } + return format::decode_dd_region(bytes, meta, win_base, expected_doc_count, workspace_, + &workspace_->docids); +} + +Status SniiPostingCursor::load_flat_chunk() { + DCHECK(workspace_ != nullptr); + Slice dd_bytes; + Slice prx_bytes; + if (entry_.kind == format::DictEntryKind::kInline) { + dd_bytes = Slice(entry_.frq_bytes.data(), static_cast(flat_dd_len_)); + prx_bytes = Slice(entry_.prx_bytes); + } else { + RETURN_IF_ERROR(read_context_->posting_cache_->resolve( + PostingStream::kDocs, flat_dd_abs_, flat_dd_len_, &workspace_->docs_scratch, + &dd_bytes, + read_context_->memory_reporter_ == nullptr + ? nullptr + : &workspace_->docs_scratch_reservation)); + if (term_has_positions_) { + RETURN_IF_ERROR(read_context_->posting_cache_->resolve( + PostingStream::kPrx, flat_prx_abs_, flat_prx_len_, &workspace_->prx_scratch, + &prx_bytes, + read_context_->memory_reporter_ == nullptr + ? nullptr + : &workspace_->prx_scratch_reservation)); + } + } + + RETURN_IF_ERROR(workspace_->reserve_docids(entry_.df)); + RETURN_IF_ERROR(decode_dd(dd_bytes, entry_.dd_meta, /*win_base=*/0, entry_.df)); + format::PrxDecodedShape prx_shape; + if (term_has_positions_) { + RETURN_IF_ERROR(decode_prx(prx_bytes, &prx_shape)); + } + if (workspace_->docids.size() != entry_.df || workspace_->docids.empty() || + workspace_->docids.back() >= index_->stats().doc_count) { + return posting_corruption("decoded docids differ from flat entry shape", source_ordinal_); + } + if (term_has_positions_ && + (prx_shape.total_docs != entry_.df || + prx_shape.total_positions != workspace_->positions_flat.size() || + prx_shape.has_zero_frequency || + workspace_->position_offsets.size() != static_cast(entry_.df) + 1 || + workspace_->position_offsets.empty() || workspace_->position_offsets.front() != 0 || + workspace_->position_offsets.back() != workspace_->positions_flat.size())) { + return posting_corruption("dd/prx document shape mismatch", source_ordinal_); + } + decoded_docs_ = entry_.df; + if (term_has_positions_) { + decoded_total_freq_ = prx_shape.total_positions; + decoded_max_freq_ = prx_shape.max_frequency; + } + flat_loaded_ = true; + return Status::OK(); +} + +Status SniiPostingCursor::load_windowed_chunk() { + DCHECK(workspace_ != nullptr); + format::WindowMeta meta; + RETURN_IF_ERROR(workspace_->prelude.window(next_window_, &meta)); + reader::WindowAbsRange range; + RETURN_IF_ERROR(reader::windowed_window_range( + *index_, entry_, frq_base_, prx_base_, workspace_->prelude, next_window_, + /*want_positions=*/term_has_positions_, /*want_freq=*/false, &range)); + + Slice dd_bytes; + Slice prx_bytes; + RETURN_IF_ERROR(read_context_->posting_cache_->resolve( + PostingStream::kDocs, range.dd_off, range.dd_len, &workspace_->docs_scratch, &dd_bytes, + read_context_->memory_reporter_ == nullptr ? nullptr + : &workspace_->docs_scratch_reservation)); + if (term_has_positions_) { + RETURN_IF_ERROR(read_context_->posting_cache_->resolve( + PostingStream::kPrx, range.prx_off, range.prx_len, &workspace_->prx_scratch, + &prx_bytes, + read_context_->memory_reporter_ == nullptr ? nullptr + : &workspace_->prx_scratch_reservation)); + } + RETURN_IF_ERROR(workspace_->reserve_docids(meta.doc_count)); + RETURN_IF_ERROR(decode_dd(dd_bytes, + format::FrqRegionMeta {.zstd = meta.dd_zstd, + .uncomp_len = meta.dd_uncomp_len, + .disk_len = meta.dd_disk_len, + .crc = meta.crc_dd, + .verify_crc = meta.verify_crc}, + meta.win_base, meta.doc_count)); + format::PrxDecodedShape prx_shape; + if (term_has_positions_) { + RETURN_IF_ERROR(decode_prx(prx_bytes, &prx_shape)); + } + if (workspace_->docids.size() != meta.doc_count || workspace_->docids.empty() || + workspace_->docids.back() != meta.last_docid || + (next_window_ != 0 && workspace_->docids.front() <= meta.win_base)) { + return posting_corruption("window docid shape or last docid mismatch", source_ordinal_); + } + if (term_has_positions_ && + (prx_shape.total_docs != meta.doc_count || + prx_shape.total_positions != workspace_->positions_flat.size() || + prx_shape.has_zero_frequency || + workspace_->position_offsets.size() != static_cast(meta.doc_count) + 1 || + workspace_->position_offsets.empty() || workspace_->position_offsets.front() != 0 || + workspace_->position_offsets.back() != workspace_->positions_flat.size())) { + return posting_corruption("dd/prx document shape mismatch", source_ordinal_); + } + if (term_has_positions_ && entry_.term_stats_present && + prx_shape.max_frequency != meta.max_freq) { + return posting_corruption("window max frequency mismatch", source_ordinal_); + } + if (decoded_docs_ > entry_.df || meta.doc_count > entry_.df - decoded_docs_) { + return posting_corruption("decoded document count exceeds df", source_ordinal_); + } + decoded_docs_ += meta.doc_count; + if (term_has_positions_) { + if (prx_shape.total_positions > + std::numeric_limits::max() - decoded_total_freq_) { + return posting_corruption("total term frequency overflow", source_ordinal_); + } + decoded_total_freq_ += prx_shape.total_positions; + decoded_max_freq_ = std::max(decoded_max_freq_, prx_shape.max_frequency); + } + ++next_window_; + return Status::OK(); +} + +Status SniiPostingCursor::load_next_chunk(bool* loaded) { + DCHECK(loaded != nullptr); + DCHECK(workspace_ != nullptr); + *loaded = false; + // docids/positions_flat/position_offsets are intentionally NOT cleared here: + // both decode paths size them exactly (resize_uninitialized + post-decode + // shape checks), and clearing first would turn every warm re-decode into a + // cold 0->n grow whose std::vector tail value-initialization memsets the + // whole chunk. On the not-loaded path the stale contents are never read. + + if (shape_ == Shape::kFlat) { + if (flat_loaded_) { + return Status::OK(); + } + RETURN_IF_ERROR(load_flat_chunk()); + *loaded = true; + return Status::OK(); + } + if (next_window_ >= workspace_->prelude.n_windows()) { + return Status::OK(); + } + RETURN_IF_ERROR(load_windowed_chunk()); + *loaded = true; + return Status::OK(); +} + +Status SniiPostingCursor::finish_source() { + if (decoded_docs_ != entry_.df) { + return posting_corruption("decoded document count differs from df", source_ordinal_); + } + if (entry_.term_stats_present && + (decoded_total_freq_ != entry_.ttf_delta || decoded_max_freq_ != entry_.max_freq)) { + return posting_corruption("decoded term statistics mismatch", source_ordinal_); + } + exhausted_ = true; + term_lease_.reset(); + return Status::OK(); +} + +Status SniiPostingCursor::map_decoded_chunk() { + DCHECK(workspace_ != nullptr); + std::vector& runs = workspace_->destination_runs; + std::vector& docids = workspace_->docids; + std::vector& frequencies = workspace_->frequencies; + runs.clear(); + frequencies.clear(); + const size_t document_count = docids.size(); + const size_t max_run_count = + std::min(document_count, rowid_conversion_->destination_segment_doc_counts().size()); + RETURN_IF_ERROR( + workspace_->reserve_remapped(document_count, max_run_count, term_has_positions_)); + + auto append_live_document = [&](uint32_t segment, uint32_t docid, size_t live_docs) { + if (runs.empty() || runs.back().destination_segment != segment) { + if (!runs.empty()) { + runs.back().document_end = static_cast(live_docs); + } + runs.push_back({.destination_segment = segment}); + } + docids[live_docs] = docid; + }; + + if (!source_has_deletions_) { + if (term_has_positions_) { + // Frequencies are adjacent position-offset deltas; fill them in + // their own pass so the remap loop below stays a pure gather. + const uint32_t* offsets = workspace_->position_offsets.data(); + for (size_t ordinal = 0; ordinal < document_count; ++ordinal) { + frequencies.push_back(offsets[ordinal + 1] - offsets[ordinal]); + } + } + // The gather walks source_mapping_ at monotonically increasing but + // sparse indexes the hardware prefetcher cannot follow; the future + // lookup indexes are already decoded, so prefetch them explicitly. + // Without deletions every mapping entry is live, so the destination + // segment can never equal the uint32 sentinel. + constexpr size_t kMapPrefetchDistance = 16; + uint32_t current_segment = std::numeric_limits::max(); + for (size_t ordinal = 0; ordinal < document_count; ++ordinal) { + if (ordinal + kMapPrefetchDistance < document_count) { + __builtin_prefetch(&source_mapping_[docids[ordinal + kMapPrefetchDistance]]); + } + const auto [segment, docid] = source_mapping_[docids[ordinal]]; + if (segment != current_segment) { + if (!runs.empty()) { + runs.back().document_end = static_cast(ordinal); + } + runs.push_back({.destination_segment = segment}); + current_segment = segment; + } + docids[ordinal] = docid; + } + } else if (!term_has_positions_) { + constexpr uint32_t kDeleted = std::numeric_limits::max(); + size_t live_docs = 0; + for (uint32_t source_docid : docids) { + const auto [segment, docid] = source_mapping_[source_docid]; + if (segment != kDeleted) { + append_live_document(segment, docid, live_docs++); + } + } + docids.resize(live_docs); + } else { + constexpr uint32_t kDeleted = std::numeric_limits::max(); + size_t write_position = 0; + size_t live_docs = 0; + workspace_->position_offsets[0] = 0; + for (size_t ordinal = 0; ordinal < docids.size(); ++ordinal) { + const auto [segment, docid] = source_mapping_[docids[ordinal]]; + if (segment == kDeleted) { + continue; + } + const size_t begin = workspace_->position_offsets[ordinal]; + const size_t end = workspace_->position_offsets[ordinal + 1]; + const uint32_t frequency = static_cast(end - begin); + if (write_position != begin) { + std::copy(workspace_->positions_flat.begin() + begin, + workspace_->positions_flat.begin() + end, + workspace_->positions_flat.begin() + write_position); + } + write_position += frequency; + append_live_document(segment, docid, live_docs); + frequencies.push_back(frequency); + workspace_->position_offsets[++live_docs] = static_cast(write_position); + } + docids.resize(live_docs); + workspace_->positions_flat.resize(write_position); + workspace_->position_offsets.resize(live_docs + 1); + } + + if (!runs.empty()) { + runs.back().document_end = static_cast(docids.size()); + } + next_destination_run_ = 0; + return Status::OK(); +} + +void SniiPostingCursor::emit_next_mapped_run(RemappedPostingChunk* chunk) { + DCHECK(chunk != nullptr); + DCHECK(workspace_ != nullptr); + DCHECK_LT(next_destination_run_, workspace_->destination_runs.size()); + + const size_t run_ordinal = next_destination_run_++; + const DestinationPostingRun& run = workspace_->destination_runs[run_ordinal]; + const size_t document_begin = + run_ordinal == 0 ? 0 : workspace_->destination_runs[run_ordinal - 1].document_end; + const size_t document_end = run.document_end; + DCHECK_LT(document_begin, document_end); + DCHECK_LE(document_end, workspace_->docids.size()); + const size_t document_count = document_end - document_begin; + + chunk->destination_segment = run.destination_segment; + chunk->destination_docids = + std::span(workspace_->docids).subspan(document_begin, document_count); + if (!term_has_positions_) { + return; + } + + chunk->freqs = std::span(workspace_->frequencies) + .subspan(document_begin, document_count); + chunk->position_offsets = std::span(workspace_->position_offsets) + .subspan(document_begin, document_count + 1); + const size_t position_begin = chunk->position_offsets.front(); + const size_t position_end = chunk->position_offsets.back(); + DCHECK_LE(position_begin, position_end); + DCHECK_LE(position_end, workspace_->positions_flat.size()); + chunk->positions_flat = std::span(workspace_->positions_flat) + .subspan(position_begin, position_end - position_begin); +} + +Status SniiPostingCursor::next_chunk(RemappedPostingChunk* chunk, bool* has_chunk) { + if (chunk == nullptr || has_chunk == nullptr) { + return Status::Error( + "posting_cursor: null chunk output"); + } + *chunk = {}; + *has_chunk = false; + if (!failed_.ok()) { + return failed_; + } + if (!initialized_) { + return Status::Error( + "posting_cursor: next_chunk before init"); + } + if (exhausted_) { + return Status::OK(); + } + + if (next_destination_run_ < workspace_->destination_runs.size()) { + emit_next_mapped_run(chunk); + *has_chunk = true; + return Status::OK(); + } + + while (true) { + bool loaded = false; + const Status status = load_next_chunk(&loaded); + if (!status.ok()) { + return poison(status); + } + if (!loaded) { + *chunk = {}; + const Status finish_status = finish_source(); + if (!finish_status.ok()) { + return poison(finish_status); + } + return Status::OK(); + } + const Status map_status = map_decoded_chunk(); + if (!map_status.ok()) { + return poison(map_status); + } + if (!workspace_->destination_runs.empty()) { + emit_next_mapped_run(chunk); + *has_chunk = true; + return Status::OK(); + } + *chunk = {}; + } +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/posting_cursor.h b/be/src/storage/index/snii/compaction/posting_cursor.h new file mode 100644 index 00000000000000..356f08d7a9326a --- /dev/null +++ b/be/src/storage/index/snii/compaction/posting_cursor.h @@ -0,0 +1,258 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/compaction/region_reader.h" +#include "storage/index/snii/compaction/rowid_conversion.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_decode_stats.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::compaction { + +struct DestinationPostingRun { + uint32_t destination_segment = 0; + uint32_t document_end = 0; +}; + +// One destination-homogeneous slice of a decoded physical chunk. Positioned +// runs retain the decoder's offset base; positions_flat covers exactly +// [position_offsets.front(), position_offsets.back()). +struct RemappedPostingChunk { + uint32_t destination_segment = 0; + std::span destination_docids; + std::span freqs; + std::span position_offsets; + std::span positions_flat; +}; + +// Checks [region.offset, region.offset + region.length) against the source +// file without ever evaluating the potentially-overflowing addition. +Status validate_posting_region(const format::RegionRef& region, uint64_t file_size); +bool posting_entry_has_positions(const format::DictEntry& entry); + +class SniiPostingCursor; + +// Persistent read/decode state for every term decoded from one source logical +// index. Docs and positions share one bounded physical chunk cache while their +// decoder vectors and prelude workspace survive term cursor destruction. +class SniiPostingReadContext { +public: + static constexpr size_t kMaxReadAheadBudgetBytes = + 2 * SequentialRegionReader::kDefaultChunkBytes; + static constexpr size_t kMaxRetainedDecoderWorkspaceBytes = 64ULL << 10; + + SniiPostingReadContext(const reader::LogicalIndexReader* index, + size_t total_read_ahead_budget_bytes, + writer::MemoryReporter* memory_reporter = nullptr) + : index_(index), + total_read_ahead_budget_bytes_(total_read_ahead_budget_bytes), + memory_reporter_(memory_reporter) {} + + SniiPostingReadContext(const SniiPostingReadContext&) = delete; + SniiPostingReadContext& operator=(const SniiPostingReadContext&) = delete; + SniiPostingReadContext(SniiPostingReadContext&&) = delete; + SniiPostingReadContext& operator=(SniiPostingReadContext&&) = delete; + + Status init(); + + const reader::LogicalIndexReader* index() const { return index_; } + bool initialized() const { return initialized_; } + size_t total_read_ahead_budget_bytes() const { return total_read_ahead_budget_bytes_; } + uint64_t docs_read_calls() const; + uint64_t prx_read_calls() const; + uint64_t docs_buffer_hits() const; + uint64_t prx_buffer_hits() const; + uint64_t physical_read_ranges() const; + uint64_t physical_read_bytes() const; + size_t resident_read_ahead_capacity_bytes() const; + size_t decoder_workspace_capacity_bytes() const; + size_t retained_decoder_workspace_limit_bytes() const { + return std::min(total_read_ahead_budget_bytes_, kMaxRetainedDecoderWorkspaceBytes); + } + +private: + friend class SniiPostingCursor; + + struct DecoderWorkspace final : public format::PrxCsrAllocationGate { + format::FrqPreludeReader prelude; + writer::MemoryReporter::Reservation docs_scratch_reservation; + writer::MemoryReporter::Reservation prx_scratch_reservation; + writer::MemoryReporter::Reservation docids_reservation; + writer::MemoryReporter::Reservation positions_reservation; + writer::MemoryReporter::Reservation position_offsets_reservation; + writer::MemoryReporter::Reservation destination_runs_reservation; + writer::MemoryReporter::Reservation frequencies_reservation; + writer::MemoryReporter::Reservation decompressed_reservation; + std::vector docs_scratch; + std::vector prx_scratch; + std::vector decompressed; + std::vector docids; + std::vector positions_flat; + std::vector position_offsets; + std::vector destination_runs; + std::vector frequencies; + bool reservations_enabled = false; + + size_t capacity_bytes() const; + void init_memory_reporter(writer::MemoryReporter* memory_reporter); + Status reserve_docids(size_t count); + Status reserve_remapped(size_t document_count, size_t run_count, bool retain_frequencies); + Status reserve_csr(std::vector* pos_flat, size_t position_count, + std::vector* pos_off, size_t offset_count) override; + Status reserve_decompression(size_t bytes, std::vector** buffer) override; + void release_large_buffers(size_t retained_capacity_limit_bytes); + }; + + class TermLease { + public: + ~TermLease(); + + TermLease(const TermLease&) = delete; + TermLease& operator=(const TermLease&) = delete; + + private: + friend class SniiPostingReadContext; + explicit TermLease(SniiPostingReadContext* context) : context_(context) {} + SniiPostingReadContext* context_ = nullptr; + }; + + Status acquire_term(bool has_docs_range, bool has_prx_range, + const format::RegionRef& docs_range, const format::RegionRef& prx_range, + std::unique_ptr* lease); + Status validate_next_range(const format::RegionRef& range, bool has_previous, + uint64_t previous_end, const char* stream, uint64_t* end) const; + Status poison(Status status); + Status poison_active_term(Status status, TermLease* lease); + void release_term(); + const Status& failed_status() const { return failed_; } + + const reader::LogicalIndexReader* index_ = nullptr; + size_t total_read_ahead_budget_bytes_ = 0; + writer::MemoryReporter* memory_reporter_ = nullptr; + format::RegionRef posting_region_; + std::unique_ptr posting_cache_; + DecoderWorkspace decoder_workspace_; + + uint64_t last_docs_end_ = 0; + uint64_t last_prx_end_ = 0; + bool has_docs_range_ = false; + bool has_prx_range_ = false; + bool term_active_ = false; + bool initialized_ = false; + Status failed_ = Status::OK(); +}; + +// Sequential positions-posting decoder for one term in one source SNII index. +// It accepts all v1 physical shapes (inline, slim POD-ref and windowed POD-ref), +// validates the decoded doc/frequency/position stream, applies a validated +// row-id conversion, and yields surviving rows in destination-order chunks. +// +// The borrowed read context and row-id capability must outlive the cursor. A +// cursor is single-use and single-threaded. Only one cursor may hold a context +// term lease at a time. +class SniiPostingCursor { +public: + SniiPostingCursor(SniiPostingReadContext* read_context, format::DictEntry entry, + uint64_t frq_base, uint64_t prx_base, uint32_t source_ordinal, + const ValidatedRowIdConversion* rowid_conversion) + : read_context_(read_context), + index_(read_context == nullptr ? nullptr : read_context->index()), + entry_(std::move(entry)), + frq_base_(frq_base), + prx_base_(prx_base), + source_ordinal_(source_ordinal), + rowid_conversion_(rowid_conversion), + term_has_positions_(posting_entry_has_positions(entry_)) {} + + // Validates index capability, posting locators and fixed entry geometry, + // then acquires the source context's exclusive term lease. Payload decoding + // is lazy: corrupt DD/PRX frames surface from next_chunk(). + Status init(); + + // Returns the next non-empty remapped chunk. Its spans, including an exact + // base-relative positions slice, remain valid until the next call. + // Decoder/corruption failures poison the cursor and are returned unchanged + // by subsequent calls. + Status next_chunk(RemappedPostingChunk* chunk, bool* has_chunk); + bool has_positions() const { return term_has_positions_; } + +private: + enum class Shape : uint8_t { kFlat, kWindowed }; + + Status validate_entry_geometry(); + Status prepare_flat_ranges(); + Status prepare_windowed_ranges(); + Status prepare_windowed(); + Status load_next_chunk(bool* loaded); + Status load_flat_chunk(); + Status load_windowed_chunk(); + Status decode_dd(Slice bytes, const format::FrqRegionMeta& meta, uint64_t win_base, + uint32_t expected_doc_count); + Status decode_prx(Slice bytes, format::PrxDecodedShape* shape); + Status map_decoded_chunk(); + void emit_next_mapped_run(RemappedPostingChunk* chunk); + Status finish_source(); + Status poison(Status status); + + SniiPostingReadContext* read_context_ = nullptr; + const reader::LogicalIndexReader* index_ = nullptr; + format::DictEntry entry_; + uint64_t frq_base_ = 0; + uint64_t prx_base_ = 0; + uint32_t source_ordinal_ = 0; + const ValidatedRowIdConversion* rowid_conversion_ = nullptr; + std::span> source_mapping_; + bool source_has_deletions_ = false; + bool term_has_positions_ = true; + + Shape shape_ = Shape::kFlat; + std::unique_ptr term_lease_; + SniiPostingReadContext::DecoderWorkspace* workspace_ = nullptr; + + uint64_t flat_dd_abs_ = 0; + uint64_t flat_dd_len_ = 0; + uint64_t flat_prx_abs_ = 0; + uint64_t flat_prx_len_ = 0; + uint32_t next_window_ = 0; + bool flat_loaded_ = false; + + uint64_t decoded_docs_ = 0; + uint64_t decoded_total_freq_ = 0; + uint32_t decoded_max_freq_ = 0; + size_t next_destination_run_ = 0; + + bool initialized_ = false; + bool exhausted_ = false; + Status failed_ = Status::OK(); +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/posting_run_merger.cpp b/be/src/storage/index/snii/compaction/posting_run_merger.cpp new file mode 100644 index 00000000000000..bdd5b31c8742cd --- /dev/null +++ b/be/src/storage/index/snii/compaction/posting_run_merger.cpp @@ -0,0 +1,463 @@ +// 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 "storage/index/snii/compaction/posting_run_merger.h" + +#include +#include +#include +#include +#include + +#include "common/check.h" + +namespace doris::snii::compaction { + +namespace { + +Status invalid_source(std::string_view reason) { + return Status::Error("posting_run_merger: {}", reason); +} + +Status merge_corruption(std::string_view reason) { + return Status::Error("posting_run_merger: {}", + reason); +} + +bool posting_after(uint32_t lhs_segment, uint32_t lhs_docid, uint32_t rhs_segment, + uint32_t rhs_docid) { + return lhs_segment > rhs_segment || (lhs_segment == rhs_segment && lhs_docid > rhs_docid); +} + +#ifdef BE_TEST +std::atomic posting_run_frontier_update_counter {0}; +std::atomic posting_run_frontier_comparison_counter {0}; +std::atomic posting_run_document_counter {0}; +std::atomic posting_run_emitted_run_counter {0}; +std::atomic posting_run_boundary_search_counter {0}; +std::atomic posting_run_shape_scan_document_counter {0}; +std::atomic posting_run_legacy_fill_call_counter {0}; +std::atomic posting_run_copied_document_counter {0}; +#endif + +size_t lower_bound_docid(std::span docids, size_t begin, uint32_t target) { +#ifdef BE_TEST + posting_run_boundary_search_counter.fetch_add(1, std::memory_order_relaxed); +#endif + const size_t size = docids.size(); + if (begin >= size || docids[begin] >= target) { + return begin; + } + // Gallop before the binary search: when many sources interleave (a full + // compaction merging dozens of sorted segments), destination runs shrink to + // one or two documents, and a plain binary search over the whole remaining + // chunk costs O(log chunk) comparisons per emitted run. Doubling probes + // resolve those short runs in O(1) while keeping O(log run) for long runs. + size_t less = begin; + size_t probe = 1; + while (less + probe < size && docids[less + probe] < target) { + less += probe; + probe *= 2; + } + size_t low = less + 1; + size_t high = std::min(less + probe, size); + while (low < high) { + const size_t middle = low + (high - low) / 2; + if (docids[middle] < target) { + low = middle + 1; + } else { + high = middle; + } + } + return low; +} + +} // namespace + +void MergedPostingRuns::ActivePostingChunk::refresh_frontier() { + DCHECK_LT(ordinal, chunk.destination_docids.size()); + frontier_segment = chunk.destination_segment; + frontier_docid = chunk.destination_docids[ordinal]; +} + +Status MergedPostingRuns::ActivePostingChunk::validate_and_refresh_frontier( + bool retain_positions, std::span destination_doc_counts) { + if (chunk.destination_docids.empty() || ordinal >= chunk.destination_docids.size()) { + return merge_corruption("destination posting run is empty"); + } + if (chunk.destination_segment >= destination_doc_counts.size()) { + return merge_corruption("destination posting run segment is out of range"); + } + if (retain_positions) { + if (chunk.freqs.size() != chunk.destination_docids.size() || + chunk.position_offsets.size() != chunk.destination_docids.size() + 1) { + return merge_corruption("positioned posting run has an invalid shape"); + } + const uint32_t position_begin = chunk.position_offsets.front(); + const uint32_t position_end = chunk.position_offsets.back(); + if (position_end < position_begin || + position_end - position_begin != chunk.positions_flat.size()) { + return merge_corruption("positioned posting run has invalid offsets"); + } + } else if (!chunk.freqs.empty() || !chunk.position_offsets.empty() || + !chunk.positions_flat.empty()) { + return merge_corruption("docs-only posting run has positioned payload"); + } + + uint32_t previous_offset = retain_positions ? chunk.position_offsets.front() : 0; + for (size_t document = 0; document < chunk.destination_docids.size(); ++document) { + const uint32_t docid = chunk.destination_docids[document]; + if (docid >= destination_doc_counts[chunk.destination_segment]) { + return merge_corruption("destination posting run document is out of range"); + } + if (document > 0 && docid <= chunk.destination_docids[document - 1]) { + return merge_corruption("destination posting run is not strictly monotone"); + } + if (retain_positions) { + const uint32_t next_offset = chunk.position_offsets[document + 1]; + if (next_offset < previous_offset || + next_offset - previous_offset != chunk.freqs[document]) { + return merge_corruption("positioned posting run offsets differ from frequencies"); + } + previous_offset = next_offset; + } +#ifdef BE_TEST + posting_run_shape_scan_document_counter.fetch_add(1, std::memory_order_relaxed); +#endif + } + if (has_previous_chunk_posting && + !posting_after(chunk.destination_segment, chunk.destination_docids.front(), + previous_chunk_segment, previous_chunk_docid)) { + return merge_corruption("source posting chunks are not globally monotone"); + } + previous_chunk_segment = chunk.destination_segment; + previous_chunk_docid = chunk.destination_docids.back(); + has_previous_chunk_posting = true; + refresh_frontier(); + return Status::OK(); +} + +bool MergedPostingRuns::FrontierBefore::operator()(size_t lhs, size_t rhs) const { +#ifdef BE_TEST + posting_run_frontier_comparison_counter.fetch_add(1, std::memory_order_relaxed); +#endif + const ActivePostingChunk& lhs_chunk = (*active_chunks)[lhs]; + const ActivePostingChunk& rhs_chunk = (*active_chunks)[rhs]; + if (lhs_chunk.frontier_segment != rhs_chunk.frontier_segment) { + return lhs_chunk.frontier_segment < rhs_chunk.frontier_segment; + } + if (lhs_chunk.frontier_docid != rhs_chunk.frontier_docid) { + return lhs_chunk.frontier_docid < rhs_chunk.frontier_docid; + } + return lhs < rhs; +} + +MergedPostingRuns::MergedPostingRuns(std::vector> cursors, + bool retain_positions, bool counts_as_semantic_token, + std::span destination_doc_counts, + std::span destination_semantic_token_counts) + : cursors_(std::move(cursors)), + active_frontier_(FrontierBefore {.active_chunks = &active_chunks_}), + retain_positions_(retain_positions), + counts_as_semantic_token_(counts_as_semantic_token), + destination_doc_counts_(destination_doc_counts), + destination_semantic_token_counts_(destination_semantic_token_counts) {} + +Status MergedPostingRuns::init() { + if (initialized_) { + return invalid_source("source initialized twice"); + } + if (cursors_.empty() || destination_doc_counts_.empty()) { + return invalid_source("source or destination set is empty"); + } + if (counts_as_semantic_token_ && + destination_semantic_token_counts_.size() != destination_doc_counts_.size()) { + return invalid_source("semantic token counters differ from destination count"); + } + active_chunks_.resize(cursors_.size()); + for (size_t cursor_ordinal = 0; cursor_ordinal < cursors_.size(); ++cursor_ordinal) { + bool has_chunk = false; + RETURN_IF_ERROR(cursors_[cursor_ordinal]->next_chunk(&active_chunks_[cursor_ordinal].chunk, + &has_chunk)); + if (has_chunk) { + RETURN_IF_ERROR(active_chunks_[cursor_ordinal].validate_and_refresh_frontier( + retain_positions_, destination_doc_counts_)); + } + } + active_frontier_.build(cursors_.size(), [this](size_t source) { + return !active_chunks_[source].chunk.destination_docids.empty(); + }); + initialized_ = true; + return Status::OK(); +} + +bool MergedPostingRuns::empty() const { + DCHECK(initialized_); + return active_frontier_.empty(); +} + +uint32_t MergedPostingRuns::next_destination() const { + DCHECK(initialized_); + DCHECK(!active_destination_.has_value()); + DCHECK(!pending_source_.has_value()); + return front_segment(); +} + +Status MergedPostingRuns::begin_destination(uint32_t destination) { + if (!initialized_ || active_destination_.has_value() || pending_source_.has_value() || + active_frontier_.empty()) { + return invalid_source("cannot begin destination"); + } + if (front_segment() != destination) { + return invalid_source("destination differs from frontier"); + } + active_destination_ = destination; + return Status::OK(); +} + +Status MergedPostingRuns::next_run(uint32_t max_docs, writer::PostingRunView* run, bool* has_run) { + if (max_docs == 0 || run == nullptr || has_run == nullptr) { + return invalid_source("invalid next_run arguments"); + } + if (!initialized_ || !active_destination_.has_value()) { + return invalid_source("source has no active destination"); + } + *run = {}; + *has_run = false; + RETURN_IF_ERROR(settle_pending_run()); + if (active_frontier_.empty() || front_segment() != *active_destination_) { + active_destination_.reset(); + return Status::OK(); + } + RETURN_IF_ERROR(select_front_run(max_docs, run)); + *has_run = true; + return Status::OK(); +} + +Status MergedPostingRuns::fill(uint32_t target_docs, writer::TermPostingBuffer* out, + bool* exhausted) { + if (target_docs == 0 || out == nullptr || exhausted == nullptr) { + return invalid_source("invalid fill arguments"); + } + if (!initialized_ || !active_destination_.has_value()) { + return invalid_source("source has no active destination"); + } + if (!out->empty()) { + return invalid_source("output must be empty"); + } +#ifdef BE_TEST + posting_run_legacy_fill_call_counter.fetch_add(1, std::memory_order_relaxed); +#endif + + while (out->document_count() < target_docs) { + writer::PostingRunView run; + bool has_run = false; + RETURN_IF_ERROR(next_run(static_cast(target_docs - out->document_count()), &run, + &has_run)); + if (!has_run) { + *exhausted = true; + return Status::OK(); + } + const size_t position_count = run.positions_flat.size(); + writer::MutableTermPostingSpan destination; + RETURN_IF_ERROR(out->grow_uninitialized(run.docids.size(), retain_positions_, + position_count, &destination)); + std::ranges::copy(run.docids, destination.docids.begin()); + if (retain_positions_) { + std::ranges::copy(run.freqs, destination.freqs.begin()); + std::ranges::copy(run.positions_flat, destination.positions_flat.begin()); + } +#ifdef BE_TEST + posting_run_copied_document_counter.fetch_add(run.docids.size(), std::memory_order_relaxed); +#endif + } + + RETURN_IF_ERROR(settle_pending_run()); + *exhausted = active_frontier_.empty() || front_segment() != *active_destination_; + if (*exhausted) { + active_destination_.reset(); + } + return Status::OK(); +} + +uint32_t MergedPostingRuns::front_segment() const { + return active_chunks_[active_frontier_.winner()].frontier_segment; +} + +Status MergedPostingRuns::select_front_run(size_t max_docs, writer::PostingRunView* run) { + DCHECK_GT(max_docs, 0); + DCHECK(!pending_source_.has_value()); + const size_t cursor_ordinal = active_frontier_.winner(); + std::optional> next_frontier; + const size_t runner_up = active_frontier_.runner_up(); + if (runner_up != IndexedWinnerTree::kNoSource) { + const ActivePostingChunk& next = active_chunks_[runner_up]; + next_frontier = std::pair(next.frontier_segment, next.frontier_docid); + } + + ActivePostingChunk& active = active_chunks_[cursor_ordinal]; + RETURN_IF_ERROR(select_run(&active, max_docs, next_frontier, run)); + pending_source_ = cursor_ordinal; + return Status::OK(); +} + +Status MergedPostingRuns::select_run(ActivePostingChunk* active, size_t max_docs, + std::optional> next_frontier, + writer::PostingRunView* run) { + DCHECK(active != nullptr); + const auto docids = active->chunk.destination_docids; + if (active->chunk.destination_segment != *active_destination_) { + return merge_corruption("posting run differs from the active destination"); + } + const size_t begin = active->ordinal; + size_t end = begin + std::min(max_docs, docids.size() - begin); + if (next_frontier.has_value() && next_frontier->first == active->chunk.destination_segment) { + end = std::min(end, lower_bound_docid(docids, begin, next_frontier->second)); + } + if (end == begin) { + return merge_corruption("destination postings contain an equal merge frontier"); + } + if (has_previous_posting_ && !posting_after(active->frontier_segment, active->frontier_docid, + previous_segment_, previous_docid_)) { + return merge_corruption("destination postings are duplicated or not globally monotone"); + } + + const size_t document_count = end - begin; + size_t position_begin = 0; + size_t position_count = 0; + if (retain_positions_) { + const size_t position_base = active->chunk.position_offsets.front(); + const size_t absolute_position_begin = active->chunk.position_offsets[begin]; + const size_t absolute_position_end = active->chunk.position_offsets[end]; + DCHECK_GE(absolute_position_begin, position_base); + DCHECK_GE(absolute_position_end, absolute_position_begin); + DCHECK_LE(absolute_position_end - position_base, active->chunk.positions_flat.size()); + position_begin = absolute_position_begin - position_base; + position_count = absolute_position_end - absolute_position_begin; + } + + run->docids = docids.subspan(begin, document_count); + run->freqs = retain_positions_ ? active->chunk.freqs.subspan(begin, document_count) + : std::span {}; + run->position_offsets = + retain_positions_ ? active->chunk.position_offsets.subspan(begin, document_count + 1) + : std::span {}; + run->positions_flat = + retain_positions_ ? active->chunk.positions_flat.subspan(position_begin, position_count) + : std::span {}; + + if (counts_as_semantic_token_) { + uint64_t& token_count = destination_semantic_token_counts_[*active_destination_]; + if (position_count > std::numeric_limits::max() - token_count) { + return merge_corruption("semantic token count overflows uint64"); + } + token_count += position_count; + } + previous_segment_ = active->chunk.destination_segment; + previous_docid_ = docids[end - 1]; + has_previous_posting_ = true; + active->ordinal = end; +#ifdef BE_TEST + posting_run_document_counter.fetch_add(document_count, std::memory_order_relaxed); + posting_run_emitted_run_counter.fetch_add(1, std::memory_order_relaxed); +#endif + return Status::OK(); +} + +Status MergedPostingRuns::settle_pending_run() { + if (!pending_source_.has_value()) { + return Status::OK(); + } + const size_t cursor_ordinal = *pending_source_; + RETURN_IF_ERROR(advance_front_source(cursor_ordinal, &active_chunks_[cursor_ordinal])); + pending_source_.reset(); + return Status::OK(); +} + +Status MergedPostingRuns::advance_front_source(size_t cursor_ordinal, ActivePostingChunk* active) { + DCHECK(active != nullptr); + bool has_chunk = active->ordinal < active->chunk.destination_docids.size(); + if (!has_chunk) { + active->chunk = {}; + active->ordinal = 0; + RETURN_IF_ERROR(cursors_[cursor_ordinal]->next_chunk(&active->chunk, &has_chunk)); + } + if (has_chunk) { + if (active->ordinal == 0) { + RETURN_IF_ERROR(active->validate_and_refresh_frontier(retain_positions_, + destination_doc_counts_)); + } else { + active->refresh_frontier(); + } + } + active_frontier_.update(cursor_ordinal, has_chunk); +#ifdef BE_TEST + posting_run_frontier_update_counter.fetch_add(1, std::memory_order_relaxed); +#endif + return Status::OK(); +} + +#ifdef BE_TEST +namespace testing { + +void reset_posting_run_merge_counters() { + posting_run_frontier_update_counter.store(0, std::memory_order_relaxed); + posting_run_frontier_comparison_counter.store(0, std::memory_order_relaxed); + posting_run_document_counter.store(0, std::memory_order_relaxed); + posting_run_emitted_run_counter.store(0, std::memory_order_relaxed); + posting_run_boundary_search_counter.store(0, std::memory_order_relaxed); + posting_run_shape_scan_document_counter.store(0, std::memory_order_relaxed); + posting_run_legacy_fill_call_counter.store(0, std::memory_order_relaxed); + posting_run_copied_document_counter.store(0, std::memory_order_relaxed); +} + +uint64_t posting_run_frontier_updates() { + return posting_run_frontier_update_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_frontier_comparisons() { + return posting_run_frontier_comparison_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_documents() { + return posting_run_document_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_emitted_runs() { + return posting_run_emitted_run_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_boundary_searches() { + return posting_run_boundary_search_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_shape_scan_documents() { + return posting_run_shape_scan_document_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_legacy_fill_calls() { + return posting_run_legacy_fill_call_counter.load(std::memory_order_relaxed); +} + +uint64_t posting_run_copied_documents() { + return posting_run_copied_document_counter.load(std::memory_order_relaxed); +} + +} // namespace testing +#endif + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/posting_run_merger.h b/be/src/storage/index/snii/compaction/posting_run_merger.h new file mode 100644 index 00000000000000..7792a862d35f22 --- /dev/null +++ b/be/src/storage/index/snii/compaction/posting_run_merger.h @@ -0,0 +1,116 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/compaction/indexed_winner_tree.h" +#include "storage/index/snii/compaction/posting_cursor.h" +#include "storage/index/snii/writer/posting_window_emitter.h" +#include "storage/index/snii/writer/term_posting_source.h" + +namespace doris::snii::compaction { + +// K-way merge of destination-homogeneous cursor chunks. next_run() borrows the +// selected cursor workspace without copying; the view remains valid until a +// later next_run() call advances that cursor. +class MergedPostingRuns final : public writer::TermPostingSource { + struct ActivePostingChunk { + RemappedPostingChunk chunk; + size_t ordinal = 0; + uint32_t frontier_segment = 0; + uint32_t frontier_docid = 0; + uint32_t previous_chunk_segment = 0; + uint32_t previous_chunk_docid = 0; + bool has_previous_chunk_posting = false; + + void refresh_frontier(); + Status validate_and_refresh_frontier(bool retain_positions, + std::span destination_doc_counts); + }; + + struct FrontierBefore { + const std::vector* active_chunks = nullptr; + + bool operator()(size_t lhs, size_t rhs) const; + }; + +public: + MergedPostingRuns(std::vector> cursors, + bool retain_positions, bool counts_as_semantic_token, + std::span destination_doc_counts, + std::span destination_semantic_token_counts); + + Status init(); + bool empty() const; + uint32_t next_destination() const; + Status begin_destination(uint32_t destination); + Status next_run(uint32_t max_docs, writer::PostingRunView* run, bool* has_run); + + // Transitional streamed-session adapter. Compaction tests use next_run() + // directly; the assembler integration removes this handoff in the next + // milestone. + Status fill(uint32_t target_docs, writer::TermPostingBuffer* out, bool* exhausted) override; + +private: + uint32_t front_segment() const; + Status select_front_run(size_t max_docs, writer::PostingRunView* run); + Status select_run(ActivePostingChunk* active, size_t max_docs, + std::optional> next_frontier, + writer::PostingRunView* run); + Status settle_pending_run(); + Status advance_front_source(size_t cursor_ordinal, ActivePostingChunk* active); + + std::vector> cursors_; + std::vector active_chunks_; + IndexedWinnerTree active_frontier_; + bool retain_positions_ = true; + bool counts_as_semantic_token_ = false; + std::span destination_doc_counts_; + std::span destination_semantic_token_counts_; + std::optional active_destination_; + std::optional pending_source_; + uint32_t previous_segment_ = 0; + uint32_t previous_docid_ = 0; + bool has_previous_posting_ = false; + bool initialized_ = false; +}; + +#ifdef BE_TEST +namespace testing { + +void reset_posting_run_merge_counters(); +uint64_t posting_run_frontier_updates(); +uint64_t posting_run_frontier_comparisons(); +uint64_t posting_run_documents(); +uint64_t posting_run_emitted_runs(); +uint64_t posting_run_boundary_searches(); +uint64_t posting_run_shape_scan_documents(); +uint64_t posting_run_legacy_fill_calls(); +uint64_t posting_run_copied_documents(); + +} // namespace testing +#endif + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/region_reader.cpp b/be/src/storage/index/snii/compaction/region_reader.cpp new file mode 100644 index 00000000000000..3832b8505523ef --- /dev/null +++ b/be/src/storage/index/snii/compaction/region_reader.cpp @@ -0,0 +1,274 @@ +// 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 "storage/index/snii/compaction/region_reader.h" + +#include +#include + +namespace doris::snii::compaction { + +namespace { + +Status reserve_read_buffer(size_t target, std::vector* buffer, + writer::MemoryReporter::Reservation* reservation) { + if (reservation == nullptr) return Status::OK(); + if (buffer->capacity() >= target) { + DCHECK_EQ(reservation->bytes(), buffer->capacity()); + return Status::OK(); + } + writer::MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(reservation->prepare_replacement(target, &replacement)); + buffer->reserve(target); + DCHECK_EQ(buffer->capacity(), target); + *reservation = std::move(replacement); + return Status::OK(); +} + +} // namespace + +size_t SharedAlignedRegionCache::stream_index(PostingStream stream) { + const size_t index = static_cast(stream); + DCHECK_LT(index, kStreamCount); + return index; +} + +Status SharedAlignedRegionCache::init() { + if (initialized_) { + return Status::Error( + "shared_region_cache: init called twice"); + } + if (reader_ == nullptr) { + return Status::Error( + "shared_region_cache: null reader"); + } + if (total_budget_bytes_ < kSlotCount) { + return Status::Error( + "shared_region_cache: budget must hold two chunks"); + } + if (region_len_ > std::numeric_limits::max() - region_off_ || + region_off_ > reader_->size() || region_len_ > reader_->size() - region_off_) { + return Status::Error( + "shared_region_cache: region outside source file"); + } + + block_bytes_ = total_budget_bytes_ / kSlotCount; + for (size_t slot_index = 0; slot_index < slots_.size(); ++slot_index) { + Slot& slot = slots_[slot_index]; + if (memory_reporter_ != nullptr) { + slot_reservations_[slot_index] = memory_reporter_->make_reservation(); + RETURN_IF_ERROR(slot_reservations_[slot_index].set_bytes(block_bytes_)); + slot.bytes.reserve(block_bytes_); + DCHECK_EQ(slot.bytes.capacity(), block_bytes_); + } + slot.bytes.resize(block_bytes_); + } + if (resident_capacity_bytes() > total_budget_bytes_) { + return Status::Error( + "shared_region_cache: allocator exceeded cache budget"); + } + initialized_ = true; + return Status::OK(); +} + +size_t SharedAlignedRegionCache::resident_capacity_bytes() const { + size_t capacity = 0; + for (const Slot& slot : slots_) { + capacity += slot.bytes.capacity(); + } + return capacity; +} + +uint64_t SharedAlignedRegionCache::read_calls(PostingStream stream) const { + return read_calls_[stream_index(stream)]; +} + +uint64_t SharedAlignedRegionCache::buffer_hits(PostingStream stream) const { + return buffer_hits_[stream_index(stream)]; +} + +void SharedAlignedRegionCache::unpin(PostingStream stream) { + const size_t index = stream_index(stream); + const int8_t slot_index = stream_slots_[index]; + if (slot_index < 0) { + return; + } + Slot& slot = slots_[static_cast(slot_index)]; + DCHECK_GT(slot.pins, 0); + --slot.pins; + stream_slots_[index] = -1; +} + +Status SharedAlignedRegionCache::read_physical(PostingStream stream, uint64_t abs_off, size_t len, + std::vector* out) { + RETURN_IF_ERROR(reader_->read_at(abs_off, len, out)); + ++physical_read_ranges_; + physical_read_bytes_ += len; + ++read_calls_[stream_index(stream)]; + return Status::OK(); +} + +Status SharedAlignedRegionCache::resolve(PostingStream stream, uint64_t abs_off, uint64_t len, + std::vector* scratch, Slice* out, + writer::MemoryReporter::Reservation* scratch_reservation) { + if (scratch == nullptr || out == nullptr) { + return Status::Error("shared_region_cache: null out"); + } + if (!initialized_) { + return Status::Error( + "shared_region_cache: resolve before init"); + } + *out = Slice(); + if (abs_off < region_off_) { + return Status::Error( + "shared_region_cache: window outside region"); + } + const uint64_t relative_off = abs_off - region_off_; + if (relative_off > region_len_ || len > region_len_ - relative_off) { + return Status::Error( + "shared_region_cache: window outside region"); + } + if (len > std::numeric_limits::max()) { + return Status::Error( + "shared_region_cache: window length out of range"); + } + + unpin(stream); + const size_t want = static_cast(len); + if (want == 0) { + return Status::OK(); + } + + const uint64_t aligned_relative_off = (relative_off / block_bytes_) * block_bytes_; + const uint64_t block_off = region_off_ + aligned_relative_off; + const size_t block_len = static_cast( + std::min(block_bytes_, region_len_ - aligned_relative_off)); + const size_t offset_in_block = static_cast(relative_off - aligned_relative_off); + const bool fits_one_block = want <= block_len - offset_in_block; + if (!fits_one_block) { + RETURN_IF_ERROR(reserve_read_buffer(want, scratch, scratch_reservation)); + RETURN_IF_ERROR(read_physical(stream, abs_off, want, scratch)); + if (scratch->size() != want) { + return Status::Error( + "shared_region_cache: short exact read"); + } + *out = Slice(scratch->data(), want); + return Status::OK(); + } + + const size_t stream_id = stream_index(stream); + for (size_t slot_index = 0; slot_index < slots_.size(); ++slot_index) { + Slot& slot = slots_[slot_index]; + if (slot.valid && slot.offset == block_off && offset_in_block + want <= slot.bytes.size()) { + ++slot.pins; + stream_slots_[stream_id] = static_cast(slot_index); + ++buffer_hits_[stream_id]; + *out = Slice(slot.bytes.data() + offset_in_block, want); + return Status::OK(); + } + } + + size_t slot_index = 0; + while (slot_index < slots_.size() && slots_[slot_index].pins != 0) { + ++slot_index; + } + DCHECK_LT(slot_index, slots_.size()); + Slot& slot = slots_[slot_index]; + slot.valid = false; + RETURN_IF_ERROR(read_physical(stream, block_off, block_len, &slot.bytes)); + if (slot.bytes.size() != block_len) { + return Status::Error( + "shared_region_cache: short chunk read"); + } + slot.offset = block_off; + slot.pins = 1; + slot.valid = true; + stream_slots_[stream_id] = static_cast(slot_index); + *out = Slice(slot.bytes.data() + offset_in_block, want); + return Status::OK(); +} + +Status SequentialRegionReader::resolve(uint64_t abs_off, uint64_t len, + std::vector* scratch, Slice* out) { + if (scratch == nullptr || out == nullptr) { + return Status::Error("region_reader: null out"); + } + if (reader_ == nullptr) { + return Status::Error("region_reader: null reader"); + } + *out = Slice(); + // Subtraction-based bounds avoid wrapping region_offset+region_length or + // abs_off+len at UINT64_MAX. + if (region_len_ > std::numeric_limits::max() - region_off_ || abs_off < region_off_) { + return Status::Error( + "region_reader: window outside region"); + } + const uint64_t relative_off = abs_off - region_off_; + if (relative_off > region_len_ || len > region_len_ - relative_off) { + return Status::Error( + "region_reader: window outside region"); + } + if (len > std::numeric_limits::max()) { + return Status::Error( + "region_reader: window length out of range"); + } + const size_t want = static_cast(len); + if (want == 0) { + return Status::OK(); + } + + // 1. Buffered hit: zero-copy slice, no file read. + if (!buf_.empty() && abs_off >= buf_off_ && abs_off - buf_off_ + want <= buf_.size()) { + ++buffer_hits_; + *out = Slice(buf_.data() + (abs_off - buf_off_), want); + return Status::OK(); + } + + // 3. Oversized or backward miss: one exact range read into the caller's + // scratch, keeping the buffered chunk (and the forward stream position) + // intact. + const bool backward = !buf_.empty() && abs_off < buf_off_; + if (want > chunk_bytes_ || backward) { + RETURN_IF_ERROR(reader_->read_at(abs_off, want, scratch)); + ++read_calls_; + if (scratch->size() != want) { + return Status::Error( + "region_reader: short read"); + } + *out = Slice(scratch->data(), want); + return Status::OK(); + } + + // 2. Forward miss: refill the chunk starting at the window, clamped to the + // region end so read-ahead never reads past the region (whose tail may abut + // other file sections or EOF). want <= chunk and the range check above + // guarantee the window fits the refilled chunk. + const uint64_t remaining = region_len_ - relative_off; + const size_t fill = static_cast(std::min(chunk_bytes_, remaining)); + RETURN_IF_ERROR(reader_->read_at(abs_off, fill, &buf_)); + ++read_calls_; + if (buf_.size() != fill) { + buf_.clear(); + return Status::Error( + "region_reader: short chunk read"); + } + buf_off_ = abs_off; + *out = Slice(buf_.data(), want); + return Status::OK(); +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/region_reader.h b/be/src/storage/index/snii/compaction/region_reader.h new file mode 100644 index 00000000000000..63304883cd8ed7 --- /dev/null +++ b/be/src/storage/index/snii/compaction/region_reader.h @@ -0,0 +1,156 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/writer/memory_reporter.h" + +// SequentialRegionReader -- chunked sequential read-ahead over ONE contiguous +// byte region of a source file (T2.3, compaction index-merge fast path). +// +// The merge walks a source segment's posting region in ascending offset order +// (the writer laid the per-term [prx][frq] spans out in term order, and the +// term cursor replays terms in that same order), so per-window read_at calls +// would issue thousands of tiny reads over an already-sequential byte stream. +// This reader amortizes them: resolve() serves a window from the buffered +// chunk when possible and only touches the file on a miss. +// +// resolve() preference order (documented contract, pinned by UT): +// 1. window fully inside the buffered chunk -> zero-copy Slice into the +// buffer, NO file read; +// 2. forward miss with len <= chunk_bytes -> ONE chunk read starting at the +// window (clamped to the region end, never past it) and a slice of it; +// 3. oversized (len > chunk_bytes) or backward window -> ONE exact range +// read into *scratch, leaving the buffered chunk untouched (a rare +// backward probe must not thrash the forward stream). +// A window not fully inside [region_offset, region_offset+region_length) is +// Corruption -- posting locators were already validated against the region by +// LogicalIndexReader::resolve_*_window, so an out-of-region request here means +// a caller bug or corrupt state, never a legal miss. +// +// The returned Slice is valid until the NEXT resolve() call (buffer path) or +// until *scratch is next modified (fallback path); callers decode immediately. +// Single-threaded, borrowed FileReader must outlive the region reader. +namespace doris::snii::compaction { + +enum class PostingStream : uint8_t { kDocs = 0, kPrx = 1 }; + +// Two logical monotone posting streams share two aligned physical chunks. A +// stream pins at most its current chunk, so resolving the other stream cannot +// invalidate an outstanding Slice. Requests crossing an aligned chunk use the +// caller's per-stream scratch and do not evict either pinned chunk. +class SharedAlignedRegionCache { +public: + SharedAlignedRegionCache(io::FileReader* reader, uint64_t region_offset, uint64_t region_length, + size_t total_budget_bytes, + writer::MemoryReporter* memory_reporter = nullptr) + : reader_(reader), + region_off_(region_offset), + region_len_(region_length), + total_budget_bytes_(total_budget_bytes), + memory_reporter_(memory_reporter) {} + + Status init(); + Status resolve(PostingStream stream, uint64_t abs_off, uint64_t len, + std::vector* scratch, Slice* out, + writer::MemoryReporter::Reservation* scratch_reservation = nullptr); + + uint64_t physical_read_ranges() const { return physical_read_ranges_; } + uint64_t physical_read_bytes() const { return physical_read_bytes_; } + uint64_t read_calls(PostingStream stream) const; + uint64_t buffer_hits(PostingStream stream) const; + size_t resident_capacity_bytes() const; + +private: + static constexpr size_t kStreamCount = 2; + static constexpr size_t kSlotCount = 2; + + struct Slot { + std::vector bytes; + uint64_t offset = 0; + uint8_t pins = 0; + bool valid = false; + }; + + static size_t stream_index(PostingStream stream); + void unpin(PostingStream stream); + Status read_physical(PostingStream stream, uint64_t abs_off, size_t len, + std::vector* out); + + io::FileReader* reader_ = nullptr; + uint64_t region_off_ = 0; + uint64_t region_len_ = 0; + size_t total_budget_bytes_ = 0; + size_t block_bytes_ = 0; + writer::MemoryReporter* memory_reporter_ = nullptr; + std::array slot_reservations_; + std::array slots_; + std::array stream_slots_ {-1, -1}; + std::array read_calls_ {}; + std::array buffer_hits_ {}; + uint64_t physical_read_ranges_ = 0; + uint64_t physical_read_bytes_ = 0; + bool initialized_ = false; +}; + +class SequentialRegionReader { +public: + // Default chunk: large enough to amortize per-window read overhead over + // slim terms, small enough that k sources x 1 chunk stays negligible next + // to the merge's memory-precheck budget. Configurable per instance (the + // compaction wiring exposes a config knob in T2.6). + static constexpr size_t kDefaultChunkBytes = 4ULL << 20; + + SequentialRegionReader(io::FileReader* reader, uint64_t region_offset, uint64_t region_length, + size_t chunk_bytes = kDefaultChunkBytes) + : reader_(reader), + region_off_(region_offset), + region_len_(region_length), + chunk_bytes_(chunk_bytes == 0 ? kDefaultChunkBytes : chunk_bytes) {} + + // Resolves the absolute byte window [abs_off, abs_off+len) per the + // contract above. len == 0 yields an empty slice without touching the + // file. + Status resolve(uint64_t abs_off, uint64_t len, std::vector* scratch, Slice* out); + + // Observability for tests/profiling: physical reads issued vs windows + // served straight from the buffered chunk. + uint64_t read_calls() const { return read_calls_; } + uint64_t buffer_hits() const { return buffer_hits_; } + +private: + io::FileReader* reader_ = nullptr; + uint64_t region_off_ = 0; + uint64_t region_len_ = 0; + size_t chunk_bytes_ = kDefaultChunkBytes; + + std::vector buf_; // buffered chunk; empty until the first fill + uint64_t buf_off_ = 0; // absolute offset of buf_[0] (valid when buf_ non-empty) + + uint64_t read_calls_ = 0; + uint64_t buffer_hits_ = 0; +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/rowid_conversion.cpp b/be/src/storage/index/snii/compaction/rowid_conversion.cpp new file mode 100644 index 00000000000000..e05b80572d20c9 --- /dev/null +++ b/be/src/storage/index/snii/compaction/rowid_conversion.cpp @@ -0,0 +1,240 @@ +// 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 "storage/index/snii/compaction/rowid_conversion.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" + +namespace doris::snii::compaction { +namespace { + +constexpr uint32_t kDeleted = std::numeric_limits::max(); + +struct HeapEntry { + uint64_t destination_ordinal = 0; + size_t source_ordinal = 0; + size_t source_docid = 0; +}; + +struct HeapEntryGreater { + bool operator()(const HeapEntry& lhs, const HeapEntry& rhs) const { + if (lhs.destination_ordinal != rhs.destination_ordinal) { + return lhs.destination_ordinal > rhs.destination_ordinal; + } + return lhs.source_ordinal > rhs.source_ordinal; + } +}; + +uint64_t destination_ordinal(const std::pair& destination, + const std::vector& destination_segment_prefixes) { + return destination_segment_prefixes[destination.first] + destination.second; +} + +} // namespace + +Status validate_rowid_conversion(const RowIdConversionMap& conversion, + const std::vector& source_segment_doc_counts, + const std::vector& destination_segment_doc_counts) { + if (conversion.size() != source_segment_doc_counts.size()) { + return Status::InvalidArgument( + fmt::format("SNII rowid conversion source segment count mismatch: conversion={}, " + "doc_counts={}", + conversion.size(), source_segment_doc_counts.size())); + } + if (destination_segment_doc_counts.size() > + static_cast(std::numeric_limits::max())) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion destination segment count {} exceeds uint32 encoding", + destination_segment_doc_counts.size())); + } + + std::vector destination_segment_prefixes; + destination_segment_prefixes.reserve(destination_segment_doc_counts.size() + 1); + destination_segment_prefixes.push_back(0); + uint64_t destination_doc_count = 0; + for (size_t destination_segment = 0; + destination_segment < destination_segment_doc_counts.size(); ++destination_segment) { + const uint64_t segment_doc_count = destination_segment_doc_counts[destination_segment]; + if (segment_doc_count > std::numeric_limits::max() - destination_doc_count) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion destination prefix sum overflows uint64 at segment {}", + destination_segment)); + } + destination_doc_count += segment_doc_count; + destination_segment_prefixes.push_back(destination_doc_count); + } + + for (size_t source = 0; source < conversion.size(); ++source) { + const auto& source_conversion = conversion[source]; + if (source_conversion.size() != source_segment_doc_counts[source]) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion source doc count mismatch at source {}: " + "conversion={}, doc_count={}", + source, source_conversion.size(), source_segment_doc_counts[source])); + } + + bool has_previous = false; + uint64_t previous_ordinal = 0; + for (size_t source_docid = 0; source_docid < source_conversion.size(); ++source_docid) { + const auto& destination = source_conversion[source_docid]; + const bool segment_deleted = destination.first == kDeleted; + const bool row_deleted = destination.second == kDeleted; + if (segment_deleted != row_deleted) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion entry is partially deleted at source {} doc {}: " + "destination=({}, {})", + source, source_docid, destination.first, destination.second)); + } + if (segment_deleted) { + continue; + } + if (destination.first >= destination_segment_doc_counts.size()) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion destination segment {} is out of range at " + "source {} doc {} (segment_count={})", + destination.first, source, source_docid, + destination_segment_doc_counts.size())); + } + if (destination.second >= destination_segment_doc_counts[destination.first]) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion destination row {} is out of range at source {} " + "doc {} (destination segment {} has {} docs)", + destination.second, source, source_docid, destination.first, + destination_segment_doc_counts[destination.first])); + } + + const uint64_t ordinal = destination_ordinal(destination, destination_segment_prefixes); + if (has_previous && ordinal <= previous_ordinal) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion source {} is not strictly increasing at doc {}: " + "destination ordinal {} follows {}", + source, source_docid, ordinal, previous_ordinal)); + } + previous_ordinal = ordinal; + has_previous = true; + } + } + + std::priority_queue, HeapEntryGreater> heap; + auto push_next_live = [&](size_t source, size_t source_docid) { + const auto& source_conversion = conversion[source]; + while (source_docid < source_conversion.size() && + source_conversion[source_docid].first == kDeleted) { + ++source_docid; + } + if (source_docid < source_conversion.size()) { + heap.push({destination_ordinal(source_conversion[source_docid], + destination_segment_prefixes), + source, source_docid}); + } + }; + + for (size_t source = 0; source < conversion.size(); ++source) { + push_next_live(source, 0); + } + + uint64_t expected_ordinal = 0; + while (!heap.empty()) { + const HeapEntry entry = heap.top(); + heap.pop(); + if (entry.destination_ordinal < expected_ordinal) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion has duplicate destination ordinal {} at source {} " + "doc {}", + entry.destination_ordinal, entry.source_ordinal, entry.source_docid)); + } + if (entry.destination_ordinal > expected_ordinal) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion is missing destination ordinal {} before ordinal {} " + "at source {} doc {}", + expected_ordinal, entry.destination_ordinal, entry.source_ordinal, + entry.source_docid)); + } + ++expected_ordinal; + push_next_live(entry.source_ordinal, entry.source_docid + 1); + } + + if (expected_ordinal != destination_doc_count) { + return Status::InvalidArgument(fmt::format( + "SNII rowid conversion is missing destination ordinal {}: covered {} of {} " + "destination docs", + expected_ordinal, expected_ordinal, destination_doc_count)); + } + return Status::OK(); +} + +ValidatedRowIdConversion::ValidatedRowIdConversion( + const RowIdConversionMap* conversion, std::vector source_segment_doc_counts, + std::vector destination_segment_doc_counts) + : conversion_(conversion), + source_segment_doc_counts_(std::move(source_segment_doc_counts)), + destination_segment_doc_counts_(std::move(destination_segment_doc_counts)) { + DORIS_CHECK(conversion_ != nullptr); + source_has_deletions_.reserve(conversion_->size()); + for (const auto& source : *conversion_) { + const bool has_deletions = + std::ranges::any_of(source, [](const std::pair& mapping) { + return mapping.first == kDeleted; + }); + source_has_deletions_.push_back(static_cast(has_deletions)); + } +} + +Status ValidatedRowIdConversion::create(const RowIdConversionMap* conversion, + std::span source_segment_doc_counts, + std::span destination_segment_doc_counts, + std::unique_ptr* out) { + if (out == nullptr) { + return Status::InvalidArgument("SNII validated rowid conversion has null out parameter"); + } + out->reset(); + if (conversion == nullptr) { + return Status::InvalidArgument("SNII validated rowid conversion has null conversion"); + } + + std::vector source_counts(source_segment_doc_counts.begin(), + source_segment_doc_counts.end()); + std::vector destination_counts(destination_segment_doc_counts.begin(), + destination_segment_doc_counts.end()); + RETURN_IF_ERROR(validate_rowid_conversion(*conversion, source_counts, destination_counts)); + out->reset(new ValidatedRowIdConversion(conversion, std::move(source_counts), + std::move(destination_counts))); + return Status::OK(); +} + +std::span> ValidatedRowIdConversion::source_mapping( + size_t source_ordinal) const { + DCHECK_LT(source_ordinal, conversion_->size()); + return std::span>((*conversion_)[source_ordinal]); +} + +bool ValidatedRowIdConversion::source_has_deletions(size_t source_ordinal) const { + DCHECK_LT(source_ordinal, source_has_deletions_.size()); + return source_has_deletions_[source_ordinal] != 0; +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/rowid_conversion.h b/be/src/storage/index/snii/compaction/rowid_conversion.h new file mode 100644 index 00000000000000..ac6a339ba01f4f --- /dev/null +++ b/be/src/storage/index/snii/compaction/rowid_conversion.h @@ -0,0 +1,85 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris::snii::compaction { + +using RowIdConversionMap = std::vector>>; + +// Capability proving that a complete row-id conversion has passed the global +// shape, bounds, monotonicity and destination-coverage validation. Construction +// is restricted to create(), so merge plans cannot accidentally accept an +// unvalidated conversion. The conversion map is borrowed and must outlive this +// token and every merge plan prepared from it. +class ValidatedRowIdConversion { +public: + ValidatedRowIdConversion(const ValidatedRowIdConversion&) = delete; + ValidatedRowIdConversion& operator=(const ValidatedRowIdConversion&) = delete; + + static Status create(const RowIdConversionMap* conversion, + std::span source_segment_doc_counts, + std::span destination_segment_doc_counts, + std::unique_ptr* out); + + size_t source_segment_count() const { return source_segment_doc_counts_.size(); } + const std::vector& source_segment_doc_counts() const { + return source_segment_doc_counts_; + } + const std::vector& destination_segment_doc_counts() const { + return destination_segment_doc_counts_; + } + bool source_has_deletions(size_t source_ordinal) const; + std::span> source_mapping(size_t source_ordinal) const; + +private: + ValidatedRowIdConversion(const RowIdConversionMap* conversion, + std::vector source_segment_doc_counts, + std::vector destination_segment_doc_counts); + + const RowIdConversionMap* conversion_ = nullptr; + std::vector source_segment_doc_counts_; + std::vector destination_segment_doc_counts_; + std::vector source_has_deletions_; +}; + +// Validates the complete row-id conversion before an SNII index fast merge +// writes any destination bytes. Each inner conversion vector belongs to one +// source segment and is indexed by source docid. A deleted source doc must be +// represented by exactly (UINT32_MAX, UINT32_MAX); every other pair is a live +// (destination segment, destination docid). +// +// In addition to shape and bounds, the function proves that: +// * each source stream is strictly increasing in global destination order; +// * all source streams together contain every destination doc exactly once. +// +// Completeness is checked by a k-way merge of the monotonic source streams. +// Memory is O(source segments + destination segments), never O(output rows). +Status validate_rowid_conversion(const RowIdConversionMap& conversion, + const std::vector& source_segment_doc_counts, + const std::vector& destination_segment_doc_counts); + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/snii_index_compaction.cpp b/be/src/storage/index/snii/compaction/snii_index_compaction.cpp new file mode 100644 index 00000000000000..5bba7f822ae01b --- /dev/null +++ b/be/src/storage/index/snii/compaction/snii_index_compaction.cpp @@ -0,0 +1,557 @@ +// 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 "storage/index/snii/compaction/snii_index_compaction.h" + +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "common/logging.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/snii/compaction/eligibility.h" +#include "storage/index/snii/compaction/posting_run_merger.h" +#include "storage/index/snii/compaction/term_cursor.h" +#include "storage/index/snii/compaction/term_merge_frontier.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::compaction { + +namespace { + +Status invalid_plan(std::string_view reason) { + return Status::Error("snii_compaction: {}", reason); +} + +Status merge_corruption(std::string_view reason) { + return Status::Error("snii_compaction: {}", + reason); +} + +bool is_well_formed_common_gram(std::string_view term) { + namespace inverted_index = segment_v2::inverted_index; + if (!term.starts_with(inverted_index::CG_V1_MARKER) || + term.size() > inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES) { + return false; + } + constexpr size_t kLengthBytes = 8; + const size_t length_offset = inverted_index::CG_V1_MARKER.size(); + if (term.size() < length_offset + kLengthBytes + 1) { + return false; + } + uint32_t left_length = 0; + for (size_t i = 0; i < kLengthBytes; ++i) { + const char digit = term[length_offset + i]; + if (!((digit >= '0' && digit <= '9') || (digit >= 'a' && digit <= 'f'))) { + return false; + } + left_length = (left_length << 4) | + static_cast(digit <= '9' ? digit - '0' : digit - 'a' + 10); + } + const size_t separator = length_offset + kLengthBytes; + if (term[separator] != ':') { + return false; + } + const std::string_view components = term.substr(separator + 1); + if (left_length > components.size()) { + return false; + } + return inverted_index::validate_common_grams_logical_term(components.substr(0, left_length), + "left term") + .ok() && + inverted_index::validate_common_grams_logical_term(components.substr(left_length), + "right term") + .ok(); +} + +template +Status reserve_tracked_vector(std::vector* values, size_t additional, + writer::MemoryReporter::Reservation* reservation) { + if (reservation == nullptr || additional == 0) return Status::OK(); + if (additional > std::numeric_limits::max() - values->size()) { + return Status::Error( + "snii_compaction: destination posting vector size overflows"); + } + const size_t required = values->size() + additional; + if (required <= values->capacity()) { + DCHECK_EQ(reservation->bytes(), values->capacity() * sizeof(T)); + return Status::OK(); + } + size_t target = std::max(64, required); + if (values->capacity() != 0 && values->capacity() <= std::numeric_limits::max() / 2) { + target = std::max(target, values->capacity() * 2); + } + if (target > static_cast(std::numeric_limits::max()) / sizeof(T)) { + return Status::Error( + "snii_compaction: destination posting reservation exceeds int64"); + } + writer::MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(reservation->prepare_replacement(target * sizeof(T), &replacement)); + values->reserve(target); + DCHECK_EQ(values->capacity(), target); + *reservation = std::move(replacement); + return Status::OK(); +} +} // namespace + +SniiPlainT2MergePlan::SniiPlainT2MergePlan( + std::vector source_indexes, + const ValidatedRowIdConversion* rowid_conversion, + std::vector destination_segment_num_rows, + std::vector destination_null_reservations, + std::vector> destination_null_docids, + SniiCompactionEligibility eligibility, + std::vector destination_norm_reservations, + std::vector> destination_encoded_norms, + std::shared_ptr memory_reporter, + std::vector> read_contexts) + : source_indexes_(std::move(source_indexes)), + rowid_conversion_(rowid_conversion), + destination_segment_num_rows_(std::move(destination_segment_num_rows)), + memory_reporter_(std::move(memory_reporter)), + destination_null_reservations_(std::move(destination_null_reservations)), + destination_null_docids_(std::move(destination_null_docids)), + destination_null_docids_taken_(destination_null_docids_.size(), false), + eligibility_(std::move(eligibility)), + destination_norm_reservations_(std::move(destination_norm_reservations)), + destination_encoded_norms_(std::move(destination_encoded_norms)), + destination_encoded_norms_taken_(destination_encoded_norms_.size(), false), + destination_semantic_token_counts_(destination_segment_num_rows_.size(), 0), + read_contexts_(std::move(read_contexts)) {} + +Status SniiPlainT2MergePlan::prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + size_t total_read_ahead_budget_bytes, + std::unique_ptr* out) { + return prepare(std::move(source_indexes), rowid_conversion, total_read_ahead_budget_bytes, + nullptr, out); +} + +Status SniiPlainT2MergePlan::prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + size_t total_read_ahead_budget_bytes, + std::shared_ptr memory_reporter, + std::unique_ptr* out) { + SniiCompactionEligibility eligibility; + eligibility.kind = SniiStreamedMergeKind::kPlainT2; + return prepare(std::move(source_indexes), rowid_conversion, eligibility, + total_read_ahead_budget_bytes, std::move(memory_reporter), out); +} + +Status SniiPlainT2MergePlan::prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + const SniiCompactionEligibility& eligibility, + size_t total_read_ahead_budget_bytes, + std::unique_ptr* out) { + return prepare(std::move(source_indexes), rowid_conversion, eligibility, + total_read_ahead_budget_bytes, nullptr, out); +} + +Status SniiPlainT2MergePlan::prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + const SniiCompactionEligibility& eligibility, + size_t total_read_ahead_budget_bytes, + std::shared_ptr memory_reporter, + std::unique_ptr* out) { + if (out == nullptr) { + return invalid_plan("null plan out parameter"); + } + out->reset(); + if (source_indexes.empty()) { + return invalid_plan("no source indexes"); + } + const std::vector& destination_segment_num_rows = + rowid_conversion.destination_segment_doc_counts(); + if (destination_segment_num_rows.empty()) { + return invalid_plan("no destination segments"); + } + if (source_indexes.size() != rowid_conversion.source_segment_count()) { + return invalid_plan("source index count differs from row-id conversion"); + } + // Division first avoids source_count * minimum overflow. Rejecting tiny + // allocations is an IO gate: the raw rebuild is cheaper than issuing a + // range request for every small posting window. + if (source_indexes.size() > total_read_ahead_budget_bytes / kMinReadAheadBudgetPerSource) { + return invalid_plan("read-ahead budget is below the per-source IO floor"); + } + const size_t per_source_read_ahead_budget = + std::min(SniiPostingReadContext::kMaxReadAheadBudgetBytes, + total_read_ahead_budget_bytes / source_indexes.size()); + DORIS_CHECK_GE(per_source_read_ahead_budget, kMinReadAheadBudgetPerSource); + DORIS_CHECK_LE(per_source_read_ahead_budget, + total_read_ahead_budget_bytes / source_indexes.size()); + + for (size_t source_ordinal = 0; source_ordinal < source_indexes.size(); ++source_ordinal) { + const reader::LogicalIndexReader* source = source_indexes[source_ordinal]; + if (source == nullptr) { + return invalid_plan("null source index"); + } + RETURN_IF_ERROR(validate_snii_source_eligibility(*source, source_ordinal, eligibility)); + if (source->stats().doc_count > std::numeric_limits::max()) { + return merge_corruption("source doc count exceeds the SNII uint32 docid domain"); + } + if (source->stats().doc_count != + rowid_conversion.source_segment_doc_counts()[source_ordinal]) { + return invalid_plan("source doc count differs from validated row-id conversion"); + } + } + + std::vector destination_null_reservations; + destination_null_reservations.reserve(destination_segment_num_rows.size()); + for (size_t destination_ordinal = 0; destination_ordinal < destination_segment_num_rows.size(); + ++destination_ordinal) { + destination_null_reservations.push_back(memory_reporter == nullptr + ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation()); + } + std::vector> destination_null_docids(destination_segment_num_rows.size()); + for (size_t source_ordinal = 0; source_ordinal < source_indexes.size(); ++source_ordinal) { + reader::NullDocidsScanMemory scan_memory; + RETURN_IF_ERROR(source_indexes[source_ordinal]->null_docids_scan_memory(&scan_memory)); + writer::MemoryReporter::Reservation source_output_reservation = + memory_reporter == nullptr ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + writer::MemoryReporter::Reservation source_frame_reservation = + memory_reporter == nullptr ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + writer::MemoryReporter::Reservation source_decode_reservation = + memory_reporter == nullptr ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + if (memory_reporter != nullptr) { + RETURN_IF_ERROR(source_output_reservation.set_bytes(scan_memory.output_bytes)); + RETURN_IF_ERROR(source_frame_reservation.set_bytes(scan_memory.frame_bytes)); + } + std::vector source_null_docids; + source_null_docids.reserve(source_indexes[source_ordinal]->stats().null_count); + if (memory_reporter != nullptr) { + DORIS_CHECK_EQ(source_null_docids.capacity() * sizeof(uint32_t), + source_output_reservation.bytes()); + } + RETURN_IF_ERROR(source_indexes[source_ordinal]->read_null_docids( + &source_null_docids, [&](uint64_t bytes) { + return memory_reporter == nullptr ? Status::OK() + : source_decode_reservation.set_bytes(bytes); + })); + source_frame_reservation.reset(); + source_decode_reservation.reset(); + const auto source_mapping = rowid_conversion.source_mapping(source_ordinal); + for (uint32_t source_docid : source_null_docids) { + DCHECK_LT(source_docid, source_mapping.size()); + const auto [destination_segment, destination_docid] = source_mapping[source_docid]; + const bool segment_deleted = + destination_segment == std::numeric_limits::max(); + const bool doc_deleted = destination_docid == std::numeric_limits::max(); + DCHECK_EQ(segment_deleted, doc_deleted); + if (!segment_deleted) { + DCHECK_LT(destination_segment, destination_null_docids.size()); + DCHECK_LT(destination_docid, destination_segment_num_rows[destination_segment]); + RETURN_IF_ERROR(reserve_tracked_vector( + &destination_null_docids[destination_segment], 1, + memory_reporter == nullptr + ? nullptr + : &destination_null_reservations[destination_segment])); + destination_null_docids[destination_segment].push_back(destination_docid); + } + } + } + for (auto& null_docids : destination_null_docids) { + std::ranges::sort(null_docids); + DCHECK(std::adjacent_find(null_docids.begin(), null_docids.end()) == null_docids.end()); + } + + std::vector destination_norm_reservations; + std::vector> destination_encoded_norms; + if (eligibility.kind == SniiStreamedMergeKind::kCommonGramsT3) { + DORIS_CHECK(eligibility.common_grams_metadata_seed.has_value()); + destination_norm_reservations.reserve(destination_segment_num_rows.size()); + destination_encoded_norms.resize(destination_segment_num_rows.size()); + for (size_t destination_ordinal = 0; + destination_ordinal < destination_segment_num_rows.size(); ++destination_ordinal) { + destination_norm_reservations.push_back(memory_reporter == nullptr + ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation()); + const uint32_t doc_count = destination_segment_num_rows[destination_ordinal]; + auto& norms = destination_encoded_norms[destination_ordinal]; + RETURN_IF_ERROR(reserve_tracked_vector( + &norms, doc_count, + memory_reporter == nullptr ? nullptr : &destination_norm_reservations.back())); + norms.resize(doc_count); + if (memory_reporter != nullptr) { + DORIS_CHECK_EQ(destination_norm_reservations.back().bytes(), norms.capacity()); + } + } + for (size_t source_ordinal = 0; source_ordinal < source_indexes.size(); ++source_ordinal) { + writer::MemoryReporter::Reservation source_norms_reservation = + memory_reporter == nullptr ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + if (memory_reporter != nullptr) { + RETURN_IF_ERROR(source_norms_reservation.set_bytes( + source_indexes[source_ordinal]->compaction_norms_cache_charge())); + } + format::NormsPodReader source_norms; + RETURN_IF_ERROR(source_indexes[source_ordinal]->open_norms(&source_norms)); + const auto source_mapping = rowid_conversion.source_mapping(source_ordinal); + // The norms POD is only CRC-self-consistent; nothing upstream ties its + // doc_count to the validated conversion, and the loop below indexes + // source_mapping by it. Reconcile loudly (once per source, cold path). + if (source_norms.doc_count() != source_mapping.size()) { + return merge_corruption("norms doc count differs from validated row-id conversion"); + } + for (uint32_t source_docid = 0; source_docid < source_norms.doc_count(); + ++source_docid) { + const auto [destination_segment, destination_docid] = source_mapping[source_docid]; + const bool deleted = destination_segment == std::numeric_limits::max(); + DCHECK_EQ(deleted, destination_docid == std::numeric_limits::max()); + if (deleted) { + continue; + } + DCHECK_LT(destination_segment, destination_encoded_norms.size()); + DCHECK_LT(destination_docid, destination_encoded_norms[destination_segment].size()); + destination_encoded_norms[destination_segment][destination_docid] = + source_norms.encoded_norm(source_docid); + } + source_indexes[source_ordinal]->release_compaction_norms(); + source_norms_reservation.reset(); + } + } + + std::vector> read_contexts; + read_contexts.reserve(source_indexes.size()); + for (const reader::LogicalIndexReader* source : source_indexes) { + auto context = std::make_unique( + source, per_source_read_ahead_budget, memory_reporter.get()); + RETURN_IF_ERROR(context->init()); + read_contexts.push_back(std::move(context)); + } + + out->reset(new SniiPlainT2MergePlan( + std::move(source_indexes), &rowid_conversion, destination_segment_num_rows, + std::move(destination_null_reservations), std::move(destination_null_docids), + eligibility, std::move(destination_norm_reservations), + std::move(destination_encoded_norms), std::move(memory_reporter), + std::move(read_contexts))); + return Status::OK(); +} + +const std::vector& SniiPlainT2MergePlan::destination_null_docids( + size_t destination_segment) const { + DORIS_CHECK_LT(destination_segment, destination_null_docids_.size()); + return destination_null_docids_[destination_segment]; +} + +writer::TrackedNullDocids SniiPlainT2MergePlan::take_destination_null_docids( + size_t destination_segment) { + DORIS_CHECK_LT(destination_segment, destination_null_docids_.size()); + DORIS_CHECK(!destination_null_docids_taken_[destination_segment]); + destination_null_docids_taken_[destination_segment] = true; + return writer::TrackedNullDocids(std::move(destination_null_reservations_[destination_segment]), + std::move(destination_null_docids_[destination_segment])); +} + +const std::vector& SniiPlainT2MergePlan::destination_encoded_norms( + size_t destination_segment) const { + DORIS_CHECK(eligibility_.kind == SniiStreamedMergeKind::kCommonGramsT3); + DORIS_CHECK_LT(destination_segment, destination_encoded_norms_.size()); + return destination_encoded_norms_[destination_segment]; +} + +writer::TrackedEncodedNorms SniiPlainT2MergePlan::take_destination_encoded_norms( + size_t destination_segment) { + DORIS_CHECK(eligibility_.kind == SniiStreamedMergeKind::kCommonGramsT3); + DORIS_CHECK_LT(destination_segment, destination_encoded_norms_.size()); + DORIS_CHECK(!destination_encoded_norms_taken_[destination_segment]); + destination_encoded_norms_taken_[destination_segment] = true; + return writer::TrackedEncodedNorms( + std::move(destination_norm_reservations_[destination_segment]), + std::move(destination_encoded_norms_[destination_segment])); +} + +format::IndexConfig SniiPlainT2MergePlan::destination_index_config() const { + return eligibility_.kind == SniiStreamedMergeKind::kCommonGramsT3 + ? format::IndexConfig::kDocsPositionsScoring + : format::IndexConfig::kDocsPositions; +} + +std::optional +SniiPlainT2MergePlan::destination_common_grams_metadata(size_t destination_segment) const { + if (eligibility_.kind == SniiStreamedMergeKind::kPlainT2) { + return std::nullopt; + } + DORIS_CHECK(eligibility_.common_grams_metadata_seed.has_value()); + DORIS_CHECK_LT(destination_segment, destination_segment_num_rows_.size()); + auto metadata = *eligibility_.common_grams_metadata_seed; + metadata.scoring_doc_count = destination_segment_num_rows_[destination_segment]; + metadata.scoring_token_count = 0; + return metadata; +} + +Status SniiPlainT2MergePlan::poison(Status status) { + DORIS_CHECK(!status.ok()); + if (failed_.ok()) { + failed_ = std::move(status); + } + return failed_; +} + +Status SniiPlainT2MergePlan::take_front_source(TermMergeFrontier* frontier, CurrentTerm* current) { + SniiSegmentTermCursor* source = frontier->front(); + const uint32_t source_ordinal = source->source_ordinal(); + DCHECK_LT(source_ordinal, source_indexes_.size()); + const uint64_t frq_base = source->frq_base(); + const uint64_t prx_base = source->prx_base(); + format::DictEntry entry = source->take_entry(); + if (current->posting_cursors.empty()) { + current->term = std::move(entry.term); + } + + const bool source_has_positions = posting_entry_has_positions(entry); + if (!current->common_gram && !source_has_positions) { + return merge_corruption("ordinary term has docs-only posting shape"); + } + if (!current->has_positions.has_value()) { + current->has_positions = source_has_positions; + } else if (*current->has_positions != source_has_positions) { + return merge_corruption("same term has inconsistent position shape across sources"); + } + + auto cursor = std::make_unique(read_contexts_[source_ordinal].get(), + std::move(entry), frq_base, prx_base, + source_ordinal, rowid_conversion_); + RETURN_IF_ERROR(cursor->init()); + current->posting_cursors.push_back(std::move(cursor)); + return frontier->advance_front(); +} + +Status SniiPlainT2MergePlan::take_current_term(TermMergeFrontier* frontier, CurrentTerm* current) { + DCHECK(frontier != nullptr); + DCHECK(current != nullptr); + DCHECK(!frontier->empty()); + const std::string_view group_term = frontier->front()->term(); + current->posting_cursors.reserve(source_indexes_.size()); + if (eligibility_.kind == SniiStreamedMergeKind::kCommonGramsT3) { + current->common_gram = segment_v2::inverted_index::is_internal_term_key(group_term); + if (current->common_gram && !is_well_formed_common_gram(group_term)) { + return merge_corruption("CommonGrams source contains an unknown internal term marker"); + } + current->counts_as_semantic_token = !current->common_gram; + } + + do { + RETURN_IF_ERROR(take_front_source(frontier, current)); + } while (!frontier->empty() && frontier->front()->term() == current->term); + DCHECK(current->has_positions.has_value()); + return Status::OK(); +} + +Status SniiPlainT2MergePlan::write_current_term( + CurrentTerm current, std::span sessions) { + MergedPostingRuns posting_source(std::move(current.posting_cursors), *current.has_positions, + current.counts_as_semantic_token, + destination_segment_num_rows_, + destination_semantic_token_counts_); + RETURN_IF_ERROR(posting_source.init()); + while (!posting_source.empty()) { + const uint32_t destination = posting_source.next_destination(); + DCHECK_LT(destination, sessions.size()); + RETURN_IF_ERROR(posting_source.begin_destination(destination)); + writer::StreamedTermPostings postings {.term = current.term, + .retain_positions = *current.has_positions, + .source = &posting_source}; + RETURN_IF_ERROR(sessions[destination]->push_term(std::move(postings))); + } + return Status::OK(); +} + +Status SniiPlainT2MergePlan::merge_terms( + std::span sessions) { + std::vector> term_cursors; + std::vector term_cursor_ptrs; + term_cursors.reserve(source_indexes_.size()); + term_cursor_ptrs.reserve(source_indexes_.size()); + for (size_t source_ordinal = 0; source_ordinal < source_indexes_.size(); ++source_ordinal) { + term_cursors.push_back(std::make_unique( + source_indexes_[source_ordinal], static_cast(source_ordinal), + memory_reporter_.get())); + term_cursor_ptrs.push_back(term_cursors.back().get()); + } + TermMergeFrontier term_frontier; + RETURN_IF_ERROR(term_frontier.init(std::move(term_cursor_ptrs))); + + while (!term_frontier.empty()) { + CurrentTerm current; + RETURN_IF_ERROR(take_current_term(&term_frontier, ¤t)); + RETURN_IF_ERROR(write_current_term(std::move(current), sessions)); + } + + if (eligibility_.kind == SniiStreamedMergeKind::kCommonGramsT3) { + for (size_t destination_ordinal = 0; destination_ordinal < sessions.size(); + ++destination_ordinal) { + RETURN_IF_ERROR(sessions[destination_ordinal]->set_semantic_token_count( + destination_semantic_token_counts_[destination_ordinal])); + } + } + for (writer::SniiStreamedIndexSession* session : sessions) { + RETURN_IF_ERROR(session->finish()); + } + return Status::OK(); +} + +Status SniiPlainT2MergePlan::execute(std::span sessions) { + const auto abort_sessions = [&sessions](const Status& cause) { + DCHECK(!cause.ok()); + for (writer::SniiStreamedIndexSession* session : sessions) { + if (session != nullptr) { + session->abort(cause); + } + } + }; + if (!failed_.ok()) { + abort_sessions(failed_); + return failed_; + } + if (executed_) { + const Status status = invalid_plan("merge plan executed twice"); + abort_sessions(status); + return poison(status); + } + if (sessions.size() != destination_segment_num_rows_.size()) { + const Status status = invalid_plan("destination session count mismatch"); + abort_sessions(status); + return poison(status); + } + for (writer::SniiStreamedIndexSession* session : sessions) { + if (session == nullptr) { + const Status status = invalid_plan("null destination session"); + abort_sessions(status); + return poison(status); + } + } + executed_ = true; + const Status status = merge_terms(sessions); + if (!status.ok()) { + abort_sessions(status); + return poison(status); + } + return Status::OK(); +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/snii_index_compaction.h b/be/src/storage/index/snii/compaction/snii_index_compaction.h new file mode 100644 index 00000000000000..e012ae4e294111 --- /dev/null +++ b/be/src/storage/index/snii/compaction/snii_index_compaction.h @@ -0,0 +1,146 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/compaction/eligibility.h" +#include "storage/index/snii/compaction/posting_cursor.h" +#include "storage/index/snii/compaction/rowid_conversion.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" + +namespace doris::snii::compaction { + +class TermMergeFrontier; + +// Prepared, one-shot merge of one plain positions-only logical index across +// source segments. prepare() performs every O(1)-metadata, row-id and NULL +// preflight before the caller creates destination streamed sessions. execute() +// then performs exactly one dictionary/posting pass and seals every session. +// +// Source readers and the validated row-id conversion token are borrowed and +// must remain stable through execute(). The plan owns the per-source read +// contexts, destination row counts and remapped NULL docids. Its aggregate +// read-ahead allocation never exceeds the explicit prepare() budget. +class SniiPlainT2MergePlan { +public: + // Below this aggregate-per-source budget, two posting streams would issue + // tiny range reads and turn a memory fallback into an IO regression. + static constexpr size_t kMinReadAheadBudgetPerSource = 64U << 10; + + SniiPlainT2MergePlan(const SniiPlainT2MergePlan&) = delete; + SniiPlainT2MergePlan& operator=(const SniiPlainT2MergePlan&) = delete; + SniiPlainT2MergePlan(SniiPlainT2MergePlan&&) = delete; + SniiPlainT2MergePlan& operator=(SniiPlainT2MergePlan&&) = delete; + + static Status prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + size_t total_read_ahead_budget_bytes, + std::unique_ptr* out); + static Status prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + size_t total_read_ahead_budget_bytes, + std::shared_ptr memory_reporter, + std::unique_ptr* out); + static Status prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + const SniiCompactionEligibility& eligibility, + size_t total_read_ahead_budget_bytes, + std::unique_ptr* out); + static Status prepare(std::vector source_indexes, + const ValidatedRowIdConversion& rowid_conversion, + const SniiCompactionEligibility& eligibility, + size_t total_read_ahead_budget_bytes, + std::shared_ptr memory_reporter, + std::unique_ptr* out); + + const std::vector& destination_null_docids(size_t destination_segment) const; + writer::TrackedNullDocids take_destination_null_docids(size_t destination_segment); + const std::vector& destination_encoded_norms(size_t destination_segment) const; + writer::TrackedEncodedNorms take_destination_encoded_norms(size_t destination_segment); + format::IndexConfig destination_index_config() const; + std::optional + destination_common_grams_metadata(size_t destination_segment) const; + format::CommonGramsPostingPolicy destination_common_grams_posting_policy() const { + return eligibility_.common_grams_posting_policy; + } + size_t destination_segment_count() const { return destination_segment_num_rows_.size(); } + + // Sessions must correspond one-for-one with destination segments and must + // already carry the doc_count and destination_null_docids() returned by this + // plan. Any failure is terminal: the first error is sticky, sessions remain + // unsealable, and the whole unpublished compaction output must be discarded. + Status execute(std::span sessions); + +private: + struct CurrentTerm { + std::string term; + std::vector> posting_cursors; + std::optional has_positions; + bool common_gram = false; + bool counts_as_semantic_token = false; + }; + + SniiPlainT2MergePlan( + std::vector source_indexes, + const ValidatedRowIdConversion* rowid_conversion, + std::vector destination_segment_num_rows, + std::vector destination_null_reservations, + std::vector> destination_null_docids, + SniiCompactionEligibility eligibility, + std::vector destination_norm_reservations, + std::vector> destination_encoded_norms, + std::shared_ptr memory_reporter, + std::vector> read_contexts); + + Status take_front_source(TermMergeFrontier* frontier, CurrentTerm* current); + Status take_current_term(TermMergeFrontier* frontier, CurrentTerm* current); + Status write_current_term(CurrentTerm current, + std::span sessions); + Status merge_terms(std::span sessions); + Status poison(Status status); + + std::vector source_indexes_; + const ValidatedRowIdConversion* rowid_conversion_ = nullptr; + std::vector destination_segment_num_rows_; + std::shared_ptr memory_reporter_; + // Reservations precede their vectors so physical memory is destroyed first. + std::vector destination_null_reservations_; + std::vector> destination_null_docids_; + std::vector destination_null_docids_taken_; + SniiCompactionEligibility eligibility_; + // Reservations precede their vectors so physical memory is destroyed first. + std::vector destination_norm_reservations_; + std::vector> destination_encoded_norms_; + std::vector destination_encoded_norms_taken_; + std::vector destination_semantic_token_counts_; + std::vector> read_contexts_; + bool executed_ = false; + Status failed_ = Status::OK(); +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/term_cursor.cpp b/be/src/storage/index/snii/compaction/term_cursor.cpp new file mode 100644 index 00000000000000..1c25f1b57ff06a --- /dev/null +++ b/be/src/storage/index/snii/compaction/term_cursor.cpp @@ -0,0 +1,164 @@ +// 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 "storage/index/snii/compaction/term_cursor.h" + +#include + +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/format/sampled_term_index.h" + +namespace doris::snii::compaction { + +namespace { + +Status add_entry_memory(uint64_t bytes, uint64_t* total) { + if (bytes > std::numeric_limits::max() - *total) { + return Status::Error( + "term_cursor: decoded dictionary entries exceed uint64 memory accounting"); + } + *total += bytes; + return Status::OK(); +} + +Status decoded_entries_memory_bytes(const std::vector& entries, uint64_t* out) { + if (entries.capacity() > std::numeric_limits::max() / sizeof(format::DictEntry)) { + return Status::Error( + "term_cursor: decoded dictionary entry slots exceed uint64 memory accounting"); + } + uint64_t bytes = entries.capacity() * sizeof(format::DictEntry); + for (const format::DictEntry& entry : entries) { + RETURN_IF_ERROR(add_entry_memory(format::std_string_heap_bytes(entry.term), &bytes)); + RETURN_IF_ERROR(add_entry_memory(entry.frq_bytes.capacity(), &bytes)); + RETURN_IF_ERROR(add_entry_memory(entry.prx_bytes.capacity(), &bytes)); + } + *out = bytes; + return Status::OK(); +} + +} // namespace + +uint64_t big_endian_term_prefix(std::string_view term) { + uint64_t prefix = 0; + for (size_t i = 0; i < term.size() && i < sizeof(prefix); ++i) { + prefix |= static_cast(static_cast(term[i])) << (56 - i * 8); + } + return prefix; +} + +Status SniiSegmentTermCursor::next(bool* has_term) { + if (has_term == nullptr) { + return Status::Error("term_cursor: null has_term"); + } + *has_term = false; + if (!failed_.ok()) { + return failed_; + } + if (index_ == nullptr) { + return Status::Error("term_cursor: null index"); + } + if (exhausted_) { + return Status::OK(); + } + + if (started_) { + ++pos_; + } else { + started_ = true; + } + // Cross a block boundary (or start): materialize the next DICT block. The + // loop form also tolerates a (format-legal but writer-never-produced) + // empty block. + while (pos_ >= entries_.size()) { + if (next_block_ >= index_->n_dict_blocks()) { + exhausted_ = true; + entries_.clear(); + return Status::OK(); + } + + writer::MemoryReporter::Reservation decode_reservation = + memory_reporter_ == nullptr ? writer::MemoryReporter::Reservation() + : memory_reporter_->make_reservation(); + if (memory_reporter_ != nullptr) { + reader::DictBlockScanMemory memory; + Status st = index_->dict_block_scan_memory(next_block_, &memory); + if (!st.ok()) { + failed_ = st; + return failed_; + } + st = decode_reservation.set_bytes(memory.decode_bytes); + if (!st.ok()) { + failed_ = st; + return failed_; + } + + previous_entries_reservation_.reset(); + previous_entries_reservation_ = std::move(entries_reservation_); + entries_reservation_ = memory_reporter_->make_reservation(); + st = entries_reservation_.set_bytes(memory.entries_bytes); + if (!st.ok()) { + failed_ = st; + return failed_; + } + } + const Status st = index_->decode_dict_block(next_block_, &entries_, &frq_base_, &prx_base_); + if (!st.ok()) { + failed_ = st; + return failed_; + } + if (memory_reporter_ != nullptr) { + uint64_t actual_entries_bytes = 0; + Status memory_status = decoded_entries_memory_bytes(entries_, &actual_entries_bytes); + if (memory_status.ok()) { + memory_status = entries_reservation_.set_bytes(actual_entries_bytes); + } + if (!memory_status.ok()) { + std::vector().swap(entries_); + entries_reservation_.reset(); + failed_ = memory_status; + return failed_; + } + } + ++next_block_; + pos_ = 0; + } + + const format::DictEntry& e = entries_[pos_]; + // Hidden bigram / sentinel gate: classify by the FULL marker so a user term + // that merely begins with a raw 0x1F byte passes through (design ruling on + // the 0x1F-prefix corner). Any full-marker term aborts this column's merge + // -- the caller must fall back to rebuild (v1 excludes legacy bigrams). + if (format::is_phrase_bigram_term(e.term)) { + failed_ = Status::Error( + "term_cursor: source dictionary contains a legacy phrase-bigram/sentinel term; " + "column must fall back to index rebuild (src_ord={})", + source_ordinal_); + return failed_; + } + if (has_prev_ && e.term <= prev_term_) { + failed_ = Status::Error( + "term_cursor: dictionary term order violated (src_ord={})", source_ordinal_); + return failed_; + } + prev_term_ = e.term; + term_prefix_ = big_endian_term_prefix(e.term); + has_prev_ = true; + *has_term = true; + return Status::OK(); +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/term_cursor.h b/be/src/storage/index/snii/compaction/term_cursor.h new file mode 100644 index 00000000000000..2c73e8442c7f17 --- /dev/null +++ b/be/src/storage/index/snii/compaction/term_cursor.h @@ -0,0 +1,124 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/writer/memory_reporter.h" + +// SniiSegmentTermCursor -- pull-model full-dictionary scan over ONE source +// segment's logical index (T2.3, compaction index-merge fast path). +// +// The cursor walks the source's DICT blocks in ordinal order and yields every +// DictEntry in lexicographic term order, one block resident at a time (never +// the whole vocabulary). Entries are passed through UNINTERPRETED: the locator +// (inline bytes / slim pod_ref / windowed pod_ref) and the per-block +// kNoTermStats flag (DictEntry::term_stats_present) reach the downstream +// decoder exactly as the reader produced them -- the merge pump (T2.4) decides +// how to decode, and kNoTermStats inputs must have their ttf/max_freq recomputed +// from the actual freq stream (correctness invariant 2), so this layer must not +// synthesize stats. +// +// Hidden-term gate (base-drift addendum ruling): the v1 merge fast path does +// NOT merge legacy phrase-bigram postings. Any dictionary term carrying the +// FULL 0x1F bigram marker (hidden bigram pair or the bare-marker sentinel, +// classified by format::is_phrase_bigram_term -- NOT by a raw leading 0x1F +// byte, which a legitimate user term may begin with) makes next() return +// INVERTED_INDEX_NOT_SUPPORTED so the caller aborts THIS column's merge and +// falls back to a full rebuild. The error is deliberately raised from next() +// (not swallowed by skipping) because silently dropping hidden postings would +// change phrase semantics on the merged output. +namespace doris::snii::compaction { + +uint64_t big_endian_term_prefix(std::string_view term); + +class SniiSegmentTermCursor { +public: + // `index` is borrowed and must outlive the cursor. `source_ordinal` is the + // caller's stable id for this source segment (frontier tie-break + docid + // remapping key downstream). + SniiSegmentTermCursor(const reader::LogicalIndexReader* index, uint32_t source_ordinal, + writer::MemoryReporter* memory_reporter = nullptr) + : index_(index), + source_ordinal_(source_ordinal), + memory_reporter_(memory_reporter), + previous_entries_reservation_(memory_reporter == nullptr + ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation()), + entries_reservation_(memory_reporter == nullptr + ? writer::MemoryReporter::Reservation() + : memory_reporter->make_reservation()) {} + + // Advances to the next dictionary term. *has_term=false once the + // dictionary is exhausted (an empty index yields it on the first call). + // Errors: + // INVERTED_INDEX_NOT_SUPPORTED -- hidden bigram/sentinel term (see above); + // Corruption -- the dictionary violated strict lexicographic order. + // After an error the cursor is poisoned and keeps returning the error. + Status next(bool* has_term); + + // Accessors for the CURRENT term; valid only after next() returned + // *has_term=true and, for term()/entry(), before take_entry(). + const std::string& term() const { return entries_[pos_].term; } + uint64_t term_prefix() const { return term_prefix_; } + const format::DictEntry& entry() const { return entries_[pos_]; } + // Moves the current entry out (term + locator + any inline posting bytes). + // The cursor stays positioned, but term()/entry() must not be used again + // until the next next() call. + format::DictEntry take_entry() { return std::move(entries_[pos_]); } + + // frq/prx bases of the DICT block owning the current entry -- required to + // resolve pod_ref locators against the source's posting region. + uint64_t frq_base() const { return frq_base_; } + uint64_t prx_base() const { return prx_base_; } + uint32_t source_ordinal() const { return source_ordinal_; } + +private: + const reader::LogicalIndexReader* index_ = nullptr; + uint32_t source_ordinal_ = 0; + writer::MemoryReporter* memory_reporter_ = nullptr; + // Keep the preceding block's charge for one additional block transition. + // The frontier advances a source before the current term's posting cursor + // has released the DictEntry moved out of that block. + writer::MemoryReporter::Reservation previous_entries_reservation_; + writer::MemoryReporter::Reservation entries_reservation_; + + uint32_t next_block_ = 0; // next DICT block ordinal to decode + std::vector entries_; // current block, materialized + size_t pos_ = 0; // current entry within entries_ + uint64_t frq_base_ = 0; + uint64_t prx_base_ = 0; + uint64_t term_prefix_ = 0; + + bool started_ = false; // first next() must not pre-increment pos_ + bool exhausted_ = false; + Status failed_ = Status::OK(); // sticky error (poisoned cursor) + // Strict-order guard across block boundaries: DICT blocks and the merge + // frontier both assume a strictly increasing term sequence; a violation + // means a corrupt dictionary and must fail the merge, not scramble output. + std::string prev_term_; + bool has_prev_ = false; +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/term_merge_frontier.cpp b/be/src/storage/index/snii/compaction/term_merge_frontier.cpp new file mode 100644 index 00000000000000..8c5bd6ebc8c539 --- /dev/null +++ b/be/src/storage/index/snii/compaction/term_merge_frontier.cpp @@ -0,0 +1,103 @@ +// 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 "storage/index/snii/compaction/term_merge_frontier.h" + +#include + +#include "common/check.h" + +namespace doris::snii::compaction { + +bool TermMergeFrontier::Before::operator()(size_t lhs, size_t rhs) const { + const uint64_t lhs_prefix = (*cursors)[lhs]->term_prefix(); + const uint64_t rhs_prefix = (*cursors)[rhs]->term_prefix(); + if (lhs_prefix != rhs_prefix) { + return lhs_prefix < rhs_prefix; + } + const std::string& lhs_term = (*cursors)[lhs]->term(); + const std::string& rhs_term = (*cursors)[rhs]->term(); + if (lhs_term != rhs_term) { + return lhs_term < rhs_term; + } + return (*cursors)[lhs]->source_ordinal() < (*cursors)[rhs]->source_ordinal(); +} + +Status TermMergeFrontier::init(std::vector cursors) { + if (initialized_) { + return Status::Error( + "term_merge_frontier: init called twice"); + } + initialized_ = true; + cursors_ = std::move(cursors); + std::vector live(cursors_.size(), 0); + for (size_t source = 0; source < cursors_.size(); ++source) { + SniiSegmentTermCursor* cursor = cursors_[source]; + if (cursor == nullptr) { + failed_ = Status::Error( + "term_merge_frontier: null cursor"); + return failed_; + } + bool has_term = false; + const Status status = cursor->next(&has_term); + if (!status.ok()) { + failed_ = status; + return failed_; + } + live[source] = has_term; + } + frontier_.build(cursors_.size(), [&live](size_t source) { return live[source] != 0; }); + return Status::OK(); +} + +bool TermMergeFrontier::empty() const { + DCHECK(initialized_); + DCHECK(failed_.ok()); + return frontier_.empty(); +} + +SniiSegmentTermCursor* TermMergeFrontier::front() const { + DCHECK(initialized_); + DCHECK(failed_.ok()); + return cursors_[frontier_.winner()]; +} + +Status TermMergeFrontier::advance_front() { + if (!initialized_) { + return Status::Error( + "term_merge_frontier: advance before init"); + } + if (!failed_.ok()) { + return failed_; + } + if (frontier_.empty()) { + return Status::Error( + "term_merge_frontier: advance on empty frontier"); + } + + const size_t source = frontier_.winner(); + bool has_term = false; + const Status status = cursors_[source]->next(&has_term); + if (!status.ok()) { + failed_ = status; + return failed_; + } + frontier_.update(source, has_term); + return Status::OK(); +} + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/compaction/term_merge_frontier.h b/be/src/storage/index/snii/compaction/term_merge_frontier.h new file mode 100644 index 00000000000000..8bb36a3a849dca --- /dev/null +++ b/be/src/storage/index/snii/compaction/term_merge_frontier.h @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/compaction/indexed_winner_tree.h" +#include "storage/index/snii/compaction/term_cursor.h" + +namespace doris::snii::compaction { + +// K-way term frontier. Each source cursor is a dense winner-tree leaf. +// The caller consumes front()->entry() before advance_front(); advancing a +// source updates one leaf-to-root path and never materializes an intermediate +// merged-term group. +class TermMergeFrontier { + struct Before { + const std::vector* cursors = nullptr; + + bool operator()(size_t lhs, size_t rhs) const; + }; + +public: + TermMergeFrontier() : frontier_(Before {.cursors = &cursors_}) {} + + Status init(std::vector cursors); + + bool empty() const; + SniiSegmentTermCursor* front() const; + Status advance_front(); + +private: + std::vector cursors_; + IndexedWinnerTree frontier_; + bool initialized_ = false; + Status failed_ = Status::OK(); +}; + +} // namespace doris::snii::compaction diff --git a/be/src/storage/index/snii/encoding/byte_sink.cpp b/be/src/storage/index/snii/encoding/byte_sink.cpp new file mode 100644 index 00000000000000..7d91a6552f6658 --- /dev/null +++ b/be/src/storage/index/snii/encoding/byte_sink.cpp @@ -0,0 +1,60 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "storage/index/snii/encoding/byte_sink.h" + +#include "storage/index/snii/encoding/varint.h" + +namespace doris::snii { + +void ByteSink::put_fixed16(uint16_t v) { + for (int i = 0; i < 2; ++i) buf_.push_back(static_cast(v >> (8 * i))); +} + +void ByteSink::put_fixed32(uint32_t v) { + for (int i = 0; i < 4; ++i) buf_.push_back(static_cast(v >> (8 * i))); +} + +void ByteSink::put_fixed64(uint64_t v) { + for (int i = 0; i < 8; ++i) buf_.push_back(static_cast(v >> (8 * i))); +} + +void ByteSink::put_varint32(uint32_t v) { + while (v >= 0x80) { + buf_.push_back(static_cast(v) | 0x80); + v >>= 7; + } + buf_.push_back(static_cast(v)); +} + +void ByteSink::put_varint64(uint64_t v) { + while (v >= 0x80) { + buf_.push_back(static_cast(v) | 0x80); + v >>= 7; + } + buf_.push_back(static_cast(v)); +} + +void ByteSink::put_zigzag(int64_t v) { + put_varint64(zigzag_encode(v)); +} + +void ByteSink::put_bytes(Slice s) { + buf_.insert(buf_.end(), s.data(), s.data() + s.size()); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/byte_sink.h b/be/src/storage/index/snii/encoding/byte_sink.h new file mode 100644 index 00000000000000..e96ae112e20a31 --- /dev/null +++ b/be/src/storage/index/snii/encoding/byte_sink.h @@ -0,0 +1,71 @@ +// 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. + +#pragma once + +#include +#include + +#include "storage/index/snii/common/slice.h" + +namespace doris::snii { + +// append-only write cursor: all section serialization goes through this; manual byte assembly is forbidden. +// All multi-byte fixed-width fields are little-endian. +class ByteSink { +public: + void put_u8(uint8_t v) { buf_.push_back(v); } + void put_fixed16(uint16_t v); + void put_fixed32(uint32_t v); + void put_fixed64(uint64_t v); + void put_varint32(uint32_t v); + void put_varint64(uint64_t v); + void put_zigzag(int64_t v); + void put_bytes(Slice s); + + size_t size() const { return buf_.size(); } + const std::vector& buffer() const { return buf_; } + Slice view() const { return Slice(buf_); } + + // Reserves capacity for `additional` MORE bytes on top of the current size(), + // so a caller about to append a known-length run pays at most one reallocation + // instead of the geometric-growth reallocs a byte-at-a-time put_u8 loop would + // trigger. The argument is RELATIVE (absolute target = size() + additional): a + // sink REUSED across encodes (e.g. the shared PFOR-run `out`) keeps accumulating + // correctly rather than no-op'ing on a repeated per-run reserve of the same + // value. Affects only the backing buffer's capacity -- the emitted bytes are + // unchanged. + void reserve(size_t additional) { buf_.reserve(buf_.size() + additional); } + + // Resets the cursor to empty while RETAINING the backing capacity, so a sink can + // be reused across many small encodes (e.g. per-window region/prx scratch in the + // windowed posting builder) without re-allocating each time -- this avoids the + // cumulative small-allocation churn that fragments the heap arena and inflates + // peak RSS during the merge of a high-df term split into thousands of windows. + void clear() { buf_.clear(); } + + // Moves the backing buffer OUT to the caller (the sink is left empty), so an encoded + // section can be handed off without the copy (+ copy-induced capacity slack) that + // reading buffer() and copy-assigning would incur. Use only when the sink is not + // reused afterward (a stack-local about to die, or one that is clear()'d next). + std::vector take() { return std::move(buf_); } + +private: + std::vector buf_; +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/byte_source.cpp b/be/src/storage/index/snii/encoding/byte_source.cpp new file mode 100644 index 00000000000000..93619aa97064e7 --- /dev/null +++ b/be/src/storage/index/snii/encoding/byte_source.cpp @@ -0,0 +1,179 @@ +// 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 "storage/index/snii/encoding/byte_source.h" + +#include + +#include "storage/index/snii/encoding/varint.h" + +namespace doris::snii { + +Status ByteSource::get_u8(uint8_t* v) { + if (remaining() < 1) + return Status::Error("get_u8 overrun"); + *v = s_[pos_++]; + return Status::OK(); +} + +Status ByteSource::get_fixed16(uint16_t* v) { + if (remaining() < 2) + return Status::Error( + "get_fixed16 overrun"); + uint16_t r = 0; + for (int i = 0; i < 2; ++i) r |= static_cast(s_[pos_ + i]) << (8 * i); + pos_ += 2; + *v = r; + return Status::OK(); +} + +Status ByteSource::get_fixed32(uint32_t* v) { + if (remaining() < 4) + return Status::Error( + "get_fixed32 overrun"); + uint32_t r = 0; + for (int i = 0; i < 4; ++i) r |= static_cast(s_[pos_ + i]) << (8 * i); + pos_ += 4; + *v = r; + return Status::OK(); +} + +Status ByteSource::get_fixed64(uint64_t* v) { + if (remaining() < 8) + return Status::Error( + "get_fixed64 overrun"); + uint64_t r = 0; + for (int i = 0; i < 8; ++i) r |= static_cast(s_[pos_ + i]) << (8 * i); + pos_ += 8; + *v = r; + return Status::OK(); +} + +Status ByteSource::get_varint64(uint64_t* v) { + const uint8_t* p = s_.data() + pos_; + const uint8_t* next = nullptr; + RETURN_IF_ERROR(decode_varint64(p, s_.data() + s_.size(), v, &next)); + pos_ = static_cast(next - s_.data()); + return Status::OK(); +} + +Status ByteSource::get_varint32(uint32_t* v) { + uint64_t tmp; + RETURN_IF_ERROR(get_varint64(&tmp)); + if (tmp > 0xFFFFFFFFu) + return Status::Error("varint32 overflow"); + *v = static_cast(tmp); + return Status::OK(); +} + +Status ByteSource::decode_delta_run(size_t count, std::vector* out) { + if (out == nullptr) { + return Status::Error( + "byte_source: null delta run output"); + } + if (count > std::numeric_limits::max() - out->size()) { + return Status::Error( + "byte_source: delta run size overflow"); + } + const uint8_t* const begin = s_.data(); + const uint8_t* const end = begin + s_.size(); + const uint8_t* p = begin + pos_; + const size_t original_size = out->size(); + out->reserve(out->size() + count); + uint32_t running = 0; + const auto fail = [&](const char* message) { + out->resize(original_size); + return Status::Error(message); + }; + for (size_t i = 0; i < count; ++i) { + if (p >= end) { + return fail("byte_source: delta run past end"); + } + uint32_t b = *p++; + uint32_t delta = b & 0x7FU; + if (b >= 0x80) { // multi-byte (rare for position deltas) + uint32_t shift = 7; + for (;;) { + if (p >= end) { + return fail("byte_source: delta run past end"); + } + b = *p++; + if (shift == 28 && (b & 0xF0U) != 0) { + return fail("byte_source: delta varint32 overflow"); + } + delta |= (b & 0x7FU) << shift; + if ((b & 0x80) == 0) { + break; + } + if (shift == 28) { + return fail("byte_source: delta varint32 overflow"); + } + shift += 7; + } + } + if (delta > std::numeric_limits::max() - running) { + return fail("byte_source: delta prefix sum overflow"); + } + running += delta; + out->push_back(running); + } + pos_ = static_cast(p - begin); + return Status::OK(); +} + +Status ByteSource::skip_varints(size_t count) { + const uint8_t* const begin = s_.data(); + const uint8_t* const end = begin + s_.size(); + const uint8_t* p = begin + pos_; + // Each varint ends at the first byte whose continuation bit (0x80) is clear. + // Scanning for `count` such terminators skips the values with one branch per + // byte -- no shift/accumulate/store and no per-value bounds Status. (A SIMD + // bulk terminator-count was tried and reverted: the skipped position runs + // between selected docs are almost always 1-3 varints -- far below a 16-byte + // block -- so the vector path never amortized, and the larger body stopped + // this function from inlining into the CSR reader, a net CPU regression on + // the httplogs/agentlogs phrase-prefix profiles.) + for (size_t k = 0; k < count; ++k) { + while (p < end && (*p & 0x80) != 0) { + ++p; + } + if (p >= end) { + return Status::Error( + "byte_source: varint skip past end"); + } + ++p; // consume the terminator byte + } + pos_ = static_cast(p - begin); + return Status::OK(); +} + +Status ByteSource::get_zigzag(int64_t* v) { + uint64_t tmp; + RETURN_IF_ERROR(get_varint64(&tmp)); + *v = zigzag_decode(tmp); + return Status::OK(); +} + +Status ByteSource::get_bytes(size_t n, Slice* out) { + if (remaining() < n) + return Status::Error("get_bytes overrun"); + *out = s_.subslice(pos_, n); + pos_ += n; + return Status::OK(); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/byte_source.h b/be/src/storage/index/snii/encoding/byte_source.h new file mode 100644 index 00000000000000..0490e6949b8de9 --- /dev/null +++ b/be/src/storage/index/snii/encoding/byte_source.h @@ -0,0 +1,91 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" + +namespace doris::snii { + +// Slice read cursor: all section deserialization goes through this; any overrun returns Corruption. +class ByteSource { +public: + explicit ByteSource(Slice s) : s_(s) {} + + Status get_u8(uint8_t* v); + Status get_fixed16(uint16_t* v); + Status get_fixed32(uint32_t* v); + Status get_fixed64(uint64_t* v); + Status get_varint32(uint32_t* v); + Status get_varint64(uint64_t* v); + + // Single-byte fast-path varint32. The CSR position reader decodes one count + // header (`pos_count`) per doc in a window -- for every doc, selected or + // skipped -- and that per-doc read, routed through the out-of-line + // get_varint32 -> get_varint64 -> decode_varint64 chain, dominates the loop + // once positions themselves are skipped/inlined. Almost all counts are < 128 + // (one byte), so inline that case here and fall back to the full decoder for + // multi-byte values and bounds handling. + Status get_varint32_fast(uint32_t* v) { + const size_t p = pos_; + if (p < s_.size()) { + const uint8_t b0 = s_[p]; + if (b0 < 0x80) { + *v = b0; + pos_ = p + 1; + return Status::OK(); + } + } + return get_varint32(v); + } + // Advances past `count` LEB128 varints WITHOUT decoding their values -- just + // scans continuation bytes. Cheaper than get_varint* per value when the + // decoded value is unused (e.g. skipping a non-selected doc's position + // deltas in a CSR window, where the vast majority of docs in a window are + // not in the candidate set). Returns Corruption on truncation. + Status skip_varints(size_t count); + + // Decodes `count` LEB128 varints, treats them as ASCENDING deltas + // (running prefix sum starting at 0), and APPENDS the running values to + // `out`. A tight inline decoder -- no per-value get_varint32/get_varint64/ + // decode_varint64 call chain, no per-value Status, and a single-byte fast + // path (position deltas are almost always < 128). This is the hot loop of + // the CSR position reader for candidate (selected) docs. Returns Corruption + // on truncation, a >32-bit value, or a uint32 prefix-sum overflow. Failure + // leaves the cursor and output vector unchanged. + Status decode_delta_run(size_t count, std::vector* out); + Status get_zigzag(int64_t* v); + Status get_bytes(size_t n, Slice* out); + + size_t remaining() const { return s_.size() - pos_; } + size_t position() const { return pos_; } + bool eof() const { return pos_ == s_.size(); } + + // Returns a sub-view starting at absolute offset start with length len (used by framer etc. to rewind over the CRC coverage region). + Slice slice_from(size_t start, size_t len) const { return s_.subslice(start, len); } + +private: + Slice s_; + size_t pos_ = 0; +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/crc32c.cpp b/be/src/storage/index/snii/encoding/crc32c.cpp new file mode 100644 index 00000000000000..39d7c6f58fe487 --- /dev/null +++ b/be/src/storage/index/snii/encoding/crc32c.cpp @@ -0,0 +1,278 @@ +// 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 "storage/index/snii/encoding/crc32c.h" + +// T21 test-seam implementation. Production crc32c()/crc32c_extend() (in the header) +// delegate to the bundled Google crc32c thirdparty, which already runs a +// runtime-dispatched, hardware-accelerated and interleaved CRC32C. The reference +// sub-paths defined here -- portable slice-by-8, serial SSE4.2 hardware, and the +// 3-way interleaved SSE4.2 hardware algorithm T21 specifies -- exist ONLY so that +// unit tests can prove, byte-for-byte across all sizes and alignments, that the +// production path equals the canonical CRC32C (same bit-reflected Castagnoli +// polynomial) and that the hardware path is engaged. All of them are compiled out +// of release builds by the BE_TEST gate, so production pays nothing: no extra code, +// no static slice-by-8 table, no startup CPUID probe. Every function here is a pure +// function with no shared mutable state, so it is trivially thread-safe. +#ifdef BE_TEST + +#include +#include +#include + +#if defined(__x86_64__) || defined(_M_X64) +#define SNII_CRC32C_X86 1 +#include // __get_cpuid, bit_SSE4_2 +#include // _mm_crc32_u8/u32/u64 (SSE4.2) +#endif + +namespace doris::snii { +namespace { + +// Bit-reflected Castagnoli polynomial (CRC32C / iSCSI). Identical to the constant +// the removed in-tree implementation used, so every reference path below yields +// the same on-disk checksum value as the production library. +constexpr uint32_t kPoly = 0x82F63B78U; + +// Below this length the 3-way path's lane setup and GF(2) shift-combine outweigh +// the throughput win, so crc32c_hw3 falls back to the serial hardware path. Small +// buffers (inline prx windows, small pod_ref regions) therefore never pay the +// combine cost. Exposed to tests via crc32c_interleave_threshold(). +constexpr size_t kInterleaveThreshold = 1024; + +// Builds the slice-by-8 lookup tables. Column 0 is the classic byte table; each +// successive column folds in one more byte of look-ahead, letting the inner loop +// consume 8 bytes per iteration with 8 table reads + XORs instead of 8 dependent +// shift/lookup steps. The checksum value is identical to the byte-at-a-time loop. +std::array, 8> make_slice8_table() { + std::array, 8> t {}; + for (uint32_t i = 0; i < 256; ++i) { + uint32_t c = i; + for (int k = 0; k < 8; ++k) { + c = (c & 1) ? (kPoly ^ (c >> 1)) : (c >> 1); + } + t[0][i] = c; + } + for (uint32_t i = 0; i < 256; ++i) { + uint32_t c = t[0][i]; + for (int s = 1; s < 8; ++s) { + c = t[0][c & 0xFF] ^ (c >> 8); + t[s][i] = c; + } + } + return t; +} + +const std::array, 8> kSlice8 = make_slice8_table(); + +inline uint32_t load_le32(const uint8_t* p) { + return static_cast(p[0]) | (static_cast(p[1]) << 8) | + (static_cast(p[2]) << 16) | (static_cast(p[3]) << 24); +} + +// Pure software slice-by-8 (used as the portable path and the hardware fallback). +// Operates on the raw (pre/post-inversion applied by the callers) CRC register. +uint32_t crc32c_slice8(uint32_t crc, const uint8_t* p, size_t n) { + while (n >= 8) { + crc ^= load_le32(p); + const uint32_t hi = load_le32(p + 4); + crc = kSlice8[7][crc & 0xFF] ^ kSlice8[6][(crc >> 8) & 0xFF] ^ + kSlice8[5][(crc >> 16) & 0xFF] ^ kSlice8[4][crc >> 24] ^ kSlice8[3][hi & 0xFF] ^ + kSlice8[2][(hi >> 8) & 0xFF] ^ kSlice8[1][(hi >> 16) & 0xFF] ^ kSlice8[0][hi >> 24]; + p += 8; + n -= 8; + } + while (n--) { + crc = kSlice8[0][(crc ^ *p++) & 0xFF] ^ (crc >> 8); + } + return crc; +} + +#if SNII_CRC32C_X86 +// Serial hardware CRC32C via the SSE4.2 crc32 instruction. The intrinsics operate +// on the same bit-reflected Castagnoli polynomial as the tables, so the result is +// byte-identical. This TU is compiled without -msse4.2, so gate the intrinsics +// behind a function-level target attribute and a runtime CPUID check. This is the +// authoritative serial hardware path that crc32c_hw3 reuses for each lane and tail. +__attribute__((target("sse4.2"))) uint32_t crc32c_hw_serial(uint32_t crc, const uint8_t* p, + size_t n) { + while (n >= 8) { + uint64_t v; + std::memcpy(&v, p, sizeof(v)); // unaligned-safe; x86 folds to a plain load + crc = static_cast(_mm_crc32_u64(crc, v)); + p += 8; + n -= 8; + } + if (n >= 4) { + crc = _mm_crc32_u32(crc, load_le32(p)); + p += 4; + n -= 4; + } + while (n--) { + crc = _mm_crc32_u8(crc, *p++); + } + return crc; +} + +// GF(2) 32x32 bit-matrix helpers (zlib crc32_combine style). A matrix column mat[i] +// is the image of the i-th unit CRC register; gf2_matrix_times sums (XORs) the +// columns selected by the set bits of vec. +uint32_t gf2_matrix_times(const uint32_t* mat, uint32_t vec) { + uint32_t sum = 0; + while (vec != 0) { + if (vec & 1) { + sum ^= *mat; + } + vec >>= 1; + ++mat; + } + return sum; +} + +void gf2_matrix_square(uint32_t* square, const uint32_t* mat) { + for (int n = 0; n < 32; ++n) { + square[n] = gf2_matrix_times(mat, mat[n]); + } +} + +// crc32c_shift(crc, bytes): advance the raw CRC32C register as if `bytes` zero +// bytes were appended, i.e. crc . x^(8*bytes) mod P in the bit-reflected domain. +// This is the linear operator the 3-way combine applies to a lane's partial CRC so +// it lines up with the following lanes -- a pure table/matrix computation with no +// PCLMULQDQ, hence no extra CPUID gate. bytes == 0 returns crc unchanged. +uint32_t crc32c_shift(uint32_t crc, size_t bytes) { + uint32_t even[32]; // operator for 2^k zero bits, doubled each round + uint32_t odd[32]; // operator for 2^(k-1) zero bits + + // odd = operator for a single zero bit: column 0 is the polynomial, columns + // 1..31 shift the register right by one (bit i maps to bit i-1). + odd[0] = kPoly; + uint32_t row = 1; + for (int n = 1; n < 32; ++n) { + odd[n] = row; + row <<= 1; + } + gf2_matrix_square(even, odd); // even = two zero bits + gf2_matrix_square(odd, even); // odd = four zero bits + + size_t len = bytes; + do { + gf2_matrix_square(even, odd); // first pass: even = one zero byte (8 bits) + if (len & 1) { + crc = gf2_matrix_times(even, crc); + } + len >>= 1; + if (len == 0) { + break; + } + gf2_matrix_square(odd, even); + if (len & 1) { + crc = gf2_matrix_times(odd, crc); + } + len >>= 1; + } while (len != 0); + return crc; +} + +// 3-way interleaved hardware CRC32C (rocksdb/folly/Intel style). Splits the buffer +// into three equal, 8-byte-aligned lanes processed by independent _mm_crc32_u64 +// accumulators (breaking the ~3-cycle loop-carried dependency of the serial path), +// then stitches them with the GF(2) shift-combine and finishes the remainder +// serially. Byte-identical to crc32c_hw_serial / crc32c_slice8. Below the +// threshold it defers to the serial path so small buffers skip the combine cost. +uint32_t crc32c_hw3(uint32_t crc, const uint8_t* p, size_t n) { + if (n < kInterleaveThreshold) { + return crc32c_hw_serial(crc, p, n); + } + // 8-byte-aligned lane length keeps every lane on the u64 fast path. n >= 1024 + // guarantees L >= 336 > 0, so 3L <= n and the tail (n - 3L) is well defined. + const size_t lane = (n / 3) & ~static_cast(7); + const uint32_t crc_a = crc32c_hw_serial(crc, p, lane); // seeded lane + const uint32_t crc_b = crc32c_hw_serial(0, p + lane, lane); // raw lane + const uint32_t crc_c = crc32c_hw_serial(0, p + 2 * lane, lane); // raw lane + // crc(seed, A||B) == shift(crc(seed, A), |B|) ^ crc(0, B), applied twice. + uint32_t comb = crc32c_shift(crc_a, lane) ^ crc_b; + comb = crc32c_shift(comb, lane) ^ crc_c; + return crc32c_hw_serial(comb, p + 3 * lane, n - 3 * lane); // serial tail +} + +bool detect_sse42() { + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; + if (!__get_cpuid(1, &eax, &ebx, &ecx, &edx)) { + return false; + } + return (ecx & bit_SSE4_2) != 0; +} + +const bool kHasSse42 = detect_sse42(); +#endif // SNII_CRC32C_X86 + +} // namespace + +namespace detail { + +// Portable software path. Always available; the canonical scalar reference for +// every other path and for the on-disk checksum value. +uint32_t crc32c_slice8_extend(uint32_t crc, Slice data) { + crc = ~crc; + crc = crc32c_slice8(crc, data.data(), data.size()); + return ~crc; +} + +// Serial hardware path (falls back to slice8 without SSE4.2). +uint32_t crc32c_hw_serial_extend(uint32_t crc, Slice data) { + crc = ~crc; +#if SNII_CRC32C_X86 + if (kHasSse42) { + crc = crc32c_hw_serial(crc, data.data(), data.size()); + return ~crc; + } +#endif + crc = crc32c_slice8(crc, data.data(), data.size()); + return ~crc; +} + +// 3-way interleaved hardware path (falls back to slice8 without SSE4.2; internally +// falls back to the serial path below the interleave threshold). +uint32_t crc32c_hw3_extend(uint32_t crc, Slice data) { + crc = ~crc; +#if SNII_CRC32C_X86 + if (kHasSse42) { + crc = crc32c_hw3(crc, data.data(), data.size()); + return ~crc; + } +#endif + crc = crc32c_slice8(crc, data.data(), data.size()); + return ~crc; +} + +size_t crc32c_interleave_threshold() { + return kInterleaveThreshold; +} + +bool crc32c_has_hw() { +#if SNII_CRC32C_X86 + return kHasSse42; +#else + return false; +#endif +} + +} // namespace detail +} // namespace doris::snii + +#endif // BE_TEST diff --git a/be/src/storage/index/snii/encoding/crc32c.h b/be/src/storage/index/snii/encoding/crc32c.h new file mode 100644 index 00000000000000..775a781c3d39f0 --- /dev/null +++ b/be/src/storage/index/snii/encoding/crc32c.h @@ -0,0 +1,72 @@ +// 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. + +#pragma once + +#include + +#include +#include + +#include "storage/index/snii/common/slice.h" + +namespace doris::snii { + +// CRC32C (Castagnoli, polynomial 0x1EDC6F41). Used to checksum the tail of each +// format block. Thin inline adapter over Doris's bundled Google crc32c thirdparty +// (crc32c::Extend / crc32c::Crc32c). That library computes the same canonical +// CRC32C (same reflected polynomial, same standard pre/post inversion), so every +// on-disk checksum stays byte-identical to the previous in-tree slice-by-8 / +// SSE4.2 implementation -- this is an implementation swap, not a format change. +// The leading :: keeps the crc32c namespace distinct from crc32c() below. +inline uint32_t crc32c_extend(uint32_t crc, Slice data) { + return ::crc32c::Extend(crc, data.data(), data.size()); +} + +inline uint32_t crc32c(Slice data) { + return ::crc32c::Crc32c(data.data(), data.size()); +} + +#ifdef BE_TEST +// T21 test seam. The production crc32c()/crc32c_extend() above delegate to the +// bundled Google crc32c thirdparty (see commit d0416bb4129), which already runs a +// runtime-dispatched, hardware-accelerated and interleaved CRC32C -- so T21's +// "hardware interleaved CRC" goal is already met (and exceeded: that library adds +// a PCLMULQDQ fold a hand-rolled 3-way _mm_crc32_u64 lacks). Rather than regress +// that reuse, the reference sub-paths below let unit tests prove, byte-for-byte +// across all sizes/alignments, that the production path equals the canonical +// CRC32C and that the hardware path is engaged: +// * crc32c_slice8_extend -- portable software slice-by-8 (always available); +// * crc32c_hw_serial_extend -- serial SSE4.2 _mm_crc32 hardware path; +// * crc32c_hw3_extend -- 3-way interleaved SSE4.2 hardware path with a +// GF(2) shift-combine and a 1024-byte fall-back to +// the serial path (the algorithm T21 specifies). +// hw_serial/hw3 fall back to slice8 when SSE4.2 is absent. Each *_extend applies +// the standard ~crc pre/post inversion, so *_extend(0, d) == crc32c(d). The whole +// seam plus its static slice-by-8 table and startup CPUID probe are compiled out +// of release builds by this BE_TEST gate, so production carries no extra code. +// Pure functions with no shared mutable state (CONCURRENCY: N/A). +namespace detail { +uint32_t crc32c_slice8_extend(uint32_t crc, Slice data); +uint32_t crc32c_hw_serial_extend(uint32_t crc, Slice data); +uint32_t crc32c_hw3_extend(uint32_t crc, Slice data); +size_t crc32c_interleave_threshold(); +bool crc32c_has_hw(); +} // namespace detail +#endif // BE_TEST + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/pfor.cpp b/be/src/storage/index/snii/encoding/pfor.cpp new file mode 100644 index 00000000000000..dcc620fcc54c8f --- /dev/null +++ b/be/src/storage/index/snii/encoding/pfor.cpp @@ -0,0 +1,452 @@ +// 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 "storage/index/snii/encoding/pfor.h" + +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" + +namespace doris::snii { +namespace { + +// Unaligned little-endian 64-bit load from a raw byte pointer (single +// instruction on x86; memcpy is the portable, UB-free spelling the compiler +// folds to a mov). +inline uint64_t load_u64_le(const uint8_t* p) { + uint64_t v; + std::memcpy(&v, p, sizeof(v)); +#if defined(__BIG_ENDIAN__) || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) + v = __builtin_bswap64(v); +#endif + return v; +} + +// TEST-ONLY seam backing doris::snii::testing::pfor_width_evals(). value_width() +// is the SINGLE per-value bit-width evaluation point on the encode path, so this +// counter equals the number of values processed per run -- the deterministic +// signal that the histogram path scans each value exactly once (vs the former +// O(maxw*n) re-scan). Compiled out entirely in non-test builds so production pays +// nothing; a relaxed atomic under BE_TEST keeps it data-race-free under TSAN. +#ifdef BE_TEST +std::atomic g_width_evals {0}; +#endif + +// Number of significant bits of v (its minimal bit_width); value_width(0) == 0 is +// preserved (matching the former bits_for). clz(0) is undefined behaviour, so +// v == 0 is mapped explicitly here -- never via clz(v | 1), which would mis-score +// 0 as width 1 and corrupt the histogram. +inline uint8_t value_width(uint32_t v) { +#ifdef BE_TEST + g_width_evals.fetch_add(1, std::memory_order_relaxed); +#endif + return v ? static_cast(32 - __builtin_clz(v)) : 0; +} + +// Choose the bit_width that minimizes total bytes (packed + exceptions), with +// exception cost estimated at ~6 bytes each. A single O(n) pass builds a bit-width +// histogram (recording each value's width into widths[] for the encoder to reuse), +// then an O(maxw) suffix-sum gives the exception count per candidate width -- +// replacing the former O(maxw*n) re-scan. The cost formula and the ascending-w +// strict-'<' tie-break are kept identical, so the chosen width (and hence every +// encoded byte) is unchanged. +uint8_t choose_width(const uint32_t* v, size_t n, uint8_t* widths) { + // hist[b] = #values whose bit-width == b. n is capped at kFrqBaseUnit (256) on + // the production path; uint32_t buckets keep a direct caller with a larger run + // correct without affecting the chosen width. + uint32_t hist[33] = {0}; + uint8_t maxw = 0; + for (size_t i = 0; i < n; ++i) { + const uint8_t b = value_width(v[i]); + widths[i] = b; + ++hist[b]; + maxw = std::max(b, maxw); + } + // suffix[k] = #values whose bit-width >= k, so the exceptions for candidate + // width w (values needing more than w bits) are exactly suffix[w + 1]. + size_t suffix[34] = {0}; + for (int k = 32; k >= 0; --k) { + suffix[k] = suffix[k + 1] + hist[k]; + } + uint8_t best = maxw; + size_t best_cost = SIZE_MAX; + for (uint8_t w = 0; w <= maxw; ++w) { + const size_t exc = suffix[w + 1]; + const size_t cost = (static_cast(w) * n + 7) / 8 + exc * 6; + if (cost < best_cost) { + best_cost = cost; + best = w; + } + } + return best; +} + +uint32_t low_mask(uint8_t w) { + return (w >= 32) ? 0xFFFFFFFFU : ((1U << w) - 1U); +} + +// Bit-pack the low w bits of each value, writing 0 at exception positions +// (widths[i] > w) instead of the value's low bits. This is byte-identical to +// packing a copy in which those slots were pre-zeroed (the former `low[i] = 0` +// placeholder), so the encoder no longer materializes that copy. +void bitpack_masked(const uint32_t* v, const uint8_t* widths, size_t n, uint8_t w, ByteSink* out) { + if (w == 0) { + return; + } + // Pre-size for the exact packed byte count (ceil(w*n/8)) so the per-byte put_u8 + // loop below never reallocates mid-pack. ByteSink::reserve is RELATIVE to out's + // current size, and `out` is reused across the PFOR runs of one region, so this + // accumulates run-to-run instead of no-op'ing. Capacity-only: the packed bytes + // and their values are byte-identical to the un-reserved path. + const size_t packed = (static_cast(w) * n + 7) / 8; + out->reserve(packed); + const uint32_t mask = low_mask(w); + uint64_t acc = 0; + int filled = 0; + for (size_t i = 0; i < n; ++i) { + const uint32_t lo = (widths[i] > w) ? 0U : (v[i] & mask); + acc |= static_cast(lo) << filled; + filled += w; + while (filled >= 8) { + out->put_u8(static_cast(acc)); + acc >>= 8; + filled -= 8; + } + } + if (filled > 0) { + out->put_u8(static_cast(acc)); + } +} + +void bitunpack_tail(const uint8_t* base, size_t packed, size_t n, uint8_t w, size_t i, + uint64_t mask, uint32_t* out) { + for (; i < n; ++i) { + const size_t bit_off = static_cast(w) * i; + const size_t byte_off = bit_off >> 3; + uint64_t word = 0; + for (size_t b = byte_off; b < packed && b < byte_off + 8; ++b) { + word |= static_cast(base[b]) << ((b - byte_off) * 8); + } + out[i] = static_cast((word >> (bit_off & 7)) & mask); + } +} + +void bitunpack_w1(const uint8_t* base, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 8 <= n; i += 8, ++byte) { + const uint8_t v = base[byte]; + out[i] = v & 1U; + out[i + 1] = (v >> 1) & 1U; + out[i + 2] = (v >> 2) & 1U; + out[i + 3] = (v >> 3) & 1U; + out[i + 4] = (v >> 4) & 1U; + out[i + 5] = (v >> 5) & 1U; + out[i + 6] = (v >> 6) & 1U; + out[i + 7] = (v >> 7) & 1U; + } + if (i < n) { + const uint8_t v = base[byte]; + for (uint8_t bit = 0; i < n; ++i, ++bit) { + out[i] = (v >> bit) & 1U; + } + } +} + +void bitunpack_w2(const uint8_t* base, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 4 <= n; i += 4, ++byte) { + const uint8_t v = base[byte]; + out[i] = v & 3U; + out[i + 1] = (v >> 2) & 3U; + out[i + 2] = (v >> 4) & 3U; + out[i + 3] = (v >> 6) & 3U; + } + if (i < n) { + const uint8_t v = base[byte]; + for (uint8_t shift = 0; i < n; ++i, shift += 2) { + out[i] = (v >> shift) & 3U; + } + } +} + +void bitunpack_w3(const uint8_t* base, size_t packed, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 8 <= n; i += 8, byte += 3) { + const uint32_t b0 = base[byte]; + const uint32_t b1 = base[byte + 1]; + const uint32_t b2 = base[byte + 2]; + out[i] = b0 & 7U; + out[i + 1] = (b0 >> 3) & 7U; + out[i + 2] = ((b0 >> 6) | (b1 << 2)) & 7U; + out[i + 3] = (b1 >> 1) & 7U; + out[i + 4] = (b1 >> 4) & 7U; + out[i + 5] = ((b1 >> 7) | (b2 << 1)) & 7U; + out[i + 6] = (b2 >> 2) & 7U; + out[i + 7] = (b2 >> 5) & 7U; + } + bitunpack_tail(base, packed, n, 3, i, 7U, out); +} + +void bitunpack_w4(const uint8_t* base, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 2 <= n; i += 2, ++byte) { + const uint8_t v = base[byte]; + out[i] = v & 15U; + out[i + 1] = (v >> 4) & 15U; + } + if (i < n) { + out[i] = base[byte] & 15U; + } +} + +void bitunpack_w5(const uint8_t* base, size_t packed, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 8 <= n; i += 8, byte += 5) { + const uint32_t b0 = base[byte]; + const uint32_t b1 = base[byte + 1]; + const uint32_t b2 = base[byte + 2]; + const uint32_t b3 = base[byte + 3]; + const uint32_t b4 = base[byte + 4]; + out[i] = b0 & 31U; + out[i + 1] = ((b0 >> 5) | (b1 << 3)) & 31U; + out[i + 2] = (b1 >> 2) & 31U; + out[i + 3] = ((b1 >> 7) | (b2 << 1)) & 31U; + out[i + 4] = ((b2 >> 4) | (b3 << 4)) & 31U; + out[i + 5] = (b3 >> 1) & 31U; + out[i + 6] = ((b3 >> 6) | (b4 << 2)) & 31U; + out[i + 7] = (b4 >> 3) & 31U; + } + bitunpack_tail(base, packed, n, 5, i, 31U, out); +} + +void bitunpack_w6(const uint8_t* base, size_t packed, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 4 <= n; i += 4, byte += 3) { + const uint32_t b0 = base[byte]; + const uint32_t b1 = base[byte + 1]; + const uint32_t b2 = base[byte + 2]; + out[i] = b0 & 63U; + out[i + 1] = ((b0 >> 6) | (b1 << 2)) & 63U; + out[i + 2] = ((b1 >> 4) | (b2 << 4)) & 63U; + out[i + 3] = (b2 >> 2) & 63U; + } + bitunpack_tail(base, packed, n, 6, i, 63U, out); +} + +void bitunpack_w7(const uint8_t* base, size_t packed, size_t n, uint32_t* out) { + size_t i = 0; + size_t byte = 0; + for (; i + 8 <= n; i += 8, byte += 7) { + const uint32_t b0 = base[byte]; + const uint32_t b1 = base[byte + 1]; + const uint32_t b2 = base[byte + 2]; + const uint32_t b3 = base[byte + 3]; + const uint32_t b4 = base[byte + 4]; + const uint32_t b5 = base[byte + 5]; + const uint32_t b6 = base[byte + 6]; + out[i] = b0 & 127U; + out[i + 1] = ((b0 >> 7) | (b1 << 1)) & 127U; + out[i + 2] = ((b1 >> 6) | (b2 << 2)) & 127U; + out[i + 3] = ((b2 >> 5) | (b3 << 3)) & 127U; + out[i + 4] = ((b3 >> 4) | (b4 << 4)) & 127U; + out[i + 5] = ((b4 >> 3) | (b5 << 5)) & 127U; + out[i + 6] = ((b5 >> 2) | (b6 << 6)) & 127U; + out[i + 7] = (b6 >> 1) & 127U; + } + bitunpack_tail(base, packed, n, 7, i, 127U, out); +} + +void bitunpack_w8(const uint8_t* base, size_t n, uint32_t* out) { + for (size_t i = 0; i < n; ++i) { + out[i] = base[i]; + } +} + +void bitunpack_generic(const uint8_t* base, size_t packed, size_t n, uint8_t w, uint32_t* out) { + const uint64_t mask = low_mask(w); + size_t i = 0; + if (packed >= 8) { + const size_t last_safe_byte = packed - 8; + for (; i < n; ++i) { + const size_t bit_off = static_cast(w) * i; + const size_t byte_off = bit_off >> 3; + if (byte_off > last_safe_byte) { + break; + } + out[i] = static_cast((load_u64_le(base + byte_off) >> (bit_off & 7)) & mask); + } + } + bitunpack_tail(base, packed, n, w, i, mask, out); +} + +Status bitunpack(ByteSource* src, size_t n, uint8_t w, uint32_t* out) { + if (w == 0) { + std::memset(out, 0, n * sizeof(uint32_t)); + return Status::OK(); + } + // Pull the packed run once and unpack from the contiguous slice; this keeps + // the hot decode path free of per-byte ByteSource calls. + const size_t packed = (static_cast(w) * n + 7) / 8; + Slice buf; + RETURN_IF_ERROR(src->get_bytes(packed, &buf)); + const uint8_t* base = buf.data(); + + switch (w) { + case 1: + bitunpack_w1(base, n, out); + break; + case 2: + bitunpack_w2(base, n, out); + break; + case 3: + bitunpack_w3(base, packed, n, out); + break; + case 4: + bitunpack_w4(base, n, out); + break; + case 5: + bitunpack_w5(base, packed, n, out); + break; + case 6: + bitunpack_w6(base, packed, n, out); + break; + case 7: + bitunpack_w7(base, packed, n, out); + break; + case 8: + bitunpack_w8(base, n, out); + break; + default: + bitunpack_generic(base, packed, n, w, out); + break; + } + return Status::OK(); +} + +} // namespace + +namespace testing { +// Test-only op-count seam; see pfor.h. Reports/resets the per-value bit-width +// evaluation counter, and is a no-op in non-test builds where the counter is +// compiled out. +uint64_t pfor_width_evals() { +#ifdef BE_TEST + return g_width_evals.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_pfor_width_evals() { +#ifdef BE_TEST + g_width_evals.store(0, std::memory_order_relaxed); +#endif +} +} // namespace testing + +void pfor_encode(const uint32_t* values, size_t n, ByteSink* out) { + // n is hard-capped at kFrqBaseUnit (256) by encode_pfor_runs, so the stack + // buffer is used on every production call; the heap fallback only guards a + // direct caller passing a larger run. + uint8_t widths_stack[256]; + std::vector widths_heap; + uint8_t* widths = widths_stack; + if (n > sizeof(widths_stack)) { + widths_heap.resize(n); + widths = widths_heap.data(); + } + + const uint8_t w = choose_width(values, n, widths); + out->put_u8(w); + + // Count exceptions for the varint header by reusing the cached per-value + // widths -- no further bit-width evaluation. + uint32_t n_exc = 0; + for (size_t i = 0; i < n; ++i) { + n_exc += (widths[i] > w); + } + out->put_varint32(n_exc); + + // Pack the low w bits, writing 0 at exception slots (byte-identical to the + // former zeroed-copy approach) so no separate `low` buffer is materialized. + bitpack_masked(values, widths, n, w, out); + + // Exception table: (index_delta, full_value) in ascending index order. + uint32_t prev = 0; + for (size_t i = 0; i < n; ++i) { + if (widths[i] > w) { + out->put_varint32(static_cast(i) - prev); + out->put_varint32(values[i]); + prev = static_cast(i); + } + } +} + +Status pfor_decode(ByteSource* src, size_t n, uint32_t* out) { + uint8_t w; + RETURN_IF_ERROR(src->get_u8(&w)); + uint32_t n_exc; + RETURN_IF_ERROR(src->get_varint32(&n_exc)); + RETURN_IF_ERROR(bitunpack(src, n, w, out)); + uint32_t idx = 0; + for (uint32_t i = 0; i < n_exc; ++i) { + uint32_t d, val; + RETURN_IF_ERROR(src->get_varint32(&d)); + RETURN_IF_ERROR(src->get_varint32(&val)); + idx += d; + if (idx >= n) { + return Status::Error( + "pfor exception index out of range"); + } + out[idx] = val; + } + return Status::OK(); +} + +Status pfor_skip(ByteSource* src, size_t n) { + uint8_t w = 0; + RETURN_IF_ERROR(src->get_u8(&w)); + uint32_t n_exc = 0; + RETURN_IF_ERROR(src->get_varint32(&n_exc)); + const size_t packed = (static_cast(w) * n + 7) / 8; + Slice unused; + RETURN_IF_ERROR(src->get_bytes(packed, &unused)); + uint32_t idx = 0; + for (uint32_t i = 0; i < n_exc; ++i) { + uint32_t d = 0; + uint32_t val = 0; + RETURN_IF_ERROR(src->get_varint32(&d)); + RETURN_IF_ERROR(src->get_varint32(&val)); + idx += d; + if (idx >= n) { + return Status::Error( + "pfor exception index out of range"); + } + } + return Status::OK(); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/pfor.h b/be/src/storage/index/snii/encoding/pfor.h new file mode 100644 index 00000000000000..709d35748249f1 --- /dev/null +++ b/be/src/storage/index/snii/encoding/pfor.h @@ -0,0 +1,53 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" + +namespace doris::snii { + +// PFOR integer block encoder/decoder (unsigned uint32 array). +// Encoded layout: [u8 bit_width][varint n_exceptions][bit-packed low +// bits][exception table]. Selects the bit_width that minimizes total byte size; +// values exceeding it go into the exception table (index_delta, full_value). +// delta/zigzag is handled by the upper layer (.frq window); PFOR only processes +// unsigned integer arrays. +void pfor_encode(const uint32_t* values, size_t n, ByteSink* out); +Status pfor_decode(ByteSource* src, size_t n, uint32_t* out); +Status pfor_skip(ByteSource* src, size_t n); + +} // namespace doris::snii + +// Test-only instrumentation seam (mirrors the dict-block decode-counter pattern). +// pfor_width_evals() returns a process-global count of per-value bit-width +// evaluations performed by pfor_encode since the last reset -- one per +// value_width() call, the single evaluation point on the encode path. Deterministic +// perf tests assert it equals the number of encoded values per run, proving the +// histogram path scans each value exactly once (vs the former O(maxw*n) re-scan). +// Compiled to a no-op in non-test builds; reset between tests. +namespace doris::snii::testing { + +uint64_t pfor_width_evals(); +void reset_pfor_width_evals(); + +} // namespace doris::snii::testing diff --git a/be/src/storage/index/snii/encoding/section_framer.cpp b/be/src/storage/index/snii/encoding/section_framer.cpp new file mode 100644 index 00000000000000..e67b618758dae8 --- /dev/null +++ b/be/src/storage/index/snii/encoding/section_framer.cpp @@ -0,0 +1,59 @@ +// 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 "storage/index/snii/encoding/section_framer.h" + +#include "storage/index/snii/encoding/crc32c.h" + +namespace doris::snii { + +void SectionFramer::write(ByteSink& sink, uint8_t section_type, Slice payload) { + // Single-copy framing: write [type][varint64 len][payload] straight into the + // target sink, then crc exactly those bytes. view() is taken AFTER the payload + // and BEFORE the crc, so subslice([start, framed_len)) is over a settled, + // contiguous buffer with no pending realloc/aliasing. Byte-identical to the + // former temp-ByteSink assembly, minus one heap alloc + one payload copy. + const size_t start = sink.size(); + sink.put_u8(section_type); + sink.put_varint64(payload.size()); + sink.put_bytes(payload); + const size_t framed_len = sink.size() - start; + const uint32_t crc = crc32c(sink.view().subslice(start, framed_len)); + sink.put_fixed32(crc); +} + +Status SectionFramer::read(ByteSource& src, FramedSection* out) { + size_t start = src.position(); + uint8_t type; + RETURN_IF_ERROR(src.get_u8(&type)); + uint64_t len; + RETURN_IF_ERROR(src.get_varint64(&len)); + Slice payload; + RETURN_IF_ERROR(src.get_bytes(static_cast(len), &payload)); + size_t framed_len = src.position() - start; + uint32_t stored; + RETURN_IF_ERROR(src.get_fixed32(&stored)); + if (crc32c(src.slice_from(start, framed_len)) != stored) { + return Status::Error( + "section crc mismatch"); + } + out->type = type; + out->payload = payload; + return Status::OK(); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/section_framer.h b/be/src/storage/index/snii/encoding/section_framer.h new file mode 100644 index 00000000000000..9a248381ea97b6 --- /dev/null +++ b/be/src/storage/index/snii/encoding/section_framer.h @@ -0,0 +1,44 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" + +namespace doris::snii { + +// A framed section: type + payload view. +struct FramedSection { + uint8_t type = 0; + Slice payload; +}; + +// Unified section framing: [u8 type][varint64 len][payload][fixed32 crc32c(type+len+payload)]. +// All full-format sections reuse this encode/checksum path to avoid ad-hoc hand-assembly. +// Unknown optional sections are dispatched by the caller based on type; read still verifies the CRC and skips the payload. +class SectionFramer { +public: + static void write(ByteSink& sink, uint8_t section_type, Slice payload); + static Status read(ByteSource& src, FramedSection* out); +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/varint.cpp b/be/src/storage/index/snii/encoding/varint.cpp new file mode 100644 index 00000000000000..a53a08de6b8d3b --- /dev/null +++ b/be/src/storage/index/snii/encoding/varint.cpp @@ -0,0 +1,73 @@ +// 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 "storage/index/snii/encoding/varint.h" + +namespace doris::snii { + +size_t varint_len(uint64_t v) { + size_t n = 1; + while (v >= 0x80) { + v >>= 7; + ++n; + } + return n; +} + +size_t encode_varint64(uint64_t v, uint8_t* out) { + size_t i = 0; + while (v >= 0x80) { + out[i++] = static_cast(v) | 0x80; + v >>= 7; + } + out[i++] = static_cast(v); + return i; +} + +size_t encode_varint32(uint32_t v, uint8_t* out) { + return encode_varint64(v, out); +} + +Status decode_varint64(const uint8_t* p, const uint8_t* end, uint64_t* v, const uint8_t** next) { + uint64_t result = 0; + int shift = 0; + while (p < end) { + uint8_t b = *p++; + result |= static_cast(b & 0x7F) << shift; + if ((b & 0x80) == 0) { + *v = result; + *next = p; + return Status::OK(); + } + shift += 7; + if (shift >= 64) + return Status::Error( + "varint64 overflow"); + } + return Status::Error("varint truncated"); +} + +Status decode_varint32(const uint8_t* p, const uint8_t* end, uint32_t* v, const uint8_t** next) { + uint64_t tmp; + RETURN_IF_ERROR(decode_varint64(p, end, &tmp, next)); + if (tmp > 0xFFFFFFFFu) + return Status::Error("varint32 overflow"); + *v = static_cast(tmp); + return Status::OK(); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/varint.h b/be/src/storage/index/snii/encoding/varint.h new file mode 100644 index 00000000000000..978812ea6e66f6 --- /dev/null +++ b/be/src/storage/index/snii/encoding/varint.h @@ -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. + +#pragma once + +#include +#include + +#include "common/status.h" + +namespace doris::snii { + +// LEB128 variable-length integer encoding + zigzag. out buffer must be >=10 bytes; returns number of bytes written. +size_t varint_len(uint64_t v); +size_t encode_varint32(uint32_t v, uint8_t* out); +size_t encode_varint64(uint64_t v, uint8_t* out); + +// Decode a varint from the range [p, end); on success *next points to the next byte after the consumed input. +Status decode_varint32(const uint8_t* p, const uint8_t* end, uint32_t* v, const uint8_t** next); +Status decode_varint64(const uint8_t* p, const uint8_t* end, uint64_t* v, const uint8_t** next); + +inline uint64_t zigzag_encode(int64_t v) { + return (static_cast(v) << 1) ^ static_cast(v >> 63); +} +inline int64_t zigzag_decode(uint64_t v) { + return static_cast(v >> 1) ^ -static_cast(v & 1); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/zstd_codec.cpp b/be/src/storage/index/snii/encoding/zstd_codec.cpp new file mode 100644 index 00000000000000..6c01c0cb2ebdef --- /dev/null +++ b/be/src/storage/index/snii/encoding/zstd_codec.cpp @@ -0,0 +1,55 @@ +// 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 "storage/index/snii/encoding/zstd_codec.h" + +#include + +#include + +#include "storage/index/snii/common/uninitialized_buffer.h" + +namespace doris::snii { + +Status zstd_compress(Slice input, int level, std::vector* out) { + size_t bound = ZSTD_compressBound(input.size()); + out->resize(bound); + size_t n = ZSTD_compress(out->data(), bound, input.data(), input.size(), level); + if (ZSTD_isError(n)) { + return Status::Error(std::string("zstd compress: ") + + ZSTD_getErrorName(n)); + } + out->resize(n); + return Status::OK(); +} + +Status zstd_decompress(Slice input, size_t expected_uncomp_len, std::vector* out) { + // Sized then fully overwritten by ZSTD_decompress (length-checked below). + resize_uninitialized(*out, expected_uncomp_len); + size_t n = ZSTD_decompress(out->data(), expected_uncomp_len, input.data(), input.size()); + if (ZSTD_isError(n)) { + return Status::Error( + std::string("zstd decompress: ") + ZSTD_getErrorName(n)); + } + if (n != expected_uncomp_len) { + return Status::Error( + "zstd decompressed length mismatch"); + } + return Status::OK(); +} + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/encoding/zstd_codec.h b/be/src/storage/index/snii/encoding/zstd_codec.h new file mode 100644 index 00000000000000..c8291ce024fca7 --- /dev/null +++ b/be/src/storage/index/snii/encoding/zstd_codec.h @@ -0,0 +1,33 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" + +namespace doris::snii { + +// Thin ZSTD wrapper. Used for compressing large payloads such as .prx windows. Decompression requires the caller to supply the original uncompressed length (from the block header). +Status zstd_compress(Slice input, int level, std::vector* out); +Status zstd_decompress(Slice input, size_t expected_uncomp_len, std::vector* out); + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/format/bootstrap_header.cpp b/be/src/storage/index/snii/format/bootstrap_header.cpp new file mode 100644 index 00000000000000..6cdc1869385bec --- /dev/null +++ b/be/src/storage/index/snii/format/bootstrap_header.cpp @@ -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. + +#include "storage/index/snii/format/bootstrap_header.h" + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" + +namespace doris::snii::format { + +namespace { + +// Number of bytes covered by header_checksum: everything except the trailing +// crc32c. +constexpr size_t kChecksumCoverage = kBootstrapHeaderSize - 4; + +// Writes all fixed fields except the trailing checksum. Field order is the +// on-disk contract; reuse ByteSink fixed-width primitives, never hand-assemble +// bytes. +void encode_fields(const BootstrapHeader& header, ByteSink* sink) { + sink->put_fixed32(header.magic); + sink->put_fixed32((static_cast(header.min_reader_version) << 16) | + header.format_version); + sink->put_fixed32(header.flags); + sink->put_fixed32(kBootstrapHeaderSize); // header_length is always derived + sink->put_u8(header.tail_pointer_size); +} + +} // namespace + +Status encode_bootstrap_header(const BootstrapHeader& header, ByteSink* sink) { + if (sink == nullptr) { + return Status::Error("bootstrap_header: null sink"); + } + ByteSink fields; + encode_fields(header, &fields); + const uint32_t checksum = crc32c(fields.view()); + sink->put_bytes(fields.view()); + sink->put_fixed32(checksum); + return Status::OK(); +} + +Status decode_bootstrap_header(Slice data, BootstrapHeader* out) { + if (out == nullptr) { + return Status::Error("bootstrap_header: null out"); + } + // Reject any size other than the exact fixed header: short input is + // truncation, longer input means stray trailing bytes the parser would + // otherwise ignore. + if (data.size() != kBootstrapHeaderSize) { + return Status::Error( + "bootstrap_header: wrong header size"); + } + + ByteSource src(data); + uint32_t magic = 0; + uint32_t version_pair = 0; + uint32_t flags = 0; + uint32_t header_length = 0; + uint8_t tail_pointer_size = 0; + uint32_t stored_checksum = 0; + RETURN_IF_ERROR(src.get_fixed32(&magic)); + RETURN_IF_ERROR(src.get_fixed32(&version_pair)); + RETURN_IF_ERROR(src.get_fixed32(&flags)); + RETURN_IF_ERROR(src.get_fixed32(&header_length)); + RETURN_IF_ERROR(src.get_u8(&tail_pointer_size)); + RETURN_IF_ERROR(src.get_fixed32(&stored_checksum)); + + if (magic != kContainerMagic) { + return Status::Error( + "bootstrap_header: bad container magic"); + } + const uint32_t computed = crc32c(data.subslice(0, kChecksumCoverage)); + if (computed != stored_checksum) { + return Status::Error( + "bootstrap_header: checksum mismatch"); + } + + const auto min_reader_version = static_cast((version_pair >> 16) & 0xFFFFu); + const auto format_version = static_cast(version_pair & 0xFFFFu); + if (format_version != kFormatVersion) { + return Status::Error( + "bootstrap_header: unsupported container format_version"); + } + if (min_reader_version > kFormatVersion) { + return Status::Error( + "bootstrap_header: container requires a newer reader version"); + } + + out->magic = magic; + out->format_version = format_version; + out->min_reader_version = min_reader_version; + out->flags = flags; + out->header_length = header_length; + out->tail_pointer_size = tail_pointer_size; + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/bootstrap_header.h b/be/src/storage/index/snii/format/bootstrap_header.h new file mode 100644 index 00000000000000..6a8c0e8083d3ac --- /dev/null +++ b/be/src/storage/index/snii/format/bootstrap_header.h @@ -0,0 +1,71 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +// Fixed container header at the very start of a {rowset_id}_{seg_id}.idx file. +// Identifies the SNII container and carries basic compatibility info so a +// reader can fail fast before touching any streamed section or the tail meta +// region. +// +// On-disk layout (all multi-byte fields little-endian, fixed width; NOT framed +// by SectionFramer because it must be parseable without prior knowledge of the +// file): +// u32 magic == kContainerMagic +// u16 format_version == kFormatVersion +// u16 min_reader_version readers with kFormatVersion < this MUST refuse to +// read u32 flags container-level feature flags u32 +// header_length total bytes of this header including the checksum u8 +// tail_pointer_size size of the fixed tail pointer at EOF (hint for the +// reader) u32 header_checksum crc32c over all preceding header bytes +struct BootstrapHeader { + uint32_t magic = kContainerMagic; + uint16_t format_version = kFormatVersion; + uint16_t min_reader_version = kMinReaderVersion; + uint32_t flags = 0; + uint32_t header_length = 0; + uint8_t tail_pointer_size = 0; +}; + +// Total fixed on-disk size of the header, including the trailing crc32c. +inline constexpr uint32_t kBootstrapHeaderSize = + 4 /*magic*/ + 2 /*format_version*/ + 2 /*min_reader_version*/ + 4 /*flags*/ + + 4 /*header_length*/ + 1 /*tail_pointer_size*/ + 4 /*header_checksum*/; + +// Serializes the header to sink: writes header_length = kBootstrapHeaderSize +// and appends a crc32c over all preceding bytes. The caller's header_length +// field is ignored on input (it is always derived). Returns OK. +Status encode_bootstrap_header(const BootstrapHeader& header, ByteSink* sink); + +// Parses and validates a bootstrap header from the front of data. +// - too short / trailing bytes beyond the fixed header -> kCorruption +// - magic != kContainerMagic -> kCorruption +// - checksum mismatch -> kCorruption +// - format_version != kFormatVersion -> kUnsupported +// - min_reader_version > kFormatVersion -> kUnsupported +Status decode_bootstrap_header(Slice data, BootstrapHeader* out); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/bsbf.cpp b/be/src/storage/index/snii/format/bsbf.cpp new file mode 100644 index 00000000000000..d214a99e745ae9 --- /dev/null +++ b/be/src/storage/index/snii/format/bsbf.cpp @@ -0,0 +1,257 @@ +// 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 "storage/index/snii/format/bsbf.h" + +#include + +#include "storage/index/snii/encoding/crc32c.h" + +#if defined(__x86_64__) || defined(_M_X64) +#include +#define SNII_BSBF_X86 1 +#endif + +#define XXH_INLINE_ALL +#include "xxhash.h" + +namespace doris::snii::format { + +const uint32_t kBsbfSalt[kBsbfBitsSetPerBlock] = {0x47b6137bU, 0x44974d91U, 0x8824ad5bU, + 0xa2b7289dU, 0x705495c7U, 0x2df1424bU, + 0x9efc4947U, 0x5c6bfb31U}; + +namespace { + +void store_le32(uint8_t* p, uint32_t v) { + p[0] = static_cast(v); + p[1] = static_cast(v >> 8); + p[2] = static_cast(v >> 16); + p[3] = static_cast(v >> 24); +} +uint32_t load_le32(const uint8_t* p) { + return static_cast(p[0]) | (static_cast(p[1]) << 8) | + (static_cast(p[2]) << 16) | (static_cast(p[3]) << 24); +} + +bool cpu_has_avx2() { +#if defined(SNII_BSBF_X86) + static const bool v = __builtin_cpu_supports("avx2"); + return v; +#else + return false; +#endif +} + +// --- scalar kernels --- +inline void masks_scalar(uint32_t key, uint32_t m[8]) { + for (int i = 0; i < 8; ++i) m[i] = 1u << ((key * kBsbfSalt[i]) >> 27); +} +bool block_contains_scalar(uint64_t hash, const uint8_t* block) { + const uint32_t* w = reinterpret_cast(block); // LE + uint32_t m[8]; + masks_scalar(static_cast(hash), m); + for (int i = 0; i < 8; ++i) + if ((load_le32(reinterpret_cast(w + i)) & m[i]) != m[i]) return false; + return true; +} +void insert_scalar(uint32_t* words, uint32_t block, uint32_t key) { + uint32_t m[8]; + masks_scalar(key, m); + for (int i = 0; i < 8; ++i) words[block * 8 + i] |= m[i]; +} +bool find_scalar(const uint32_t* words, uint32_t block, uint32_t key) { + uint32_t m[8]; + masks_scalar(key, m); + for (int i = 0; i < 8; ++i) + if ((words[block * 8 + i] & m[i]) != m[i]) return false; + return true; +} + +#if defined(SNII_BSBF_X86) +// --- AVX2 kernels: a 256-bit block is one YMM register --- +__attribute__((target("avx2"))) __m256i mask_avx2(uint32_t key) { + const __m256i salt = + _mm256_setr_epi32(static_cast(kBsbfSalt[0]), static_cast(kBsbfSalt[1]), + static_cast(kBsbfSalt[2]), static_cast(kBsbfSalt[3]), + static_cast(kBsbfSalt[4]), static_cast(kBsbfSalt[5]), + static_cast(kBsbfSalt[6]), static_cast(kBsbfSalt[7])); + const __m256i prod = _mm256_mullo_epi32(_mm256_set1_epi32(static_cast(key)), salt); + const __m256i shifts = _mm256_srli_epi32(prod, 27); // top 5 bits -> 0..31 + return _mm256_sllv_epi32(_mm256_set1_epi32(1), shifts); +} +__attribute__((target("avx2"))) bool block_contains_avx2(uint64_t hash, const uint8_t* block) { + const __m256i m = mask_avx2(static_cast(hash)); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(block)); + return _mm256_testc_si256(b, m) != 0; // (~b & m) == 0 -> b contains m +} +__attribute__((target("avx2"))) void insert_avx2(uint32_t* words, uint32_t block, uint32_t key) { + __m256i* p = reinterpret_cast<__m256i*>(words + block * 8); + _mm256_storeu_si256(p, _mm256_or_si256(_mm256_loadu_si256(p), mask_avx2(key))); +} +__attribute__((target("avx2"))) bool find_avx2(const uint32_t* words, uint32_t block, + uint32_t key) { + const __m256i m = mask_avx2(key); + const __m256i b = _mm256_loadu_si256(reinterpret_cast(words + block * 8)); + return _mm256_testc_si256(b, m) != 0; +} +#endif + +} // namespace + +uint64_t bsbf_hash(std::string_view term) { + return XXH64(term.data(), term.size(), /*seed=*/0); +} + +uint32_t bsbf_optimal_num_bytes(uint32_t ndv, double fpp) { + // Parquet OptimalNumOfBits, then >>3 for bytes. + const double m = -8.0 * ndv / std::log(1 - std::pow(fpp, 1.0 / 8)); + uint32_t num_bits; + if (m < 0 || m > static_cast(kBsbfMaxBytes) * 8) { + num_bits = kBsbfMaxBytes << 3; + } else { + num_bits = static_cast(m); + } + if (num_bits < (kBsbfMinBytes << 3)) num_bits = kBsbfMinBytes << 3; + // G16-g: round up to the 32-byte block only, NOT to the next power of 2. + // The block index is fastrange ((h >> 32) * num_blocks >> 32, bsbf.h), + // which is uniform for ANY block count -- the old power-of-2 rounding was + // a sizing convention, not an addressing requirement, and cost up to ~2x + // bitset bytes (e.g. a 31M-term segment: 40.3 MB optimal -> 64 MB rounded) + // while silently over-delivering on fpp. Exact sizing keeps the fpp at + // the kBsbfFpp design point. + const uint32_t block_bits = kBsbfBytesPerBlock << 3; + num_bits = (num_bits + block_bits - 1) / block_bits * block_bits; + if (num_bits > (kBsbfMaxBytes << 3)) num_bits = kBsbfMaxBytes << 3; + return num_bits >> 3; +} + +bool bsbf_block_contains(uint64_t hash, const uint8_t block[kBsbfBytesPerBlock]) { +#if defined(SNII_BSBF_X86) + if (cpu_has_avx2()) return block_contains_avx2(hash, block); +#endif + return block_contains_scalar(hash, block); +} + +Status BsbfBuilder::create(uint32_t ndv, double fpp, BsbfBuilder* out) { + if (out == nullptr) return Status::Error("bsbf: null out"); + if (!(fpp > 0.0 && fpp < 1.0)) + return Status::Error("bsbf: fpp out of (0,1)"); + if (ndv == 0) ndv = 1; + out->num_bytes_ = bsbf_optimal_num_bytes(ndv, fpp); + out->num_blocks_ = out->num_bytes_ / kBsbfBytesPerBlock; + out->ndv_ = ndv; + out->words_.assign(out->num_bytes_ / 4, 0u); + return Status::OK(); +} + +void BsbfBuilder::insert(uint64_t hash) { + const uint32_t block = bsbf_block_index(hash, num_blocks_); + const uint32_t key = static_cast(hash); +#if defined(SNII_BSBF_X86) + if (cpu_has_avx2()) { + insert_avx2(words_.data(), block, key); + return; + } +#endif + insert_scalar(words_.data(), block, key); +} + +bool BsbfBuilder::maybe_contains(uint64_t hash) const { + const uint32_t block = bsbf_block_index(hash, num_blocks_); + const uint32_t key = static_cast(hash); +#if defined(SNII_BSBF_X86) + if (cpu_has_avx2()) return find_avx2(words_.data(), block, key); +#endif + return find_scalar(words_.data(), block, key); +} + +Status BsbfBuilder::serialize(ByteSink* sink) const { + if (sink == nullptr) + return Status::Error("bsbf: null sink"); + if (num_bytes_ == 0) + return Status::Error("bsbf: not built"); + uint8_t hdr[kBsbfHeaderSize] = {0}; + hdr[0] = 'B'; + hdr[1] = 'S'; + hdr[2] = 'B'; + hdr[3] = 'F'; + hdr[4] = 1; // version + hdr[5] = 0; // hash strategy: XXH64 seed 0 + hdr[6] = 0; // index strategy: fastrange + hdr[7] = 0; // pad + store_le32(hdr + 8, num_bytes_); + store_le32(hdr + 12, num_blocks_); + store_le32(hdr + 16, ndv_); + store_le32(hdr + 20, crc32c(Slice(hdr, 20))); // header crc over [0,20) + const uint8_t* bits = reinterpret_cast(words_.data()); + store_le32(hdr + 24, crc32c(Slice(bits, num_bytes_))); // bitset crc + sink->put_bytes(Slice(hdr, kBsbfHeaderSize)); + sink->put_bytes(Slice(bits, num_bytes_)); // contiguous, uncompressed, LE + return Status::OK(); +} + +Status BsbfHeader::parse(Slice h, uint64_t section_base, BsbfHeader* out) { + if (out == nullptr) return Status::Error("bsbf: null out"); + if (h.size() < kBsbfHeaderSize) + return Status::Error("bsbf: short header"); + const uint8_t* p = h.data(); + if (p[0] != 'B' || p[1] != 'S' || p[2] != 'B' || p[3] != 'F') + return Status::Error("bsbf: bad magic"); + if (p[4] != 1) + return Status::Error("bsbf: bad version"); + if (p[5] != 0) + return Status::Error( + "bsbf: unsupported hash strategy"); + if (p[6] != 0) + return Status::Error( + "bsbf: unsupported index strategy"); + if (crc32c(Slice(p, 20)) != load_le32(p + 20)) + return Status::Error( + "bsbf: header crc mismatch"); + const uint32_t nb = load_le32(p + 8); + const uint32_t nblk = load_le32(p + 12); + // G16-g: num_bytes only needs 32-byte block alignment -- the fastrange + // block index is uniform for any block count. Power-of-2 sizes (all + // pre-G16-g filters) remain valid as a subset. + if (nb < kBsbfMinBytes || nb > kBsbfMaxBytes || nb % kBsbfBytesPerBlock != 0) + return Status::Error( + "bsbf: num_bytes out of range or not block-aligned"); + if (nblk != nb / kBsbfBytesPerBlock) + return Status::Error( + "bsbf: num_blocks mismatch"); + out->num_bytes = nb; + out->num_blocks = nblk; + out->bitset_crc = load_le32(p + 24); + out->bitset_base = section_base + kBsbfHeaderSize; + return Status::OK(); +} + +Status bsbf_probe(io::FileReader* reader, const BsbfHeader& header, uint64_t hash, + bool* maybe_present) { + if (reader == nullptr || maybe_present == nullptr) + return Status::Error("bsbf: null arg"); + std::vector blk; + RETURN_IF_ERROR(reader->read_at(header.block_offset(hash), kBsbfBytesPerBlock, &blk)); + if (blk.size() < kBsbfBytesPerBlock) + return Status::Error( + "bsbf: short block read"); + *maybe_present = bsbf_block_contains(hash, blk.data()); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/bsbf.h b/be/src/storage/index/snii/format/bsbf.h new file mode 100644 index 00000000000000..7877073fedd10b --- /dev/null +++ b/be/src/storage/index/snii/format/bsbf.h @@ -0,0 +1,135 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/io/file_reader.h" + +// Block-split bloom filter (BSBF) -- Apache Parquet split-block spec, with an +// S3-native on-demand single-block probe that none of the reference implementations +// (Apache Parquet, Doris storage, Doris format/parquet) ship. +// +// BIT FORMAT IS PARQUET-CANONICAL (interoperable with Apache Parquet / Doris +// format/parquet for the bitset bytes): +// - 256-bit (32-byte) blocks, 8 bits set per block. +// - key = XXH64(term, seed=0); high 32 bits select the block via FASTRANGE +// `block = ((hash>>32) * num_blocks) >> 32` (no power-of-2 requirement); low 32 +// bits select 8 in-block positions `1 << ((key * SALT[i]) >> 27)`. +// - num_bytes via Parquet OptimalNumOfBytes: power of 2 in [32, 128 MiB]. +// +// SNII WRAPPER (NOT Parquet's variable thrift header): a FIXED 28-byte header, then +// the contiguous, uncompressed, little-endian bitset. Because the header size is a +// constant, the bitset start is a constant offset (`section_base + 28`) and block i +// is at `section_base + 28 + i*32` -- so a single 32-byte block can be range-read on +// demand WITHOUT parsing a variable-length header and WITHOUT loading the whole blob. +namespace doris::snii::format { + +constexpr uint32_t kBsbfBytesPerBlock = 32; // 256-bit block +constexpr uint32_t kBsbfBitsSetPerBlock = 8; // 8 uint32 words / block +constexpr uint32_t kBsbfMinBytes = 32; +constexpr uint32_t kBsbfMaxBytes = 128u * 1024 * 1024; // Parquet kMaximumBloomFilterBytes +constexpr uint32_t kBsbfHeaderSize = 28; // FIXED (constant bitset offset) +// L0/L1 tiering threshold (the "fast-reject absent terms" design): a bsbf section whose total +// size is <= this is loaded WHOLE into the resident reader at open (L0 -> free +// in-memory probe, no per-lookup round); larger filters stay L1 (header-only, probed +// one 32-byte block on demand). 256 KiB fits in a single cloud FileCache block. +constexpr uint32_t kBsbfResidentMaxBytes = 256u * 1024; + +// Canonical Parquet/Doris split-block SALT (8 odd 32-bit constants). +extern const uint32_t kBsbfSalt[kBsbfBitsSetPerBlock]; + +// XXH64(term, seed=0) -- the Parquet-canonical key (NOT XXH3, NOT Doris murmur). +uint64_t bsbf_hash(std::string_view term); + +// Parquet OptimalNumOfBytes(ndv, fpp): power of 2 in [32, 128 MiB]. +uint32_t bsbf_optimal_num_bytes(uint32_t ndv, double fpp); + +// Fastrange block index from a 64-bit hash and the block count. +inline uint32_t bsbf_block_index(uint64_t hash, uint32_t num_blocks) { + return static_cast(((hash >> 32) * num_blocks) >> 32); +} + +// Pure 32-byte-block kernel: does `block` contain the key's 8 bits? SIMD (AVX2) +// accelerated at runtime when available, scalar otherwise. Returns true => the term +// MAY be present (could be a false positive); false => DEFINITELY ABSENT. +bool bsbf_block_contains(uint64_t hash, const uint8_t block[kBsbfBytesPerBlock]); + +// In-memory builder + serializer. +class BsbfBuilder { +public: + BsbfBuilder() = default; + + // Sizes the filter for `ndv` distinct keys at target `fpp`. fpp in (0,1). + static Status create(uint32_t ndv, double fpp, BsbfBuilder* out); + + // Insert a key / term. SIMD-accelerated. + void insert(uint64_t hash); + void insert_term(std::string_view term) { insert(bsbf_hash(term)); } + + // In-memory probe over the resident bitset (build/warm path). SIMD-accelerated. + bool maybe_contains(uint64_t hash) const; + bool maybe_contains_term(std::string_view term) const { + return maybe_contains(bsbf_hash(term)); + } + + // Serialize [28-byte header][contiguous LE bitset] into `sink`. The header carries + // magic/version/hash+index strategy/num_bytes/num_blocks/ndv + header & bitset + // crc32c. The bitset is Parquet-canonical bytes. + Status serialize(ByteSink* sink) const; + + uint32_t num_bytes() const { return num_bytes_; } + uint32_t num_blocks() const { return num_blocks_; } + size_t resident_capacity_bytes() const { return words_.capacity() * sizeof(uint32_t); } + +private: + std::vector words_; // num_bytes_/4, blocks of 8 words + uint32_t num_bytes_ = 0; + uint32_t num_blocks_ = 0; + uint32_t ndv_ = 0; +}; + +// Resident header (28 bytes), parsed once at open. Validates magic/version/crc/bounds. +struct BsbfHeader { + uint32_t num_bytes = 0; + uint32_t num_blocks = 0; + uint32_t bitset_crc = 0; // stored crc32c of the bitset body (for L0 verification) + uint64_t bitset_base = 0; // absolute file offset of block 0 = section_base + 28 + + // Parse a 28-byte header located at `section_base` in the file. The bitset_base + // is set to section_base + kBsbfHeaderSize. + static Status parse(Slice header28, uint64_t section_base, BsbfHeader* out); + + // Absolute file offset of the 32-byte block this hash maps to. + uint64_t block_offset(uint64_t hash) const { + return bitset_base + + static_cast(bsbf_block_index(hash, num_blocks)) * kBsbfBytesPerBlock; + } +}; + +// On-demand probe: read EXACTLY ONE 32-byte block via `reader`, then test. No whole +// blob load, no deep copy. *maybe_present=false means DEFINITELY ABSENT. +Status bsbf_probe(io::FileReader* reader, const BsbfHeader& header, uint64_t hash, + bool* maybe_present); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/core_metadata.cpp b/be/src/storage/index/snii/format/core_metadata.cpp new file mode 100644 index 00000000000000..f9007cc96768b3 --- /dev/null +++ b/be/src/storage/index/snii/format/core_metadata.cpp @@ -0,0 +1,285 @@ +// 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 "storage/index/snii/format/core_metadata.h" + +#include +#include +#include +#include + +#include "gen_cpp/snii.pb.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" + +namespace doris::snii::format { +namespace { + +using segment_v2::inverted_index::CommonGramsCoverage; +using segment_v2::inverted_index::CommonGramsSegmentMetadata; +using segment_v2::inverted_index::PlainTermKeyVersion; +using segment_v2::inverted_index::ScoringCoverage; +using segment_v2::inverted_index::validate_common_grams_segment_metadata; +using segment_v2::inverted_index::validate_snii_scoring_metadata; + +Status corrupted(std::string_view message) { + return Status::Error(message); +} + +Status unsupported(std::string_view message) { + return Status::Error(message); +} + +Status validate_index_config(uint32_t value, IndexConfig* out) { + switch (value) { + case static_cast(IndexConfig::kDocsOnly): + case static_cast(IndexConfig::kDocsPositions): + case static_cast(IndexConfig::kDocsPositionsScoring): + *out = static_cast(value); + return Status::OK(); + default: + return unsupported("core metadata: unsupported index config"); + } +} + +Status validate_posting_policy(uint32_t value, CommonGramsPostingPolicy* out) { + switch (value) { + case 0: + *out = CommonGramsPostingPolicy::kNone; + return Status::OK(); + case 1: + *out = CommonGramsPostingPolicy::kHybridV1; + return Status::OK(); + default: + return unsupported("core metadata: unsupported CommonGrams posting policy"); + } +} + +Status validate_plain_term_key_version(uint32_t value) { + switch (value) { + case static_cast(PlainTermKeyVersion::kLegacyRaw): + case static_cast(PlainTermKeyVersion::kEscapedV1): + case static_cast(PlainTermKeyVersion::kRawNoInternal): + return Status::OK(); + default: + return unsupported("core metadata: unsupported plain-term key version"); + } +} + +Status validate_common_grams_coverage(uint32_t value) { + switch (value) { + case static_cast(CommonGramsCoverage::kNone): + case static_cast(CommonGramsCoverage::kComplete): + case static_cast(CommonGramsCoverage::kMixed): + return Status::OK(); + default: + return unsupported("core metadata: unsupported CommonGrams coverage"); + } +} + +Status validate_scoring_coverage(uint32_t value) { + switch (value) { + case static_cast(ScoringCoverage::kNone): + case static_cast(ScoringCoverage::kComplete): + return Status::OK(); + default: + return unsupported("core metadata: unsupported scoring coverage"); + } +} + +void encode_region_ref(const RegionRef& ref, doris::snii::SniiRegionRefPB* out) { + out->set_offset(ref.offset); + out->set_length(ref.length); +} + +void encode_common_grams(const CommonGramsSegmentMetadata& metadata, + doris::snii::SniiCommonGramsMetadataPB* out) { + out->set_plain_term_key_version(static_cast(metadata.plain_term_key_version)); + out->set_common_grams_coverage(static_cast(metadata.common_grams_coverage)); + out->set_common_grams_semantics_version(metadata.common_grams_semantics_version); + out->set_common_grams_key_version(metadata.common_grams_key_version); + out->set_common_grams_dictionary_identity(metadata.common_grams_dictionary_identity); + out->set_base_analyzer_fingerprint(metadata.base_analyzer_fingerprint); + out->set_common_grams_fingerprint(metadata.common_grams_fingerprint); + out->set_scoring_coverage(static_cast(metadata.scoring_coverage)); + out->set_scoring_stats_version(metadata.scoring_stats_version); + out->set_norm_semantics_version(metadata.norm_semantics_version); + out->set_scoring_doc_count(metadata.scoring_doc_count); + out->set_scoring_token_count(metadata.scoring_token_count); +} + +Status decode_region_ref(const doris::snii::SniiRegionRefPB& input, RegionRef* out) { + if (!input.has_offset() || !input.has_length()) { + return corrupted("core metadata: missing region reference field"); + } + *out = {.offset = input.offset(), .length = input.length()}; + return Status::OK(); +} + +Status decode_common_grams(const doris::snii::SniiCommonGramsMetadataPB& input, + CommonGramsSegmentMetadata* out) { + if (!input.has_plain_term_key_version() || !input.has_common_grams_coverage() || + !input.has_common_grams_semantics_version() || !input.has_common_grams_key_version() || + !input.has_common_grams_dictionary_identity() || !input.has_base_analyzer_fingerprint() || + !input.has_common_grams_fingerprint() || !input.has_scoring_coverage() || + !input.has_scoring_stats_version() || !input.has_norm_semantics_version() || + !input.has_scoring_doc_count() || !input.has_scoring_token_count()) { + return corrupted("core metadata: missing CommonGrams metadata field"); + } + RETURN_IF_ERROR(validate_plain_term_key_version(input.plain_term_key_version())); + RETURN_IF_ERROR(validate_common_grams_coverage(input.common_grams_coverage())); + RETURN_IF_ERROR(validate_scoring_coverage(input.scoring_coverage())); + *out = {.plain_term_key_version = + static_cast(input.plain_term_key_version()), + .common_grams_coverage = + static_cast(input.common_grams_coverage()), + .common_grams_semantics_version = input.common_grams_semantics_version(), + .common_grams_key_version = input.common_grams_key_version(), + .common_grams_dictionary_identity = input.common_grams_dictionary_identity(), + .base_analyzer_fingerprint = input.base_analyzer_fingerprint(), + .common_grams_fingerprint = input.common_grams_fingerprint(), + .scoring_coverage = static_cast(input.scoring_coverage()), + .scoring_stats_version = input.scoring_stats_version(), + .norm_semantics_version = input.norm_semantics_version(), + .scoring_doc_count = input.scoring_doc_count(), + .scoring_token_count = input.scoring_token_count()}; + return validate_common_grams_segment_metadata(*out); +} + +Status decode_core_pb(const doris::snii::SniiCoreMetadataPB& input, CoreMetadata* out) { + if (!input.has_index_config() || !input.has_stats() || !input.has_section_refs()) { + return corrupted("core metadata: missing required field"); + } + RETURN_IF_ERROR(validate_index_config(input.index_config(), &out->index_config)); + + const auto& stats = input.stats(); + if (!stats.has_doc_count() || !stats.has_indexed_doc_count() || !stats.has_term_count() || + !stats.has_sum_total_term_freq() || !stats.has_null_count()) { + return corrupted("core metadata: missing statistics field"); + } + out->stats = {.doc_count = stats.doc_count(), + .indexed_doc_count = stats.indexed_doc_count(), + .term_count = stats.term_count(), + .sum_total_term_freq = stats.sum_total_term_freq(), + .null_count = stats.null_count()}; + + const auto& refs = input.section_refs(); + if (!refs.has_dict_region() || !refs.has_posting_region() || !refs.has_norms() || + !refs.has_null_bitmap() || !refs.has_bsbf()) { + return corrupted("core metadata: missing section reference"); + } + RETURN_IF_ERROR(decode_region_ref(refs.dict_region(), &out->section_refs.dict_region)); + RETURN_IF_ERROR(decode_region_ref(refs.posting_region(), &out->section_refs.posting_region)); + RETURN_IF_ERROR(decode_region_ref(refs.norms(), &out->section_refs.norms)); + RETURN_IF_ERROR(decode_region_ref(refs.null_bitmap(), &out->section_refs.null_bitmap)); + RETURN_IF_ERROR(decode_region_ref(refs.bsbf(), &out->section_refs.bsbf)); + + if (input.has_common_grams()) { + CommonGramsSegmentMetadata common_grams; + RETURN_IF_ERROR(decode_common_grams(input.common_grams(), &common_grams)); + out->common_grams_metadata = std::move(common_grams); + } + + RETURN_IF_ERROR(validate_posting_policy(input.common_grams_posting_policy(), + &out->common_grams_posting_policy)); + if (out->common_grams_posting_policy == CommonGramsPostingPolicy::kHybridV1 && + (!out->common_grams_metadata.has_value() || + out->common_grams_metadata->common_grams_coverage != CommonGramsCoverage::kMixed)) { + return corrupted("core metadata: hybrid policy requires mixed CommonGrams metadata"); + } + const bool has_scoring_tier = out->index_config == IndexConfig::kDocsPositionsScoring; + if (has_scoring_tier) { + if (out->section_refs.norms.length == 0) { + return corrupted("core metadata: scoring index requires a norms region"); + } + } + if (has_scoring_tier || + (out->common_grams_metadata.has_value() && + out->common_grams_metadata->scoring_coverage == ScoringCoverage::kComplete)) { + RETURN_IF_ERROR(validate_snii_scoring_metadata( + out->common_grams_metadata ? &*out->common_grams_metadata : nullptr, + out->stats.doc_count, out->stats.sum_total_term_freq, has_scoring_tier, + has_positions(out->index_config), out->section_refs.norms.length != 0)); + } + return Status::OK(); +} + +} // namespace + +Status encode_core_metadata(const CoreMetadata& metadata, ByteSink* out) { + if (out == nullptr) { + return Status::Error("core metadata: null output"); + } + + doris::snii::SniiCoreMetadataPB core; + core.set_index_config(static_cast(metadata.index_config)); + auto* stats = core.mutable_stats(); + stats->set_doc_count(metadata.stats.doc_count); + stats->set_indexed_doc_count(metadata.stats.indexed_doc_count); + stats->set_term_count(metadata.stats.term_count); + stats->set_sum_total_term_freq(metadata.stats.sum_total_term_freq); + stats->set_null_count(metadata.stats.null_count); + auto* refs = core.mutable_section_refs(); + encode_region_ref(metadata.section_refs.dict_region, refs->mutable_dict_region()); + encode_region_ref(metadata.section_refs.posting_region, refs->mutable_posting_region()); + encode_region_ref(metadata.section_refs.norms, refs->mutable_norms()); + encode_region_ref(metadata.section_refs.null_bitmap, refs->mutable_null_bitmap()); + encode_region_ref(metadata.section_refs.bsbf, refs->mutable_bsbf()); + if (metadata.common_grams_metadata.has_value()) { + encode_common_grams(*metadata.common_grams_metadata, core.mutable_common_grams()); + } + if (metadata.common_grams_posting_policy != CommonGramsPostingPolicy::kNone) { + core.set_common_grams_posting_policy( + static_cast(metadata.common_grams_posting_policy)); + } + + CoreMetadata validated; + RETURN_IF_ERROR(decode_core_pb(core, &validated)); + const size_t size = core.ByteSizeLong(); + if (size > static_cast(std::numeric_limits::max())) { + return corrupted("core metadata: protobuf payload exceeds INT_MAX"); + } + std::string payload(size, '\0'); + if (!core.SerializeToArray(payload.data(), static_cast(size))) { + return corrupted("core metadata: protobuf serialization failed"); + } + SectionFramer::write(*out, static_cast(SectionType::kCoreMetadataPB), Slice(payload)); + return Status::OK(); +} + +Status decode_core_metadata(Slice framed_bytes, CoreMetadata* out) { + if (out == nullptr) { + return Status::Error("core metadata: null output"); + } + *out = {}; + ByteSource source(framed_bytes); + FramedSection section; + RETURN_IF_ERROR(SectionFramer::read(source, §ion)); + if (!source.eof() || section.type != static_cast(SectionType::kCoreMetadataPB)) { + return corrupted("core metadata: invalid frame"); + } + if (section.payload.size() > static_cast(std::numeric_limits::max())) { + return corrupted("core metadata: protobuf payload exceeds INT_MAX"); + } + doris::snii::SniiCoreMetadataPB core; + if (!core.ParseFromArray(section.payload.data(), static_cast(section.payload.size()))) { + return corrupted("core metadata: protobuf parsing failed"); + } + return decode_core_pb(core, out); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/core_metadata.h b/be/src/storage/index/snii/format/core_metadata.h new file mode 100644 index 00000000000000..ce761901a145ec --- /dev/null +++ b/be/src/storage/index/snii/format/core_metadata.h @@ -0,0 +1,62 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/stats_block.h" + +namespace doris::snii::format { + +struct RegionRef { + uint64_t offset = 0; + uint64_t length = 0; +}; + +struct SectionRefs { + RegionRef dict_region; + RegionRef posting_region; + RegionRef norms; + RegionRef null_bitmap; + RegionRef bsbf; +}; + +enum class CommonGramsPostingPolicy : uint8_t { + kNone = 0, + kDocsOnlyV1 = 1, + kHybridV1 = kDocsOnlyV1, +}; + +struct CoreMetadata { + IndexConfig index_config = IndexConfig::kDocsOnly; + StatsBlock stats; + SectionRefs section_refs; + std::optional common_grams_metadata; + CommonGramsPostingPolicy common_grams_posting_policy = CommonGramsPostingPolicy::kNone; +}; + +Status encode_core_metadata(const CoreMetadata& metadata, ByteSink* out); +Status decode_core_metadata(Slice framed_bytes, CoreMetadata* out); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/dict_block.cpp b/be/src/storage/index/snii/format/dict_block.cpp new file mode 100644 index 00000000000000..56416924f8abd0 --- /dev/null +++ b/be/src/storage/index/snii/format/dict_block.cpp @@ -0,0 +1,604 @@ +// 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 "storage/index/snii/format/dict_block.h" + +#include +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/sampled_term_index.h" // std_string_heap_bytes + +namespace doris::snii::format { + +namespace { + +constexpr size_t kFooterBytes = sizeof(uint32_t); // trailing crc32c +constexpr size_t kNAnchorsBytes = sizeof(uint32_t); // n_anchors u32 +constexpr size_t kAnchorOffBytes = sizeof(uint32_t); // per-anchor offset u32 + +size_t estimate_statless_entry_upper_bound(const DictEntry& e, IndexTier tier) { + size_t body = 0; + body += varint_len(static_cast(e.term.size())); + body += varint_len(static_cast(e.term.size())); + body += e.term.size(); + body += 1; // flags + body += varint_len(e.df); + + const bool tier_has_stats = tier >= IndexTier::kT2; + if (e.kind == DictEntryKind::kInline) { + body += varint_len(static_cast(e.frq_bytes.size())) + e.frq_bytes.size(); + body += varint_len(e.inline_dd_disk_len); + body += 1 + varint_len(e.dd_meta.uncomp_len); // win_mode + DD metadata + if (tier_has_stats) { + body += varint_len(e.freq_meta.uncomp_len); + body += varint_len(static_cast(e.prx_bytes.size())) + e.prx_bytes.size(); + } + } else { + body += varint_len(e.frq_off_delta) + varint_len(e.frq_len); + if (e.enc == DictEntryEnc::kWindowed) { + body += varint_len(e.prelude_len) + varint_len(e.frq_docs_len); + } else { + body += varint_len(e.frq_docs_len); + body += 1 + varint_len(e.dd_meta.uncomp_len) + sizeof(uint32_t); + if (tier_has_stats) { + body += varint_len(e.freq_meta.uncomp_len) + sizeof(uint32_t); + } + } + if (tier_has_stats) { + body += varint_len(e.prx_off_delta) + varint_len(e.prx_len); + } + } + return varint_len(static_cast(body)) + body; +} + +// Estimate the encoded upper-bound byte size of one entry (no actual encoding; used by +// estimated_bytes). Take the maximum varint width of each variable-length field plus payload bytes +// to guarantee an upper bound. +size_t estimate_entry_bytes(const DictEntry& e, IndexTier tier, bool term_stats) { + size_t body = 0; + body += varint_len(static_cast(e.term.size())); // prefix_len upper bound + body += varint_len(static_cast(e.term.size())); // suffix_len upper bound + body += e.term.size(); // suffix bytes upper bound + body += 1; // flags + body += 10; // df upper bound + if (term_stats) { + body += 10; // ttf_delta + body += 10; // max_freq + } + if (e.kind == DictEntryKind::kInline) { + body += 10 + e.frq_bytes.size(); + body += 10 + e.prx_bytes.size(); + } else { + body += 10 * 5; // frq_off/frq_len/prelude/prx_off/prx_len upper bound + } + const size_t legacy_estimate = varint_len(static_cast(body)) + body; + if (term_stats) { + return legacy_estimate; + } + return std::max(legacy_estimate, estimate_statless_entry_upper_bound(e, tier)); +} + +} // namespace + +// ---- DictBlockBuilder ---- + +DictBlockBuilder::DictBlockBuilder(IndexTier tier, bool has_positions, uint64_t frq_base, + uint64_t prx_base, uint32_t anchor_interval, bool term_stats) + : tier_(tier), + has_positions_(has_positions), + term_stats_(term_stats), + frq_base_(frq_base), + prx_base_(prx_base), + anchor_interval_(anchor_interval == 0 ? 1 : anchor_interval) {} + +void DictBlockBuilder::add_entry(const DictEntry& entry) { + if (is_anchor(n_entries_)) { + ++n_anchors_; + } + entries_est_ += estimate_entry_bytes(entry, tier_, term_stats_); + entries_.push_back(entry); + ++n_entries_; +} + +void DictBlockBuilder::add_entry(DictEntry&& entry) { + if (is_anchor(n_entries_)) { + ++n_anchors_; + } + // estimate_entry_bytes reads `entry`, so it MUST run before the move below: + // sizing a moved-from (empty) entry would undercount entries_est_ and split + // blocks incorrectly. finish() output is unaffected either way -- it depends + // only on the entries actually queued, not on how they were appended. + entries_est_ += estimate_entry_bytes(entry, tier_, term_stats_); + entries_.push_back(std::move(entry)); + ++n_entries_; +} + +size_t DictBlockBuilder::estimated_bytes() const { + size_t header = varint_len(static_cast(n_entries_)) + 2; // +ver +flags + header += varint_len(frq_base_); + if (has_positions_) { + header += varint_len(prx_base_); + } + const size_t anchors = n_anchors_ * kAnchorOffBytes + kNAnchorsBytes; + return header + entries_est_ + anchors + kFooterBytes; +} + +void DictBlockBuilder::encode_covered(ByteSink* sink) const { + // header. + sink->put_varint64(static_cast(n_entries_)); + sink->put_u8(kDictBlockFormatVer); + sink->put_u8(static_cast((has_positions_ ? dict_block_flags::kHasPositions : 0U) | + (term_stats_ ? 0U : dict_block_flags::kNoTermStats))); + sink->put_varint64(frq_base_); + if (has_positions_) { + sink->put_varint64(prx_base_); + } + + // entries: anchor entries use prev_term="" and record their byte offset within the block. + std::vector anchor_offsets; + anchor_offsets.reserve(n_anchors_); + ByteSink entry_body_scratch; + std::string_view prev; + for (uint32_t i = 0; i < n_entries_; ++i) { + const bool anchor = is_anchor(i); + if (anchor) { + anchor_offsets.push_back(static_cast(sink->size())); + } + const std::string_view prev_term = anchor ? std::string_view {} : prev; + // finish() is void and entry encoding into an in-memory ByteSink cannot fail; + // explicitly discard the (now [[nodiscard]] Status) return. + static_cast(encode_dict_entry(entries_[i], prev_term, tier_, sink, term_stats_, + &entry_body_scratch)); + prev = entries_[i].term; + } + + // anchor_offsets[] + n_anchors. + for (uint32_t off : anchor_offsets) { + sink->put_fixed32(off); + } + sink->put_fixed32(static_cast(anchor_offsets.size())); +} + +void DictBlockBuilder::finish(ByteSink* sink) const { + ByteSink body; // header + entries + anchor_offsets + n_anchors (crc covered region) + encode_covered(&body); + + // Write the entire block (including crc footer) to sink. + sink->put_bytes(body.view()); + sink->put_fixed32(crc32c(body.view())); +} + +std::vector DictBlockBuilder::finish_owned() const { + ByteSink block; + encode_covered(&block); + const uint32_t checksum = crc32c(block.view()); + block.reserve(kFooterBytes); + block.put_fixed32(checksum); + return block.take(); +} + +// ---- DictBlockReader ---- + +namespace { + +// Verify the block length is sufficient and validate the trailing crc; return a Slice of the covered region (excluding crc footer). +Status verify_crc(Slice block, Slice* covered) { + if (block.size() < kFooterBytes + kNAnchorsBytes) { + return Status::Error( + "dict_block: block too short to contain footer"); + } + const size_t covered_len = block.size() - kFooterBytes; + *covered = block.subslice(0, covered_len); + + ByteSource crc_src(block.subslice(covered_len, kFooterBytes)); + uint32_t stored = 0; + RETURN_IF_ERROR(crc_src.get_fixed32(&stored)); + if (crc32c(*covered) != stored) { + return Status::Error( + "dict_block: crc32c checksum mismatch"); + } + return Status::OK(); +} + +// Read and verify that block_flags is consistent with has_positions. +Status check_flags(uint8_t flags, bool has_positions) { + const bool flag_pos = (flags & dict_block_flags::kHasPositions) != 0; + if (flag_pos != has_positions) { + return Status::Error( + "dict_block: has_positions inconsistent with block_flags"); + } + return Status::OK(); +} + +} // namespace + +Status DictBlockReader::open(Slice block, IndexTier tier, bool has_positions, + DictBlockReader* out) { + if (out == nullptr) { + return Status::Error("dict_block: out is null"); + } + *out = DictBlockReader {}; + + // Decode instrumentation seam: one increment per block materialization (the + // CRC verify + anchor parse below, preceded by a zstd decompress for a + // compressed block). A dict-block cache eliminates repeats of exactly this. + testing::note_dict_block_decode(); + + Slice covered; + RETURN_IF_ERROR(verify_crc(block, &covered)); + out->block_ = covered; + out->tier_ = tier; + out->has_positions_ = has_positions; + + // header. + ByteSource src(covered); + uint64_t n_entries = 0; + RETURN_IF_ERROR(src.get_varint64(&n_entries)); + uint8_t ver = 0; + uint8_t flags = 0; + RETURN_IF_ERROR(src.get_u8(&ver)); + RETURN_IF_ERROR(src.get_u8(&flags)); + if (ver != kDictBlockFormatVer) { + return Status::Error( + "dict_block: unsupported entry_format_ver"); + } + RETURN_IF_ERROR(check_flags(flags, has_positions)); + out->term_stats_ = (flags & dict_block_flags::kNoTermStats) == 0; + RETURN_IF_ERROR(src.get_varint64(&out->frq_base_)); + if (has_positions) { + RETURN_IF_ERROR(src.get_varint64(&out->prx_base_)); + } + + out->n_entries_ = static_cast(n_entries); + out->entries_begin_ = src.position(); + + // The anchor table is at the tail of covered: [... anchor_offsets[n] n_anchors(u32)]. + if (covered.size() < kNAnchorsBytes) { + return Status::Error( + "dict_block: missing n_anchors"); + } + ByteSource na_src(covered.subslice(covered.size() - kNAnchorsBytes, kNAnchorsBytes)); + uint32_t n_anchors = 0; + RETURN_IF_ERROR(na_src.get_fixed32(&n_anchors)); + + const size_t anchor_table_bytes = static_cast(n_anchors) * kAnchorOffBytes; + if (covered.size() < kNAnchorsBytes + anchor_table_bytes || + out->entries_begin_ + anchor_table_bytes + kNAnchorsBytes > covered.size()) { + return Status::Error( + "dict_block: anchor table out of range"); + } + const size_t anchor_table_begin = covered.size() - kNAnchorsBytes - anchor_table_bytes; + + ByteSource at_src(covered.subslice(anchor_table_begin, anchor_table_bytes)); + out->anchor_offsets_.resize(n_anchors); + out->anchor_terms_.resize(n_anchors); + for (uint32_t i = 0; i < n_anchors; ++i) { + uint32_t off = 0; + RETURN_IF_ERROR(at_src.get_fixed32(&off)); + if (off >= anchor_table_begin) { + return Status::Error( + "dict_block: anchor offset out of range"); + } + // Anchor offsets must be strictly monotonically increasing, and the first anchor must be exactly the start of the entries region (entry 0 is always an anchor). + // Otherwise scan_from_anchor's segment-length computation seg_end-seg_begin would underflow as size_t and cause an out-of-range read, + // guarding against non-monotonic offset tables with a re-stamped crc (remote on-demand read / cache misalignment scenarios). + if (i == 0) { + if (off != out->entries_begin_) { + return Status::Error( + "dict_block: first anchor offset is not the start of entries"); + } + } else if (off <= out->anchor_offsets_[i - 1]) { + return Status::Error( + "dict_block: anchor offsets are not strictly increasing"); + } + out->anchor_offsets_[i] = off; + // Anchor entries are encoded with prev_term="" and can be decoded independently to retrieve their term. + ByteSource e_src(covered.subslice(off, anchor_table_begin - off)); + DictEntry probe; + RETURN_IF_ERROR(decode_dict_entry(&e_src, std::string_view {}, tier, &probe, + (flags & dict_block_flags::kNoTermStats) == 0)); + out->anchor_terms_[i] = std::move(probe.term); + } + return Status::OK(); +} + +size_t DictBlockReader::heap_bytes() const { + size_t bytes = anchor_offsets_.capacity() * sizeof(uint32_t) + + anchor_terms_.capacity() * sizeof(std::string); + for (const auto& term : anchor_terms_) { + bytes += std_string_heap_bytes(term); + } + return bytes; +} + +bool DictBlockReader::locate_anchor(std::string_view target, size_t* anchor_idx) const { + if (anchor_terms_.empty()) { + return false; + } + if (target < std::string_view(anchor_terms_.front())) { + return false; + } + // The last anchor_term <= target. + size_t lo = 0; + size_t hi = anchor_terms_.size(); // open interval + while (lo + 1 < hi) { + const size_t mid = lo + (hi - lo) / 2; + if (std::string_view(anchor_terms_[mid]) <= target) { + lo = mid; + } else { + hi = mid; + } + } + *anchor_idx = lo; + return true; +} + +Status DictBlockReader::decode_all(std::vector* out) const { + if (out == nullptr) { + return Status::Error("dict_block: out is null"); + } + out->clear(); + out->reserve(n_entries_); + for (size_t a = 0; a < anchor_offsets_.size(); ++a) { + const size_t seg_begin = anchor_offsets_[a]; + const bool is_last = a + 1 == anchor_offsets_.size(); + const size_t seg_end = is_last ? (block_.size() - kNAnchorsBytes - + anchor_offsets_.size() * kAnchorOffBytes) + : anchor_offsets_[a + 1]; + if (seg_end < seg_begin || seg_end > block_.size()) { + return Status::Error( + "dict_block: anchor segment range invalid"); + } + ByteSource src(block_.subslice(seg_begin, seg_end - seg_begin)); + std::string prev; // first entry of a segment is an anchor (prev_term="") + while (!src.eof()) { + DictEntry e; + RETURN_IF_ERROR( + decode_dict_entry(&src, std::string_view(prev), tier_, &e, term_stats_)); + prev = e.term; + out->push_back(std::move(e)); + } + } + if (out->size() != n_entries_) { + return Status::Error( + "dict_block: decoded entry count mismatch"); + } + return Status::OK(); +} + +Status DictBlockReader::scan_from_anchor(size_t anchor_idx, std::string_view target, bool* found, + DictEntry* out) const { + // Byte range of this anchor segment: [anchor_offset, next anchor offset or anchor table start). + const size_t seg_begin = anchor_offsets_[anchor_idx]; + const bool is_last = anchor_idx + 1 == anchor_offsets_.size(); + const size_t seg_end = + is_last ? (block_.size() - kNAnchorsBytes - anchor_offsets_.size() * kAnchorOffBytes) + : anchor_offsets_[anchor_idx + 1]; + + // Fallback: open() has already verified anchor monotonicity; this additionally guards against seg_end block_.size()) { + return Status::Error( + "dict_block: anchor segment range invalid"); + } + ByteSource src(block_.subslice(seg_begin, seg_end - seg_begin)); + std::string prev; // the first entry in the segment is an anchor, prev_term="" + while (!src.eof()) { + // Key-first: decode only the (front-coded) term key, then decide whether + // the body is worth materializing. Non-matching entries skip their body + // entirely -- with anchor_interval=16 that turns ~16 body decodes per + // lookup into 1 (the matched entry) or 0 (a miss). + DictEntry e; + size_t body_start = 0; + uint64_t entry_total = 0; + RETURN_IF_ERROR( + decode_dict_entry_key(&src, std::string_view(prev), &e, &body_start, &entry_total)); + if (e.term == target) { + RETURN_IF_ERROR( + decode_dict_entry_rest(&src, tier_, body_start, entry_total, &e, term_stats_)); + *found = true; + *out = std::move(e); + return Status::OK(); + } + if (std::string_view(e.term) > target) { + *found = false; // already past target; entries are sorted so it does not exist + return Status::OK(); + } + // Before target: skip the body but keep the key as the front-coding base. + RETURN_IF_ERROR(skip_dict_entry_body(&src, body_start, entry_total)); + prev = std::move(e.term); + } + *found = false; + return Status::OK(); +} + +Status DictBlockReader::find_term(std::string_view target, bool* found, DictEntry* out) const { + if (found == nullptr || out == nullptr) { + return Status::Error("dict_block: found / out is null"); + } + *found = false; + size_t anchor_idx = 0; + if (!locate_anchor(target, &anchor_idx)) { + return Status::OK(); + } + return scan_from_anchor(anchor_idx, target, found, out); +} + +Status DictBlockReader::visit_prefix_range(std::string_view prefix, + const std::function& accept_key, + const std::function& on_hit, + bool* prefix_exhausted) const { + if (!on_hit || prefix_exhausted == nullptr) { + return Status::Error( + "dict_block: null visit_prefix_range args"); + } + *prefix_exhausted = false; + if (anchor_offsets_.empty()) { + return Status::OK(); // empty block: nothing to enumerate + } + + // Anchor-jump: start at the anchor segment that may contain prefix. Earlier + // segments hold only terms < prefix (every term is < the next anchor term + // <= prefix), so they are skipped without any decode. An empty prefix or one + // sorting before the first anchor starts at anchor 0. + size_t anchor_idx = 0; + if (!prefix.empty()) { + locate_anchor(prefix, &anchor_idx); // false leaves anchor_idx at 0 + } + + // Scan from the chosen anchor to the end of the entries region: the prefix + // range may span several anchor segments. Anchor entries are encoded with + // prefix_len=0, so a single running `prev` reconstructs every term correctly + // even as the scan crosses segment boundaries. + const size_t seg_begin = anchor_offsets_[anchor_idx]; + const size_t entries_end = + block_.size() - kNAnchorsBytes - anchor_offsets_.size() * kAnchorOffBytes; + if (entries_end < seg_begin || entries_end > block_.size()) { + return Status::Error( + "dict_block: entries region range invalid"); + } + + ByteSource src(block_.subslice(seg_begin, entries_end - seg_begin)); + std::string prev; // the chosen anchor segment starts at an anchor (prev="") + while (!src.eof()) { + DictEntry e; + size_t body_start = 0; + uint64_t entry_total = 0; + RETURN_IF_ERROR( + decode_dict_entry_key(&src, std::string_view(prev), &e, &body_start, &entry_total)); + const std::string_view t(e.term); + if (t < prefix) { + // Still before the range (only reachable inside the anchor segment + // that straddles prefix): skip the body, keep the front-coding base. + RETURN_IF_ERROR(skip_dict_entry_body(&src, body_start, entry_total)); + prev = std::move(e.term); + continue; + } + const bool has_prefix = t.size() >= prefix.size() && t.starts_with(prefix); + if (!has_prefix) { + *prefix_exhausted = true; // sorted: no further matches here or later + return Status::OK(); + } + if (accept_key && !accept_key(t)) { + // Key-only rejection: never pay for this entry's body. + RETURN_IF_ERROR(skip_dict_entry_body(&src, body_start, entry_total)); + prev = std::move(e.term); + continue; + } + // Accepted: materialize the body and hand the entry to the visitor. + RETURN_IF_ERROR( + decode_dict_entry_rest(&src, tier_, body_start, entry_total, &e, term_stats_)); + prev = e.term; // copy the key before the entry is moved into on_hit + bool stop = false; + RETURN_IF_ERROR(on_hit(std::move(e), &stop)); + if (stop) { + return Status::OK(); + } + } + return Status::OK(); +} + +Status DictBlockReader::visit_term_range(std::string_view lower_inclusive, + std::optional upper_exclusive, + const std::function& accept_key, + const std::function& on_hit, + bool* range_exhausted) const { + if (!on_hit || range_exhausted == nullptr) { + return Status::Error( + "dict_block: null visit_term_range args"); + } + *range_exhausted = false; + if (upper_exclusive.has_value() && *upper_exclusive <= lower_inclusive) { + *range_exhausted = true; + return Status::OK(); + } + if (anchor_offsets_.empty()) { + return Status::OK(); + } + + size_t anchor_idx = 0; + if (!lower_inclusive.empty()) { + locate_anchor(lower_inclusive, &anchor_idx); + } + const size_t seg_begin = anchor_offsets_[anchor_idx]; + const size_t entries_end = + block_.size() - kNAnchorsBytes - anchor_offsets_.size() * kAnchorOffBytes; + if (entries_end < seg_begin || entries_end > block_.size()) { + return Status::Error( + "dict_block: entries region range invalid"); + } + + ByteSource src(block_.subslice(seg_begin, entries_end - seg_begin)); + std::string prev; + while (!src.eof()) { + DictEntry entry; + size_t body_start = 0; + uint64_t entry_total = 0; + RETURN_IF_ERROR(decode_dict_entry_key(&src, std::string_view(prev), &entry, &body_start, + &entry_total)); + const std::string_view term(entry.term); + if (term < lower_inclusive) { + RETURN_IF_ERROR(skip_dict_entry_body(&src, body_start, entry_total)); + prev = std::move(entry.term); + continue; + } + if (upper_exclusive.has_value() && term >= *upper_exclusive) { + *range_exhausted = true; + return Status::OK(); + } + if (accept_key && !accept_key(term)) { + RETURN_IF_ERROR(skip_dict_entry_body(&src, body_start, entry_total)); + prev = std::move(entry.term); + continue; + } + RETURN_IF_ERROR( + decode_dict_entry_rest(&src, tier_, body_start, entry_total, &entry, term_stats_)); + prev = entry.term; + bool stop = false; + RETURN_IF_ERROR(on_hit(std::move(entry), &stop)); + if (stop) { + return Status::OK(); + } + } + return Status::OK(); +} + +} // namespace doris::snii::format + +namespace doris::snii::testing { +namespace { +std::atomic& dict_decode_atomic() { + static std::atomic counter {0}; + return counter; +} +} // namespace + +uint64_t dict_decode_counter() { + return dict_decode_atomic().load(std::memory_order_relaxed); +} + +void reset_dict_decode_counter() { + dict_decode_atomic().store(0, std::memory_order_relaxed); +} + +void note_dict_block_decode() { + dict_decode_atomic().fetch_add(1, std::memory_order_relaxed); +} + +} // namespace doris::snii::testing diff --git a/be/src/storage/index/snii/format/dict_block.h b/be/src/storage/index/snii/format/dict_block.h new file mode 100644 index 00000000000000..7a299a1f162f8e --- /dev/null +++ b/be/src/storage/index/snii/format/dict_block.h @@ -0,0 +1,232 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" + +// DICT block —— a positioning unit mapping term → postings read plan, and also +// the unit for remote on-demand fetching, caching, and CRC checksum +// verification (see docs/design/SNII-design-spec.source.md "DICT block" and +// "dict lookup flow summary" sections). +// +// Byte layout (strictly implemented; multi-byte fixed-width fields are +// little-endian, variable-length integers use LEB128): +// header: +// n_entries varint +// entry_format_ver u8 # = kDictBlockFormatVer +// block_flags u8 # bit0 = has_positions (consistency check +// against the value passed to reader) frq_base varint64 prx_base +// varint64 # present only when has_positions is set +// entries[n_entries] # variable-length DictEntry, front-coded in +// lexicographic order anchor_offsets[n_anchors] # u32 * n_anchors, byte +// offset of each anchor entry within the block n_anchors u32 crc32c +// u32 # covers [header .. n_anchors], detects corruption (sole CRC +// layer) +// +// Anchor rule: every anchor_interval entries, one "term anchor" is forced — +// that entry is encoded with prev_term="" (prefix_len=0, storing the full +// term), and its byte offset is recorded in anchor_offsets; non-anchor entries +// use the preceding entry's term as prev_term for front coding. The reader can +// start from any anchor and scan independently without needing earlier terms, +// enabling anchor binary search + local scan for exact term lookup. +namespace doris::snii::format { + +// DICT block entry_format_ver: self-describing version of the DictEntry +// encoding. Reader rejects a mismatch so a query-only run cannot silently read +// an older dict-entry layout as the current one. +inline constexpr uint8_t kDictBlockFormatVer = 2; + +// block_flags bit definitions. +namespace dict_block_flags { +inline constexpr uint8_t kHasPositions = 1U << 0; // whether to write prx_base / .prx fields +// G16-f: entries omit the ttf_delta/max_freq varints (freq-dropped index -- +// the stats serve only BM25 scoring). Self-describing per block; absent on +// pre-G16-f blocks, whose entries always carry the stats on tier>=T2. +inline constexpr uint8_t kNoTermStats = 1U << 1; +// bit1-7 reserved +} // namespace dict_block_flags + +// DICT block writer: entries are added in lexicographic order via add_entry; +// internally determines anchors and accumulates size estimates, and on finish +// serializes header + entries + anchor table + CRC in one pass. The front-coding +// base is rebuilt from a local prev inside finish(), so no prev_term is retained +// as builder state. +class DictBlockBuilder { +public: + DictBlockBuilder(IndexTier tier, bool has_positions, uint64_t frq_base, uint64_t prx_base, + uint32_t anchor_interval = 16, bool term_stats = true); + + // Append one entry (caller must guarantee lexicographic term order). + // Internally decides whether it becomes an anchor. The copy overload is kept + // for callers that must retain their entry afterwards (materialized fallback, + // tests); the move overload avoids the per-term DictEntry copy -- which for an + // inline entry is two std::vector heap allocations plus the term + // copy -- on the SPIMI build path. + void add_entry(const DictEntry& entry); + void add_entry(DictEntry&& entry); + + // Upper-bound estimate of the serialized size of the current block (including + // header + entries + anchor table + CRC footer), used by the upper layer to + // decide when to cut a new block based on target_dict_block_bytes. + size_t estimated_bytes() const; + + // Number of entries. + uint32_t n_entries() const { return n_entries_; } + + // Serialize the entire block and append it to sink. + void finish(ByteSink* sink) const; + + // Serialize the entire block into an owned buffer. This avoids copying the + // CRC-covered bytes when the caller needs ownership of the complete block. + std::vector finish_owned() const; + +private: + bool is_anchor(uint32_t index) const { return index % anchor_interval_ == 0; } + void encode_covered(ByteSink* sink) const; + + IndexTier tier_; + bool has_positions_; + bool term_stats_ = true; // false: entries omit ttf/max_freq (kNoTermStats) + uint64_t frq_base_; + uint64_t prx_base_; + uint32_t anchor_interval_; + + uint32_t n_entries_ = 0; + std::vector entries_; + size_t entries_est_ = 0; // accumulated byte estimate for the entries section + size_t n_anchors_ = 0; // number of anchors +}; + +// DICT block reader: on open, verifies the CRC and parses the header / anchor +// table; find_term uses anchor binary search + local scan to locate a +// DictEntry. Holds a byte view of the block (non-owning); lifetime is managed +// by the caller. +class DictBlockReader { +public: + DictBlockReader() = default; + + // Parse and verify the entire block. CRC mismatch / truncation / invalid + // structure → Corruption; has_positions in the header inconsistent with the + // supplied argument → InvalidArgument. + static Status open(Slice block, IndexTier tier, bool has_positions, DictBlockReader* out); + + // Anchor binary search + local scan to locate target. Hit → *found=true and + // *out is filled; miss (including out-of-range, gap) → *found=false. + // Structural error → non-OK Status. + Status find_term(std::string_view target, bool* found, DictEntry* out) const; + + // Decodes EVERY entry in the block in lexicographic order into *out (each a + // self-contained DictEntry, owning its term). Used for ordered term + // enumeration (prefix / range scans). Resets the front-coding base at each + // anchor segment. Retained as the golden reference; the prefix path now + // streams via visit_prefix_range. + Status decode_all(std::vector* out) const; + + // Streams the entries of this block whose term lies in [prefix, prefix+) in + // lexicographic order, materializing only the bodies a caller keeps (T07): + // 1) anchor-jump to the anchor segment containing prefix (segments whose + // terms are all < prefix are skipped without any decode); an empty + // prefix or one before the first anchor starts at anchor 0; + // 2) within range, decode each entry's term key only; term < prefix skips + // the body and continues; a term that leaves the prefix range sets + // *prefix_exhausted=true and ends the scan (sorted order guarantees no + // further matches here or in later blocks); + // 3) accept_key(term)==false skips the body (lets callers push a key-only + // predicate down so a non-match never pays for its body); + // 4) only accepted entries have their body decoded and are handed to + // on_hit, which may request an early stop via *stop. + // accept_key may be empty (treated as accept-all). prefix_exhausted is set + // false on entry and true only when a term past the range is seen. + Status visit_prefix_range(std::string_view prefix, + const std::function& accept_key, + const std::function& on_hit, + bool* prefix_exhausted) const; + + // Streams entries in [lower_inclusive, upper_exclusive). A missing upper + // bound scans to the end of the block. Like visit_prefix_range, this starts + // from the nearest anchor and only decodes accepted entry bodies. + Status visit_term_range(std::string_view lower_inclusive, + std::optional upper_exclusive, + const std::function& accept_key, + const std::function& on_hit, + bool* range_exhausted) const; + + uint64_t frq_base() const { return frq_base_; } + uint64_t prx_base() const { return prx_base_; } + uint32_t n_entries() const { return n_entries_; } + + // Resident heap held beyond sizeof(*this): the anchor_offsets_ / anchor_terms_ + // vector buffers plus each non-SSO anchor term's heap allocation. block_ is a + // NON-owning view and is deliberately NOT counted here -- its owning buffer is + // charged by the caller (the resident block's `bytes` vector, or a + // request-scoped decoded block). Summed into + // LogicalIndexReader::memory_usage() per resident block. + size_t heap_bytes() const; + +private: + // Sequentially scan from anchor anchor_idx to the end of that anchor segment, + // searching for target. + Status scan_from_anchor(size_t anchor_idx, std::string_view target, bool* found, + DictEntry* out) const; + + // Find the last anchor index where first_term(anchor) <= target; return false + // if none exists. + bool locate_anchor(std::string_view target, size_t* anchor_idx) const; + + Slice block_; // [header .. crc) full block view + IndexTier tier_ = IndexTier::kT1; + bool has_positions_ = false; + bool term_stats_ = true; // from block flags (kNoTermStats absent => true) + uint64_t frq_base_ = 0; + uint64_t prx_base_ = 0; + uint32_t n_entries_ = 0; + + size_t entries_begin_ = 0; // absolute offset of the start of the entries section + std::vector anchor_offsets_; // byte offset within the block for each anchor entry + std::vector + anchor_terms_; // full term of each anchor entry (used for binary search) +}; + +} // namespace doris::snii::format + +// Test-only instrumentation seam. dict_decode_counter() returns a process-global +// count of DICT block decodes performed by DictBlockReader::open -- i.e. the +// optional zstd decompress + CRC verify + anchor parse that turns on-disk block +// bytes into a usable reader. This is precisely the unit a dict-block cache +// eliminates on repeat, so tests assert dict_decode_counter() == unique_blocks. +// In production DICT blocks are zstd-compressed, so this equals the zstd +// decompress count. Counters use relaxed atomics; reset between tests. +namespace doris::snii::testing { + +uint64_t dict_decode_counter(); +void reset_dict_decode_counter(); +void note_dict_block_decode(); + +} // namespace doris::snii::testing diff --git a/be/src/storage/index/snii/format/dict_block_directory.cpp b/be/src/storage/index/snii/format/dict_block_directory.cpp new file mode 100644 index 00000000000000..8801418f8a3c6e --- /dev/null +++ b/be/src/storage/index/snii/format/dict_block_directory.cpp @@ -0,0 +1,131 @@ +// 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 "storage/index/snii/format/dict_block_directory.h" + +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +namespace { + +// Each block_ref has a fixed field order; reuse ByteSink varint/fixed primitives — do not hand-craft bytes manually. +// uncomp_len trails only when the kZstd flag is set, so uncompressed-block +// directories keep their compact (v1-identical) per-ref byte layout. +void encode_ref(const BlockRef& ref, ByteSink* payload) { + payload->put_varint64(ref.offset); + payload->put_varint64(ref.length); + payload->put_varint32(ref.n_entries); + payload->put_u8(ref.flags); + payload->put_fixed32(ref.checksum); + if (ref.flags & block_ref_flags::kZstd) payload->put_varint64(ref.uncomp_len); +} + +Status decode_ref(ByteSource* ps, BlockRef* ref) { + RETURN_IF_ERROR(ps->get_varint64(&ref->offset)); + RETURN_IF_ERROR(ps->get_varint64(&ref->length)); + RETURN_IF_ERROR(ps->get_varint32(&ref->n_entries)); + RETURN_IF_ERROR(ps->get_u8(&ref->flags)); + RETURN_IF_ERROR(ps->get_fixed32(&ref->checksum)); + if (ref->flags & block_ref_flags::kZstd) { + RETURN_IF_ERROR(ps->get_varint64(&ref->uncomp_len)); + } + return Status::OK(); +} + +Status decode_payload(Slice payload, std::vector* refs) { + ByteSource ps(payload); + uint32_t n_blocks = 0; + RETURN_IF_ERROR(ps.get_varint32(&n_blocks)); + // Guard against a corrupted, inflated count from untrusted bytes: each BlockRef + // needs >= 8 bytes (flags u8 + checksum u32 + >= 1 byte for each of 3 varints), + // so cap before reserve to avoid a huge allocation. + constexpr size_t kMinRefBytes = 8; + if (n_blocks > ps.remaining() / kMinRefBytes) { + return Status::Error( + "dict_block_directory: n_blocks exceeds payload capacity"); + } + refs->clear(); + refs->reserve(n_blocks); + uint64_t previous_end = 0; + for (uint32_t i = 0; i < n_blocks; ++i) { + BlockRef ref {}; + RETURN_IF_ERROR(decode_ref(&ps, &ref)); + if (ref.length == 0) { + return Status::Error( + "dict_block_directory: zero-length block"); + } + if (ref.length > std::numeric_limits::max() - ref.offset) { + return Status::Error( + "dict_block_directory: block range end overflow"); + } + const uint64_t block_end = ref.offset + ref.length; + if (i != 0 && ref.offset < previous_end) { + return Status::Error( + "dict_block_directory: blocks overlap or are out of physical order"); + } + previous_end = block_end; + refs->push_back(ref); + } + if (!ps.eof()) { + return Status::Error( + "dict_block_directory: trailing bytes in payload"); + } + return Status::OK(); +} + +} // namespace + +void DictBlockDirectoryBuilder::finish(ByteSink* sink) const { + ByteSink payload; + payload.put_varint32(static_cast(refs_.size())); + for (const auto& ref : refs_) { + encode_ref(ref, &payload); + } + SectionFramer::write(*sink, static_cast(SectionType::kDictBlockDirectory), + payload.view()); +} + +Status DictBlockDirectoryReader::open(Slice section, DictBlockDirectoryReader* out) { + ByteSource src(section); + FramedSection sec; + RETURN_IF_ERROR(SectionFramer::read(src, &sec)); + if (!src.eof()) { + return Status::Error( + "dict_block_directory: trailing framed section bytes"); + } + if (sec.type != static_cast(SectionType::kDictBlockDirectory)) { + return Status::Error( + "dict_block_directory: unexpected section type"); + } + return decode_payload(sec.payload, &out->refs_); +} + +Status DictBlockDirectoryReader::get(uint32_t ordinal, BlockRef* out) const { + if (ordinal >= refs_.size()) { + return Status::Error( + "dict_block_directory: ordinal out of range"); + } + *out = refs_[ordinal]; + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/dict_block_directory.h b/be/src/storage/index/snii/format/dict_block_directory.h new file mode 100644 index 00000000000000..909b32f29d1374 --- /dev/null +++ b/be/src/storage/index/snii/format/dict_block_directory.h @@ -0,0 +1,95 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +namespace doris::snii::format { + +// BlockRef.flags bit definitions. +namespace block_ref_flags { +// bit0: the on-disk block bytes are zstd(uncompressed_block). When set, the +// directory also stores uncomp_len, and the reader zstd-decompresses the fetched +// [offset, offset+length) range to uncomp_len before parsing the dict block. The +// block-level crc32c (and BlockRef.checksum) cover the UNCOMPRESSED bytes, so a +// zstd block shrinks the bytes fetched from S3 while keeping the same integrity +// guarantees after decompression in RAM. +inline constexpr uint8_t kZstd = 1u << 0; +} // namespace block_ref_flags + +// Physical location and checksum info for a single DICT block. Aligned with SampledTermIndex by ordinal: +// SampledTermIndex[i]'s first_term corresponds to DictBlockDirectory[i] (see design spec +// "sampled dict index"). The read path issues a single range read over [offset, offset+length). +struct BlockRef { + uint64_t offset = 0; // absolute byte offset of the block within the container + uint64_t length = 0; // ON-DISK byte length of the block (compressed when kZstd) + uint32_t n_entries = 0; // number of DictEntry records within this block + uint8_t flags = 0; // block-level flags (block_ref_flags::*) + uint32_t checksum = 0; // crc32c of the block's UNCOMPRESSED content (verified after read) + uint64_t uncomp_len = 0; // uncompressed block byte length (stored only when kZstd set) +}; + +// DICT block directory: block ordinal → physical location mapping. +// +// on-disk layout (framed by SectionFramer with a unified type+len+crc32c wrapper): +// [u8 type=kDictBlockDirectory][varint64 payload_len][payload][fixed32 crc32c] +// payload = varint32 n_blocks +// then n_blocks × block_ref{ +// varint64 offset, varint64 length, varint32 n_entries, +// u8 flags, fixed32 checksum } +// Section-level crc detects truncation/corruption; block_ref.checksum is the per-block crc. +class DictBlockDirectoryBuilder { +public: + void add(const BlockRef& ref) { refs_.push_back(ref); } + + // Encodes as a kDictBlockDirectory framed section (with embedded crc32c) and appends to sink. + void finish(ByteSink* sink) const; + +private: + std::vector refs_; +}; + +// Reads and verifies a kDictBlockDirectory framed section; provides ordinal → BlockRef lookup. +// After parsing, all block_refs reside in the reader (entering the searcher cache along with meta). +class DictBlockDirectoryReader { +public: + // Verifies the section crc and deserializes all block_refs. + // crc mismatch / truncation / trailing bytes → kCorruption; wrong section type → kInvalidArgument. + static Status open(Slice section, DictBlockDirectoryReader* out); + + uint32_t n_blocks() const { return static_cast(refs_.size()); } + + // Resident heap held beyond sizeof(*this): the refs_ vector buffer. BlockRef + // is trivially copyable (no per-element heap), so the vector buffer is the + // whole charge. Summed into LogicalIndexReader::memory_usage(). + size_t heap_bytes() const { return refs_.capacity() * sizeof(BlockRef); } + + // Returns the ordinal-th block_ref; ordinal >= n_blocks → kNotFound. + Status get(uint32_t ordinal, BlockRef* out) const; + +private: + std::vector refs_; +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/dict_entry.cpp b/be/src/storage/index/snii/format/dict_entry.cpp new file mode 100644 index 00000000000000..e9ee72912f461b --- /dev/null +++ b/be/src/storage/index/snii/format/dict_entry.cpp @@ -0,0 +1,383 @@ +// 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 "storage/index/snii/format/dict_entry.h" + +#include +#include + +#include "common/check.h" +#include "storage/index/snii/common/slice.h" + +namespace doris::snii::format { + +namespace { + +// Body-decode counter seam (T07). Incremented once per decode_dict_entry_rest +// call so tests can assert a key-first scan only materializes the bodies it +// actually produces. Relaxed atomic: read by tests only, never a sync point. +std::atomic& body_decode_atomic() { + static std::atomic counter {0}; + return counter; +} + +// Pure-function assembly / parsing of flags bits; avoids a long inline if-else +// chain. +uint8_t pack_flags(const DictEntry& e) { + uint8_t f = 0; + if (e.kind == DictEntryKind::kInline) f |= dict_flags::kKind; + if (e.enc == DictEntryEnc::kWindowed) f |= dict_flags::kEnc; + if (e.has_sb) f |= dict_flags::kHasSb; + // bit3 has_champion / bit4 offsets_ref are always 0 in v1. + return f; +} + +void apply_flags(uint8_t f, DictEntry* e) { + e->kind = (f & dict_flags::kKind) ? DictEntryKind::kInline : DictEntryKind::kPodRef; + e->enc = (f & dict_flags::kEnc) ? DictEntryEnc::kWindowed : DictEntryEnc::kSlim; + e->has_sb = (f & dict_flags::kHasSb) != 0; +} + +// Length of the longest common prefix between term and prev_term. +uint32_t common_prefix_len(std::string_view term, std::string_view prev) { + uint32_t n = 0; + const uint32_t lim = static_cast(std::min(term.size(), prev.size())); + while (n < lim && term[n] == prev[n]) ++n; + return n; +} + +bool tier_has_stats(IndexTier tier) { + return tier >= IndexTier::kT2; +} + +// ---- Encode entry body (excluding entry_len and trailing crc) ---- + +void write_term_key(const DictEntry& e, std::string_view prev, ByteSink* sink) { + const uint32_t prefix = common_prefix_len(e.term, prev); + const std::string_view suffix = std::string_view(e.term).substr(prefix); + sink->put_varint32(prefix); + sink->put_varint32(static_cast(suffix.size())); + sink->put_bytes(Slice(suffix)); +} + +void write_stats(const DictEntry& e, IndexTier tier, bool term_stats, ByteSink* sink) { + sink->put_varint32(e.df); + // G16-f: term_stats == false (freq-dropped index, block header flag) omits + // the ttf_delta/max_freq varints -- BM25-only fields, dead without freq. + if (!tier_has_stats(tier) || !term_stats) return; + sink->put_varint64(e.ttf_delta); + sink->put_varint64(e.max_freq); +} + +// Per-window codec mode byte shared by slim/inline single-window regions. +uint8_t pack_win_mode(const DictEntry& e) { + uint8_t mode = 0; + if (e.dd_meta.zstd) mode |= 1u << 0; // dd_zstd + if (e.freq_meta.zstd) mode |= 1u << 1; // freq_zstd + return mode; +} + +// Writes the slim/inline region codec metadata (dd always; freq when tier>=T2). +// store_crc=false (INLINE entries, format v2) omits the redundant per-region +// crc32c: the inline bytes already sit inside the dict block, whose own +// block-level crc32c covers them. POD-ref entries pass store_crc=true (their +// regions live in the separately-fetched .frq POD, uncovered by the block crc). +void write_region_meta(const DictEntry& e, IndexTier tier, bool store_crc, ByteSink* sink) { + sink->put_u8(pack_win_mode(e)); + sink->put_varint64(e.dd_meta.uncomp_len); + if (store_crc) sink->put_fixed32(e.dd_meta.crc); + if (!tier_has_stats(tier)) return; + sink->put_varint64(e.freq_meta.uncomp_len); + if (store_crc) sink->put_fixed32(e.freq_meta.crc); +} + +void write_pod_ref(const DictEntry& e, IndexTier tier, ByteSink* sink) { + sink->put_varint64(e.frq_off_delta); + sink->put_varint64(e.frq_len); + if (e.enc == DictEntryEnc::kWindowed) { + sink->put_varint64(e.prelude_len); + sink->put_varint64(e.frq_docs_len); + } else { + sink->put_varint64(e.frq_docs_len); // slim pod_ref: dd region on-disk length + // POD-ref regions live in the .frq POD (not covered by the block crc): keep + // crc. + write_region_meta(e, tier, /*store_crc=*/true, sink); + } + if (!tier_has_stats(tier)) return; + sink->put_varint64(e.prx_off_delta); + sink->put_varint64(e.prx_len); +} + +void write_inline(const DictEntry& e, IndexTier tier, ByteSink* sink) { + sink->put_varint64(static_cast(e.frq_bytes.size())); + sink->put_bytes(Slice(e.frq_bytes)); + sink->put_varint64(e.inline_dd_disk_len); + // INLINE bytes are covered by the dict block crc32c: omit the redundant + // per-region crc. + write_region_meta(e, tier, /*store_crc=*/false, sink); + if (!tier_has_stats(tier)) return; + sink->put_varint64(static_cast(e.prx_bytes.size())); + sink->put_bytes(Slice(e.prx_bytes)); +} + +void write_body(const DictEntry& e, std::string_view prev, IndexTier tier, bool term_stats, + ByteSink* sink) { + write_term_key(e, prev, sink); + sink->put_u8(pack_flags(e)); + write_stats(e, tier, term_stats, sink); + if (e.kind == DictEntryKind::kInline) { + write_inline(e, tier, sink); + } else { + write_pod_ref(e, tier, sink); + } +} + +// ---- Decode entry body ---- + +Status read_term_key(ByteSource* src, std::string_view prev, DictEntry* out) { + uint32_t prefix = 0; + uint32_t suffix_len = 0; + RETURN_IF_ERROR(src->get_varint32(&prefix)); + RETURN_IF_ERROR(src->get_varint32(&suffix_len)); + if (prefix > prev.size()) { + return Status::Error( + "dict_entry: prefix_len exceeds prev_term length"); + } + Slice suffix; + RETURN_IF_ERROR(src->get_bytes(suffix_len, &suffix)); + out->term.assign(prev.substr(0, prefix)); + out->term.append(reinterpret_cast(suffix.data()), suffix.size()); + return Status::OK(); +} + +Status read_stats(ByteSource* src, IndexTier tier, bool term_stats, DictEntry* out) { + RETURN_IF_ERROR(src->get_varint32(&out->df)); + out->term_stats_present = tier_has_stats(tier) && term_stats; + if (!out->term_stats_present) return Status::OK(); + RETURN_IF_ERROR(src->get_varint64(&out->ttf_delta)); + RETURN_IF_ERROR(src->get_varint64(&out->max_freq)); + return Status::OK(); +} + +// Reads the slim/inline region codec metadata (mode/uncomp/[crc]) and fills the +// dd/freq region disk_len from the supplied total/split lengths. has_crc=false +// (INLINE entries, format v2) means no per-region crc was stored: the on-disk +// crc field is absent and region decode must skip crc verification (verify_crc= +// false) since the dict block's own crc32c already covers the inline bytes. +Status read_region_meta(ByteSource* src, IndexTier tier, bool has_crc, uint64_t dd_disk_len, + uint64_t freq_disk_len, DictEntry* out) { + uint8_t mode = 0; + RETURN_IF_ERROR(src->get_u8(&mode)); + if ((mode & ~0x3u) != 0) { + return Status::Error( + "dict_entry: unknown win_mode bits"); + } + out->dd_meta.zstd = (mode & (1u << 0)) != 0; + out->dd_meta.disk_len = dd_disk_len; + out->dd_meta.verify_crc = has_crc; + RETURN_IF_ERROR(src->get_varint64(&out->dd_meta.uncomp_len)); + if (has_crc) RETURN_IF_ERROR(src->get_fixed32(&out->dd_meta.crc)); + if (!tier_has_stats(tier)) { + if (mode & (1u << 1)) { + return Status::Error( + "dict_entry: freq mode set without freq tier"); + } + return Status::OK(); + } + out->freq_meta.zstd = (mode & (1u << 1)) != 0; + out->freq_meta.disk_len = freq_disk_len; + out->freq_meta.verify_crc = has_crc; + RETURN_IF_ERROR(src->get_varint64(&out->freq_meta.uncomp_len)); + if (has_crc) RETURN_IF_ERROR(src->get_fixed32(&out->freq_meta.crc)); + return Status::OK(); +} + +Status read_pod_ref(ByteSource* src, IndexTier tier, DictEntry* out) { + RETURN_IF_ERROR(src->get_varint64(&out->frq_off_delta)); + RETURN_IF_ERROR(src->get_varint64(&out->frq_len)); + if (out->enc == DictEntryEnc::kWindowed) { + RETURN_IF_ERROR(src->get_varint64(&out->prelude_len)); + RETURN_IF_ERROR(src->get_varint64(&out->frq_docs_len)); + if (out->prelude_len == 0 || out->prelude_len > out->frq_docs_len || + out->frq_docs_len > out->frq_len) { + return Status::Error( + "dict_entry: invalid windowed docs prefix"); + } + } else { + RETURN_IF_ERROR(src->get_varint64(&out->frq_docs_len)); + if (out->frq_docs_len > out->frq_len) { + return Status::Error( + "dict_entry: frq_docs_len exceeds frq_len"); + } + RETURN_IF_ERROR(read_region_meta(src, tier, /*has_crc=*/true, out->frq_docs_len, + out->frq_len - out->frq_docs_len, out)); + } + if (!tier_has_stats(tier)) return Status::OK(); + RETURN_IF_ERROR(src->get_varint64(&out->prx_off_delta)); + RETURN_IF_ERROR(src->get_varint64(&out->prx_len)); + return Status::OK(); +} + +Status read_byte_blob(ByteSource* src, std::vector* out) { + uint64_t len = 0; + RETURN_IF_ERROR(src->get_varint64(&len)); + Slice bytes; + RETURN_IF_ERROR(src->get_bytes(static_cast(len), &bytes)); + out->assign(bytes.data(), bytes.data() + bytes.size()); + return Status::OK(); +} + +Status read_inline(ByteSource* src, IndexTier tier, DictEntry* out) { + RETURN_IF_ERROR(read_byte_blob(src, &out->frq_bytes)); + RETURN_IF_ERROR(src->get_varint64(&out->inline_dd_disk_len)); + if (out->inline_dd_disk_len > out->frq_bytes.size()) { + return Status::Error( + "dict_entry: inline_dd_disk_len exceeds frq_bytes"); + } + const uint64_t freq_disk_len = + static_cast(out->frq_bytes.size()) - out->inline_dd_disk_len; + // INLINE entries store no per-region crc (covered by the block crc): + // has_crc=false. + RETURN_IF_ERROR(read_region_meta(src, tier, /*has_crc=*/false, out->inline_dd_disk_len, + freq_disk_len, out)); + if (!tier_has_stats(tier)) return Status::OK(); + RETURN_IF_ERROR(read_byte_blob(src, &out->prx_bytes)); + return Status::OK(); +} + +Status read_locator(ByteSource* src, IndexTier tier, DictEntry* out) { + if (out->kind == DictEntryKind::kInline) return read_inline(src, tier, out); + return read_pod_ref(src, tier, out); +} + +// Read entry_len (= body length) and verify that src has enough remaining +// bytes. +Status read_entry_len(ByteSource* src, uint64_t* total) { + RETURN_IF_ERROR(src->get_varint64(total)); + if (*total > src->remaining()) { + return Status::Error( + "dict_entry: entry_len out of range"); + } + return Status::OK(); +} + +} // namespace + +Status encode_dict_entry(const DictEntry& entry, std::string_view prev_term, IndexTier tier, + ByteSink* sink, bool term_stats) { + ByteSink body_scratch; + return encode_dict_entry(entry, prev_term, tier, sink, term_stats, &body_scratch); +} + +Status encode_dict_entry(const DictEntry& entry, std::string_view prev_term, IndexTier tier, + ByteSink* sink, bool term_stats, ByteSink* body_scratch) { + if (sink == nullptr || body_scratch == nullptr || sink == body_scratch) { + return Status::Error( + "dict_entry: sink and body_scratch must be non-null and distinct"); + } + + // Serialize the body into a temporary buffer first to obtain the exact + // length, then write entry_len + body. CRC verification is done uniformly at + // the DICT block level (covering block header + all entries + anchor table); + // CRC is not repeated at the entry level, to keep slim/inline low-frequency + // terms maximally compact (spec §DICT block/§dict entry). + body_scratch->clear(); + write_body(entry, prev_term, tier, term_stats, body_scratch); + sink->put_varint64(static_cast(body_scratch->size())); + sink->put_bytes(body_scratch->view()); + return Status::OK(); +} + +Status decode_dict_entry_key(ByteSource* src, std::string_view prev_term, DictEntry* out, + size_t* body_start, uint64_t* entry_total) { + if (src == nullptr || out == nullptr) { + return Status::Error("dict_entry: src / out is null"); + } + *out = DictEntry {}; + + uint64_t total = 0; + RETURN_IF_ERROR(read_entry_len(src, &total)); + *body_start = src->position(); + *entry_total = total; + + return read_term_key(src, prev_term, out); +} + +Status decode_dict_entry_rest(ByteSource* src, IndexTier tier, size_t body_start, + uint64_t entry_total, DictEntry* out, bool term_stats) { + if (src == nullptr || out == nullptr) { + return Status::Error("dict_entry: src / out is null"); + } + // Body-decode seam: count exactly the entries whose body we materialize. + body_decode_atomic().fetch_add(1, std::memory_order_relaxed); + + uint8_t flags = 0; + RETURN_IF_ERROR(src->get_u8(&flags)); + apply_flags(flags, out); + RETURN_IF_ERROR(read_stats(src, tier, term_stats, out)); + RETURN_IF_ERROR(read_locator(src, tier, out)); + + // The body must consume exactly entry_len bytes; otherwise the structure is + // inconsistent with the tier. + const size_t consumed = src->position() - body_start; + if (consumed != static_cast(entry_total)) { + return Status::Error( + "dict_entry: body length does not match entry_len"); + } + return Status::OK(); +} + +Status skip_dict_entry_body(ByteSource* src, size_t body_start, uint64_t entry_total) { + if (src == nullptr) { + return Status::Error("dict_entry: src is null"); + } + const size_t consumed = src->position() - body_start; + if (consumed > static_cast(entry_total)) { + return Status::Error( + "dict_entry: term key overruns entry body"); + } + const size_t advance = static_cast(entry_total) - consumed; + Slice unused; + return src->get_bytes(advance, &unused); +} + +Status decode_dict_entry(ByteSource* src, std::string_view prev_term, IndexTier tier, + DictEntry* out, bool term_stats) { + size_t body_start = 0; + uint64_t entry_total = 0; + RETURN_IF_ERROR(decode_dict_entry_key(src, prev_term, out, &body_start, &entry_total)); + return decode_dict_entry_rest(src, tier, body_start, entry_total, out, term_stats); +} + +Status skip_dict_entry(ByteSource* src) { + if (src == nullptr) + return Status::Error("dict_entry: src is null"); + uint64_t total = 0; + RETURN_IF_ERROR(read_entry_len(src, &total)); + Slice unused; + return src->get_bytes(static_cast(total), &unused); +} + +uint64_t dict_entry_body_decode_count() { + return body_decode_atomic().load(std::memory_order_relaxed); +} + +void reset_dict_entry_counters() { + body_decode_atomic().store(0, std::memory_order_relaxed); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/dict_entry.h b/be/src/storage/index/snii/format/dict_entry.h new file mode 100644 index 00000000000000..df427080067f01 --- /dev/null +++ b/be/src/storage/index/snii/format/dict_entry.h @@ -0,0 +1,184 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" + +// DictEntry —— on-disk encoding/decoding of a dict entry. +// +// Byte layout (see docs/design/SNII-design-spec.source.md "dict entry" +// section): +// entry_len varint # byte length of entry body, allowing reader to skip +// unknown extensions or fast-skip entries +// --- entry body begins here, covered by entry_len --- +// prefix_len varint # length of shared prefix with prev_term +// suffix_len varint # number of suffix bytes +// suffix u8[] # suffix bytes that differ from prev_term +// flags u8 # bit0 kind / bit1 enc / bit2 has_sb / bit3 +// has_champion(=0) / bit4 offsets_ref(=0) df varint ttf_delta varint +// # only when tier>=T2 max_freq varint # only when tier>=T2 locator: +// pod_ref: frq_off_delta varint, frq_len varint, +// [prelude_len varint, frq_docs_len varint when enc=windowed] +// # docs-only prefix [prelude][dd-block]; windowed entries +// carry # per-window region metadata in the prelude. +// [frq_docs_len varint, slim region meta when enc=slim]: +// # frq_docs_len == dd region on-disk length; the docs-only +// prefix # [frq_off, frq_off+frq_docs_len) a docid-only reader +// fetches # without the freq region. win_mode u8 (bit0 +// dd_zstd, bit1 freq_zstd) dd_uncomp_len varint, crc_dd u32 +// [freq_uncomp_len varint, crc_freq u32 when tier>=T2] +// # The single slim window is [dd_region][freq_region]; +// dd_disk_len # = frq_docs_len, freq_disk_len = frq_len - +// frq_docs_len. +// [prx_off_delta varint, prx_len varint when tier>=T2] +// inline: frq_len varint, frq_bytes u8[], # frq_bytes = +// [dd_region][freq_region] +// slim region meta (as above, sans frq_docs_len which == dd disk +// len +// carried as inline_dd_disk_len varint), +// [prx_len varint, prx_bytes u8[] when tier>=T2] +// --- entry body ends --- +// +// CRC verification is performed at the DICT block level (covering block header +// + all entries + anchor offset table), no per-entry CRC to keep slim/inline +// low-frequency terms compact (spec §DICT block line 330/348). tier and +// positions capability are provided by Core metadata (not stored redundantly +// inside entries): when tier>=T2, ttf_delta / max_freq and .prx locator/bytes +// are written. +namespace doris::snii::format { + +// Dict entry: inline or pod-ref (two states), self-described length, supports +// intra-block front coding. +struct DictEntry { + // term key (front coding relative to prev_term is applied during + // encode/decode; full term stored here). + std::string term; + + // flags. + DictEntryKind kind = DictEntryKind::kPodRef; + DictEntryEnc enc = DictEntryEnc::kSlim; + bool has_sb = false; + + // term stats. + uint32_t df = 0; + uint64_t ttf_delta = 0; // only when tier>=T2 AND the block carries term stats + // G16-f: false when decoded from a kNoTermStats block (freq-dropped index): + // ttf_delta/max_freq above are then meaningless defaults, NOT real zeros. + // Consumers that need them (stats provider / BM25) must check this flag. + bool term_stats_present = true; + uint64_t max_freq = 0; // only when tier>=T2 + + // pod_ref locator. + uint64_t frq_off_delta = 0; + uint64_t frq_len = 0; + uint64_t prelude_len = 0; // only when enc=windowed + uint64_t frq_docs_len = 0; // pod_ref docs-only prefix length + uint64_t prx_off_delta = 0; // only when tier>=T2 + uint64_t prx_len = 0; // only when tier>=T2 + + // slim/inline single-window region codecs. The window is + // [dd_region][freq_region] (no self-describing header). dd_meta drives the + // docs-only decode; freq_meta the scoring decode (only when tier>=T2). For + // slim pod_ref dd_meta.disk_len == frq_docs_len; for inline it is stored as + // inline_dd_disk_len. + FrqRegionMeta dd_meta; + FrqRegionMeta freq_meta; // only when tier>=T2 + uint64_t inline_dd_disk_len = 0; // only for inline: dd region on-disk length + + // inline payload. + std::vector frq_bytes; // = [dd_region][freq_region] + std::vector prx_bytes; // only when tier>=T2 +}; + +// Encodes an entry into sink (appending) using the layout above, with front +// coding relative to prev_term. tier determines whether optional fields are +// written. term_stats == false (G16-f: freq-dropped indexes, declared by the +// block header's kNoTermStats flag) omits the ttf_delta/max_freq varints -- +// they serve only BM25 scoring, dead on an index that dropped freq; df stays. +// Region metadata (freq/prx locators) remains tier-conditioned. +Status encode_dict_entry(const DictEntry& entry, std::string_view prev_term, IndexTier tier, + ByteSink* sink, bool term_stats = true); + +// Same encoding with caller-owned body scratch. The scratch is cleared before +// use while retaining capacity, allowing block builders to avoid one allocation +// per dictionary entry. sink and body_scratch must be distinct objects. +Status encode_dict_entry(const DictEntry& entry, std::string_view prev_term, IndexTier tier, + ByteSink* sink, bool term_stats, ByteSink* body_scratch); + +// Decodes one entry from the current position of src; term is reconstructed +// from prev_term + suffix. Verifies the trailing CRC; out-of-range / CRC +// mismatch / invalid prefix_len all return Corruption. term_stats must match +// the writer's choice (the dict block header flag carries it). +Status decode_dict_entry(ByteSource* src, std::string_view prev_term, IndexTier tier, + DictEntry* out, bool term_stats = true); + +// Skips one entry using only entry_len (does not parse internal fields or +// verify CRC). +Status skip_dict_entry(ByteSource* src); + +// ---- Key-first decode primitives (T07) ---- +// +// decode_dict_entry is split into a "key" stage and a "rest" (body) stage so a +// caller scanning many entries can decide on the (front-coded) term key alone +// whether the entry's body is worth materializing. decode_dict_entry below is +// re-expressed as key + rest and stays byte-for-byte identical in output. + +// Reads only entry_len + the front-coded term key. On return src is positioned +// at the start of the entry body (the flags byte). out is reset to defaults and +// out->term is reconstructed from prev_term + suffix; no other field is touched. +// *body_start receives the absolute src position of the body (right after the +// entry_len varint) and *entry_total the body byte length (already bounds-checked +// against src->remaining()). Used to drive key-first scans (find_term, prefix +// streaming) that skip non-matching bodies. +Status decode_dict_entry_key(ByteSource* src, std::string_view prev_term, DictEntry* out, + size_t* body_start, uint64_t* entry_total); + +// Continues from the position decode_dict_entry_key left at: reads +// flags/stats/locator into *out and verifies the body consumed exactly +// entry_total bytes (body_start anchors the consumed count). Increments the +// body-decode counter seam (see dict_entry_body_decode_count) at its top so +// tests can assert how many entry bodies a scan actually materialized. +Status decode_dict_entry_rest(ByteSource* src, IndexTier tier, size_t body_start, + uint64_t entry_total, DictEntry* out, bool term_stats = true); + +// Skips the remaining body bytes after decode_dict_entry_key, advancing src to +// the next entry without parsing flags/stats/locator. advance = entry_total - +// (src.position() - body_start); the term key already consumed by the key stage +// is accounted for. +Status skip_dict_entry_body(ByteSource* src, size_t body_start, uint64_t entry_total); + +// Test-only instrumentation seam: dict_entry_body_decode_count() returns a +// process-global count of decode_dict_entry_rest calls -- i.e. how many entry +// bodies (flags/stats/locator + any inline byte copies) a scan materialized. +// Key-first find_term / prefix streaming drive this toward the number of +// produced hits instead of the number of entries walked. Counters use relaxed +// atomics; reset between tests. +uint64_t dict_entry_body_decode_count(); +void reset_dict_entry_counters(); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/format_constants.h b/be/src/storage/index/snii/format/format_constants.h new file mode 100644 index 00000000000000..fa8401defb4fa3 --- /dev/null +++ b/be/src/storage/index/snii/format/format_constants.h @@ -0,0 +1,137 @@ +// 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. + +#pragma once + +#include + +// SNII container and per-section on-disk contract constants. +// Once published, these values are format semantics; changes require bumping +// format_version and maintaining a compatibility policy. All multi-byte +// fixed-width fields are little-endian; variable-length integers use LEB128 +// (see snii/encoding/varint.h). +namespace doris::snii::format { + +// ---- Container-level magic / version ---- +// "SNII" reads as 0x49494E53 in little-endian. +inline constexpr uint32_t kContainerMagic = 0x49494E53u; // 'S''N''I''I' +inline constexpr uint32_t kTailMagic = 0x4C494154u; // 'T''A''I''L' +inline constexpr uint16_t kFormatVersion = 1; +inline constexpr uint16_t kMinReaderVersion = 1; + +// ---- SectionFramer type ids for standalone metadata blobs ---- +enum class SectionType : uint8_t { + kSampledTermIndex = 2, + kDictBlockDirectory = 3, + kCoreMetadataPB = 6, + // G13: zstd-compressed carriers for the two large metadata blobs + // (they are highly compressible sorted string/offset tables and dominate the + // metadata group fetched serially at open). Payload = varint64 + // uncomp_len followed by zstd(original full frame), where "original full + // frame" is the byte-exact kSampledTermIndex / kDictBlockDirectory frame + // (type+len+payload+crc32c) used by the raw layout. Decompression + // therefore reproduces the raw frame verbatim and the sub-module readers + // (which re-verify the inner crc) stay unchanged. The writer emits these ONLY + // when the raw frame reaches kMetaSectionCompressMinBytes AND compression + // shrinks it; otherwise it emits the raw frame. + kSampledTermIndexZstd = 11, + kDictBlockDirectoryZstd = 12, + // Per-document one-byte BM25 norms. This must remain distinct from + // Core metadata so a corrupt section reference cannot reinterpret valid + // collection statistics as document norms. + kNormsPod = 14, +}; + +// ---- Logical index postings storage content configuration (fixed per logical +// index, not per-term) ---- Determines whether to write freq / positions / +// norms+stats. +enum class IndexConfig : uint8_t { + kDocsOnly = 0, // docid only: term/match filtering + kDocsPositions = 1, // docid+positions (+freq only when the caller keeps + // it -- SniiIndexInput::write_freq, G16-c): MATCH_PHRASE + kDocsPositionsScoring = 2, // + norms + stats: phrase + BM25 + kPositionsOffsets = 3, // reserved (highlight/RAG), not implemented in this release +}; + +// term stats / postings capability tiers: only tier>=kT2 writes +// ttf_delta/max_freq and .prx. +enum class IndexTier : uint8_t { + kT1 = 1, // docs-only + kT2 = 2, // docs-positions + kT3 = 3, // docs-positions-scoring +}; + +inline constexpr IndexTier tier_of(IndexConfig cfg) { + return cfg == IndexConfig::kDocsOnly ? IndexTier::kT1 + : cfg == IndexConfig::kDocsPositions ? IndexTier::kT2 + : IndexTier::kT3; // scoring / offsets +} +inline constexpr bool has_positions(IndexConfig cfg) { + return cfg != IndexConfig::kDocsOnly; +} +inline constexpr bool has_scoring(IndexConfig cfg) { + return cfg == IndexConfig::kDocsPositionsScoring; +} + +// ---- DictEntry flags bit definitions ---- +namespace dict_flags { +inline constexpr uint8_t kKind = 1u << 0; // 0=pod_ref / 1=inline +inline constexpr uint8_t kEnc = 1u << 1; // 0=slim / 1=windowed +inline constexpr uint8_t kHasSb = 1u << 2; // posting prelude includes sub-block directory +inline constexpr uint8_t kHasChampion = 1u << 3; // v1 always 0 +inline constexpr uint8_t kOffsetsRef = 1u << 4; // v1 always 0 +// bit5-7 reserved +} // namespace dict_flags + +enum class DictEntryKind : uint8_t { kPodRef = 0, kInline = 1 }; +enum class DictEntryEnc : uint8_t { kSlim = 0, kWindowed = 1 }; + +// ---- .prx window codec (codec byte bit0-5) ---- +// kRaw : plaintext varint payload (doc_count, per-doc pos_count + position +// deltas). kZstd : zstd-compressed plaintext payload (legacy reader still +// supported). kPfor : doc_count + per-doc pos_count (varint), then position +// deltas bit-packed +// as PFOR runs (kFrqBaseUnit each). No entropy coding -> far cheaper +// build CPU than zstd while staying competitive on size for ascending +// deltas. +enum class PrxCodec : uint8_t { + kRaw = 0, + kZstd = 1, + kPfor = 2 /* bit7 cont-reserved */ +}; + +// ---- Build-time parameters (not format semantics; may be tuned against real +// metrics) ---- +inline constexpr uint32_t kFrqBaseUnit = 256; // window base unit +inline constexpr uint32_t kSlimDfThreshold = 512; // df < this → slim +inline constexpr uint32_t kDefaultInlineThreshold = 256; // slim encoded bytes ≤ this → inline +// Adaptive window sizing (design #4): high-df windowed terms use larger windows +// to cut prelude rows + per-window header/crc overhead. Windows remain a whole +// multiple of kFrqBaseUnit so .prx alignment and win_base/last_docid semantics +// are preserved. A term whose df >= kAdaptiveWindowDfThreshold splits into +// kAdaptiveWindowDocs-sized windows instead of kFrqBaseUnit-sized ones. +inline constexpr uint32_t kAdaptiveWindowDfThreshold = 8192; // df >= this -> larger windows +inline constexpr uint32_t kAdaptiveWindowDocs = 1024; // larger window size (4 * base unit) +inline constexpr uint32_t kDefaultTargetDictBlockBytes = 64 * 1024; +// G13: SampledTermIndex / DictBlockDirectory metadata frames +// at or above this raw size are emitted zstd-compressed (kSampledTermIndexZstd / +// kDictBlockDirectoryZstd); smaller ones stay raw -- compression overhead is not +// worth it below a few KB. A build-time parameter, not format semantics: readers +// accept both layouts regardless of the value. +inline constexpr size_t kMetaSectionCompressMinBytes = 4 * 1024; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/frq_pod.cpp b/be/src/storage/index/snii/format/frq_pod.cpp new file mode 100644 index 00000000000000..f2d835eb019f36 --- /dev/null +++ b/be/src/storage/index/snii/format/frq_pod.cpp @@ -0,0 +1,439 @@ +// 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 "storage/index/snii/format/frq_pod.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/common/uninitialized_buffer.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/prx_decode_stats.h" + +namespace doris::snii::testing { +void note_frq_dd_validation_doc_visits(uint64_t count); +void note_frq_dd_materialized_values(uint64_t count); +void note_frq_raw_region_copy_bytes(uint64_t count); +} // namespace doris::snii::testing + +namespace doris::snii::format { + +namespace { + +// Auto-compression threshold: use raw when a region is smaller than this byte +// count (zstd gain is negligible and metadata overhead is relatively large). +inline constexpr size_t kAutoZstdMinBytes = 512; +// Default zstd level for auto mode. +inline constexpr int kDefaultZstdLevel = 3; +// Maximum decompressed byte size for a single region. Guards against a +// corrupted uncomp_len read from S3 that inflated to a huge value: sanity-check +// before allocating/decompressing to avoid GB-scale allocations. Windows are +// 256-doc aligned and normally far smaller than this. +inline constexpr uint32_t kMaxRegionUncompBytes = 256u * 1024 * 1024; +// Maximum doc count per .frq window (guards against a corrupted n). Window +// baseline is 256, practical combined cap is 2048, so this is a loose but +// astronomically-large-number-blocking upper bound. +inline constexpr uint32_t kMaxWindowDocs = 1u << 24; + +// Encode a uint32 array into multiple PFOR runs, each of 256 (kFrqBaseUnit) +// elements. n / run count is not written: the number of runs is derived from +// total length n and kFrqBaseUnit, and the decoder computes it the same way. +void encode_pfor_runs(std::span values, ByteSink* out) { + size_t n = values.size(); + for (size_t off = 0; off < n; off += kFrqBaseUnit) { + size_t run = (n - off < kFrqBaseUnit) ? (n - off) : kFrqBaseUnit; + pfor_encode(values.data() + off, run, out); + } +} + +// Decode n uint32 values from source (multiple PFOR runs of 256 each). +Status decode_pfor_runs(ByteSource* src, size_t n, std::vector* out) { + // Sized then fully overwritten by pfor_decode below; no zero-fill needed. + resize_uninitialized(*out, n); + for (size_t off = 0; off < n; off += kFrqBaseUnit) { + size_t run = (n - off < kFrqBaseUnit) ? (n - off) : kFrqBaseUnit; + RETURN_IF_ERROR(pfor_decode(src, run, out->data() + off)); + } + return Status::OK(); +} + +// Verifies docids are ascending and the first entry is not below win_base. +Status validate_docs(std::span docs, uint64_t win_base) { + if (docs.empty()) return Status::OK(); +#ifdef BE_TEST + ::doris::snii::testing::note_frq_dd_validation_doc_visits(1); +#endif + if (static_cast(docs.front()) < win_base) { + return Status::Error("frq: first docid below win_base"); + } + for (size_t i = 1; i < docs.size(); ++i) { +#ifdef BE_TEST + ::doris::snii::testing::note_frq_dd_validation_doc_visits(1); +#endif + if (docs[i] < docs[i - 1]) { + return Status::Error( + "frq: docids must be ascending"); + } + } + return Status::OK(); +} + +// Decision: given level and plaintext length, determine whether to compress. +bool should_compress(int level, size_t plain_len) { + if (level == 0) return false; // force raw + if (level > 0) return true; // force zstd + return plain_len >= kAutoZstdMinBytes; // auto +} + +// Encodes one region's plaintext into raw or zstd, appends the on-disk bytes to +// out, and fills meta (mode/uncomp_len/disk_len/crc). The region carries no +// header. +Status emit_region(Slice plain, int level, ByteSink* out, FrqRegionMeta* meta) { + if (out == nullptr || meta == nullptr) { + return Status::Error("frq: null region out"); + } + meta->uncomp_len = plain.size(); + if (should_compress(level, plain.size())) { + // zstd needs its own buffer: the compressed bytes differ from `plain`. + std::vector disk; + meta->zstd = true; + RETURN_IF_ERROR(zstd_compress(plain, level > 0 ? level : kDefaultZstdLevel, &disk)); + meta->disk_len = static_cast(disk.size()); + meta->crc = crc32c(Slice(disk)); + out->put_bytes(Slice(disk)); + return Status::OK(); + } + // Raw: the on-disk bytes ARE `plain` (a view over the caller's contiguous + // ByteSink), so crc and emit straight from it -- no temp `disk` alloc/copy. + // disk_len MUST stay == plain.size(): open_region enforces uncomp_len == + // disk_len for raw regions. Byte-identical to the former disk.assign() path + // (disk == plain, so crc32c(disk) == crc32c(plain), put_bytes(disk) == same + // bytes). + meta->zstd = false; + meta->disk_len = static_cast(plain.size()); + meta->crc = crc32c(plain); +#ifdef BE_TEST + ::doris::snii::testing::note_frq_raw_region_copy_bytes(plain.size()); +#endif + out->put_bytes(plain); + return Status::OK(); +} + +void finish_raw_region(ByteSink* out, size_t begin, FrqRegionMeta* meta) { + const size_t length = out->size() - begin; + const Slice appended = length == 0 ? Slice() : Slice(out->buffer().data() + begin, length); + meta->zstd = false; + meta->uncomp_len = static_cast(length); + meta->disk_len = static_cast(length); + meta->crc = crc32c(appended); +} + +void rollback_raw_region(ByteSink* out, size_t begin) { + ByteSink restored; + if (begin != 0) { + restored.put_bytes(Slice(out->buffer().data(), begin)); + } + *out = std::move(restored); +} + +Status append_raw_dd_region(std::span docs, uint64_t win_base, ByteSink* out, + FrqRegionMeta* meta) { + if (!docs.empty() && static_cast(docs.front()) < win_base) { + return Status::Error("frq: first docid below win_base"); + } + + const size_t begin = out->size(); + out->put_varint32(static_cast(docs.size())); + std::array deltas; + uint64_t previous = win_base; + for (size_t offset = 0; offset < docs.size(); offset += kFrqBaseUnit) { + const size_t count = std::min(docs.size() - offset, static_cast(kFrqBaseUnit)); + for (size_t i = 0; i < count; ++i) { + const uint32_t doc = docs[offset + i]; + if (static_cast(doc) < previous) { + rollback_raw_region(out, begin); + return Status::Error( + "frq: docids must be ascending"); + } + deltas[i] = static_cast(static_cast(doc) - previous); + previous = doc; + } + pfor_encode(deltas.data(), count, out); + } + finish_raw_region(out, begin, meta); + return Status::OK(); +} + +// Materializes a region's plaintext (raw borrows the view; zstd decompresses) +// and verifies its crc + slice length against meta. +Status open_region(Slice disk, const FrqRegionMeta& meta, std::vector* holder, + PrxCsrAllocationGate* allocation_gate, Slice* plain) { + if (disk.size() != static_cast(meta.disk_len)) { + return Status::Error( + "frq: region slice length mismatch"); + } + if (meta.uncomp_len > kMaxRegionUncompBytes) { + return Status::Error( + "frq: region uncomp_len exceeds sane cap"); + } + // Inline entries (verify_crc=false) carry no per-region crc: their on-disk + // bytes are covered by the enclosing dict block's block-level crc32c, so the + // region crc would be redundant. POD-ref regions keep their own crc check. + if (meta.verify_crc && crc32c(disk) != meta.crc) { + return Status::Error( + "frq: region crc mismatch"); + } + if (!meta.zstd) { + if (meta.uncomp_len != meta.disk_len) { + return Status::Error( + "frq: raw region length inconsistent"); + } + *plain = disk; + return Status::OK(); + } + if (allocation_gate != nullptr) { + RETURN_IF_ERROR(allocation_gate->reserve_decompression(static_cast(meta.uncomp_len), + &holder)); + DCHECK(holder != nullptr); + } + RETURN_IF_ERROR(zstd_decompress(disk, static_cast(meta.uncomp_len), holder)); + *plain = Slice(*holder); + return Status::OK(); +} + +} // namespace + +Status build_dd_region(std::span docids_ascending, uint64_t win_base, + int zstd_level_or_neg_for_auto, ByteSink* out, FrqRegionMeta* meta) { + if (out == nullptr || meta == nullptr) { + return Status::Error("frq: null dd region out"); + } + if (zstd_level_or_neg_for_auto == 0) { + return append_raw_dd_region(docids_ascending, win_base, out, meta); + } + RETURN_IF_ERROR(validate_docs(docids_ascending, win_base)); + ByteSink plain; // VInt n ++ PFOR_runs(doc_delta) + std::vector dd(docids_ascending.size()); +#ifdef BE_TEST + ::doris::snii::testing::note_frq_dd_materialized_values(dd.size()); +#endif + uint64_t prev = win_base; + for (size_t i = 0; i < docids_ascending.size(); ++i) { + dd[i] = static_cast(static_cast(docids_ascending[i]) - prev); + prev = docids_ascending[i]; + } + plain.put_varint32(static_cast(docids_ascending.size())); + encode_pfor_runs(dd, &plain); + return emit_region(plain.view(), zstd_level_or_neg_for_auto, out, meta); +} + +Status build_dd_region_from_deltas(std::span doc_deltas, + int zstd_level_or_neg_for_auto, ByteSink* out, + FrqRegionMeta* meta) { + if (out == nullptr || meta == nullptr) { + return Status::Error("frq: null dd region out"); + } + if (zstd_level_or_neg_for_auto != 0) { + return Status::Error( + "frq: direct doc-delta encoding requires raw level 0"); + } + + const size_t begin = out->size(); + out->put_varint32(static_cast(doc_deltas.size())); + encode_pfor_runs(doc_deltas, out); + finish_raw_region(out, begin, meta); + return Status::OK(); +} + +} // namespace doris::snii::format + +namespace doris::snii::testing { +#ifdef BE_TEST +namespace { +std::atomic validation_visits {0}; +std::atomic materialized_values {0}; +std::atomic raw_copy_bytes {0}; +} // namespace +#endif + +void note_frq_dd_validation_doc_visits(uint64_t count) { +#ifdef BE_TEST + validation_visits.fetch_add(count, std::memory_order_relaxed); +#endif +} +void note_frq_dd_materialized_values(uint64_t count) { +#ifdef BE_TEST + materialized_values.fetch_add(count, std::memory_order_relaxed); +#endif +} +void note_frq_raw_region_copy_bytes(uint64_t count) { +#ifdef BE_TEST + raw_copy_bytes.fetch_add(count, std::memory_order_relaxed); +#endif +} +void reset_frq_raw_encode_work() { +#ifdef BE_TEST + validation_visits.store(0, std::memory_order_relaxed); + materialized_values.store(0, std::memory_order_relaxed); + raw_copy_bytes.store(0, std::memory_order_relaxed); +#endif +} +uint64_t frq_dd_validation_doc_visits() { +#ifdef BE_TEST + return validation_visits.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t frq_dd_materialized_values() { +#ifdef BE_TEST + return materialized_values.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t frq_raw_region_copy_bytes() { +#ifdef BE_TEST + return raw_copy_bytes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +} // namespace doris::snii::testing + +namespace doris::snii::format { + +Status build_freq_region(std::span freqs, int zstd_level_or_neg_for_auto, + ByteSink* out, FrqRegionMeta* meta) { + if (out == nullptr || meta == nullptr) { + return Status::Error("frq: null freq region out"); + } + if (zstd_level_or_neg_for_auto == 0) { + const size_t begin = out->size(); + encode_pfor_runs(freqs, out); + finish_raw_region(out, begin, meta); + return Status::OK(); + } + ByteSink plain; + encode_pfor_runs(freqs, &plain); + return emit_region(plain.view(), zstd_level_or_neg_for_auto, out, meta); +} + +namespace { + +Status decode_dd_region_impl(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + const uint32_t* expected_doc_count, + PrxCsrAllocationGate* allocation_gate, std::vector* docids) { + if (docids == nullptr) + return Status::Error("frq: null docids out"); + std::vector holder; + Slice plain; + RETURN_IF_ERROR(open_region(dd_disk, meta, &holder, allocation_gate, &plain)); + ByteSource src(plain); + uint32_t n = 0; + RETURN_IF_ERROR(src.get_varint32(&n)); + if (n > kMaxWindowDocs) + return Status::Error( + "frq: doc count exceeds sane cap"); + if (expected_doc_count != nullptr && n != *expected_doc_count) { + return Status::Error( + "frq: encoded doc count differs from metadata"); + } + RETURN_IF_ERROR(decode_pfor_runs(&src, n, docids)); + if (!src.eof()) { + return Status::Error( + "frq: trailing bytes after dd region payload"); + } + uint64_t cur = win_base; + if (n != 0 && cur > std::numeric_limits::max()) { + return Status::Error( + "frq: window base exceeds uint32 docid range"); + } + for (uint32_t i = 0; i < n; ++i) { + const uint32_t delta = (*docids)[i]; + if (i != 0 && delta == 0) { + return Status::Error( + "frq: zero docid delta"); + } + if (delta > std::numeric_limits::max() - cur) { + return Status::Error( + "frq: docid accumulation overflow"); + } + cur += delta; + (*docids)[i] = static_cast(cur); + } + return Status::OK(); +} + +} // namespace + +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + std::vector* docids) { + return decode_dd_region_impl(dd_disk, meta, win_base, nullptr, nullptr, docids); +} + +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + uint32_t expected_doc_count, std::vector* docids) { + return decode_dd_region_impl(dd_disk, meta, win_base, &expected_doc_count, nullptr, docids); +} + +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + uint32_t expected_doc_count, PrxCsrAllocationGate* allocation_gate, + std::vector* docids) { + if (allocation_gate == nullptr) { + return Status::Error( + "frq: allocation gate must be non-null"); + } + return decode_dd_region_impl(dd_disk, meta, win_base, &expected_doc_count, allocation_gate, + docids); +} + +Status decode_freq_region(Slice freq_disk, const FrqRegionMeta& meta, size_t doc_count, + std::vector* freqs) { + if (freqs == nullptr) + return Status::Error("frq: null freqs out"); + std::vector holder; + Slice plain; + RETURN_IF_ERROR(open_region(freq_disk, meta, &holder, nullptr, &plain)); + if (doc_count == 0) { + if (meta.uncomp_len != 0) { + return Status::Error( + "frq: empty freq region expected"); + } + freqs->clear(); + return Status::OK(); + } + ByteSource src(plain); + RETURN_IF_ERROR(decode_pfor_runs(&src, doc_count, freqs)); + if (!src.eof()) { + return Status::Error( + "frq: trailing bytes after freq region payload"); + } + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/frq_pod.h b/be/src/storage/index/snii/format/frq_pod.h new file mode 100644 index 00000000000000..4a3ca5d95b137c --- /dev/null +++ b/be/src/storage/index/snii/format/frq_pod.h @@ -0,0 +1,150 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +// .frq region codec (FrqPod): doc-delta (dd) and freq postings, columnar + PFOR +// (see docs/design SNII "frq design" and the read-byte-optimizations +// design 1.6). +// +// PHASE D (posting-level dd/freq grouping): windows are NO LONGER +// self-describing. A windowed .frq payload is laid out as +// [prelude][dd-block][freq-block] +// where the dd-block concatenates every window's dd_region and the freq-block +// concatenates every window's freq_region. Each region is independently encoded +// (raw or zstd, chosen by size) and the per-window codec metadata (mode, +// lengths, crc, offsets) is hoisted into the frq_prelude rows -- the region +// bytes carry NO header. This makes the docs-only prefix ([prelude][dd-block]) +// ONE contiguous run a docid-only / phrase reader can fetch in a single range, +// skipping the freq-block entirely. +// +// dd_region plaintext = VInt n ++ PFOR_runs(doc_delta) # n = doc count +// dd[0] = first_docid - win_base; dd[i] = docid[i] - docid[i-1]; win_base is +// the previous window's last docid (first window = 0). +// freq_region plaintext = PFOR_runs(freq) # present iff +// has_freq PFOR runs are segmented at 256 docs (kFrqBaseUnit); a partial +// segment writes the remainder. Variable-length integers reuse +// snii/encoding/varint; PFOR reuses snii/encoding/pfor; crc32c covers each +// region's ON-DISK bytes. +namespace doris::snii::format { + +class PrxCsrAllocationGate; + +// Codec metadata for ONE encoded region (dd or freq), hoisted into the prelude. +// The region's on-disk bytes are pure payload (no header); these fields drive +// the decode. crc covers the on-disk (disk_len) bytes. +struct FrqRegionMeta { + bool zstd = false; // true => disk bytes are zstd(plaintext); false => raw + uint64_t uncomp_len = 0; // plaintext byte length (== disk_len when raw) + uint64_t disk_len = 0; // on-disk byte length of this region + uint32_t crc = 0; // crc32c of the on-disk (disk_len) bytes + // When false, decode_*_region SKIPS the per-region crc check (and the writer + // omits the 4-byte crc from the dict entry). Set false for INLINE entries: + // their region bytes live inside the dict block, whose own block-level crc32c + // already covers them, so a per-region crc is fully redundant. POD-ref + // regions (slim/windowed) live in the separately-fetched .frq POD -- their + // crc stays. + bool verify_crc = true; +}; + +// Encodes a window's dd_region plaintext (VInt n ++ PFOR_runs(doc_delta)) into +// raw or zstd (per zstd_level_or_neg_for_auto), APPENDS the on-disk bytes to +// out, and fills meta (mode/uncomp_len/disk_len/crc). The region carries no +// header. docids_ascending: ascending docids in this window (single doc or +// empty allowed). win_base: previous window's last docid (first window = 0); +// requires docids[0] >= win_base. zstd_level_or_neg_for_auto: <0 auto (zstd +// when large enough, else raw); 0 force +// raw; >0 force zstd at that level. +// Non-ascending docids / first_docid < win_base / null out returns +// InvalidArgument. +Status build_dd_region(std::span docids_ascending, uint64_t win_base, + int zstd_level_or_neg_for_auto, ByteSink* out, FrqRegionMeta* meta); + +// Trusted writer fast path for a window whose doc deltas have already been +// produced and validated by the posting stream. Appends the same +// VInt n ++ PFOR_runs(doc_delta) bytes as build_dd_region(..., level=0), without +// reconstructing absolute docids or scanning the deltas. Only raw level 0 is +// supported; any other level returns InvalidArgument without modifying out or +// meta. +Status build_dd_region_from_deltas(std::span doc_deltas, + int zstd_level_or_neg_for_auto, ByteSink* out, + FrqRegionMeta* meta); + +// Vector convenience overload (forwards a span view; no copy of the elements). +inline Status build_dd_region(const std::vector& docids_ascending, uint64_t win_base, + int zstd_level_or_neg_for_auto, ByteSink* out, FrqRegionMeta* meta) { + return build_dd_region(std::span(docids_ascending), win_base, + zstd_level_or_neg_for_auto, out, meta); +} + +// Encodes a window's freq_region plaintext (PFOR_runs(freq)) into raw or zstd, +// APPENDS the on-disk bytes to out, and fills meta. Empty freqs yields a +// zero-length region. Null out returns InvalidArgument. +Status build_freq_region(std::span freqs, int zstd_level_or_neg_for_auto, + ByteSink* out, FrqRegionMeta* meta); + +// Vector convenience overload (forwards a span view; no copy of the elements). +inline Status build_freq_region(const std::vector& freqs, int zstd_level_or_neg_for_auto, + ByteSink* out, FrqRegionMeta* meta) { + return build_freq_region(std::span(freqs), zstd_level_or_neg_for_auto, out, + meta); +} + +// Decodes a dd_region from its on-disk slice (exactly disk_len bytes) + meta + +// win_base, reconstructing ascending docids. Verifies meta.crc against the +// slice. crc mismatch / wrong slice length / truncation / decompression / +// oversized count all return a non-OK Status. The freq region is irrelevant +// here (docs-only path). +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + std::vector* docids); +// Compaction variant: validates the encoded count before resizing `docids`, so +// a corrupt frame cannot grow a pre-reserved destination buffer past its +// metadata-derived hard budget. +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + uint32_t expected_doc_count, std::vector* docids); +// Compaction variant that charges a zstd decompression buffer before allocation +// and may reuse the caller's bounded decoder workspace. +Status decode_dd_region(Slice dd_disk, const FrqRegionMeta& meta, uint64_t win_base, + uint32_t expected_doc_count, PrxCsrAllocationGate* allocation_gate, + std::vector* docids); + +// Decodes a freq_region from its on-disk slice (exactly disk_len bytes) + meta, +// producing doc_count freqs. Verifies meta.crc. doc_count == 0 yields empty +// freqs (and requires a zero-length region). crc mismatch / wrong slice length +// / etc. return a non-OK Status. +Status decode_freq_region(Slice freq_disk, const FrqRegionMeta& meta, size_t doc_count, + std::vector* freqs); + +} // namespace doris::snii::format + +// Test-only work counters for the level-0 writer path. They expose redundant +// work that a direct raw encoder must eliminate without making it part of the +// production codec API. +namespace doris::snii::testing { +void reset_frq_raw_encode_work(); +uint64_t frq_dd_validation_doc_visits(); +uint64_t frq_dd_materialized_values(); +uint64_t frq_raw_region_copy_bytes(); +} // namespace doris::snii::testing diff --git a/be/src/storage/index/snii/format/frq_prelude.cpp b/be/src/storage/index/snii/format/frq_prelude.cpp new file mode 100644 index 00000000000000..9cec931244d7d4 --- /dev/null +++ b/be/src/storage/index/snii/format/frq_prelude.cpp @@ -0,0 +1,642 @@ +// 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 "storage/index/snii/format/frq_prelude.h" + +#include +#include +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" + +namespace doris::snii::format { + +namespace { + +// Anti-DoS: a segment holds at most ~15M docs (>=1 doc/window), so 1<<24 +// windows is a generous ceiling that still prevents multi-GB allocations from a +// crafted N. (crc32c is not a MAC and cannot defend a re-stamped inflated count.) +constexpr uint64_t kMaxWindows = 1ull << 24; + +uint64_t ceil_div(uint64_t a, uint64_t b) { + return (a + b - 1) / b; +} + +uint8_t make_flags(const FrqPreludeColumns& cols) { + uint8_t flags = 0; + if (cols.has_freq) flags |= frq_prelude_flags::kHasFreq; + if (cols.has_prx) flags |= frq_prelude_flags::kHasPrx; + return flags; +} + +uint8_t make_win_mode(const WindowMeta& m, bool has_freq) { + uint8_t mode = 0; + if (m.dd_zstd) mode |= frq_win_mode::kDdZstd; + if (has_freq && m.freq_zstd) mode |= frq_win_mode::kFreqZstd; + return mode; +} + +Status checked_add_u64(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (rhs > std::numeric_limits::max() - lhs) { + return Status::Error(message); + } + *out = lhs + rhs; + return Status::OK(); +} + +Status checked_u32(uint64_t value, const char* message, uint32_t* out) { + if (value > std::numeric_limits::max()) { + return Status::Error(message); + } + *out = static_cast(value); + return Status::OK(); +} + +Status validate_window_doc_count(bool first_window, uint64_t win_base, uint64_t last_docid, + uint64_t doc_count) { + uint64_t first_docid = 0; + if (!first_window) { + RETURN_IF_ERROR(checked_add_u64(win_base, 1, "frq_prelude: window base exceeds docid range", + &first_docid)); + } + if (last_docid < first_docid) { + return Status::Error( + "frq_prelude: invalid window docid range"); + } + const uint64_t width = last_docid - first_docid + 1; + if (doc_count > width) { + return Status::Error( + "frq_prelude: doc_count exceeds window width"); + } + return Status::OK(); +} + +// Validates builder input: non-null sink, group_size>=1, sane count, and +// non-decreasing absolute last_docid across windows. +Status validate_input(const FrqPreludeColumns& cols, ByteSink* out) { + if (out == nullptr) + return Status::Error("frq_prelude: null sink"); + if (cols.group_size == 0) { + return Status::Error( + "frq_prelude: group_size must be >= 1"); + } + if (cols.windows.size() > kMaxWindows) { + return Status::Error( + "frq_prelude: window count exceeds cap"); + } + for (size_t w = 1; w < cols.windows.size(); ++w) { + if (cols.windows[w].last_docid < cols.windows[w - 1].last_docid) { + return Status::Error( + "frq_prelude: last_docid not monotonic"); + } + } + return Status::OK(); +} + +// Encodes one window row into a per-block sink. last_docid_delta is the row's +// absolute last_docid minus prev_last (the previous window's absolute last). +// +// dd_off/freq_off/prx_off are NOT serialized: each was a pure running prefix sum +// of the per-window disk/prx lengths, which the reader reconstructs (see +// decode_window_row / RunningOffsets). dd_uncomp_len/freq_uncomp_len are written +// ONLY when the region's zstd win_mode bit is set; a raw region's uncomp_len is +// defined to equal its disk_len (open_region enforces uncomp_len == disk_len on +// raw region bytes), so the reader derives it without a stored field. +void encode_window_row(const WindowMeta& m, bool has_freq, bool has_prx, uint64_t prev_last, + ByteSink* block) { + const uint8_t win_mode = make_win_mode(m, has_freq); + block->put_varint64(static_cast(m.last_docid) - prev_last); + block->put_varint64(m.doc_count); + block->put_u8(win_mode); + block->put_varint64(m.dd_disk_len); + if ((win_mode & frq_win_mode::kDdZstd) != 0) { + block->put_varint64(m.dd_uncomp_len); + } + block->put_fixed32(m.crc_dd); + if (has_freq) { + block->put_varint64(m.freq_disk_len); + if ((win_mode & frq_win_mode::kFreqZstd) != 0) { + block->put_varint64(m.freq_uncomp_len); + } + block->put_fixed32(m.crc_freq); + } + if (has_prx) { + block->put_varint64(m.prx_len); + } + block->put_varint64(m.max_freq); + block->put_u8(m.max_norm); +} + +// One super-block's serialized window block plus its directory fields. +struct SuperBlock { + ByteSink block; + uint64_t last_docid = 0; // absolute last docid of this super-block's last window +}; + +// Builds every super-block's window block (row-encoded) and records the running +// absolute last docid at each super-block boundary. +std::vector encode_super_blocks(const FrqPreludeColumns& cols) { + const uint32_t g = cols.group_size; + const size_t n = cols.windows.size(); + std::vector blocks; + blocks.reserve(static_cast(ceil_div(n, g))); + uint64_t prev_last = 0; // previous window's absolute last docid (chains across blocks) + for (size_t start = 0; start < n; start += g) { + const size_t end = std::min(n, start + g); + SuperBlock sb; + for (size_t w = start; w < end; ++w) { + encode_window_row(cols.windows[w], cols.has_freq, cols.has_prx, prev_last, &sb.block); + prev_last = cols.windows[w].last_docid; + } + sb.last_docid = prev_last; + blocks.push_back(std::move(sb)); + } + return blocks; +} + +// Serializes the super_block_dir (one row per super-block) into dir_sink, using +// each block's byte length to compute its offset within the window_dir region. +void encode_super_block_dir(const std::vector& blocks, ByteSink* dir_sink) { + uint64_t prev_last = 0; + uint64_t block_off = 0; + for (const SuperBlock& sb : blocks) { + dir_sink->put_varint64(sb.last_docid - prev_last); + dir_sink->put_varint64(block_off); + dir_sink->put_varint64(sb.block.size()); + prev_last = sb.last_docid; + block_off += sb.block.size(); + } +} + +} // namespace + +Status build_frq_prelude(const FrqPreludeColumns& cols, ByteSink* out) { + RETURN_IF_ERROR(validate_input(cols, out)); + + const std::vector blocks = encode_super_blocks(cols); + ByteSink dir_sink; + encode_super_block_dir(blocks, &dir_sink); + + // covered = header + super_block_dir (the crc covers exactly this region). + ByteSink covered; + covered.put_u8(make_flags(cols)); + covered.put_varint64(cols.windows.size()); + covered.put_varint64(cols.group_size); + covered.put_varint64(blocks.size()); + covered.put_varint64(dir_sink.size()); + covered.put_bytes(dir_sink.view()); + + out->put_bytes(covered.view()); + out->put_fixed32(crc32c(covered.view())); + for (const SuperBlock& sb : blocks) out->put_bytes(sb.block.view()); + return Status::OK(); +} + +namespace { + +// Decoded header fields shared between parse phases. +struct Header { + bool has_freq = false; + bool has_prx = false; + uint64_t n = 0; + uint64_t group_size = 0; + uint64_t n_super = 0; + uint64_t sbdir_len = 0; +}; + +// Verifies the trailing crc covers [start of buffer .. end of super_block_dir]. +// covered_len = header bytes (up to and including sbdir_len) + sbdir_len. +Status verify_covered_crc(Slice prelude, size_t header_end, uint64_t sbdir_len) { + const size_t covered = header_end + static_cast(sbdir_len); + if (covered + sizeof(uint32_t) > prelude.size()) { + return Status::Error( + "frq_prelude: buffer too short for crc region"); + } + uint32_t stored = 0; + ByteSource crc_src(prelude.subslice(covered, sizeof(uint32_t))); + RETURN_IF_ERROR(crc_src.get_fixed32(&stored)); + if (crc32c(prelude.subslice(0, covered)) != stored) { + return Status::Error( + "frq_prelude: crc32c mismatch"); + } + return Status::OK(); +} + +// Parses + validates the header (counts capped before any later reserve). +Status parse_header(ByteSource* src, Header* h) { + uint8_t flags = 0; + RETURN_IF_ERROR(src->get_u8(&flags)); + h->has_freq = (flags & frq_prelude_flags::kHasFreq) != 0; + h->has_prx = (flags & frq_prelude_flags::kHasPrx) != 0; + RETURN_IF_ERROR(src->get_varint64(&h->n)); + RETURN_IF_ERROR(src->get_varint64(&h->group_size)); + RETURN_IF_ERROR(src->get_varint64(&h->n_super)); + RETURN_IF_ERROR(src->get_varint64(&h->sbdir_len)); + if (h->n > kMaxWindows || h->n_super > kMaxWindows) { + return Status::Error( + "frq_prelude: window count exceeds sane cap"); + } + if (h->group_size == 0) { + return Status::Error( + "frq_prelude: group_size is zero"); + } + if (h->n_super != ceil_div(h->n, h->group_size)) { + return Status::Error( + "frq_prelude: n_super inconsistent with N/G"); + } + return Status::OK(); +} + +// One super-block directory row. +struct SbDirRow { + uint64_t last_docid = 0; + uint64_t block_off = 0; + uint64_t block_len = 0; +}; + +// Decodes the super_block_dir region into absolute-last-docid rows, validating +// monotonic last docids and contiguous, in-bounds block offsets. +Status decode_super_block_dir(Slice dir, const Header& h, std::vector* rows, + uint64_t* window_region_len) { + ByteSource src(dir); + rows->clear(); + rows->reserve(static_cast(h.n_super)); + uint64_t prev_last = 0; + uint64_t expect_off = 0; + for (uint64_t s = 0; s < h.n_super; ++s) { + SbDirRow r; + uint64_t ldd = 0; + RETURN_IF_ERROR(src.get_varint64(&ldd)); + RETURN_IF_ERROR(src.get_varint64(&r.block_off)); + RETURN_IF_ERROR(src.get_varint64(&r.block_len)); + RETURN_IF_ERROR(checked_add_u64( + prev_last, ldd, "frq_prelude: super-block last_docid overflow", &r.last_docid)); + uint32_t checked_last = 0; + RETURN_IF_ERROR(checked_u32(r.last_docid, "frq_prelude: super-block last_docid exceeds u32", + &checked_last)); + if (r.last_docid < prev_last || r.block_off != expect_off) { + return Status::Error( + "frq_prelude: super-block dir inconsistent"); + } + expect_off += r.block_len; + prev_last = r.last_docid; + rows->push_back(r); + } + if (!src.eof()) { + return Status::Error( + "frq_prelude: super-block dir has trailing bytes"); + } + *window_region_len = expect_off; + return Status::OK(); +} + +// Validates a per-window codec mode byte against the known bits. +Status check_win_mode(uint8_t mode, bool has_freq) { + if ((mode & ~frq_win_mode::kKnownBits) != 0) { + return Status::Error( + "frq_prelude: unknown win_mode bits"); + } + if (!has_freq && (mode & frq_win_mode::kFreqZstd) != 0) { + return Status::Error( + "frq_prelude: freq mode set without has_freq"); + } + return Status::OK(); +} + +// Running per-block byte offsets, chained across windows AND across super-blocks +// with the same lifetime as prev_last: dd/freq are prefix sums of the per-window +// on-disk region lengths, prx the prefix sum of prx lengths. They replace the +// three offset columns the row used to serialize (each was exactly this sum), so +// the reader reproduces dd_off/freq_off/prx_off bit-identically to the old +// explicit fields. +struct RunningOffsets { + uint64_t dd = 0; + uint64_t freq = 0; + uint64_t prx = 0; +}; + +// Decodes one window row, advancing prev_last to this window's absolute last and +// the running offsets past this window's dd/freq/prx regions. +Status decode_window_row(ByteSource* src, bool has_freq, bool has_prx, bool first_window, + uint64_t* prev_last, RunningOffsets* run, WindowMeta* m) { + uint64_t ldd = 0, doc_count = 0; + RETURN_IF_ERROR(src->get_varint64(&ldd)); + RETURN_IF_ERROR(src->get_varint64(&doc_count)); + uint8_t mode = 0; + RETURN_IF_ERROR(src->get_u8(&mode)); + RETURN_IF_ERROR(check_win_mode(mode, has_freq)); + m->dd_zstd = (mode & frq_win_mode::kDdZstd) != 0; + m->freq_zstd = has_freq && (mode & frq_win_mode::kFreqZstd) != 0; + + // dd region: read disk_len, derive dd_off from the running dd-block offset, + // then advance it (overflow-guarded). uncomp_len is stored only for a zstd + // region; a raw region's uncomp_len == disk_len by contract. + RETURN_IF_ERROR(src->get_varint64(&m->dd_disk_len)); + m->dd_off = run->dd; + RETURN_IF_ERROR(checked_add_u64(run->dd, m->dd_disk_len, + "frq_prelude: dd-block offset overflow", &run->dd)); + if (m->dd_zstd) { + RETURN_IF_ERROR(src->get_varint64(&m->dd_uncomp_len)); + } else { + m->dd_uncomp_len = m->dd_disk_len; + } + RETURN_IF_ERROR(src->get_fixed32(&m->crc_dd)); + if (has_freq) { + RETURN_IF_ERROR(src->get_varint64(&m->freq_disk_len)); + m->freq_off = run->freq; + RETURN_IF_ERROR(checked_add_u64(run->freq, m->freq_disk_len, + "frq_prelude: freq-block offset overflow", &run->freq)); + if (m->freq_zstd) { + RETURN_IF_ERROR(src->get_varint64(&m->freq_uncomp_len)); + } else { + m->freq_uncomp_len = m->freq_disk_len; + } + RETURN_IF_ERROR(src->get_fixed32(&m->crc_freq)); + } + if (has_prx) { + RETURN_IF_ERROR(src->get_varint64(&m->prx_len)); + m->prx_off = run->prx; + RETURN_IF_ERROR(checked_add_u64(run->prx, m->prx_len, "frq_prelude: prx offset overflow", + &run->prx)); + } + uint64_t max_freq = 0; + RETURN_IF_ERROR(src->get_varint64(&max_freq)); + RETURN_IF_ERROR(src->get_u8(&m->max_norm)); + uint64_t last_docid = 0; + RETURN_IF_ERROR(checked_add_u64(*prev_last, ldd, "frq_prelude: window last_docid overflow", + &last_docid)); + RETURN_IF_ERROR(validate_window_doc_count(first_window, *prev_last, last_docid, doc_count)); + m->win_base = *prev_last; + RETURN_IF_ERROR( + checked_u32(last_docid, "frq_prelude: window last_docid exceeds u32", &m->last_docid)); + RETURN_IF_ERROR( + checked_u32(doc_count, "frq_prelude: window doc_count exceeds u32", &m->doc_count)); + RETURN_IF_ERROR( + checked_u32(max_freq, "frq_prelude: window max_freq exceeds u32", &m->max_freq)); + *prev_last = last_docid; + return Status::OK(); +} + +// Decodes one super-block's window block (<=G rows) into the global window list, +// seeding win_base from prev_last and re-checking the recorded sb last docid. +Status decode_one_block(Slice block, const Header& h, uint64_t sb_last_docid, size_t row_count, + uint64_t* prev_last, RunningOffsets* run, + std::vector* windows) { + ByteSource src(block); + for (size_t i = 0; i < row_count; ++i) { + WindowMeta m; + RETURN_IF_ERROR(decode_window_row(&src, h.has_freq, h.has_prx, windows->empty(), prev_last, + run, &m)); + windows->push_back(m); + } + if (!src.eof()) { + return Status::Error( + "frq_prelude: window block has trailing bytes"); + } + if (*prev_last != sb_last_docid) { + return Status::Error( + "frq_prelude: window block last docid mismatch"); + } + return Status::OK(); +} + +// Decodes all window blocks pointed to by the super_block_dir. +Status decode_all_blocks(Slice window_region, const Header& h, const std::vector& dir, + std::vector* windows) { + windows->clear(); + windows->reserve(static_cast(h.n)); + uint64_t prev_last = 0; + // dd/freq/prx running offsets chain across ALL super-blocks (not reset per + // block), the same lifetime as prev_last, so the derived per-window offsets are + // continuous over the whole dd-block / freq-block / prx span. + RunningOffsets run; + for (size_t s = 0; s < dir.size(); ++s) { + const SbDirRow& r = dir[s]; + if (r.block_off + r.block_len > window_region.size() || + r.block_off + r.block_len < r.block_off) { + return Status::Error( + "frq_prelude: window block out of region"); + } + const uint64_t already = static_cast(windows->size()); + const uint64_t rows = std::min(h.group_size, h.n - already); + Slice block = window_region.subslice(static_cast(r.block_off), + static_cast(r.block_len)); + RETURN_IF_ERROR(decode_one_block(block, h, r.last_docid, static_cast(rows), + &prev_last, &run, windows)); + } + if (windows->size() != h.n) { + return Status::Error( + "frq_prelude: decoded window count mismatch"); + } + return Status::OK(); +} + +// Sums the per-window dd/freq on-disk lengths into the dd-block / freq-block +// lengths, guarding each running sum against u64 overflow. The dd_off/freq_off +// contiguity cross-checks the old prelude ran here are now tautological -- the +// reader DERIVES dd_off/freq_off as these very prefix sums (decode_window_row), +// so `m.dd_off == running-sum` holds by construction and is dropped. The +// length-overflow guards and the returned block lengths are retained: they still +// bound the dd-block/freq-block range the callers' in_bounds checks fetch against, +// and mirror the same checked_add_u64 the offset derivation uses. +Status validate_region_layout(const Header& h, const std::vector& windows, + uint64_t* dd_block_len, uint64_t* freq_block_len) { + uint64_t dd_expect = 0; + uint64_t freq_expect = 0; + for (const WindowMeta& m : windows) { + // Raw regions carry uncomp_len == disk_len (derived, not stored); this + // guard stays as defensive documentation of that invariant. + if (m.dd_disk_len > m.dd_uncomp_len && !m.dd_zstd) { + return Status::Error( + "frq_prelude: raw dd region length inconsistent"); + } + if (dd_expect + m.dd_disk_len < dd_expect) { + return Status::Error( + "frq_prelude: dd block length overflow"); + } + dd_expect += m.dd_disk_len; + if (h.has_freq) { + if (freq_expect + m.freq_disk_len < freq_expect) { + return Status::Error( + "frq_prelude: freq block length overflow"); + } + freq_expect += m.freq_disk_len; + } + } + *dd_block_len = dd_expect; + *freq_block_len = freq_expect; + return Status::OK(); +} + +} // namespace + +namespace { +// TEST-ONLY seam backing testing::window_probe_count(): one increment per window +// last_docid comparison, in select_covering_windows_cursor and in locate_window's +// level-2 scan. thread_local => race-free under the shared const reader and free of +// atomic cost in the production cursor's hot loop; tests reset/read on their own thread. +thread_local uint64_t g_window_probes = 0; +inline void note_window_probe() { + ++g_window_probes; +} +} // namespace + +Status FrqPreludeReader::open(Slice prelude, FrqPreludeReader* out) { + ByteSource src(prelude); + Header h; + RETURN_IF_ERROR(parse_header(&src, &h)); + const size_t header_end = src.position(); + RETURN_IF_ERROR(verify_covered_crc(prelude, header_end, h.sbdir_len)); + + if (header_end + static_cast(h.sbdir_len) > prelude.size()) { + return Status::Error( + "frq_prelude: sbdir_len past buffer"); + } + Slice dir = prelude.subslice(header_end, static_cast(h.sbdir_len)); + std::vector rows; + uint64_t window_region_len = 0; + RETURN_IF_ERROR(decode_super_block_dir(dir, h, &rows, &window_region_len)); + + const size_t region_start = header_end + static_cast(h.sbdir_len) + sizeof(uint32_t); + if (region_start + static_cast(window_region_len) > prelude.size()) { + return Status::Error( + "frq_prelude: window region past buffer"); + } + Slice window_region = prelude.subslice(region_start, static_cast(window_region_len)); + + out->has_freq_ = h.has_freq; + out->has_prx_ = h.has_prx; + out->group_size_ = static_cast(h.group_size); + out->n_super_ = static_cast(h.n_super); + out->sb_last_docid_.clear(); + out->sb_last_docid_.reserve(rows.size()); + for (const SbDirRow& r : rows) out->sb_last_docid_.push_back(r.last_docid); + RETURN_IF_ERROR(decode_all_blocks(window_region, h, rows, &out->windows_)); + // Packed last_docid catalogue for the covering-window cursor (in-memory only; + // byte-identical to each windows_[w].last_docid, never serialized). + out->win_last_docid_.clear(); + out->win_last_docid_.reserve(out->windows_.size()); + for (const WindowMeta& m : out->windows_) { + out->win_last_docid_.push_back(m.last_docid); + } + return validate_region_layout(h, out->windows_, &out->dd_block_len_, &out->freq_block_len_); +} + +Status FrqPreludeReader::window(uint32_t w, WindowMeta* out) const { + if (out == nullptr) + return Status::Error("frq_prelude: null window out"); + if (w >= windows_.size()) { + return Status::Error( + "frq_prelude: window index out of range"); + } + *out = windows_[w]; + return Status::OK(); +} + +Status FrqPreludeReader::locate_window(uint32_t docid, bool* found, uint32_t* w) const { + if (found == nullptr || w == nullptr) { + return Status::Error("frq_prelude: null locate out"); + } + *found = false; + if (windows_.empty()) return Status::OK(); + if (docid > windows_.back().last_docid) return Status::OK(); + + // Level 1: first super-block whose absolute last docid >= docid. + const auto sb_it = std::lower_bound(sb_last_docid_.begin(), sb_last_docid_.end(), + static_cast(docid)); + const size_t sb = static_cast(sb_it - sb_last_docid_.begin()); + // Level 2: window binary search within [sb*G, min((sb+1)*G, N)). + const size_t lo = sb * group_size_; + const size_t hi = std::min(lo + group_size_, windows_.size()); + for (size_t i = lo; i < hi; ++i) { + note_window_probe(); + if (docid <= windows_[i].last_docid) { + *found = true; + *w = static_cast(i); + return Status::OK(); + } + } + return Status::OK(); // unreachable when invariants hold; defensive miss. +} + +void select_covering_windows_cursor(const uint32_t* win_last_docid, uint32_t n_windows, + const uint64_t* sb_last_docid, uint32_t n_super, + uint32_t group_size, const std::vector& candidates, + std::vector* windows) { + windows->clear(); + if (n_windows == 0) { + return; // empty-windows guard (mirrors locate_window's windows_.empty() early-out). + } + uint32_t sb = 0; // monotonic super-block cursor + uint32_t w = 0; // monotonic window cursor + uint32_t last_emitted = UINT32_MAX; // last window pushed (run-collapse dedup) + for (uint32_t d : candidates) { + const uint64_t target = d; // widen once; keeps the comparisons template-bracket free + // Level 1: first super-block whose absolute last docid >= d. sb only advances + // forward, so across the whole call it steps at most n_super times. + while (sb < n_super && sb_last_docid[sb] < target) { + ++sb; + } + if (sb == n_super) { + break; // d past the term's last docid -> this and every later candidate miss. + } + // Boundary jump: never scan windows below the current super-block's first window. + if (w < sb * group_size) { + w = sb * group_size; + } + // Level 2: first window whose absolute last docid >= d. w only advances forward, + // so across the whole call the comparisons below total at most n_windows misses + // plus one hit per candidate => probe_count <= candidates + n_windows. + while (w < n_windows) { + note_window_probe(); + if (d <= win_last_docid[w]) { + break; + } + ++w; + } + if (w == n_windows) { + break; // defensive: invariants guarantee a hit once sb < n_super. + } + if (w != last_emitted) { + windows->push_back(w); + last_emitted = w; + } + } +} + +void FrqPreludeReader::select_covering_windows(const std::vector& candidates, + std::vector* windows) const { + select_covering_windows_cursor( + win_last_docid_.data(), static_cast(win_last_docid_.size()), + sb_last_docid_.data(), static_cast(sb_last_docid_.size()), group_size_, + candidates, windows); +} + +} // namespace doris::snii::format + +namespace doris::snii::format::testing { + +uint64_t window_probe_count() { + return g_window_probes; +} + +void reset_window_probe_count() { + g_window_probes = 0; +} + +} // namespace doris::snii::format::testing diff --git a/be/src/storage/index/snii/format/frq_prelude.h b/be/src/storage/index/snii/format/frq_prelude.h new file mode 100644 index 00000000000000..642f3aed96b71c --- /dev/null +++ b/be/src/storage/index/snii/format/frq_prelude.h @@ -0,0 +1,256 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +// FrqPrelude: a TWO-LEVEL (super-block -> window) skippable directory that +// precedes a windowed .frq posting whose payload is laid out (PHASE D, design +// 1.6) with dd and freq regions GROUPED at posting level: +// windowed .frq payload = [prelude][dd-block][freq-block] +// dd-block = dd_region_0 ++ dd_region_1 ++ ... ++ dd_region_{N-1} +// freq-block = freq_region_0 ++ ... ++ freq_region_{N-1} (iff has_freq) +// Windows are NOT self-describing: each window's full codec metadata (region +// offsets, on-disk/uncompressed lengths, modes, crcs) lives in the prelude rows. +// The docs-only prefix [prelude][dd-block] is therefore ONE contiguous run a +// docid-only / phrase reader fetches in a single range, skipping the freq-block. +// +// DictEntry records prelude_len, frq_len (whole payload) and frq_docs_len +// (= prelude_len + dd_block_len) so a reader can range-fetch the prelude first, +// then fetch either the contiguous dd-block (docs-only) or both blocks (scoring). +// +// On-disk layout (strict; all multi-byte fixed fields little-endian, VInt = +// LEB128 via snii/encoding): +// header: +// u8 flags # bit0 has_freq, bit1 has_prx +// VInt N # number of .frq windows +// VInt G # windows per super-block (group_size; >=1) +// VInt n_super # = ceil(N / G); 0 when N==0 +// VInt sbdir_len # byte length of the super_block_dir region +// u32 crc32c # covers header + super_block_dir (NOT the window blocks) +// super_block_dir[n_super]: # small, resident: one row per super-block +// VInt sb_last_docid_delta # cumulative across super-blocks => absolute last +// # docid of the super-block's last window +// VInt sb_block_off # byte offset of this super-block's window block, +// # measured from the start of the window_dir region +// VInt sb_block_len # byte length of this super-block's window block +// window_dir: n_super self-contained blocks, each holding <=G window rows. +// per window row (T18 slim layout -- dd_off/freq_off/prx_off are NOT stored; +// the reader derives them as running prefix sums of the disk/prx lengths): +// VInt last_docid_delta # cumulative WITHIN the block => absolute last docid +// # (previous window's absolute last docid = win_base; +// # first window of first block: win_base = 0) +// VInt doc_count # number of docs in the window (frq_pod needs it) +// u8 win_mode # bit0 dd_zstd, bit1 freq_zstd +// VInt dd_disk_len # dd_region on-disk byte length +// [VInt dd_uncomp_len] # dd_region plaintext length; present ONLY when +// # win_mode & kDdZstd. A raw region's uncomp_len +// # == dd_disk_len (derived, not stored). +// u32 crc_dd # crc32c of the dd_region on-disk bytes +// VInt freq_disk_len # freq_region on-disk byte length (has_freq) +// [VInt freq_uncomp_len] # freq_region plaintext length; present ONLY when +// # has_freq && win_mode & kFreqZstd (raw: derived +// # == freq_disk_len). +// u32 crc_freq # crc32c of the freq_region on-disk bytes (has_freq) +// VInt prx_len # .prx payload byte length (present iff has_prx) +// VInt max_freq # window max term frequency (WAND block-max) +// u8 max_norm # window score-max norm (WAND); 0 acceptable +// +// The reader reconstructs each window's dd_off / freq_off (byte offset within the +// dd-block / freq-block) and prx_off (offset within the entry's .prx span) as the +// running prefix sums of dd_disk_len / freq_disk_len / prx_len over all windows, +// chained across super-blocks; WindowMeta still exposes those offsets, now derived. +// +// Reconstructing win_base / absolute last_docid (READER CONTRACT) is unchanged: +// the writer chains absolute last docids across windows; each row stores the delta +// of its absolute last docid from the previous window, and sb_last_docid seeds +// each block, so super-block binary search then in-block window binary search +// locate the window covering any docid without decoding the .frq blocks. +// +// The trailing crc32c covers only header + super_block_dir; every region carries +// its own crc (crc_dd / crc_freq) in the row. +namespace doris::snii::format { + +namespace frq_prelude_flags { +inline constexpr uint8_t kHasFreq = 1u << 0; +inline constexpr uint8_t kHasPrx = 1u << 1; +// Reserved extension point (T18): kSlimRows = 1u << 2 would gate the trimmed +// window-row layout (no stored dd_off/freq_off/prx_off, conditional uncomp_len) +// as a distinct on-disk path. It is NOT emitted today: the trim folds into the +// single pre-launch v1 encoding (writer/reader symmetric, no dual decode path). +// If a `lifecycle: launched` index appears before this lands, set this bit on the +// slim writer and branch the reader on it instead of unconditionally decoding slim. +} // namespace frq_prelude_flags + +// Per-window codec mode bits (win_mode byte). +namespace frq_win_mode { +inline constexpr uint8_t kDdZstd = 1u << 0; +inline constexpr uint8_t kFreqZstd = 1u << 1; +inline constexpr uint8_t kKnownBits = kDdZstd | kFreqZstd; +} // namespace frq_win_mode + +// Absolute, decoded metadata for one window (as the reader exposes it). The dd / +// freq region locators are offsets WITHIN the dd-block / freq-block respectively +// (both blocks follow the prelude). dd_off/freq_off/prx_off are DERIVED by the +// reader as running prefix sums of the disk/prx lengths (they are no longer stored +// per row; see the header layout note) -- these public members are unchanged and +// still populated, just by derivation. The reader derives the dd-block length from +// the last window's dd_off + dd_disk_len. +struct WindowMeta { + uint32_t last_docid = 0; // absolute last docid in the window + uint64_t win_base = 0; // absolute last docid of the previous window (0 for w==0) + uint32_t doc_count = 0; + + // dd_region locator (within the dd-block). + bool dd_zstd = false; + uint64_t dd_off = 0; // DERIVED: running sum of prior windows' dd_disk_len + uint64_t dd_disk_len = 0; + uint64_t dd_uncomp_len = 0; // DERIVED == dd_disk_len for raw; stored only when dd_zstd + uint32_t crc_dd = 0; + + // freq_region locator (within the freq-block); valid only when has_freq. + bool freq_zstd = false; + uint64_t freq_off = 0; // DERIVED: running sum of prior windows' freq_disk_len + uint64_t freq_disk_len = 0; + uint64_t freq_uncomp_len = 0; // DERIVED == freq_disk_len for raw; stored only when freq_zstd + uint32_t crc_freq = 0; + + uint64_t prx_off = 0; // valid only when has_prx; DERIVED: running sum of prior prx_len + uint64_t prx_len = 0; // valid only when has_prx + uint32_t max_freq = 0; + uint8_t max_norm = 0; + + // In-memory only (NOT serialized in the prelude row). When false, the dd/freq + // region decode skips crc verification -- used when these region bytes are + // covered by an enclosing crc (e.g. an INLINE entry inside its dict block). + // Windowed/slim POD-ref rows leave this true (their regions carry a crc). + bool verify_crc = true; +}; + +// Builder input: one fully-computed WindowMeta per window, in term order, plus the +// super-block grouping factor. The writer fills last_docid (absolute), doc_count, +// the region locators/crcs, prx locator, max_freq and max_norm; win_base is derived +// during build (so callers may leave it 0). group_size must be >= 1. +struct FrqPreludeColumns { + bool has_freq = true; + bool has_prx = false; + uint32_t group_size = 64; // windows per super-block (G) + std::vector windows; +}; + +// Builds the prelude bytes and appends them to out. +// Returns InvalidArgument when out is null, group_size is 0, or the windows are +// not in non-decreasing last_docid order (a window's absolute last docid must be +// >= the previous window's). +Status build_frq_prelude(const FrqPreludeColumns& cols, ByteSink* out); + +// Reads and verifies a prelude buffer, exposing two-level skip access. The reader +// parses the header + super_block_dir on open (verifying the trailing crc) and +// eagerly decodes every window block into owned WindowMeta rows (the prelude is +// small relative to the postings). It does not retain the input. +class FrqPreludeReader { +public: + // Parses + verifies the prelude. crc mismatch / truncation / inconsistent + // offsets-or-lengths / oversized counts => kCorruption. + static Status open(Slice prelude, FrqPreludeReader* out); + + uint32_t n_windows() const { return static_cast(windows_.size()); } + uint32_t n_super_blocks() const { return n_super_; } + bool has_freq() const { return has_freq_; } + bool has_prx() const { return has_prx_; } + + // Total on-disk byte length of the dd-block (== sum of dd_disk_len; the docs-only + // prefix after the prelude). 0 when there are no windows. + uint64_t dd_block_len() const { return dd_block_len_; } + // Total on-disk byte length of the freq-block (== sum of freq_disk_len). 0 when + // !has_freq or no windows. + uint64_t freq_block_len() const { return freq_block_len_; } + + // Returns the absolute WindowMeta for window w. Out-of-range => InvalidArgument. + Status window(uint32_t w, WindowMeta* out) const; + + // Locates the window covering docid via super-block binary search then window + // binary search. *found=false (with OK) when docid is past the term's last + // docid; otherwise *w is the index of the covering window (the first window + // whose absolute last_docid >= docid). + Status locate_window(uint32_t docid, bool* found, uint32_t* w) const; + + // Selects, as a monotonic two-pointer cursor, the ascending de-duplicated set of + // windows covering the ascending `candidates` (each window covering its + // (win_base, last_docid] span). Writes them to *windows (cleared first). The + // result is element-for-element identical to calling locate_window per candidate + // and collapsing equal runs, but uses O(C + N) window last_docid comparisons + // (C = candidates, N = windows) instead of O(C * group_size). Pure in-memory over + // the decoded directory; never fails. + void select_covering_windows(const std::vector& candidates, + std::vector* windows) const; + + // Packed absolute last_docid of window w (byte-identical to window(w).last_docid), + // exposed for the covering-window cursor's contiguous scan and equivalence tests. + uint32_t window_last_docid(uint32_t w) const { + DCHECK_LT(w, win_last_docid_.size()); + return win_last_docid_[w]; + } + +private: + bool has_freq_ = false; + bool has_prx_ = false; + uint32_t group_size_ = 1; + uint32_t n_super_ = 0; + uint64_t dd_block_len_ = 0; + uint64_t freq_block_len_ = 0; + // Absolute last docid at each super-block boundary (size n_super_). + std::vector sb_last_docid_; + // All windows decoded with absolute fields, in term order (size N). + std::vector windows_; + // Packed copy of each window's absolute last_docid (size N; win_last_docid_[w] == + // windows_[w].last_docid). Built in open() so the covering-window cursor scans a + // contiguous 4B/window array rather than the ~104B WindowMeta rows. In-memory only: + // never serialized; immutable after open() (same lifetime as windows_). + std::vector win_last_docid_; +}; + +// Pure cursor core (no FrqPreludeReader / IO): selects into *windows the ascending, +// de-duplicated indices of the windows covering the ascending `candidates`, given the +// packed window last_docid array (size n_windows), the super-block last_docid directory +// (size n_super) and group_size. A super-block cursor does boundary jumps while a window +// cursor advances forward only => O(C + N) window comparisons, element-for-element equal +// to per-candidate locate_window + run collapse. *windows is cleared first; n_windows == 0 +// yields an empty result. Exposed for isolated equivalence / complexity tests. +void select_covering_windows_cursor(const uint32_t* win_last_docid, uint32_t n_windows, + const uint64_t* sb_last_docid, uint32_t n_super, + uint32_t group_size, const std::vector& candidates, + std::vector* windows); + +// TEST-ONLY observability seam (mirrors the format dict-block decode counter). Counts the +// window last_docid comparisons performed by select_covering_windows_cursor and by +// locate_window's level-2 scan, so tests can assert the cursor stays O(C + N) and bounded +// by C + N regardless of group_size, while the legacy per-candidate scan grows with G. The +// counter is thread-local: race-free under the shared const reader and free of atomic cost +// in the production cursor loop; reset and read on the thread that ran the cursor. +namespace testing { +uint64_t window_probe_count(); +void reset_window_probe_count(); +} // namespace testing + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/metadata_blob.cpp b/be/src/storage/index/snii/format/metadata_blob.cpp new file mode 100644 index 00000000000000..a8d7a0f4fd938a --- /dev/null +++ b/be/src/storage/index/snii/format/metadata_blob.cpp @@ -0,0 +1,121 @@ +// 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 "storage/index/snii/format/metadata_blob.h" + +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/zstd_codec.h" + +namespace doris::snii::format { + +namespace { + +constexpr int kMetaSectionZstdLevel = 3; +constexpr uint64_t kMaxMetaSectionUncompBytes = 256ULL * 1024 * 1024; + +size_t meta_compress_min_bytes() { + const char* s = std::getenv("SNII_META_COMPRESS_MIN"); + if (s != nullptr) { + char* end = nullptr; + const unsigned long long v = std::strtoull(s, &end, 10); + if (end != s) { + return v; + } + } + return kMetaSectionCompressMinBytes; +} + +} // namespace + +Status encode_metadata_blob(Slice raw_frame, SectionType raw_type, SectionType compressed_type, + ByteSink* out) { + if (out == nullptr) { + return Status::Error("metadata_blob: null sink"); + } + + ByteSource source(raw_frame); + FramedSection raw_section; + RETURN_IF_ERROR(SectionFramer::read(source, &raw_section)); + if (!source.eof() || raw_section.type != static_cast(raw_type)) { + return Status::Error( + "metadata_blob: raw input is not exactly one frame of the expected type"); + } + + if (raw_frame.size() >= meta_compress_min_bytes()) { + std::vector compressed; + if (zstd_compress(raw_frame, kMetaSectionZstdLevel, &compressed).ok()) { + ByteSink payload; + payload.put_varint64(raw_frame.size()); + payload.put_bytes(Slice(compressed)); + if (payload.size() + 16 < raw_frame.size()) { + SectionFramer::write(*out, static_cast(compressed_type), payload.view()); + return Status::OK(); + } + } + } + out->put_bytes(raw_frame); + return Status::OK(); +} + +Status materialize_metadata_blob(Slice stored_frame, SectionType raw_type, + SectionType compressed_type, std::vector* scratch, + Slice* raw_frame) { + if (scratch == nullptr || raw_frame == nullptr) { + return Status::Error("metadata_blob: null frame out"); + } + + ByteSource source(stored_frame); + FramedSection stored_section; + RETURN_IF_ERROR(SectionFramer::read(source, &stored_section)); + if (!source.eof()) { + return Status::Error( + "metadata_blob: trailing stored frame bytes"); + } + if (stored_section.type == static_cast(raw_type)) { + *raw_frame = stored_frame; + return Status::OK(); + } + if (stored_section.type != static_cast(compressed_type)) { + return Status::Error( + "metadata_blob: unexpected stored frame type"); + } + + ByteSource payload(stored_section.payload); + uint64_t uncomp_len = 0; + RETURN_IF_ERROR(payload.get_varint64(&uncomp_len)); + if (uncomp_len == 0 || uncomp_len > kMaxMetaSectionUncompBytes) { + return Status::Error( + "metadata_blob: zstd uncomp_len out of range"); + } + Slice compressed; + RETURN_IF_ERROR(payload.get_bytes(payload.remaining(), &compressed)); + RETURN_IF_ERROR(zstd_decompress(compressed, static_cast(uncomp_len), scratch)); + ByteSource raw_source {Slice(*scratch)}; + FramedSection raw_section; + RETURN_IF_ERROR(SectionFramer::read(raw_source, &raw_section)); + if (!raw_source.eof() || raw_section.type != static_cast(raw_type)) { + return Status::Error( + "metadata_blob: decompressed bytes are not exactly one expected raw frame"); + } + *raw_frame = Slice(*scratch); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/metadata_blob.h b/be/src/storage/index/snii/format/metadata_blob.h new file mode 100644 index 00000000000000..5e098a797e08d2 --- /dev/null +++ b/be/src/storage/index/snii/format/metadata_blob.h @@ -0,0 +1,40 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +// Appends a raw metadata frame or a zstd carrier for it. The carrier payload is +// varint64 raw-frame length followed by zstd(raw frame). +Status encode_metadata_blob(Slice raw_frame, SectionType raw_type, SectionType compressed_type, + ByteSink* out); + +// Returns a raw metadata frame from a raw frame or zstd carrier. A raw frame +// remains a view into stored_frame; a materialized carrier is a view into scratch. +Status materialize_metadata_blob(Slice stored_frame, SectionType raw_type, + SectionType compressed_type, std::vector* scratch, + Slice* raw_frame); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/metadata_directory.cpp b/be/src/storage/index/snii/format/metadata_directory.cpp new file mode 100644 index 00000000000000..404afe22b6fbde --- /dev/null +++ b/be/src/storage/index/snii/format/metadata_directory.cpp @@ -0,0 +1,150 @@ +// 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 "storage/index/snii/format/metadata_directory.h" + +#include +#include +#include +#include +#include + +#include "gen_cpp/snii.pb.h" + +namespace doris::snii::format { +namespace { + +Status corrupted(std::string_view message) { + return Status::Error(message); +} + +Status unsupported(std::string_view message) { + return Status::Error(message); +} + +Status decode_blob_ref(const doris::snii::SniiBlobRefPB& input, MetadataBlobRef* out) { + if (!input.has_offset() || !input.has_length()) { + return corrupted("metadata directory: missing blob reference field"); + } + if (input.length() == 0) { + return corrupted("metadata directory: empty mandatory blob reference"); + } + *out = {.offset = input.offset(), .length = input.length()}; + return Status::OK(); +} + +Status decode_directory_pb(const doris::snii::SniiMetadataDirectoryPB& input, + std::vector* out) { + if (input.required_features_size() != 0) { + return unsupported("metadata directory: required feature is not supported"); + } + + std::vector entries; + entries.reserve(input.indexes_size()); + for (const auto& index : input.indexes()) { + if (!index.has_index_id() || !index.has_index_suffix() || !index.has_core_metadata() || + !index.has_sampled_term_index() || !index.has_dict_block_directory()) { + return corrupted("metadata directory: missing required logical field"); + } + + LogicalIndexMetadataRef entry; + entry.index_id = index.index_id(); + entry.index_suffix = index.index_suffix(); + RETURN_IF_ERROR(decode_blob_ref(index.core_metadata(), &entry.core_metadata)); + RETURN_IF_ERROR(decode_blob_ref(index.sampled_term_index(), &entry.sampled_term_index)); + RETURN_IF_ERROR(decode_blob_ref(index.dict_block_directory(), &entry.dict_block_directory)); + + for (const auto& existing : entries) { + if (existing.index_id == entry.index_id && + existing.index_suffix == entry.index_suffix) { + return corrupted("metadata directory: duplicate logical index key"); + } + } + entries.push_back(std::move(entry)); + } + *out = std::move(entries); + return Status::OK(); +} + +void encode_blob_ref(const MetadataBlobRef& input, doris::snii::SniiBlobRefPB* out) { + out->set_offset(input.offset); + out->set_length(input.length); +} + +} // namespace + +Status MetadataDirectory::decode(Slice bytes, MetadataDirectory* out) { + if (out == nullptr) { + return Status::Error("metadata directory: null output"); + } + out->entries_.clear(); + if (bytes.size() > static_cast(std::numeric_limits::max())) { + return corrupted("metadata directory: protobuf payload exceeds INT_MAX"); + } + + doris::snii::SniiMetadataDirectoryPB directory; + if (!directory.ParseFromArray(bytes.data(), static_cast(bytes.size()))) { + return corrupted("metadata directory: protobuf parsing failed"); + } + + std::vector entries; + RETURN_IF_ERROR(decode_directory_pb(directory, &entries)); + out->entries_ = std::move(entries); + return Status::OK(); +} + +const LogicalIndexMetadataRef* MetadataDirectory::find(uint64_t index_id, + std::string_view suffix) const { + for (const auto& entry : entries_) { + if (entry.index_id == index_id && entry.index_suffix == suffix) { + return &entry; + } + } + return nullptr; +} + +Status encode_metadata_directory(const std::vector& entries, + ByteSink* out) { + if (out == nullptr) { + return Status::Error("metadata directory: null output"); + } + + doris::snii::SniiMetadataDirectoryPB directory; + for (const auto& entry : entries) { + auto* index = directory.add_indexes(); + index->set_index_id(entry.index_id); + index->set_index_suffix(entry.index_suffix); + encode_blob_ref(entry.core_metadata, index->mutable_core_metadata()); + encode_blob_ref(entry.sampled_term_index, index->mutable_sampled_term_index()); + encode_blob_ref(entry.dict_block_directory, index->mutable_dict_block_directory()); + } + + std::vector validated; + RETURN_IF_ERROR(decode_directory_pb(directory, &validated)); + const size_t size = directory.ByteSizeLong(); + if (size > static_cast(std::numeric_limits::max())) { + return corrupted("metadata directory: protobuf payload exceeds INT_MAX"); + } + std::string payload(size, '\0'); + if (!directory.SerializeToArray(payload.data(), static_cast(size))) { + return corrupted("metadata directory: protobuf serialization failed"); + } + out->put_bytes(Slice(payload)); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/metadata_directory.h b/be/src/storage/index/snii/format/metadata_directory.h new file mode 100644 index 00000000000000..a25fdcfc668a14 --- /dev/null +++ b/be/src/storage/index/snii/format/metadata_directory.h @@ -0,0 +1,60 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +namespace doris::snii::format { + +struct MetadataBlobRef { + uint64_t offset = 0; + uint64_t length = 0; +}; + +struct LogicalIndexMetadataRef { + uint64_t index_id = 0; + std::string index_suffix; + MetadataBlobRef core_metadata; + MetadataBlobRef sampled_term_index; + MetadataBlobRef dict_block_directory; +}; + +class MetadataDirectory { +public: + static Status decode(Slice bytes, MetadataDirectory* out); + + const LogicalIndexMetadataRef* find(uint64_t index_id, std::string_view suffix) const; + size_t size() const { return entries_.size(); } + const std::vector& entries() const { return entries_; } + +private: + std::vector entries_; +}; + +Status encode_metadata_directory(const std::vector& entries, + ByteSink* out); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/norms_pod.cpp b/be/src/storage/index/snii/format/norms_pod.cpp new file mode 100644 index 00000000000000..b3da050c17726f --- /dev/null +++ b/be/src/storage/index/snii/format/norms_pod.cpp @@ -0,0 +1,85 @@ +// 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 "storage/index/snii/format/norms_pod.h" + +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +void NormsPodWriter::finish(ByteSink* sink) const { + finish(norms_, sink); +} + +void NormsPodWriter::finish(std::span norms, ByteSink* sink) { + // Build inner payload: [varint64 doc_count][raw norm bytes]. + ByteSink payload; + const size_t payload_size = varint_len(norms.size()) + norms.size(); + payload.reserve(payload_size); + payload.put_varint64(norms.size()); + payload.put_bytes(Slice(norms.data(), norms.size())); + // Delegate outer framing to SectionFramer to append type+len+crc32c, avoiding manual checksum assembly. + sink->reserve(1 + varint_len(payload_size) + payload_size + sizeof(uint32_t)); + SectionFramer::write(*sink, static_cast(SectionType::kNormsPod), payload.view()); +} + +Status NormsPodReader::open(Slice framed, NormsPodReader* out) { + // framer handles CRC verify, truncation detection, and payload slicing. + ByteSource src(framed); + FramedSection sec; + RETURN_IF_ERROR(SectionFramer::read(src, &sec)); + if (sec.type != static_cast(SectionType::kNormsPod)) { + return Status::Error( + "norms POD section type mismatch"); + } + if (!src.eof()) { + return Status::Error( + "norms POD trailing framed bytes"); + } + + // Parse inner payload: [varint64 doc_count][bytes]. + ByteSource payload(sec.payload); + uint64_t doc_count = 0; + RETURN_IF_ERROR(payload.get_varint64(&doc_count)); + if (payload.position() != varint_len(doc_count)) { + return Status::Error( + "norms POD non-canonical doc_count"); + } + if (doc_count > std::numeric_limits::max()) { + return Status::Error( + "norms POD doc_count overflows uint32"); + } + // doc_count must exactly equal the remaining byte count (1 byte per doc). + if (payload.remaining() != doc_count) { + return Status::Error( + "norms POD length mismatch"); + } + + Slice bytes; + RETURN_IF_ERROR(payload.get_bytes(static_cast(doc_count), &bytes)); + out->doc_count_ = static_cast(doc_count); + out->norms_ = bytes.data(); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/norms_pod.h b/be/src/storage/index/snii/format/norms_pod.h new file mode 100644 index 00000000000000..11791240615a65 --- /dev/null +++ b/be/src/storage/index/snii/format/norms_pod.h @@ -0,0 +1,89 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +namespace doris::snii::format { + +// norms POD: per logical index / field stores 1-byte encoded doc length per doc, +// used by BM25 length normalization (SniiStatsProvider::encoded_norm) for per-docid lookup. +// +// On-disk layout (the whole section is framed by SectionFramer, which adds a type+len+crc32c envelope): +// framer payload = [varint64 doc_count][bytes encoded_norm[doc_count]] +// framer envelope = [u8 type][varint64 payload_len][payload][fixed32 crc32c] +// The encoding of encoded_norm (length -> 1B) is out of scope for this module; here we only handle raw byte storage and retrieval. +class NormsPodWriter { +public: + // Appends the encoded_norm for the next docid (docid is implicit, assigned in append order starting from 0). + void add(uint8_t encoded_norm) { norms_.push_back(encoded_norm); } + + // Number of docs accumulated so far (i.e., the next docid to be assigned). + size_t count() const { return norms_.size(); } + + // Writes [doc_count][bytes] framed by SectionFramer into sink (appends; does not clear sink). + void finish(ByteSink* sink) const; + // Zero-copy source overload used by streamed compaction. + static void finish(std::span norms, ByteSink* sink); + +private: + std::vector norms_; +}; + +// Read-only view: on open, verifies the framer CRC and checks that doc_count/payload length are consistent, +// afterwards encoded_norm(docid) is O(1) direct indexing (zero-copy, borrows the underlying buffer). +class NormsPodReader { +public: + NormsPodReader() = default; + + // Parses the entire section (including the framer envelope). Returns Corruption on CRC mismatch, truncation, or length inconsistency. + // On success, *out borrows the memory pointed to by framer_payload; the caller must ensure its lifetime. + static Status open(Slice framed, NormsPodReader* out); + + uint32_t doc_count() const { return doc_count_; } + + // Precondition (hard contract): docid < doc_count(). Semantics match std::vector::operator[]: + // the caller is responsible for guaranteeing this (docid comes from trusted postings decoded internally by SNII). Asserts in debug builds; + // no check in Release (NDEBUG). Use try_encoded_norm when the docid is untrusted and needs validation. + uint8_t encoded_norm(uint32_t docid) const { + assert(docid < doc_count_); + return norms_[docid]; + } + + // Checked access: returns InvalidArgument if docid is out of range; never reads out-of-range memory. + Status try_encoded_norm(uint32_t docid, uint8_t* out) const { + if (docid >= doc_count_) + return Status::Error("norms: docid out of range"); + *out = norms_[docid]; + return Status::OK(); + } + +private: + const uint8_t* norms_ = nullptr; + uint32_t doc_count_ = 0; +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/null_bitmap.cpp b/be/src/storage/index/snii/format/null_bitmap.cpp new file mode 100644 index 00000000000000..9e7d0d298f99f8 --- /dev/null +++ b/be/src/storage/index/snii/format/null_bitmap.cpp @@ -0,0 +1,295 @@ +// 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 "storage/index/snii/format/null_bitmap.h" + +#include +#include +#include +#include + +#include "common/check.h" +// clang-format off +// CRoaring's public header defines ROARING_CONTAINER_T; its internal headers undefine it. +#include "roaring/roaring.hh" +#include "roaring/containers/array.h" +#include "roaring/containers/bitset.h" +#include "roaring/containers/run.h" +// clang-format on +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/varint.h" + +namespace doris::snii::format { + +namespace { + +constexpr uint32_t kPortableCookieNoRun = 12346; +constexpr uint16_t kPortableCookieRun = 12347; +constexpr uint32_t kMaxPortableContainers = uint32_t {1} << 16; + +struct ParsedNullBitmap { + uint32_t doc_count = 0; + Slice roaring_bytes; + uint32_t container_count = 0; +}; + +Status parse_null_bitmap(Slice framed, ParsedNullBitmap* out) { + ByteSource src(framed); + FramedSection sec; + RETURN_IF_ERROR(SectionFramer::read(src, &sec)); + if (sec.type != kNullBitmapSectionType || src.remaining() != 0) { + return Status::Error( + "null bitmap: invalid framed section"); + } + + ByteSource payload(sec.payload); + uint64_t doc_count = 0; + RETURN_IF_ERROR(payload.get_varint64(&doc_count)); + if (doc_count > std::numeric_limits::max()) { + return Status::Error( + "null bitmap doc_count overflows uint32"); + } + + uint64_t roaring_size = 0; + RETURN_IF_ERROR(payload.get_varint64(&roaring_size)); + if (roaring_size != payload.remaining()) { + return Status::Error( + "null bitmap roaring_size differs from payload"); + } + RETURN_IF_ERROR(payload.get_bytes(static_cast(roaring_size), &out->roaring_bytes)); + + ByteSource portable(out->roaring_bytes); + uint32_t cookie = 0; + RETURN_IF_ERROR(portable.get_fixed32(&cookie)); + if (static_cast(cookie) == kPortableCookieRun) { + out->container_count = (cookie >> 16) + 1; + } else if (cookie == kPortableCookieNoRun) { + RETURN_IF_ERROR(portable.get_fixed32(&out->container_count)); + if (out->container_count > kMaxPortableContainers) { + return Status::Error( + "null bitmap: portable container count out of range"); + } + } else { + return Status::Error( + "null bitmap: invalid portable cookie"); + } + + const char* data = reinterpret_cast(out->roaring_bytes.data()); + const size_t size = out->roaring_bytes.size(); + const size_t probed = roaring::api::roaring_bitmap_portable_deserialize_size(data, size); + if (probed == 0 || probed != size) { + return Status::Error( + "null bitmap: malformed roaring container"); + } + out->doc_count = static_cast(doc_count); + return Status::OK(); +} + +} // namespace + +NullBitmapWriter:: + NullBitmapWriter() // NOLINT(modernize-use-equals-default): roaring type is incomplete in the header. + : bitmap_(std::make_unique()) {} + +NullBitmapWriter::~NullBitmapWriter() = default; + +void NullBitmapWriter::add_null(uint32_t docid) { + bitmap_->add(docid); +} + +void NullBitmapWriter::add_many(std::span docids) { + bitmap_->addMany(docids.size(), docids.data()); +} + +uint32_t NullBitmapWriter::null_count() const { + return static_cast(bitmap_->cardinality()); +} + +uint64_t NullBitmapWriter::build_memory_upper_bound(std::span sorted_docids) { + if (sorted_docids.empty()) { + return 0; + } + + uint64_t container_count = 0; + uint64_t sparse_container_count = 0; + uint64_t sparse_value_count = 0; + uint64_t dense_container_count = 0; + size_t begin = 0; + while (begin < sorted_docids.size()) { + const uint32_t key = sorted_docids[begin] >> 16; + size_t end = begin + 1; + while (end < sorted_docids.size() && sorted_docids[end] >> 16 == key) { + DCHECK_GT(sorted_docids[end], sorted_docids[end - 1]); + ++end; + } + const uint64_t cardinality = end - begin; + ++container_count; + if (cardinality <= roaring::internal::DEFAULT_MAX_SIZE) { + ++sparse_container_count; + sparse_value_count += cardinality; + } else { + ++dense_container_count; + } + begin = end; + } + + // roaring_array_t uses three parallel arrays. Account a minimum allocation + // of four slots and old+replacement overlap at geometric growth. + constexpr uint64_t kTopEntryBytes = sizeof(void*) + sizeof(uint16_t) + sizeof(uint8_t); + const uint64_t top_capacity = std::max(container_count, 4); + const uint64_t top_array_peak = 3 * top_capacity * kTopEntryBytes; + + // Sparse array growth can retain the old uint16 array while allocating its + // replacement. Eight bytes per live value covers both capacity slack and + // replacement overlap. A dense container additionally covers the largest + // geometric array capacity immediately before conversion and the new 8 KiB + // bitset while both allocations are live. + constexpr uint64_t kSparseValuePeakBytes = 8; + constexpr uint64_t kBitsetBytes = + roaring::internal::BITSET_CONTAINER_SIZE_IN_WORDS * sizeof(uint64_t); + constexpr uint64_t kDenseArrayConversionCapacity = roaring::internal::DEFAULT_MAX_SIZE * 5 / 4; + constexpr uint64_t kDenseArrayConversionBytes = + kDenseArrayConversionCapacity * sizeof(uint16_t); + const uint64_t sparse_peak = + sparse_value_count * kSparseValuePeakBytes + + sparse_container_count * sizeof(roaring::internal::array_container_t); + const uint64_t dense_peak = + dense_container_count * (kDenseArrayConversionBytes + kBitsetBytes + + sizeof(roaring::internal::array_container_t) + + sizeof(roaring::internal::bitset_container_t)); + return sizeof(roaring::Roaring) + top_array_peak + sparse_peak + dense_peak; +} + +Status NullBitmapWriter::serialization_sizes(uint32_t doc_count, + NullBitmapSerializationSizes* out) const { + if (out == nullptr) { + return Status::Error( + "null bitmap: null serialization size output"); + } + const size_t roaring_bytes = bitmap_->getSizeInBytes(); + const size_t prefix_bytes = varint_len(doc_count) + varint_len(roaring_bytes); + if (roaring_bytes > std::numeric_limits::max() - prefix_bytes) { + return Status::Error( + "null bitmap: payload size overflows"); + } + const size_t payload_bytes = prefix_bytes + roaring_bytes; + const size_t envelope_bytes = 1 + varint_len(payload_bytes) + sizeof(uint32_t); + if (payload_bytes > std::numeric_limits::max() - envelope_bytes) { + return Status::Error( + "null bitmap: framed size overflows"); + } + *out = {.roaring_bytes = roaring_bytes, + .payload_bytes = payload_bytes, + .framed_bytes = envelope_bytes + payload_bytes}; + return Status::OK(); +} + +Status NullBitmapWriter::finish(uint32_t doc_count, ByteSink* sink) const { + if (sink == nullptr) { + return Status::Error("null bitmap: null output sink"); + } + NullBitmapSerializationSizes sizes; + RETURN_IF_ERROR(serialization_sizes(doc_count, &sizes)); + + // Serialize the Roaring bitmap to its portable on-disk form. + std::vector roaring_buf(sizes.roaring_bytes); + bitmap_->write(roaring_buf.data()); + + // Build inner payload: [varint64 doc_count][varint64 roaring_size][bytes]. + ByteSink payload; + payload.reserve(sizes.payload_bytes); + payload.put_varint64(doc_count); + payload.put_varint64(sizes.roaring_bytes); + payload.put_bytes( + Slice(reinterpret_cast(roaring_buf.data()), sizes.roaring_bytes)); + DORIS_CHECK_EQ(payload.size(), sizes.payload_bytes); + + // Delegate the type + len + crc32c envelope to SectionFramer. + const size_t start = sink->size(); + sink->reserve(sizes.framed_bytes); + SectionFramer::write(*sink, kNullBitmapSectionType, payload.view()); + DORIS_CHECK_EQ(sink->size() - start, sizes.framed_bytes); + return Status::OK(); +} + +NullBitmapReader:: + NullBitmapReader() // NOLINT(modernize-use-equals-default): roaring type is incomplete in the header. + : bitmap_(std::make_unique()) {} + +NullBitmapReader::~NullBitmapReader() = default; + +NullBitmapReader::NullBitmapReader(NullBitmapReader&&) noexcept = default; +NullBitmapReader& NullBitmapReader::operator=(NullBitmapReader&&) noexcept = default; + +Status NullBitmapReader::open(Slice framed, NullBitmapReader* out) { + ParsedNullBitmap parsed; + RETURN_IF_ERROR(parse_null_bitmap(framed, &parsed)); + *out->bitmap_ = + roaring::Roaring::readSafe(reinterpret_cast(parsed.roaring_bytes.data()), + parsed.roaring_bytes.size()); + out->doc_count_ = parsed.doc_count; + return Status::OK(); +} + +Status NullBitmapReader::decoded_memory_bytes(Slice framed, uint64_t* out) { + if (out == nullptr) { + return Status::Error( + "null bitmap: null decoded memory output"); + } + ParsedNullBitmap parsed; + RETURN_IF_ERROR(parse_null_bitmap(framed, &parsed)); + + constexpr uint64_t kContainerObjectBytes = + std::max({sizeof(roaring::internal::array_container_t), + sizeof(roaring::internal::bitset_container_t), + sizeof(roaring::internal::run_container_t)}); + constexpr uint64_t kContainerMetadataBytes = + sizeof(void*) + sizeof(uint16_t) + sizeof(uint8_t) + kContainerObjectBytes; + constexpr uint64_t kFixedBytes = + sizeof(roaring::Roaring) + sizeof(roaring::api::roaring_bitmap_t); + const uint64_t container_bytes = + static_cast(parsed.container_count) * kContainerMetadataBytes; + if (parsed.roaring_bytes.size() > + std::numeric_limits::max() - container_bytes - kFixedBytes) { + return Status::Error( + "null bitmap: decoded memory size overflows"); + } + *out = parsed.roaring_bytes.size() + container_bytes + kFixedBytes; + return Status::OK(); +} + +bool NullBitmapReader::is_null(uint32_t docid) const { + return bitmap_->contains(docid); +} + +uint32_t NullBitmapReader::null_count() const { + return static_cast(bitmap_->cardinality()); +} + +void NullBitmapReader::copy_to(roaring::Roaring* out) const { + *out = *bitmap_; +} + +void NullBitmapReader::append_docids(std::vector& out) const { + for (uint32_t docid : *bitmap_) { + out.push_back(docid); + } +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/null_bitmap.h b/be/src/storage/index/snii/format/null_bitmap.h new file mode 100644 index 00000000000000..a975a76d399d5d --- /dev/null +++ b/be/src/storage/index/snii/format/null_bitmap.h @@ -0,0 +1,131 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +// Forward-declare the CRoaring C++ bitmap so this header stays free of the +// (large) roaring include; the concrete type is only needed in the .cpp. +namespace roaring { +class Roaring; +} // namespace roaring + +namespace doris::snii::format { + +// SectionFramer type byte for the null-bitmap POD. There is no dedicated +// SectionType enum value yet, so we use a documented literal (0x20) outside the +// currently allocated enum range (1..9) to avoid colliding with existing types. +inline constexpr uint8_t kNullBitmapSectionType = 0x20; + +struct NullBitmapSerializationSizes { + size_t roaring_bytes = 0; + size_t payload_bytes = 0; + size_t framed_bytes = 0; +}; + +// NullBitmap POD: per logical index, a Roaring bitmap of null docids (docs whose +// value is NULL / not indexed). It decouples per-doc NULL information from the +// per-term dictionary / postings so NULL handling can pull only this side POD. +// +// On-disk layout (the whole section is framed by SectionFramer, which adds a +// type + varint64 len + payload + fixed32 crc32c envelope): +// framer payload = [varint64 doc_count][varint64 roaring_size][roaring_bytes] +// roaring_bytes is the portable CRoaring serialization (Roaring::write). +class NullBitmapWriter { +public: + NullBitmapWriter(); + ~NullBitmapWriter(); + + NullBitmapWriter(const NullBitmapWriter&) = delete; + NullBitmapWriter& operator=(const NullBitmapWriter&) = delete; + + // Marks docid as NULL (adding the same docid twice is idempotent). + void add_null(uint32_t docid); + void add_many(std::span docids); + + // Number of distinct null docids accumulated so far. + uint32_t null_count() const; + + // Conservative pre-allocation charge for constructing CRoaring from sorted + // docids. It includes top-level array growth and array-to-bitset conversion + // overlap, so the caller must retain this charge until the bitmap is destroyed. + static uint64_t build_memory_upper_bound(std::span sorted_docids); + + Status serialization_sizes(uint32_t doc_count, NullBitmapSerializationSizes* out) const; + + // Serializes [doc_count][roaring_size][roaring_bytes] framed by SectionFramer + // and appends it to sink (does not clear sink). doc_count is the total number + // of docs in the logical index (recorded so the reader can round-trip it). + Status finish(uint32_t doc_count, ByteSink* sink) const; + +private: + std::unique_ptr bitmap_; +}; + +// Read-only view: on open, SectionFramer verifies the CRC and truncation; this +// class then guards roaring_size against the remaining payload bytes before +// deserializing the Roaring bitmap (anti-DoS), so a corrupt size cannot trigger +// an oversized allocation/read. is_null() is then an O(1) membership test. +class NullBitmapReader { +public: + NullBitmapReader(); + ~NullBitmapReader(); + + NullBitmapReader(const NullBitmapReader&) = delete; + NullBitmapReader& operator=(const NullBitmapReader&) = delete; + NullBitmapReader(NullBitmapReader&&) noexcept; + NullBitmapReader& operator=(NullBitmapReader&&) noexcept; + + // Parses the entire section (framer envelope + payload). Returns Corruption on + // CRC mismatch, truncation, doc_count overflow, or an oversized roaring_size. + static Status open(Slice framed, NullBitmapReader* out); + + // Exact pre-allocation charge for CRoaring::readSafe, excluding the caller's + // framed input bytes and decoded docid output vector. + static Status decoded_memory_bytes(Slice framed, uint64_t* out); + + // True iff docid was marked NULL. docids outside the null set (including those + // >= doc_count) return false. + bool is_null(uint32_t docid) const; + + // Number of distinct null docids in the bitmap. + uint32_t null_count() const; + + // Copies the decoded bitmap into the caller-owned Roaring object. + void copy_to(roaring::Roaring* out) const; + + // Appends decoded docids in ascending order without cloning the bitmap. + void append_docids(std::vector& out) const; + + // Total doc count of the logical index, as recorded by the writer. + uint32_t doc_count() const { return doc_count_; } + +private: + std::unique_ptr bitmap_; + uint32_t doc_count_ = 0; +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/phrase_bigram.h b/be/src/storage/index/snii/format/phrase_bigram.h new file mode 100644 index 00000000000000..3d88f897e9bd6b --- /dev/null +++ b/be/src/storage/index/snii/format/phrase_bigram.h @@ -0,0 +1,33 @@ +// 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. + +#pragma once + +#include + +namespace doris::snii::format { + +inline constexpr std::string_view kPhraseBigramTermMarker = + "\x1F" + "SNII_PHRASE_BIGRAM" + "\x1F"; + +inline bool is_phrase_bigram_term(std::string_view term) { + return term.starts_with(kPhraseBigramTermMarker); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/prx_decode_stats.h b/be/src/storage/index/snii/format/prx_decode_stats.h new file mode 100644 index 00000000000000..f4c3bb093cba88 --- /dev/null +++ b/be/src/storage/index/snii/format/prx_decode_stats.h @@ -0,0 +1,106 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" + +namespace doris::snii::format { + +// Optional allocation seam for full CSR decode callers that must gate every +// retained/output allocation before touching the physical buffers. Query +// decoders leave this null; native compaction supplies its shared-budget +// workspace. Implementations must leave both CSR buffers unchanged when a +// reservation fails. +class PrxCsrAllocationGate { +public: + virtual ~PrxCsrAllocationGate() = default; + + virtual Status reserve_csr(std::vector* pos_flat, size_t position_count, + std::vector* pos_off, size_t offset_count) = 0; + virtual Status reserve_decompression(size_t bytes, std::vector** buffer) = 0; +}; + +struct PrxDecodeStats { + uint64_t raw_frames = 0; + uint64_t zstd_frames = 0; + uint64_t pfor_frames = 0; + uint64_t plaintext_bytes = 0; + uint64_t total_docs = 0; + uint64_t selected_docs = 0; + uint64_t total_positions = 0; + uint64_t selected_positions = 0; + uint64_t fetch_ns = 0; + // Inclusive successful-frame time: header/CRC validation, optional + // decompression, and payload decode. + uint64_t decode_ns = 0; + // Phrase verification excluding only the inclusive decode_ns delta. + uint64_t phrase_verify_ns = 0; + + void merge(const PrxDecodeStats& other); + [[nodiscard]] uint64_t frame_count() const { return raw_frames + zstd_frames + pfor_frames; } + [[nodiscard]] bool is_valid() const { + return selected_docs <= total_docs && selected_positions <= total_positions; + } + bool operator==(const PrxDecodeStats&) const = default; +}; + +struct PrxDecodedShape { + uint32_t total_docs = 0; + uint64_t total_positions = 0; + uint32_t max_frequency = 0; + bool has_zero_frequency = false; +}; + +// Query-plan and matcher calibration inputs are deliberately separate from PrxDecodeStats: the +// latter's 11 production counters remain a stable decode contract. +struct PhraseQueryExecutionStats { + uint64_t exact_candidate_docs = 0; + uint64_t exact_candidate_visits = 0; + uint64_t prefix_leading_candidate_docs = 0; + uint64_t prefix_tail_candidate_visits = 0; + uint64_t common_grams_candidate_queries = 0; + uint64_t common_grams_plain_plans = 0; + uint64_t common_grams_gram_plans = 0; + uint64_t common_grams_fallback_no_gram = 0; + uint64_t common_grams_fallback_incompatible = 0; + uint64_t common_grams_fallback_kill_switch = 0; + uint64_t common_grams_fallback_cost = 0; + uint64_t common_grams_fallback_base_analyzer_mismatch = 0; + uint64_t common_grams_fallback_prefix_tail_empty = 0; + uint64_t common_grams_authoritative_empty = 0; + uint64_t common_grams_plain_posting_bytes = 0; + uint64_t common_grams_gram_posting_bytes = 0; + uint64_t common_grams_plain_estimated_candidate_df = 0; + uint64_t common_grams_gram_estimated_candidate_df = 0; + uint64_t common_grams_plain_estimated_cost = 0; + uint64_t common_grams_gram_estimated_cost = 0; + uint64_t common_grams_planning_ns = 0; +}; + +struct PrxDecodeContext { + PrxDecodeStats* stats = nullptr; + PrxDecodedShape* shape = nullptr; + PhraseQueryExecutionStats* query_stats = nullptr; + PrxCsrAllocationGate* allocation_gate = nullptr; +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/prx_pod.cpp b/be/src/storage/index/snii/format/prx_pod.cpp new file mode 100644 index 00000000000000..a336bfc2b833db --- /dev/null +++ b/be/src/storage/index/snii/format/prx_pod.cpp @@ -0,0 +1,1459 @@ +// 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 "storage/index/snii/format/prx_pod.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/common/uninitialized_buffer.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +void PrxDecodeStats::merge(const PrxDecodeStats& other) { + raw_frames += other.raw_frames; + zstd_frames += other.zstd_frames; + pfor_frames += other.pfor_frames; + plaintext_bytes += other.plaintext_bytes; + total_docs += other.total_docs; + selected_docs += other.selected_docs; + total_positions += other.total_positions; + selected_positions += other.selected_positions; + fetch_ns += other.fetch_ns; + decode_ns += other.decode_ns; + phrase_verify_ns += other.phrase_verify_ns; +} + +Status validate_prx_window_limits(const PrxWindowLimits& limits) { + if (limits.max_docs == 0 || limits.max_positions == 0 || limits.max_uncomp_bytes == 0) { + return Status::Error( + "prx: window limits must be non-zero"); + } + if (limits.max_docs > kReaderPrxWindowLimits.max_docs || + limits.max_positions > kReaderPrxWindowLimits.max_positions || + limits.max_uncomp_bytes > kReaderPrxWindowLimits.max_uncomp_bytes) { + return Status::Error( + "prx: writer window limits exceed reader limits"); + } + return Status::OK(); +} + +namespace { + +using PrxClock = std::chrono::steady_clock; + +PrxClock::time_point prx_clock_now() { +#ifdef BE_TEST + testing::note_prx_clock_read(); +#endif + return PrxClock::now(); +} + +uint64_t elapsed_ns(PrxClock::time_point start) { + const auto elapsed = + std::chrono::duration_cast(prx_clock_now() - start).count(); + return std::max(1, static_cast(elapsed)); +} + +// Auto-compression threshold: use raw when payload is smaller than this (zstd +// gain is negligible and metadata overhead is relatively large). +inline constexpr size_t kAutoZstdMinBytes = 512; +// Default zstd level in auto mode. +inline constexpr int kDefaultZstdLevel = 3; +// Maximum decompressed byte size for a single .prx window. Guards against a +// corrupted uncomp_len read from S3 inflated to a huge value: sanity-check +// before allocating/decompressing to avoid GB-scale allocations. Windows are +// 256-doc aligned and normally far below this limit. +inline constexpr uint32_t kMaxWindowUncompBytes = kReaderPrxWindowLimits.max_uncomp_bytes; +// Anti-DoS cap on position count decoded from a single window before +// allocation. +inline constexpr uint32_t kMaxWindowPositions = + kReaderPrxWindowLimits.max_positions; // 64M positions/window +// Anti-DoS cap on doc count decoded from a single window before allocation. A +// corrupt doc_count is otherwise fed straight to assign()/reserve() -> +// bad_alloc. +inline constexpr uint32_t kMaxWindowDocs = kReaderPrxWindowLimits.max_docs; // 16M docs/window + +// Writer-side precondition for the FLAT builders: the per-doc partition `freqs` +// must address exactly the positions present in `flat`. If sum(freqs) overruns +// flat.size() a (positions_flat, freqs) mismatch would index flat[off+i] past +// the span end -- an out-of-bounds read on caller-supplied data. Reject it as +// InvalidArgument BEFORE any indexing so the bug surfaces as a clean Status, +// never UB. (sum < size leaves trailing positions unused, which is also a +// writer bug, so we require exact equality.) Uint64 accumulation cannot +// overflow for uint32 freqs. +Status check_flat_partition(std::span flat, std::span freqs) { + size_t sum = 0; + for (uint32_t fc : freqs) { + if (fc > flat.size() - sum) { + return Status::Error( + "prx: sum(freqs) exceeds positions_flat size"); + } + sum += fc; + } + if (sum != flat.size()) { + return Status::Error( + "prx: sum(freqs) does not match positions_flat size"); + } + return Status::OK(); +} + +Status validate_flat_positions(std::span flat, std::span freqs) { + size_t off = 0; + for (uint32_t fc : freqs) { + uint32_t previous = 0; + for (uint32_t i = 0; i < fc; ++i) { + const uint32_t position = flat[off + i]; + if (i != 0 && position < previous) { + return Status::Error( + "prx: positions within a doc must be ascending"); + } + previous = position; + } + off += fc; + } + return Status::OK(); +} + +Status validate_per_doc_window(std::span> per_doc, + const PrxWindowLimits& limits, size_t* total_positions) { + RETURN_IF_ERROR(validate_prx_window_limits(limits)); + if (per_doc.size() > limits.max_docs) { + return Status::Error( + "prx: doc count exceeds writer window limit"); + } + uint64_t total = 0; + for (const auto& positions : per_doc) { + total += positions.size(); + if (total > limits.max_positions) { + return Status::Error( + "prx: position count exceeds writer window limit"); + } + } + *total_positions = static_cast(total); + return Status::OK(); +} + +// Encode per-doc position lists into a self-describing plain payload (doc_count +// + per-doc delta stream). +Status encode_payload(std::span> per_doc, ByteSink* out) { + out->put_varint32(static_cast(per_doc.size())); + for (const auto& doc : per_doc) { + out->put_varint32(static_cast(doc.size())); + uint32_t prev = 0; + for (size_t i = 0; i < doc.size(); ++i) { + uint32_t pos = doc[i]; + if (i > 0 && pos < prev) { + return Status::Error( + "prx: positions within a doc must be ascending"); + } + out->put_varint32(i == 0 ? pos : pos - prev); + prev = pos; + } + } + return Status::OK(); +} + +// FLAT-positions encoder: identical wire output to encode_payload above, but +// reads positions from a single flat span partitioned per-doc by `freqs` (doc d +// owns the next freqs[d] entries). The public entry point has already validated +// that sum(freqs) == flat.size(). This avoids materializing a vector-of-vectors +// for the window. +Status encode_payload_flat(std::span flat, std::span freqs, + ByteSink* out) { + out->put_varint32(static_cast(freqs.size())); + size_t off = 0; + for (uint32_t fc : freqs) { + out->put_varint32(fc); + uint32_t prev = 0; + for (uint32_t i = 0; i < fc; ++i) { + const uint32_t pos = flat[off + i]; + if (i > 0 && pos < prev) { + return Status::Error( + "prx: positions within a doc must be ascending"); + } + out->put_varint32(i == 0 ? pos : pos - prev); + prev = pos; + } + off += fc; + } + return Status::OK(); +} + +// Encode a uint32 array into PFOR runs of kFrqBaseUnit (256) elements each. The +// run count is derived by the decoder from the total length, so it is not +// stored. +void encode_pfor_runs(std::span values, ByteSink* out) { + const size_t n = values.size(); + for (size_t off = 0; off < n; off += kFrqBaseUnit) { + const size_t run = (n - off < kFrqBaseUnit) ? (n - off) : kFrqBaseUnit; + pfor_encode(values.data() + off, run, out); + } +} + +// Decode n uint32 values (multiple PFOR runs of kFrqBaseUnit each) into out. +Status decode_pfor_runs(ByteSource* src, size_t n, std::vector* out) { + // Sized then fully overwritten by pfor_decode below (every [0, n) slot is + // written); no zero-fill needed beyond what std::vector mandates. + resize_uninitialized(*out, n); + for (size_t off = 0; off < n; off += kFrqBaseUnit) { + const size_t run = (n - off < kFrqBaseUnit) ? (n - off) : kFrqBaseUnit; + RETURN_IF_ERROR(pfor_decode(src, run, out->data() + off)); + } + return Status::OK(); +} + +size_t varint32_size(uint32_t value) { + size_t bytes = 1; + while (value >= 128) { + value >>= 7; + ++bytes; + } + return bytes; +} + +// Derive the per-doc position deltas ONCE into `deltas` (flat, in doc order: the +// first position of each doc is absolute, the rest are deltas within the doc), +// enforcing the ascending-position precondition after the public entry point +// validated the exact (flat, freqs) partition. The loop is identical to the delta +// derivation the old encode_pfor_payload_flat ran inline, lifted out so the auto +// path can feed BOTH the PFOR payload and (only when needed) the raw plaintext +// payload from one buffer instead of walking `flat` twice. Accumulate the exact +// raw payload size in the same pass so codec selection does not rescan deltas. +Status compute_flat_deltas(std::span flat, std::span freqs, + std::vector* const deltas, size_t* const plain_size) { +#ifdef BE_TEST + testing::note_prx_delta_materialization(); +#endif + deltas->clear(); + deltas->reserve(flat.size()); + *plain_size = varint32_size(static_cast(freqs.size())); + size_t off = 0; + for (uint32_t fc : freqs) { + *plain_size += varint32_size(fc); + uint32_t prev = 0; + for (uint32_t i = 0; i < fc; ++i) { + const uint32_t pos = flat[off + i]; + if (i > 0 && pos < prev) { + return Status::Error( + "prx: positions within a doc must be ascending"); + } + const uint32_t delta = i == 0 ? pos : pos - prev; + deltas->push_back(delta); + *plain_size += varint32_size(delta); + prev = pos; + } + off += fc; + } + return Status::OK(); +} + +// PFOR window payload (self-describing; no entropy coding): +// VInt doc_count +// VInt total_pos # sum of all pos_counts +// PFOR_runs(pos_counts) # doc_count values (bit-packed; mostly 1 -> ~1 +// bit) PFOR_runs(position_deltas) # total_pos deltas, flat across docs (first +// per +// # doc absolute, rest delta-within-doc) +// Bit-packing the per-doc pos_counts (vs one varint each) is the size win: in a +// uniform corpus most docs have freq 1, so the count column packs to ~1 bit/doc. +// Emits byte-for-byte the same payload the old encode_pfor_payload_flat produced +// (doc_count == freqs.size(), total_pos == deltas.size() == sum(freqs)), but +// reads the already-derived `deltas` instead of re-walking the positions. +void encode_pfor_payload_from_deltas(std::span freqs, + std::span deltas, ByteSink* out) { + out->put_varint32(static_cast(freqs.size())); + out->put_varint32(static_cast(deltas.size())); + encode_pfor_runs(freqs, out); + encode_pfor_runs(deltas, out); +} + +// Raw plaintext payload (self-describing per-doc boundaries): +// VInt doc_count +// per doc: VInt pos_count, then pos_count position deltas (VInt) +// Emits byte-for-byte the same payload the old encode_payload_flat produced, but +// reads the already-derived `deltas` instead of re-walking the positions and +// re-running the partition/ascending checks. +void encode_payload_from_deltas(std::span freqs, std::span deltas, + ByteSink* out) { + out->put_varint32(static_cast(freqs.size())); + size_t off = 0; + for (uint32_t fc : freqs) { + out->put_varint32(fc); + for (uint32_t i = 0; i < fc; ++i) { + out->put_varint32(deltas[off + i]); + } + off += fc; + } +} + +// Decode per-doc position lists from a PFOR payload. +Status decode_pfor_payload(Slice plain, std::vector>* out) { + ByteSource src(plain); + uint32_t doc_count = 0, total_pos = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + RETURN_IF_ERROR(src.get_varint32(&total_pos)); + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + std::vector pos_counts; + RETURN_IF_ERROR(decode_pfor_runs(&src, doc_count, &pos_counts)); + uint64_t sum = 0; + for (uint32_t d = 0; d < doc_count; ++d) sum += pos_counts[d]; + if (sum != total_pos) { + return Status::Error( + "prx: pos_count sum mismatch"); + } + std::vector deltas; + RETURN_IF_ERROR(decode_pfor_runs(&src, total_pos, &deltas)); + out->clear(); + out->reserve(doc_count); + size_t off = 0; + for (uint32_t d = 0; d < doc_count; ++d) { + std::vector doc; + doc.reserve(pos_counts[d]); + uint32_t prev = 0; + for (uint32_t i = 0; i < pos_counts[d]; ++i) { + prev = (i == 0) ? deltas[off + i] : prev + deltas[off + i]; + doc.push_back(prev); + } + off += pos_counts[d]; + out->push_back(std::move(doc)); + } + if (!src.eof()) + return Status::Error( + "prx: trailing bytes after pfor payload"); + return Status::OK(); +} + +// Writes a PFOR window: codec=pfor, payload, crc(header+payload). +void write_pfor(Slice payload, ByteSink* sink) { + // Single-copy framing: write [codec][varint len][payload] straight into the + // caller's sink, then crc exactly those bytes. view() is taken AFTER the + // payload and BEFORE the crc, so subslice([start, framed_len)) is over a + // settled, contiguous buffer with no pending realloc/aliasing. Byte-identical + // to the former temp-ByteSink assembly, minus one heap alloc + one payload copy. + const size_t start = sink->size(); + sink->put_u8(static_cast(PrxCodec::kPfor)); + sink->put_varint32(static_cast(payload.size())); + sink->put_bytes(payload); + const size_t framed_len = sink->size() - start; + const uint32_t crc = crc32c(sink->view().subslice(start, framed_len)); + sink->put_fixed32(crc); +} + +void write_raw(Slice plain, ByteSink* sink); + +// Emit a RAW frame directly from the already-derived deltas. This is used by +// the singleton fast path so choosing RAW does not allocate a temporary plain +// payload or encode/copy a losing PFOR payload first. +void write_raw_from_deltas(std::span freqs, std::span deltas, + size_t plain_size, ByteSink* sink) { + const size_t start = sink->size(); + sink->put_u8(static_cast(PrxCodec::kRaw)); + sink->put_varint32(static_cast(plain_size)); + const size_t payload_start = sink->size(); + encode_payload_from_deltas(freqs, deltas, sink); + DCHECK_EQ(sink->size() - payload_start, plain_size); + const size_t framed_len = sink->size() - start; + const uint32_t crc = crc32c(sink->view().subslice(start, framed_len)); + sink->put_fixed32(crc); +} + +Status validate_single_doc_byte_limits(std::span positions, + std::span freqs, bool auto_codec, + uint32_t max_uncomp_bytes) { + size_t offset = 0; + std::vector one_doc_deltas; + for (size_t doc = 0; doc < freqs.size(); ++doc) { + const uint32_t frequency = freqs[doc]; + const auto one_freq = freqs.subspan(doc, 1); + const auto one_doc_positions = positions.subspan(offset, frequency); + size_t plain_size = varint32_size(1) + varint32_size(frequency); + uint32_t previous = 0; + for (size_t i = 0; i < one_doc_positions.size(); ++i) { + const uint32_t position = one_doc_positions[i]; + plain_size += varint32_size(i == 0 ? position : position - previous); + previous = position; + } + if (plain_size > max_uncomp_bytes) { + if (!auto_codec) { + return Status::Error( + "prx: one document exceeds the writer window byte limit"); + } + if (one_doc_deltas.capacity() < frequency) { + std::vector().swap(one_doc_deltas); + one_doc_deltas.reserve(frequency); + } else { + one_doc_deltas.clear(); + } + previous = 0; + for (size_t i = 0; i < one_doc_positions.size(); ++i) { + const uint32_t position = one_doc_positions[i]; + one_doc_deltas.push_back(i == 0 ? position : position - previous); + previous = position; + } + ByteSink pfor_payload; + encode_pfor_payload_from_deltas(one_freq, one_doc_deltas, &pfor_payload); + if (pfor_payload.size() > max_uncomp_bytes) { + return Status::Error( + "prx: one document exceeds the writer window byte limit"); + } + } + offset += frequency; + } + return Status::OK(); +} + +size_t pfor_frame_size(size_t payload_size) { + return 1 + varint32_size(static_cast(payload_size)) + payload_size + sizeof(uint32_t); +} + +size_t raw_frame_size(size_t plain_size) { + return 1 + varint32_size(static_cast(plain_size)) + plain_size + sizeof(uint32_t); +} + +size_t zstd_frame_size(size_t plain_size, size_t compressed_size) { + return 1 + varint32_size(static_cast(plain_size)) + + varint32_size(static_cast(compressed_size)) + compressed_size + + sizeof(uint32_t); +} + +struct AutoPrxCodecChoice { + PrxCodec codec = PrxCodec::kPfor; + bool readable = false; +}; + +// Select the smallest complete reader-safe frame among the candidates already +// materialized by the zstd/fallback path. Preserve the existing codec on equal +// sizes by considering PFOR, then ZSTD, then RAW and replacing the winner only +// on a strict size reduction. Sub-threshold singleton RAW selection happens +// before PFOR materialization in build_prx_window_auto_from_flat. +AutoPrxCodecChoice select_auto_prx_codec(size_t pfor_payload_size, size_t plain_payload_size, + size_t compressed_payload_size, bool has_compressed, + uint32_t max_uncomp_bytes) { + const bool pfor_readable = pfor_payload_size <= max_uncomp_bytes; + const bool plain_readable = plain_payload_size <= max_uncomp_bytes; + AutoPrxCodecChoice choice; + size_t selected_frame_size = 0; + if (pfor_readable) { + choice = {.codec = PrxCodec::kPfor, .readable = true}; + selected_frame_size = pfor_frame_size(pfor_payload_size); + } + if (has_compressed && plain_readable) { + const size_t frame_size = zstd_frame_size(plain_payload_size, compressed_payload_size); + if (!choice.readable || frame_size < selected_frame_size) { + choice = {.codec = PrxCodec::kZstd, .readable = true}; + selected_frame_size = frame_size; + } + } + if (plain_readable) { + const size_t frame_size = raw_frame_size(plain_payload_size); + if (!choice.readable || frame_size < selected_frame_size) { + choice = {.codec = PrxCodec::kRaw, .readable = true}; + } + } + return choice; +} + +void write_zstd_compressed(Slice plain, Slice compressed, ByteSink* sink) { + // Single-copy framing (see write_pfor): assemble [codec][uncomp_len][comp_len] + // [compressed] in the caller's sink and crc that span before appending the crc. + const size_t start = sink->size(); + sink->put_u8(static_cast(PrxCodec::kZstd)); + sink->put_varint32(static_cast(plain.size())); + sink->put_varint32(static_cast(compressed.size())); + sink->put_bytes(compressed); + const size_t framed_len = sink->size() - start; + const uint32_t crc = crc32c(sink->view().subslice(start, framed_len)); + sink->put_fixed32(crc); +} + +// Shared auto-mode path for BOTH .prx builders. A single-doc/freq=1 window has a +// provably smaller RAW frame: RAW stores three payload varints, while PFOR adds +// two run headers and is always 4-5 bytes larger. Emit that RAW frame directly +// before any PFOR work. Other sub-threshold windows retain the existing PFOR +// policy so an unmeasured second payload encode cannot regress import CPU. At +// and above the zstd threshold the raw plaintext is already required for +// compression, so all materialized candidates participate in the exact complete +// frame-size comparison at no additional encoding cost. +Status build_prx_window_auto_from_flat(std::span positions_flat, + std::span freqs, int zstd_level, + const PrxWindowLimits& limits, ByteSink* sink, + PrxWindowBuildOutcome* outcome) { + if (freqs.size() == 1 && freqs.front() == 1) { + const size_t plain_size = + varint32_size(1) + varint32_size(1) + varint32_size(positions_flat.front()); + if (plain_size > limits.max_uncomp_bytes) { + return Status::Error( + "prx: one document exceeds the writer window byte limit"); + } + write_raw_from_deltas(freqs, positions_flat, plain_size, sink); + *outcome = PrxWindowBuildOutcome::kBuilt; + return Status::OK(); + } + + std::vector deltas; + size_t plain_size = 0; + RETURN_IF_ERROR(compute_flat_deltas(positions_flat, freqs, &deltas, &plain_size)); + const bool plain_readable = plain_size <= limits.max_uncomp_bytes; + + ByteSink payload; + encode_pfor_payload_from_deltas(freqs, deltas, &payload); + const bool pfor_readable = payload.size() <= limits.max_uncomp_bytes; + if (!pfor_readable && !plain_readable) { + if (freqs.size() <= 1) { + return Status::Error( + "prx: one document exceeds the writer window byte limit"); + } + RETURN_IF_ERROR(validate_single_doc_byte_limits(positions_flat, freqs, true, + limits.max_uncomp_bytes)); + *outcome = PrxWindowBuildOutcome::kNeedsSplit; + return Status::OK(); + } + if (plain_readable && (plain_size >= kAutoZstdMinBytes || !pfor_readable)) { + ByteSink plain; + encode_payload_from_deltas(freqs, deltas, &plain); + DCHECK_EQ(plain.size(), plain_size); + std::vector compressed; + const bool has_compressed = plain_size >= kAutoZstdMinBytes; + if (plain_size >= kAutoZstdMinBytes) { + testing::note_prx_raw_build(); + RETURN_IF_ERROR(zstd_compress(plain.view(), zstd_level, &compressed)); + } + const AutoPrxCodecChoice choice = + select_auto_prx_codec(payload.size(), plain.size(), compressed.size(), + has_compressed, limits.max_uncomp_bytes); + DCHECK(choice.readable); + if (choice.codec == PrxCodec::kZstd) { + write_zstd_compressed(plain.view(), Slice(compressed), sink); + } else if (choice.codec == PrxCodec::kRaw) { + write_raw(plain.view(), sink); + } else { + write_pfor(payload.view(), sink); + } + *outcome = PrxWindowBuildOutcome::kBuilt; + return Status::OK(); + } + DCHECK(pfor_readable); + write_pfor(payload.view(), sink); + *outcome = PrxWindowBuildOutcome::kBuilt; + return Status::OK(); +} + +// Decode per-doc position lists from a plain payload. +Status decode_payload(Slice plain, std::vector>* out) { + ByteSource src(plain); + uint32_t doc_count = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + out->clear(); + out->reserve(doc_count); + for (uint32_t d = 0; d < doc_count; ++d) { + uint32_t pos_count = 0; + RETURN_IF_ERROR(src.get_varint32_fast(&pos_count)); + std::vector doc; + doc.reserve(pos_count); + uint32_t prev = 0; + for (uint32_t i = 0; i < pos_count; ++i) { + uint32_t delta = 0; + RETURN_IF_ERROR(src.get_varint32(&delta)); + if (i != 0 && delta > std::numeric_limits::max() - prev) { + return Status::Error( + "prx: position accumulation overflow"); + } + prev = (i == 0) ? delta : prev + delta; + doc.push_back(prev); + } + out->push_back(std::move(doc)); + } + if (!src.eof()) + return Status::Error( + "prx: trailing bytes after payload"); + return Status::OK(); +} + +// CSR decode of a PFOR payload: all docs' positions into one flat buffer + +// per-doc offsets, with NO per-doc std::vector allocation. `pos_off` has +// doc_count+1 entries (pos_off[0]==0); doc d's positions are +// pos_flat[pos_off[d] .. pos_off[d+1]). +Status decode_pfor_payload_csr(Slice plain, std::vector* pos_flat, + std::vector* pos_off, + PrxCsrAllocationGate* allocation_gate, uint32_t* max_frequency, + bool* has_zero_frequency) { + ByteSource src(plain); + uint32_t doc_count = 0, total_pos = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + RETURN_IF_ERROR(src.get_varint32(&total_pos)); + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + if (allocation_gate != nullptr) { + RETURN_IF_ERROR(allocation_gate->reserve_csr(pos_flat, total_pos, pos_off, + static_cast(doc_count) + 1)); + } + pos_off->clear(); + pos_off->reserve(static_cast(doc_count) + 1); + RETURN_IF_ERROR(decode_pfor_runs(&src, doc_count, pos_off)); + uint64_t sum = 0; + uint32_t decoded_max_frequency = 0; + bool decoded_zero_frequency = false; + for (uint32_t d = 0; d < doc_count; ++d) { + sum += (*pos_off)[d]; + decoded_max_frequency = std::max(decoded_max_frequency, (*pos_off)[d]); + decoded_zero_frequency |= (*pos_off)[d] == 0; + } + if (sum != total_pos) + return Status::Error( + "prx: pos_count sum mismatch"); + // pos_flat is sized to total_pos by decode_pfor_runs (resize_uninitialized); + // a separate reserve is redundant. pos_off keeps its reserve (push_back below). + RETURN_IF_ERROR(decode_pfor_runs(&src, total_pos, pos_flat)); + size_t off = 0; + uint32_t next_off = 0; + for (uint32_t d = 0; d < doc_count; ++d) { + const uint32_t pos_count = (*pos_off)[d]; + (*pos_off)[d] = next_off; + uint32_t prev = 0; + for (uint32_t i = 0; i < pos_count; ++i) { + uint32_t& value = (*pos_flat)[off + i]; + if (i != 0 && value > std::numeric_limits::max() - prev) { + return Status::Error( + "prx: position accumulation overflow"); + } + prev = (i == 0) ? value : prev + value; + value = prev; + } + off += pos_count; + next_off += pos_count; + } + pos_off->push_back(next_off); + if (!src.eof()) + return Status::Error( + "prx: trailing bytes after pfor payload"); + *max_frequency = decoded_max_frequency; + *has_zero_frequency = decoded_zero_frequency; + return Status::OK(); +} + +Status validate_doc_ordinals(std::span doc_ordinals, uint32_t doc_count) { + uint32_t prev = 0; + for (size_t i = 0; i < doc_ordinals.size(); ++i) { + const uint32_t doc = doc_ordinals[i]; + if (doc >= doc_count) { + return Status::Error( + "prx: selected doc ordinal out of range"); + } + if (i != 0 && doc <= prev) { + return Status::Error( + "prx: selected doc ordinals must be strictly ascending"); + } + prev = doc; + } + return Status::OK(); +} + +struct SelectedRange { + SelectedRange(uint32_t begin_, uint32_t end_, uint32_t out_begin_) + : begin(begin_), end(end_), out_begin(out_begin_) {} + + uint32_t begin; + uint32_t end; + uint32_t out_begin; +}; + +uint32_t count_covered_pfor_runs(std::span selected, uint32_t total_pos) { + if (selected.empty() || total_pos == 0) { + return 0; + } + uint32_t runs = 0; + uint32_t next_run = 0; + for (const SelectedRange& range : selected) { + if (range.begin == range.end) { + continue; + } + const uint32_t first_run = range.begin / kFrqBaseUnit; + const uint32_t last_run = (range.end - 1) / kFrqBaseUnit; + const uint32_t counted_first = std::max(first_run, next_run); + if (counted_first <= last_run) { + runs += last_run - counted_first + 1; + next_run = last_run + 1; + } + } + return runs; +} + +bool should_decode_full_prx_positions(std::span selected, + uint32_t selected_pos_count, uint32_t total_pos) { + if (selected.empty() || total_pos == 0) { + return false; + } + if (selected_pos_count * 2 >= total_pos) { + return true; + } + const uint32_t total_runs = (total_pos + kFrqBaseUnit - 1) / kFrqBaseUnit; + const uint32_t covered_runs = count_covered_pfor_runs(selected, total_pos); + return covered_runs * 4 >= total_runs * 3; +} + +Status decode_selected_pfor_count_ranges(ByteSource* src, uint32_t doc_count, + std::span doc_ordinals, + std::vector& selected, + std::vector& pos_off, uint64_t* total_pos_count, + uint32_t* selected_pos_count, uint32_t* max_frequency, + bool* has_zero_frequency) { + selected.clear(); + selected.reserve(doc_ordinals.size()); + pos_off.clear(); + pos_off.reserve(doc_ordinals.size() + 1); + pos_off.push_back(0); + + *selected_pos_count = 0; + uint32_t delta_begin = 0; + size_t next_doc = 0; + *total_pos_count = 0; + *max_frequency = 0; + *has_zero_frequency = false; + std::array run_buf {}; + for (uint32_t run_begin = 0; run_begin < doc_count; run_begin += kFrqBaseUnit) { + const uint32_t run_len = std::min(kFrqBaseUnit, doc_count - run_begin); + RETURN_IF_ERROR(pfor_decode(src, run_len, run_buf.data())); + for (uint32_t i = 0; i < run_len; ++i) { + const uint32_t d = run_begin + i; + const uint32_t count = run_buf[i]; + *max_frequency = std::max(*max_frequency, count); + *has_zero_frequency |= count == 0; + *total_pos_count += count; + if (*total_pos_count > kMaxWindowPositions) { + return Status::Error( + "prx: pos_count sum exceeds sane cap"); + } + if (next_doc < doc_ordinals.size() && doc_ordinals[next_doc] == d) { + selected.emplace_back(delta_begin, delta_begin + count, *selected_pos_count); + *selected_pos_count += count; + pos_off.push_back(*selected_pos_count); + ++next_doc; + } + delta_begin += count; + } + } + if (next_doc != doc_ordinals.size()) { + return Status::Error( + "prx: selected doc ordinal was not decoded"); + } + return Status::OK(); +} + +Status decode_selected_pfor_positions(ByteSource* src, uint32_t total_pos, + std::span selected, bool decode_all_runs, + std::span pos_flat) { + std::array run_buf {}; + size_t range_idx = 0; + uint32_t prev = 0; + for (uint32_t run_begin = 0; run_begin < total_pos; run_begin += kFrqBaseUnit) { + const uint32_t run_len = std::min(kFrqBaseUnit, total_pos - run_begin); + const uint32_t run_end = run_begin + run_len; + while (range_idx < selected.size() && selected[range_idx].end <= run_begin) { + ++range_idx; + prev = 0; + } + if (!decode_all_runs && + (range_idx == selected.size() || selected[range_idx].begin >= run_end)) { + RETURN_IF_ERROR(pfor_skip(src, run_len)); + continue; + } + + RETURN_IF_ERROR(pfor_decode(src, run_len, run_buf.data())); + while (range_idx < selected.size() && selected[range_idx].begin < run_end) { + const SelectedRange& range = selected[range_idx]; + const uint32_t copy_begin = std::max(range.begin, run_begin); + const uint32_t copy_end = std::min(range.end, run_end); + if (copy_begin == range.begin) { + prev = 0; + } + uint32_t dst = range.out_begin + copy_begin - range.begin; + for (uint32_t off = copy_begin; off < copy_end; ++off) { + const uint32_t delta = run_buf[off - run_begin]; + if (off != range.begin && delta > std::numeric_limits::max() - prev) { + return Status::Error( + "prx: position accumulation overflow"); + } + prev = (off == range.begin) ? delta : prev + delta; + pos_flat[dst++] = prev; + } + if (copy_end < range.end) { + break; + } + ++range_idx; + prev = 0; + } + } + return Status::OK(); +} + +Status decode_pfor_payload_csr_selective(Slice plain, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off, + uint32_t* decoded_doc_count, + uint64_t* decoded_total_positions, uint32_t* max_frequency, + bool* has_zero_frequency) { + ByteSource src(plain); + uint32_t doc_count = 0, total_pos = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + RETURN_IF_ERROR(src.get_varint32(&total_pos)); + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + RETURN_IF_ERROR(validate_doc_ordinals(doc_ordinals, doc_count)); + + pos_flat->clear(); + + std::vector selected; + uint64_t sum = 0; + uint32_t selected_pos_count = 0; + RETURN_IF_ERROR(decode_selected_pfor_count_ranges(&src, doc_count, doc_ordinals, selected, + *pos_off, &sum, &selected_pos_count, + max_frequency, has_zero_frequency)); + if (sum != total_pos) { + return Status::Error( + "prx: pos_count sum mismatch"); + } + + const bool decode_all_runs = + should_decode_full_prx_positions(selected, selected_pos_count, total_pos); + pos_flat->resize(selected_pos_count); + RETURN_IF_ERROR(decode_selected_pfor_positions( + &src, total_pos, selected, decode_all_runs, + std::span(pos_flat->data(), pos_flat->size()))); + if (!src.eof()) { + return Status::Error( + "prx: trailing bytes after pfor payload"); + } + *decoded_doc_count = doc_count; + *decoded_total_positions = total_pos; + return Status::OK(); +} + +// CSR decode of a plain (raw) payload. See decode_pfor_payload_csr. +Status scan_payload_csr_shape(Slice plain, uint32_t* doc_count, uint32_t* total_positions) { + ByteSource src(plain); + RETURN_IF_ERROR(src.get_varint32(doc_count)); + if (*doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + uint64_t total_pos = 0; + for (uint32_t d = 0; d < *doc_count; ++d) { + uint32_t pos_count = 0; + RETURN_IF_ERROR(src.get_varint32_fast(&pos_count)); + total_pos += pos_count; + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + RETURN_IF_ERROR(src.skip_varints(pos_count)); + } + if (!src.eof()) { + return Status::Error( + "prx: trailing bytes after payload"); + } + *total_positions = static_cast(total_pos); + return Status::OK(); +} + +Status decode_payload_csr(Slice plain, std::vector* pos_flat, + std::vector* pos_off, PrxCsrAllocationGate* allocation_gate, + uint32_t* max_frequency, bool* has_zero_frequency) { + if (allocation_gate != nullptr) { + uint32_t preflight_doc_count = 0; + uint32_t preflight_total_positions = 0; + RETURN_IF_ERROR( + scan_payload_csr_shape(plain, &preflight_doc_count, &preflight_total_positions)); + RETURN_IF_ERROR(allocation_gate->reserve_csr(pos_flat, preflight_total_positions, pos_off, + static_cast(preflight_doc_count) + 1)); + } + ByteSource src(plain); + uint32_t doc_count = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + pos_flat->clear(); + pos_off->clear(); + pos_off->reserve(static_cast(doc_count) + 1); + pos_off->push_back(0); + uint64_t total_pos = 0; + uint32_t decoded_max_frequency = 0; + bool decoded_zero_frequency = false; + for (uint32_t d = 0; d < doc_count; ++d) { + uint32_t pos_count = 0; + RETURN_IF_ERROR(src.get_varint32_fast(&pos_count)); + decoded_max_frequency = std::max(decoded_max_frequency, pos_count); + decoded_zero_frequency |= pos_count == 0; + total_pos += pos_count; + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + // Tight inline prefix-sum decode (single-byte fast path) -- see + // decode_delta_run. Shared with the selective reader below. + RETURN_IF_ERROR(src.decode_delta_run(pos_count, pos_flat)); + pos_off->push_back(static_cast(pos_flat->size())); + } + if (!src.eof()) + return Status::Error( + "prx: trailing bytes after payload"); + *max_frequency = decoded_max_frequency; + *has_zero_frequency = decoded_zero_frequency; + return Status::OK(); +} + +Status decode_payload_csr_selective(Slice plain, std::span doc_ordinals, + std::vector* pos_flat, std::vector* pos_off, + uint32_t* decoded_doc_count, uint64_t* decoded_total_positions, + uint32_t* max_frequency, bool* has_zero_frequency) { + ByteSource src(plain); + uint32_t doc_count = 0; + RETURN_IF_ERROR(src.get_varint32(&doc_count)); + if (doc_count > kMaxWindowDocs) { + return Status::Error( + "prx: doc count exceeds sane cap"); + } + RETURN_IF_ERROR(validate_doc_ordinals(doc_ordinals, doc_count)); + pos_flat->clear(); + pos_off->clear(); + pos_off->reserve(doc_ordinals.size() + 1); + pos_off->push_back(0); + size_t next_doc = 0; + uint64_t total_pos = 0; + uint32_t decoded_max_frequency = 0; + bool decoded_zero_frequency = false; + for (uint32_t d = 0; d < doc_count; ++d) { + uint32_t pos_count = 0; + RETURN_IF_ERROR(src.get_varint32_fast(&pos_count)); + decoded_max_frequency = std::max(decoded_max_frequency, pos_count); + decoded_zero_frequency |= pos_count == 0; + total_pos += pos_count; + if (total_pos > kMaxWindowPositions) { + return Status::Error( + "prx: position count exceeds sane cap"); + } + const bool selected = next_doc < doc_ordinals.size() && doc_ordinals[next_doc] == d; + if (!selected) { + // Skip this doc's position deltas without decoding them -- the CSR + // layout is sequential so we must advance past them, but only the + // candidate (selected) docs' positions are ever used. With a sparse + // candidate set (the common phrase / phrase-prefix case after docid + // narrowing) most docs in a window are skipped, so this avoids the + // dominant varint-decode cost. + RETURN_IF_ERROR(src.skip_varints(pos_count)); + continue; + } + // Selected doc: decode its `pos_count` ascending position deltas with a + // tight inline prefix-sum decoder (single-byte fast path, no per-value + // get_varint32/Status call chain). This is the CPU hotspot for narrowed + // phrase/phrase-prefix candidate sets, where each selected doc's varint + // run dominates after non-selected docs are skipped. + RETURN_IF_ERROR(src.decode_delta_run(pos_count, pos_flat)); + pos_off->push_back(static_cast(pos_flat->size())); + ++next_doc; + } + if (!src.eof()) + return Status::Error( + "prx: trailing bytes after payload"); + *decoded_doc_count = doc_count; + *decoded_total_positions = total_pos; + *max_frequency = decoded_max_frequency; + *has_zero_frequency = decoded_zero_frequency; + return Status::OK(); +} + +// Decision: given level and plain length, determine whether to compress. +bool should_compress(int level, size_t plain_len) { + if (level == 0) return false; // force raw + if (level > 0) return true; // force zstd + return plain_len >= kAutoZstdMinBytes; // auto +} + +// Write a raw window: codec=raw, uncomp_len, crc(header+payload), payload. +void write_raw(Slice plain, ByteSink* sink) { + // Single-copy framing (see write_pfor): assemble [codec][uncomp_len][payload] + // in the caller's sink and crc that span before appending the crc. + const size_t start = sink->size(); + sink->put_u8(static_cast(PrxCodec::kRaw)); + sink->put_varint32(static_cast(plain.size())); + sink->put_bytes(plain); + const size_t framed_len = sink->size() - start; + const uint32_t crc = crc32c(sink->view().subslice(start, framed_len)); + sink->put_fixed32(crc); +} + +// Write a zstd window: codec=zstd, uncomp_len, comp_len, crc(header+payload), +// payload. +Status write_zstd(Slice plain, int level, ByteSink* sink) { + std::vector comp; + RETURN_IF_ERROR(zstd_compress(plain, level > 0 ? level : kDefaultZstdLevel, &comp)); + write_zstd_compressed(plain, Slice(comp), sink); + return Status::OK(); +} + +struct FramedPrxWindow { + uint8_t codec = 0; + uint32_t uncomp_len = 0; + Slice payload; +}; + +// Read header + payload and verify the frame CRC. +Status read_framed(ByteSource* src, FramedPrxWindow* frame) { + const size_t start = src->position(); + RETURN_IF_ERROR(src->get_u8(&frame->codec)); + if (frame->codec != static_cast(PrxCodec::kRaw) && + frame->codec != static_cast(PrxCodec::kZstd) && + frame->codec != static_cast(PrxCodec::kPfor)) { + return Status::Error("prx: unknown codec"); + } + RETURN_IF_ERROR(src->get_varint32(&frame->uncomp_len)); + if (frame->uncomp_len > kMaxWindowUncompBytes) { + return Status::Error( + "prx: uncomp_len exceeds sane window cap"); + } + size_t payload_len = frame->uncomp_len; + if (frame->codec == static_cast(PrxCodec::kZstd)) { + uint32_t comp_len = 0; + RETURN_IF_ERROR(src->get_varint32(&comp_len)); + payload_len = comp_len; + } + RETURN_IF_ERROR(src->get_bytes(payload_len, &frame->payload)); + const size_t framed_len = src->position() - start; + uint32_t stored = 0; + RETURN_IF_ERROR(src->get_fixed32(&stored)); + if (crc32c(src->slice_from(start, framed_len)) != stored) { + return Status::Error( + "prx: window crc mismatch"); + } + return Status::OK(); +} + +void initialize_frame_stats(const FramedPrxWindow& encoded, PrxDecodeStats* stats) { + if (encoded.codec == static_cast(PrxCodec::kRaw)) { + stats->raw_frames = 1; + stats->plaintext_bytes = encoded.payload.size(); + } else if (encoded.codec == static_cast(PrxCodec::kZstd)) { + stats->zstd_frames = 1; + stats->plaintext_bytes = encoded.uncomp_len; + } else { + stats->pfor_frames = 1; + stats->plaintext_bytes = encoded.uncomp_len; + } +} + +Status decode_csr_frame(const FramedPrxWindow& encoded, std::span doc_ordinals, + bool decode_all_docs, bool all_docs_selected, + std::vector* pos_flat, std::vector* pos_off, + PrxDecodeStats* stats, PrxDecodedShape* shape, + PrxCsrAllocationGate* allocation_gate) { + if (!decode_all_docs && allocation_gate != nullptr) { + return Status::Error( + "prx: selective decode cannot use an allocation gate"); + } + if (stats != nullptr) { + initialize_frame_stats(encoded, stats); + } + uint32_t total_docs = 0; + uint64_t total_positions = 0; + uint32_t max_frequency = 0; + bool has_zero_frequency = false; + + std::vector local_decompressed; + Slice plain = encoded.payload; + if (encoded.codec == static_cast(PrxCodec::kZstd)) { + std::vector* decompressed = &local_decompressed; + if (allocation_gate != nullptr) { + RETURN_IF_ERROR( + allocation_gate->reserve_decompression(encoded.uncomp_len, &decompressed)); + DCHECK(decompressed != nullptr); + DCHECK_GE(decompressed->capacity(), encoded.uncomp_len); + } + RETURN_IF_ERROR(zstd_decompress(encoded.payload, encoded.uncomp_len, decompressed)); + plain = Slice(*decompressed); + } + + if (decode_all_docs) { + if (encoded.codec == static_cast(PrxCodec::kPfor)) { + RETURN_IF_ERROR(decode_pfor_payload_csr(plain, pos_flat, pos_off, allocation_gate, + &max_frequency, &has_zero_frequency)); + } else { + RETURN_IF_ERROR(decode_payload_csr(plain, pos_flat, pos_off, allocation_gate, + &max_frequency, &has_zero_frequency)); + } + total_docs = static_cast(pos_off->size() - 1); + total_positions = pos_flat->size(); + if (!all_docs_selected) { + RETURN_IF_ERROR(validate_doc_ordinals(doc_ordinals, total_docs)); + } + } else if (encoded.codec == static_cast(PrxCodec::kPfor)) { + RETURN_IF_ERROR(decode_pfor_payload_csr_selective(plain, doc_ordinals, pos_flat, pos_off, + &total_docs, &total_positions, + &max_frequency, &has_zero_frequency)); + } else { + RETURN_IF_ERROR(decode_payload_csr_selective(plain, doc_ordinals, pos_flat, pos_off, + &total_docs, &total_positions, &max_frequency, + &has_zero_frequency)); + } + if (shape != nullptr) { + shape->total_docs = total_docs; + shape->total_positions = total_positions; + shape->max_frequency = max_frequency; + shape->has_zero_frequency = has_zero_frequency; + } + if (stats != nullptr) { + stats->total_docs = total_docs; + stats->total_positions = total_positions; + if (all_docs_selected) { + stats->selected_docs = total_docs; + stats->selected_positions = total_positions; + } else if (!decode_all_docs) { + stats->selected_docs = doc_ordinals.size(); + stats->selected_positions = pos_flat->size(); + } + } + return Status::OK(); +} + +Status read_prx_window_csr_impl(ByteSource* source, std::span doc_ordinals, + bool decode_all_docs, bool all_docs_selected, + std::vector* pos_flat, std::vector* pos_off, + PrxDecodeContext* context) { + if (source == nullptr || pos_flat == nullptr || pos_off == nullptr) { + return Status::Error("prx: null arg"); + } + const bool collect_stats = context != nullptr && context->stats != nullptr; + PrxDecodeStats frame_stats; + PrxDecodeStats* stats = collect_stats ? &frame_stats : nullptr; + + PrxClock::time_point decode_start; + if (collect_stats) { + decode_start = prx_clock_now(); + } + FramedPrxWindow encoded; + RETURN_IF_ERROR(read_framed(source, &encoded)); + RETURN_IF_ERROR(decode_csr_frame(encoded, doc_ordinals, decode_all_docs, all_docs_selected, + pos_flat, pos_off, stats, + context == nullptr ? nullptr : context->shape, + context == nullptr ? nullptr : context->allocation_gate)); + if (collect_stats) { + // Stop inclusive decode timing before the logical-selection scan. Phrase + // execution wraps this call in PhraseVerifyTimer, so that scan remains + // part of verification rather than format decode. + frame_stats.decode_ns = elapsed_ns(decode_start); + if (decode_all_docs && !all_docs_selected) { + uint64_t selected_positions = 0; + for (uint32_t ordinal : doc_ordinals) { + DCHECK_LT(static_cast(ordinal) + 1, pos_off->size()); + selected_positions += (*pos_off)[ordinal + 1] - (*pos_off)[ordinal]; + } + frame_stats.selected_docs = doc_ordinals.size(); + frame_stats.selected_positions = selected_positions; + } + // Format decode is complete. Caller-level CSR invariants are validated + // afterwards, so a later phrase error intentionally retains this work. + context->stats->merge(frame_stats); + } + return Status::OK(); +} + +} // namespace + +Status build_prx_window(std::span> per_doc_positions, + int zstd_level_or_negative_for_auto, ByteSink* sink) { + return build_prx_window(per_doc_positions, zstd_level_or_negative_for_auto, + kReaderPrxWindowLimits, sink); +} + +Status build_prx_window(std::span> per_doc_positions, + int zstd_level_or_negative_for_auto, const PrxWindowLimits& limits, + ByteSink* sink) { + if (sink == nullptr) return Status::Error("prx: null sink"); + size_t total_positions = 0; + RETURN_IF_ERROR(validate_per_doc_window(per_doc_positions, limits, &total_positions)); + // Forced legacy codecs (level 0 = raw varint, level > 0 = zstd) are kept so + // the test/legacy paths still exercise them; the auto path (< 0) now emits + // PFOR bit-packed deltas -- no entropy coding, far cheaper build CPU than + // zstd-3. + if (zstd_level_or_negative_for_auto >= 0) { + ByteSink plain; + RETURN_IF_ERROR(encode_payload(per_doc_positions, &plain)); + if (plain.size() > limits.max_uncomp_bytes) { + return Status::Error( + "prx: encoded payload exceeds writer window byte limit"); + } + Slice plain_view = plain.view(); + if (!should_compress(zstd_level_or_negative_for_auto, plain_view.size())) { + write_raw(plain_view, sink); + return Status::OK(); + } + return write_zstd(plain_view, zstd_level_or_negative_for_auto, sink); + } + // Auto mode: flatten the per-doc lists into (positions_flat, freqs) exactly as + // the former encode_pfor_payload did, then run the shared single-encode path so + // this builder stays byte-identical to build_prx_window_flat. + std::vector flat, freqs; + freqs.reserve(per_doc_positions.size()); + flat.reserve(total_positions); + for (const auto& doc : per_doc_positions) { + freqs.push_back(static_cast(doc.size())); + flat.insert(flat.end(), doc.begin(), doc.end()); + } + // G16-h: level < -1 is auto mode at zstd level |level| (-1 stays the default). + const int auto_level = zstd_level_or_negative_for_auto == -1 ? kDefaultZstdLevel + : -zstd_level_or_negative_for_auto; + PrxWindowBuildOutcome outcome = PrxWindowBuildOutcome::kBuilt; + RETURN_IF_ERROR( + build_prx_window_auto_from_flat(flat, freqs, auto_level, limits, sink, &outcome)); + if (outcome == PrxWindowBuildOutcome::kNeedsSplit) { + return Status::Error( + "prx: encoded payload exceeds writer window byte limit"); + } + return Status::OK(); +} + +Status build_prx_window_flat(std::span positions_flat, + std::span freqs, int zstd_level_or_negative_for_auto, + ByteSink* sink) { + return build_prx_window_flat(positions_flat, freqs, zstd_level_or_negative_for_auto, + kReaderPrxWindowLimits, sink); +} + +Status build_prx_window_flat(std::span positions_flat, + std::span freqs, int zstd_level_or_negative_for_auto, + const PrxWindowLimits& limits, ByteSink* sink) { + PrxWindowBuildOutcome outcome = PrxWindowBuildOutcome::kBuilt; + RETURN_IF_ERROR(try_build_prx_window_flat( + positions_flat, freqs, zstd_level_or_negative_for_auto, limits, sink, &outcome)); + if (outcome == PrxWindowBuildOutcome::kNeedsSplit) { + return Status::Error( + "prx: window exceeds writer limits"); + } + return Status::OK(); +} + +Status try_build_prx_window_flat(std::span positions_flat, + std::span freqs, + int zstd_level_or_negative_for_auto, const PrxWindowLimits& limits, + ByteSink* sink, PrxWindowBuildOutcome* outcome) { + if (sink == nullptr || outcome == nullptr) { + return Status::Error("prx: null arg"); + } + RETURN_IF_ERROR(validate_prx_window_limits(limits)); + RETURN_IF_ERROR(check_flat_partition(positions_flat, freqs)); + if (freqs.size() > limits.max_docs || positions_flat.size() > limits.max_positions) { + RETURN_IF_ERROR(validate_flat_positions(positions_flat, freqs)); + for (uint32_t frequency : freqs) { + if (frequency > limits.max_positions) { + return Status::Error( + "prx: one document exceeds the writer window position limit"); + } + } + if (freqs.size() <= 1) { + return Status::Error( + "prx: one document exceeds the writer window shape limit"); + } + RETURN_IF_ERROR(validate_single_doc_byte_limits(positions_flat, freqs, + zstd_level_or_negative_for_auto < 0, + limits.max_uncomp_bytes)); + *outcome = PrxWindowBuildOutcome::kNeedsSplit; + return Status::OK(); + } + if (zstd_level_or_negative_for_auto >= 0) { + ByteSink plain; + RETURN_IF_ERROR(encode_payload_flat(positions_flat, freqs, &plain)); + if (plain.size() > limits.max_uncomp_bytes) { + if (freqs.size() <= 1) { + return Status::Error( + "prx: one document exceeds the writer window byte limit"); + } + RETURN_IF_ERROR(validate_single_doc_byte_limits(positions_flat, freqs, false, + limits.max_uncomp_bytes)); + *outcome = PrxWindowBuildOutcome::kNeedsSplit; + return Status::OK(); + } + Slice plain_view = plain.view(); + if (!should_compress(zstd_level_or_negative_for_auto, plain_view.size())) { + write_raw(plain_view, sink); + *outcome = PrxWindowBuildOutcome::kBuilt; + return Status::OK(); + } + RETURN_IF_ERROR(write_zstd(plain_view, zstd_level_or_negative_for_auto, sink)); + *outcome = PrxWindowBuildOutcome::kBuilt; + return Status::OK(); + } + // Auto mode: shared path with a direct singleton RAW fast path, then PFOR, + // with raw plaintext materialized only for zstd or a tightened-limit fallback. + // G16-h: level < -1 is auto mode at zstd level |level| (-1 stays the default). + const int auto_level = zstd_level_or_negative_for_auto == -1 ? kDefaultZstdLevel + : -zstd_level_or_negative_for_auto; + return build_prx_window_auto_from_flat(positions_flat, freqs, auto_level, limits, sink, + outcome); +} + +Status read_prx_window(ByteSource* source, std::vector>* per_doc_positions) { + if (source == nullptr || per_doc_positions == nullptr) { + return Status::Error("prx: null arg"); + } + FramedPrxWindow frame; + RETURN_IF_ERROR(read_framed(source, &frame)); + if (frame.codec == static_cast(PrxCodec::kPfor)) { + return decode_pfor_payload(frame.payload, per_doc_positions); + } + if (frame.codec == static_cast(PrxCodec::kRaw)) { + return decode_payload(frame.payload, per_doc_positions); + } + std::vector plain; + RETURN_IF_ERROR(zstd_decompress(frame.payload, frame.uncomp_len, &plain)); + return decode_payload(Slice(plain), per_doc_positions); +} + +Status read_prx_window_csr(ByteSource* source, std::vector* pos_flat, + std::vector* pos_off) { + return read_prx_window_csr_impl(source, {}, true, true, pos_flat, pos_off, nullptr); +} + +Status read_prx_window_csr(ByteSource* source, std::vector* pos_flat, + std::vector* pos_off, PrxDecodeContext* context) { + return read_prx_window_csr_impl(source, {}, true, true, pos_flat, pos_off, context); +} + +Status read_prx_window_csr_for_selection(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off, + PrxDecodeContext* context) { + return read_prx_window_csr_impl(source, doc_ordinals, true, false, pos_flat, pos_off, context); +} + +Status read_prx_window_csr_selective(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off) { + return read_prx_window_csr_impl(source, doc_ordinals, false, false, pos_flat, pos_off, nullptr); +} + +Status read_prx_window_csr_selective(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off, PrxDecodeContext* context) { + return read_prx_window_csr_impl(source, doc_ordinals, false, false, pos_flat, pos_off, context); +} + +} // namespace doris::snii::format + +namespace doris::snii::format::testing { +namespace { +std::atomic& prx_raw_build_atomic() { + static std::atomic counter {0}; + return counter; +} + +#ifdef BE_TEST +std::atomic& prx_clock_read_atomic() { + static std::atomic counter {0}; + return counter; +} + +std::atomic& prx_delta_materialization_atomic() { + static std::atomic counter {0}; + return counter; +} +#endif +} // namespace + +uint64_t prx_raw_build_count() { + return prx_raw_build_atomic().load(std::memory_order_relaxed); +} + +void reset_prx_raw_build_count() { + prx_raw_build_atomic().store(0, std::memory_order_relaxed); +} + +void note_prx_raw_build() { + prx_raw_build_atomic().fetch_add(1, std::memory_order_relaxed); +} + +#ifdef BE_TEST +uint64_t prx_clock_read_count() { + return prx_clock_read_atomic().load(std::memory_order_relaxed); +} + +void reset_prx_clock_read_count() { + prx_clock_read_atomic().store(0, std::memory_order_relaxed); +} + +void note_prx_clock_read() { + prx_clock_read_atomic().fetch_add(1, std::memory_order_relaxed); +} + +uint64_t prx_delta_materialization_count() { + return prx_delta_materialization_atomic().load(std::memory_order_relaxed); +} + +void reset_prx_delta_materialization_count() { + prx_delta_materialization_atomic().store(0, std::memory_order_relaxed); +} + +void note_prx_delta_materialization() { + prx_delta_materialization_atomic().fetch_add(1, std::memory_order_relaxed); +} + +uint8_t select_auto_prx_codec_for_test(size_t pfor_payload_size, size_t plain_payload_size, + size_t compressed_payload_size, uint32_t max_uncomp_bytes) { + const AutoPrxCodecChoice choice = + select_auto_prx_codec(pfor_payload_size, plain_payload_size, compressed_payload_size, + plain_payload_size >= kAutoZstdMinBytes, max_uncomp_bytes); + DCHECK(choice.readable); + return static_cast(choice.codec); +} +#endif + +} // namespace doris::snii::format::testing diff --git a/be/src/storage/index/snii/format/prx_pod.h b/be/src/storage/index/snii/format/prx_pod.h new file mode 100644 index 00000000000000..fd96136d3ffe9f --- /dev/null +++ b/be/src/storage/index/snii/format/prx_pod.h @@ -0,0 +1,189 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/prx_decode_stats.h" + +// .prx position window (PrxPod): stores term position information for several +// docs within one window. +// +// Single-window on-disk byte layout (see docs/design SNII "prx design"): +// u8 codec # PrxCodec: 0=raw / 1=zstd / 2=pfor (bit7 cont-reserved) +// VInt uncomp_len # payload length (raw/pfor: on-disk payload bytes; zstd: +// plaintext) VInt comp_len # present only when codec==zstd u32 crc32c # +// covers header (codec..comp_len) + payload bytes payload # raw: varint +// plaintext; zstd: compressed; pfor: bit-packed +// +// raw/zstd plaintext payload (self-describing per-doc boundaries): +// VInt doc_count +// per doc: VInt pos_count, followed by pos_count position deltas (VInt) +// positions within a doc are ascending, stored as deltas (first absolute). +// +// pfor payload (one auto-build candidate; no entropy coding): +// VInt doc_count +// VInt total_pos # sum of all pos_counts +// PFOR_runs(pos_counts) # doc_count values +// PFOR_runs(position_deltas) # total_pos deltas, kFrqBaseUnit per run, +// # flat doc order (first per doc +// absolute) +// +// Multi-byte fixed-length fields are little-endian; variable-length integers +// reuse snii/encoding/varint. crc32c checksum at window tail detects +// corruption. +namespace doris::snii::format { + +struct PrxWindowLimits { + uint32_t max_docs; + uint32_t max_positions; + uint32_t max_uncomp_bytes; +}; + +enum class PrxWindowBuildOutcome : uint8_t { + kBuilt = 0, + kNeedsSplit = 1, +}; + +inline constexpr PrxWindowLimits kReaderPrxWindowLimits { + .max_docs = 1U << 24, + .max_positions = 1U << 26, + .max_uncomp_bytes = 256U * 1024 * 1024, +}; + +// Writer policies may tighten, but never exceed, the reader's allocation +// guards. Keeping this validation in the format layer makes it impossible for +// a caller to configure a writer that emits frames the current reader rejects. +Status validate_prx_window_limits(const PrxWindowLimits& limits); + +// Build a .prx window and append it to sink. +// per_doc_positions[d] is the position list for the d-th doc within this +// window; must be ascending (duplicates allowed). +// zstd_level_or_negative_for_auto: +// <0 → auto: use direct RAW for a single-doc/freq=1 window; otherwise use +// PFOR below the compression threshold and choose the smallest +// PFOR/ZSTD/RAW frame once raw plaintext is already required (default +// zstd level). +// 0 → force raw varint payload. +// >0 → force ZSTD with the given level. +// Non-ascending positions within a doc return InvalidArgument. +Status build_prx_window(std::span> per_doc_positions, + int zstd_level_or_negative_for_auto, ByteSink* sink); +Status build_prx_window(std::span> per_doc_positions, + int zstd_level_or_negative_for_auto, const PrxWindowLimits& limits, + ByteSink* sink); + +// Vector convenience overload (forwards a span view over the window's per-doc +// lists; the writer can pass a slice of its flat positions WITHOUT deep-copying +// the inner vectors into a fresh std::vector> per +// window). +inline Status build_prx_window(const std::vector>& per_doc_positions, + int zstd_level_or_negative_for_auto, ByteSink* sink) { + return build_prx_window(std::span>(per_doc_positions), + zstd_level_or_negative_for_auto, sink); +} + +// FLAT-positions builder: byte-identical output to build_prx_window above, but +// reads the window's positions from a single flat span partitioned per-doc by +// `freqs` (doc d owns the next freqs[d] entries; freqs.size() == doc count and +// sum(freqs) == positions_flat.size()). Lets the writer pass a subspan of the +// term's flat positions/freqs with NO vector-of-vectors materialization. +Status build_prx_window_flat(std::span positions_flat, + std::span freqs, int zstd_level_or_negative_for_auto, + ByteSink* sink); +Status build_prx_window_flat(std::span positions_flat, + std::span freqs, int zstd_level_or_negative_for_auto, + const PrxWindowLimits& limits, ByteSink* sink); + +// Writer-facing exact attempt. kNeedsSplit is returned only when the current +// multi-document window exceeds a shape or encoded-byte limit and every +// standalone document is representable, so retrying at document boundaries is +// guaranteed to resolve the limit. An unsplittable document, malformed input +// and codec failures remain errors. The sink is unchanged unless the outcome is +// kBuilt. +Status try_build_prx_window_flat(std::span positions_flat, + std::span freqs, + int zstd_level_or_negative_for_auto, const PrxWindowLimits& limits, + ByteSink* sink, PrxWindowBuildOutcome* outcome); + +// Read and verify a .prx window from source, reconstructing the per-doc +// position list. CRC mismatch / invalid codec / truncation / decompression +// failure all return a non-OK Status. +Status read_prx_window(ByteSource* source, std::vector>* per_doc_positions); + +// CSR variant of read_prx_window: decodes ALL docs' positions into one flat +// buffer `pos_flat` with per-doc offsets `pos_off` (size doc_count+1, +// pos_off[0]==0), so doc d's positions are pos_flat[pos_off[d] .. +// pos_off[d+1]). Avoids the per-doc std::vector allocation of read_prx_window +// -- both output vectors are flat uint32 buffers whose capacity a caller can +// retain (clear()) across windows/queries. +Status read_prx_window_csr(ByteSource* source, std::vector* pos_flat, + std::vector* pos_off); +Status read_prx_window_csr(ByteSource* source, std::vector* pos_flat, + std::vector* pos_off, PrxDecodeContext* context); + +// Decode every physical document while reporting the supplied ordinals as the +// logical selection. Used by the density fallback after docid narrowing. +Status read_prx_window_csr_for_selection(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off, PrxDecodeContext* context); + +// Selective CSR variant: decodes positions only for the requested local doc +// ordinals within this PRX window. `doc_ordinals` must be strictly ascending. +// The output uses the same CSR shape, but has doc_ordinals.size()+1 offsets. +Status read_prx_window_csr_selective(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off); +Status read_prx_window_csr_selective(ByteSource* source, std::span doc_ordinals, + std::vector* pos_flat, + std::vector* pos_off, PrxDecodeContext* context); + +} // namespace doris::snii::format + +// Test-only instrumentation seam. prx_raw_build_count() returns a process-global +// count of zstd-candidate plaintext payloads materialized by the auto-mode (.prx) +// window builders. Sub-threshold RAW fallback after an unreadable PFOR is not a +// zstd candidate and is intentionally not counted. In ordinary auto mode the +// small windows that dominate a Zipfian corpus skip the throwaway plaintext and +// second delta walk entirely, so tests assert the count is 0 for sub-threshold +// PFOR windows and exactly 1 per zstd candidate. +// Counters use relaxed atomics and are reset between tests; the writer's segment +// build is single-threaded, so the atomic adds introduce no data race. +namespace doris::snii::format::testing { + +uint64_t prx_raw_build_count(); +void reset_prx_raw_build_count(); +void note_prx_raw_build(); +#ifdef BE_TEST +uint64_t prx_clock_read_count(); +void reset_prx_clock_read_count(); +void note_prx_clock_read(); +uint64_t prx_delta_materialization_count(); +void reset_prx_delta_materialization_count(); +void note_prx_delta_materialization(); +uint8_t select_auto_prx_codec_for_test(size_t pfor_payload_size, size_t plain_payload_size, + size_t compressed_payload_size, uint32_t max_uncomp_bytes); +#endif + +} // namespace doris::snii::format::testing diff --git a/be/src/storage/index/snii/format/sampled_term_index.cpp b/be/src/storage/index/snii/format/sampled_term_index.cpp new file mode 100644 index 00000000000000..e60832132a1a51 --- /dev/null +++ b/be/src/storage/index/snii/format/sampled_term_index.cpp @@ -0,0 +1,189 @@ +// 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 "storage/index/snii/format/sampled_term_index.h" + +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" + +namespace doris::snii::format { + +namespace { + +// Longest common prefix length of term and prev (front coding primitive, consistent with dict_entry). +uint32_t common_prefix_len(std::string_view term, std::string_view prev) { + uint32_t n = 0; + const uint32_t lim = static_cast(std::min(term.size(), prev.size())); + while (n < lim && term[n] == prev[n]) ++n; + return n; +} + +// Write a front-coded term key (prefix_len + suffix_len + suffix). +void write_term_key(std::string_view term, std::string_view prev, ByteSink* sink) { + const uint32_t prefix = common_prefix_len(term, prev); + const std::string_view suffix = term.substr(prefix); + sink->put_varint32(prefix); + sink->put_varint32(static_cast(suffix.size())); + sink->put_bytes(Slice(suffix)); +} + +// Read a front-coded term key and reconstruct it into out from prev + suffix. +Status read_term_key(ByteSource* src, std::string_view prev, std::string* out) { + uint32_t prefix = 0; + uint32_t suffix_len = 0; + RETURN_IF_ERROR(src->get_varint32(&prefix)); + RETURN_IF_ERROR(src->get_varint32(&suffix_len)); + if (prefix > prev.size()) { + return Status::Error( + "sampled_term_index: prefix_len exceeds prev_term length"); + } + Slice suffix; + RETURN_IF_ERROR(src->get_bytes(suffix_len, &suffix)); + out->assign(prev.substr(0, prefix)); + out->append(reinterpret_cast(suffix.data()), suffix.size()); + return Status::OK(); +} + +} // namespace + +void SampledTermIndexBuilder::add_block_first_term(std::string_view first_term) { + first_terms_.emplace_back(first_term); +} + +void SampledTermIndexBuilder::finish(ByteSink* sink) { + ByteSink payload; + payload.put_varint32(static_cast(first_terms_.size())); + // min_term / max_term are written only when non-empty (== first/last sample_term). + if (!first_terms_.empty()) { + write_term_key(first_terms_.front(), std::string_view {}, &payload); + write_term_key(first_terms_.back(), std::string_view {}, &payload); + std::string_view prev {}; + for (const auto& t : first_terms_) { + write_term_key(t, prev, &payload); + prev = t; + } + } + SectionFramer::write(*sink, static_cast(SectionType::kSampledTermIndex), + payload.view()); +} + +namespace { + +// Parse n_blocks, min/max (not used directly; consumed for checksum alignment), and all sample_terms from payload. +Status parse_payload(Slice payload, std::vector* terms) { + ByteSource src(payload); + uint32_t n_blocks = 0; + RETURN_IF_ERROR(src.get_varint32(&n_blocks)); + if (n_blocks == 0) { + if (!src.eof()) { + return Status::Error( + "sampled_term_index: empty index contains trailing bytes"); + } + terms->clear(); + return Status::OK(); + } + + // min_term / max_term (do not drive binary search directly; must be consumed to verify structural alignment). + std::string min_term; + std::string max_term; + RETURN_IF_ERROR(read_term_key(&src, std::string_view {}, &min_term)); + RETURN_IF_ERROR(read_term_key(&src, std::string_view {}, &max_term)); + + std::vector out; + out.reserve(n_blocks); + std::string prev; + for (uint32_t i = 0; i < n_blocks; ++i) { + std::string term; + RETURN_IF_ERROR(read_term_key(&src, prev, &term)); + prev = term; + out.push_back(std::move(term)); + } + if (!src.eof()) { + return Status::Error( + "sampled_term_index: payload contains trailing bytes"); + } + if (out.front() != min_term || out.back() != max_term) { + return Status::Error( + "sampled_term_index: min/max inconsistent with sample_terms"); + } + *terms = std::move(out); + return Status::OK(); +} + +} // namespace + +Status SampledTermIndexReader::open(Slice section, SampledTermIndexReader* out) { + if (out == nullptr) { + return Status::Error("sampled_term_index: out is null"); + } + ByteSource src(section); + FramedSection sec; + RETURN_IF_ERROR(SectionFramer::read(src, &sec)); + if (!src.eof()) { + return Status::Error( + "sampled_term_index: trailing framed section bytes"); + } + if (sec.type != static_cast(SectionType::kSampledTermIndex)) { + return Status::Error( + "sampled_term_index: not a kSampledTermIndex section"); + } + *out = SampledTermIndexReader {}; + return parse_payload(sec.payload, &out->sample_terms_); +} + +size_t SampledTermIndexReader::heap_bytes() const { + size_t bytes = sample_terms_.capacity() * sizeof(std::string); + for (const auto& term : sample_terms_) { + bytes += std_string_heap_bytes(term); + } + return bytes; +} + +Status SampledTermIndexReader::locate(std::string_view target, bool* maybe_present, + uint32_t* block_ordinal) const { + if (maybe_present == nullptr || block_ordinal == nullptr) { + return Status::Error( + "sampled_term_index: output pointer is null"); + } + *maybe_present = false; + *block_ordinal = 0; + if (sample_terms_.empty()) { + return Status::OK(); // empty index: always out of range. + } + // target < min_term (first block's first term) -> before the first block, so it + // cannot exist in any block. NOTE: a target GREATER than the last sample term is + // NOT out of range -- sample_terms_ holds each block's FIRST term, so the LAST + // block can contain terms greater than its first term. Such a target routes to + // the last block (upper_bound -> end()), where find_term confirms presence. + if (target < std::string_view(sample_terms_.front())) { + return Status::OK(); + } + // Last sample_term <= target: step back one position after upper_bound. For a + // target past every sample term, upper_bound returns end() and idx = n-1 (the + // last block), which is correct. + auto it = std::upper_bound( + sample_terms_.begin(), sample_terms_.end(), target, + [](std::string_view t, const std::string& s) { return t < std::string_view(s); }); + const auto idx = (it - sample_terms_.begin()) - 1; // it > begin (< min excluded). + *maybe_present = true; + *block_ordinal = static_cast(idx); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/sampled_term_index.h b/be/src/storage/index/snii/format/sampled_term_index.h new file mode 100644 index 00000000000000..f2d63af88334e8 --- /dev/null +++ b/be/src/storage/index/snii/format/sampled_term_index.h @@ -0,0 +1,103 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/format_constants.h" + +// SampledTermIndex -- resident metadata for locating a query term to a candidate DICT block. +// +// Sampling granularity is per DICT block (not a fixed term count): each time the writer produces a DICT block, +// it writes the block's first_term into this index. Size grows proportionally to block count. At read time it is +// loaded into the searcher cache together with SniiLogicalIndexReader. See design spec "Sampled Term Index". +// +// On-disk layout (framed by SectionFramer, uniform type+len+crc32c): +// [u8 type=kSampledTermIndex][varint64 payload_len][payload][fixed32 crc32c] +// payload = +// n_blocks varint32 +// min_term len(varint32) + bytes # == sample_terms[0], omitted when n_blocks=0 +// max_term len(varint32) + bytes # == sample_terms[n-1], omitted when n_blocks=0 +// sample_terms[n_blocks]: # first_term of each block, in ascending order +// prefix_len varint32 # shared prefix length with the previous sample_term +// suffix_len varint32 +// suffix u8[suffix_len] +// +// Term bytes are compared as unsigned byte order (UTF-8 friendly, binary-safe). Front coding reuses +// the same prefix/suffix primitives as DictEntry; do not reimplement. +namespace doris::snii::format { + +// SSO-aware heap-byte accounting for a std::string. libstdc++ keeps up to 15 +// chars inline (SSO), so only capacity() > 15 implies a separate heap buffer of +// capacity()+1 bytes (the +1 is the NUL terminator); an SSO string owns no heap +// and contributes 0. Shared by the resident format readers' heap_bytes() charge +// helpers, which back LogicalIndexReader::memory_usage() (the searcher-cache +// charge). NOTE: the threshold 15 is libstdc++-specific; a different standard +// library needs a different SSO bound here. +inline size_t std_string_heap_bytes(const std::string& s) { + return s.capacity() > 15 ? s.capacity() + 1 : 0; +} + +// Builder: appends the first_term of each DICT block in block ordinal order (must be strictly ascending), +// and serializes the entire set into a single kSampledTermIndex framed section on finish. +class SampledTermIndexBuilder { +public: + // Appends the first_term of the next DICT block. Call order determines block ordinal order. + void add_block_first_term(std::string_view first_term); + + // Serializes and appends to sink. An empty collection (no blocks) is valid; n_blocks=0. + void finish(ByteSink* sink); + +private: + std::vector first_terms_; +}; + +// Reader: verifies the checksum and materializes all sample_terms on open; subsequent locate calls are pure in-memory binary search. +class SampledTermIndexReader { +public: + SampledTermIndexReader() = default; + + // Parses a kSampledTermIndex framed section. + // CRC mismatch / truncation / field overrun → kCorruption; type != kSampledTermIndex → kInvalidArgument. + static Status open(Slice section, SampledTermIndexReader* out); + + // Binary-search locate: returns the block ordinal of the last sample_term <= target. + // target < min_term or target > max_term (including empty index) → *maybe_present=false (out of range, term is definitely absent). + // Otherwise *maybe_present=true and *block_ordinal is the ordinal of the matching block. + Status locate(std::string_view target, bool* maybe_present, uint32_t* block_ordinal) const; + + uint32_t n_blocks() const { return static_cast(sample_terms_.size()); } + + // Resident heap held beyond sizeof(*this): the sample_terms_ vector buffer + // plus each non-SSO term's heap allocation. Summed into + // LogicalIndexReader::memory_usage() so the searcher-cache charge reflects the + // decoded sampled index (previously omitted -> under-charge -> over-commit). + size_t heap_bytes() const; + +private: + std::vector sample_terms_; +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/stats_block.h b/be/src/storage/index/snii/format/stats_block.h new file mode 100644 index 00000000000000..d1b07283ae437b --- /dev/null +++ b/be/src/storage/index/snii/format/stats_block.h @@ -0,0 +1,34 @@ +// 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. + +#pragma once + +#include + +namespace doris::snii::format { + +// Runtime counting statistics used for query planning and BM25. Core protobuf +// metadata owns their on-disk representation. +struct StatsBlock { + uint64_t doc_count = 0; // total doc count at segment level (including unindexed/NULL) + uint64_t indexed_doc_count = 0; // number of docs actually indexed (denominator for avgdl) + uint64_t term_count = 0; // number of unique terms in this index + uint64_t sum_total_term_freq = 0; // total token count across all indexed docs + uint64_t null_count = 0; // number of NULL / not-indexed docs +}; + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/tail_pointer.cpp b/be/src/storage/index/snii/format/tail_pointer.cpp new file mode 100644 index 00000000000000..69bf45c376072d --- /dev/null +++ b/be/src/storage/index/snii/format/tail_pointer.cpp @@ -0,0 +1,122 @@ +// 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 "storage/index/snii/format/tail_pointer.h" + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { + +namespace { + +// Byte widths of every fixed field, used to derive the constant on-disk size: +// u32 magic + u16 version + 2*u64 + u32 directory crc + u8 size + u32 tail crc. +constexpr size_t kMagicBytes = 4; +constexpr size_t kVersionBytes = 2; +constexpr size_t kU64Bytes = 8; +constexpr size_t kU32Bytes = 4; +constexpr size_t kSizeByteBytes = 1; + +constexpr size_t kFixedSize = + kMagicBytes + kVersionBytes + 2 * kU64Bytes + kU32Bytes + kSizeByteBytes + kU32Bytes; +// tail_checksum is the trailing u32 and covers every byte before it. +constexpr size_t kChecksumCoverage = kFixedSize - kU32Bytes; + +// Serializes the checksum-covered region in fixed field order into covered. +void serialize_covered(const TailPointer& tp, ByteSink* covered) { + covered->put_fixed32(kTailMagic); + covered->put_fixed16(kFormatVersion); + covered->put_fixed64(tp.directory_offset); + covered->put_fixed64(tp.directory_length); + covered->put_fixed32(tp.directory_crc32c); + covered->put_u8(static_cast(kFixedSize)); +} + +} // namespace + +size_t tail_pointer_size() { + return kFixedSize; +} + +Status encode_tail_pointer(const TailPointer& tp, ByteSink* sink) { + if (sink == nullptr) { + return Status::Error("tail_pointer: null sink"); + } + ByteSink covered; + serialize_covered(tp, &covered); + DORIS_CHECK_EQ(covered.size(), kChecksumCoverage); + const uint32_t tail_checksum = crc32c(covered.view()); + sink->put_bytes(covered.view()); + sink->put_fixed32(tail_checksum); + return Status::OK(); +} + +Status decode_tail_pointer(Slice last_bytes, TailPointer* out) { + if (out == nullptr) { + return Status::Error("tail_pointer: null output"); + } + // Anti-DoS / framing: the tail pointer is a fixed-size footer, so reject any + // input that is not exactly the fixed size before touching its contents. + if (last_bytes.size() != kFixedSize) { + return Status::Error( + "tail_pointer: input is not the fixed size"); + } + const Slice covered = last_bytes.subslice(0, kChecksumCoverage); + DORIS_CHECK_EQ(covered.size(), kChecksumCoverage); + ByteSource checksum_source(last_bytes.subslice(kChecksumCoverage, kU32Bytes)); + uint32_t tail_checksum = 0; + RETURN_IF_ERROR(checksum_source.get_fixed32(&tail_checksum)); + DORIS_CHECK(checksum_source.eof()); + if (tail_checksum != crc32c(covered)) { + return Status::Error( + "tail_pointer: tail_checksum mismatch"); + } + + // Only interpret fields after authenticating the complete covered region. + ByteSource src(covered); + + uint32_t magic = 0; + RETURN_IF_ERROR(src.get_fixed32(&magic)); + if (magic != kTailMagic) { + return Status::Error( + "tail_pointer: bad magic"); + } + + uint16_t tail_format_version = 0; + RETURN_IF_ERROR(src.get_fixed16(&tail_format_version)); + if (tail_format_version != kFormatVersion) { + return Status::Error( + "tail_pointer: unsupported container format_version"); + } + RETURN_IF_ERROR(src.get_fixed64(&out->directory_offset)); + RETURN_IF_ERROR(src.get_fixed64(&out->directory_length)); + RETURN_IF_ERROR(src.get_fixed32(&out->directory_crc32c)); + + uint8_t on_disk_size = 0; + RETURN_IF_ERROR(src.get_u8(&on_disk_size)); + if (on_disk_size != kFixedSize) { + return Status::Error( + "tail_pointer: embedded size mismatch"); + } + + DORIS_CHECK(src.eof()); + return Status::OK(); +} + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/format/tail_pointer.h b/be/src/storage/index/snii/format/tail_pointer.h new file mode 100644 index 00000000000000..fe2eb6e0bd19fc --- /dev/null +++ b/be/src/storage/index/snii/format/tail_pointer.h @@ -0,0 +1,67 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +namespace doris::snii::format { + +// Fixed-size entry written at the very end of a segment's .idx file. It lets a +// reader locate the raw metadata directory with a single read of the trailing +// tail_pointer_size() bytes (see design spec "fixed tail pointer"). +// +// On-disk layout (all multi-byte fields little-endian, FIXED total size so the +// reader can read exactly the last tail_pointer_size() bytes): +// [u32 magic = kTailMagic] +// [u16 format_version = kFormatVersion] +// [u64 directory_offset] +// [u64 directory_length] +// [u32 directory_crc32c] +// [u8 tail_pointer_size] (== tail_pointer_size()) +// [u32 tail_checksum] (crc32c over all preceding tail-pointer bytes) +// +// The fixed layout deliberately does NOT use the SectionFramer (which is +// variable-length): a footer needs a constant trailing size the reader knows up +// front. +struct TailPointer { + uint64_t directory_offset = 0; + uint64_t directory_length = 0; + uint32_t directory_crc32c = 0; +}; + +// Constant on-disk size of the tail pointer, so the reader knows how many +// trailing bytes to read. +size_t tail_pointer_size(); + +// Appends the fixed-layout tail-pointer bytes (magic / version / fields / size / +// tail_checksum) to sink. Returns Internal if the encoded size would not fit the +// fixed-size contract (a programming error, never expected at runtime). +Status encode_tail_pointer(const TailPointer& tp, ByteSink* sink); + +// Parses the trailing tail-pointer bytes. last_bytes must be exactly +// tail_pointer_size() bytes long. Verifies magic and tail_checksum, then fills +// out with the parsed fields. Wrong magic / checksum mismatch / wrong length -> +// Corruption. +Status decode_tail_pointer(Slice last_bytes, TailPointer* out); + +} // namespace doris::snii::format diff --git a/be/src/storage/index/snii/io/batch_range_fetcher.cpp b/be/src/storage/index/snii/io/batch_range_fetcher.cpp new file mode 100644 index 00000000000000..762c01d1c78024 --- /dev/null +++ b/be/src/storage/index/snii/io/batch_range_fetcher.cpp @@ -0,0 +1,102 @@ +// 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 "storage/index/snii/io/batch_range_fetcher.h" + +#include +#include + +namespace doris::snii::io { +namespace { + +Status checked_end(uint64_t offset, uint64_t len, uint64_t* out) { + if (len > std::numeric_limits::max() - offset) { + return Status::Error( + "batch_range_fetcher: range end overflow"); + } + *out = offset + len; + return Status::OK(); +} + +Status checked_size(uint64_t len, size_t* out) { + if (len > static_cast(std::numeric_limits::max())) { + return Status::Error( + "batch_range_fetcher: physical range too large"); + } + *out = static_cast(len); + return Status::OK(); +} + +} // namespace + +BatchRangeFetcher::BatchRangeFetcher(FileReader* reader, uint64_t coalesce_gap) + : reader_(reader), coalesce_gap_(coalesce_gap) {} + +size_t BatchRangeFetcher::add(uint64_t offset, uint64_t len) { + reqs_.push_back(Req {offset, len}); + return reqs_.size() - 1; +} + +void BatchRangeFetcher::clear() { + reqs_.clear(); + phys_.clear(); +} + +Status BatchRangeFetcher::fetch() { + if (reader_ == nullptr) + return Status::Error( + "batch_range_fetcher: null reader"); + phys_.clear(); + if (reqs_.empty()) return Status::OK(); + + std::vector order(reqs_.size()); + for (size_t i = 0; i < order.size(); ++i) order[i] = i; + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return reqs_[a].offset < reqs_[b].offset; }); + + // Sweep in offset order, merging requests into physical segments. + std::vector segs; + uint64_t cur_start = 0; + uint64_t cur_end = 0; + for (size_t k = 0; k < order.size(); ++k) { + Req& r = reqs_[order[k]]; + uint64_t r_end = 0; + RETURN_IF_ERROR(checked_end(r.offset, r.len, &r_end)); + RETURN_IF_ERROR(checked_size(r.len, &r.len_size)); + const bool disjoint = r.offset > cur_end && r.offset - cur_end > coalesce_gap_; + if (segs.empty() || disjoint) { + segs.push_back(Range {r.offset, 0}); // length finalized below + cur_start = r.offset; + cur_end = r_end; + } else { + cur_end = std::max(cur_end, r_end); + } + r.phys_idx = segs.size() - 1; + RETURN_IF_ERROR(checked_size(r.offset - cur_start, &r.sub_offset)); + RETURN_IF_ERROR(checked_size(cur_end - cur_start, &segs.back().len)); + } + + return reader_->read_batch(segs, &phys_); +} + +Slice BatchRangeFetcher::get(size_t h) const { + const Req& r = reqs_[h]; + const std::vector& buf = phys_[r.phys_idx]; + return Slice(buf.data() + r.sub_offset, r.len_size); +} + +} // namespace doris::snii::io diff --git a/be/src/storage/index/snii/io/batch_range_fetcher.h b/be/src/storage/index/snii/io/batch_range_fetcher.h new file mode 100644 index 00000000000000..1ef41c2fdc75d0 --- /dev/null +++ b/be/src/storage/index/snii/io/batch_range_fetcher.h @@ -0,0 +1,70 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_reader.h" + +namespace doris::snii::io { + +// Collects the byte ranges a query plan needs, coalesces overlapping/adjacent +// ranges into physical reads, and fetches them in a single batch (one serial +// I/O round as the test-side MeteredFileReader counts it). Callers retrieve each requested range by +// the handle returned from add(). This is the SNII read path's batching layer: +// it front-loads range planning so reads are issued concurrently rather than +// cursor-by-cursor. +class BatchRangeFetcher { +public: + // coalesce_gap: requests separated by a gap <= this many bytes are merged into + // one physical read (reads a few extra bytes to save a request). 0 merges only + // overlapping/adjacent ranges. + explicit BatchRangeFetcher(FileReader* reader, uint64_t coalesce_gap = 0); + + // Registers a desired range; returns a handle usable with get() after fetch(). + size_t add(uint64_t offset, uint64_t len); + + // Coalesces and issues one batched read; fills internal buffers. + Status fetch(); + + // Bytes for handle h (valid only after a successful fetch(), until clear()). + Slice get(size_t h) const; + + size_t pending() const { return reqs_.size(); } + void clear(); + +private: + struct Req { + uint64_t offset; + uint64_t len; + size_t len_size = 0; // validated size_t length after successful fetch() + size_t phys_idx = 0; // index into phys_ after fetch + size_t sub_offset = 0; // byte offset of this req within its physical read + }; + + FileReader* reader_; + uint64_t coalesce_gap_; + std::vector reqs_; + std::vector> phys_; // physical read buffers after fetch +}; + +} // namespace doris::snii::io diff --git a/be/src/storage/index/snii/io/file_reader.h b/be/src/storage/index/snii/io/file_reader.h new file mode 100644 index 00000000000000..46608aea602197 --- /dev/null +++ b/be/src/storage/index/snii/io/file_reader.h @@ -0,0 +1,67 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/io_metrics.h" + +namespace doris::snii::io { + +// One logical read request (offset, length). +struct Range { + uint64_t offset = 0; + size_t len = 0; +}; + +// The single physical-read primitive (a BE-internal read_at). All higher layers +// route reads through this so I/O can be accounted and backed by local files or +// object storage interchangeably. +class FileReader { +public: + virtual ~FileReader() = default; + + // Reads exactly len bytes starting at offset into *out (which is resized to + // len). Reading past EOF is an error (Corruption/IoError). + virtual Status read_at(uint64_t offset, size_t len, std::vector* out) = 0; + + // Reads a batch of ranges that may be served concurrently. The default is a + // sequential loop; backends that model concurrency (the test-side + // MeteredFileReader) or perform real parallel fetches (object storage) + // override this. + virtual Status read_batch(const std::vector& ranges, + std::vector>* outs) { + outs->resize(ranges.size()); + for (size_t i = 0; i < ranges.size(); ++i) { + RETURN_IF_ERROR(read_at(ranges[i].offset, ranges[i].len, &(*outs)[i])); + } + return Status::OK(); + } + + // Total size of the underlying object in bytes. + virtual uint64_t size() const = 0; + + // Optional live metrics. Readers that do not account I/O return nullptr. + virtual const IoMetrics* io_metrics() const { return nullptr; } +}; + +} // namespace doris::snii::io diff --git a/be/src/storage/index/snii/io/file_writer.h b/be/src/storage/index/snii/io/file_writer.h new file mode 100644 index 00000000000000..c14c61d2beaefd --- /dev/null +++ b/be/src/storage/index/snii/io/file_writer.h @@ -0,0 +1,40 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" + +namespace doris::snii::io { + +// Append-only writer (no seek-back), so the format can be produced in a single +// streaming pass compatible with S3FileWriter / StreamSinkFileWriter / packed +// writer. All container bytes are written front-to-back; back-references are +// resolved by writing metadata last. +class FileWriter { +public: + virtual ~FileWriter() = default; + + virtual Status append(Slice data) = 0; + virtual Status finalize() = 0; + virtual uint64_t bytes_written() const = 0; +}; + +} // namespace doris::snii::io diff --git a/be/src/storage/index/snii/io/io_metrics.h b/be/src/storage/index/snii/io/io_metrics.h new file mode 100644 index 00000000000000..0e1c628ede0137 --- /dev/null +++ b/be/src/storage/index/snii/io/io_metrics.h @@ -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. + +#pragma once + +#include + +namespace doris::snii::io { + +// Object-storage access metrics collected at FileReader boundaries. +struct IoMetrics { + uint64_t read_at_calls = 0; // BE-internal logical read requests issued + uint64_t serial_rounds = 0; // dependent serial I/O rounds + uint64_t range_gets = 0; // remote range GETs after cache coalescing + uint64_t remote_bytes = 0; // bytes fetched from remote + uint64_t total_request_bytes = 0; // sum of requested lengths before cache +}; + +inline IoMetrics delta(const IoMetrics& after, const IoMetrics& before) { + IoMetrics out; + out.read_at_calls = after.read_at_calls - before.read_at_calls; + out.serial_rounds = after.serial_rounds - before.serial_rounds; + out.range_gets = after.range_gets - before.range_gets; + out.remote_bytes = after.remote_bytes - before.remote_bytes; + out.total_request_bytes = after.total_request_bytes - before.total_request_bytes; + return out; +} + +} // namespace doris::snii::io diff --git a/be/src/storage/index/snii/query/bm25_scorer.cpp b/be/src/storage/index/snii/query/bm25_scorer.cpp new file mode 100644 index 00000000000000..61609980d69d19 --- /dev/null +++ b/be/src/storage/index/snii/query/bm25_scorer.cpp @@ -0,0 +1,65 @@ +// 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 "storage/index/snii/query/bm25_scorer.h" + +#include +#include + +namespace doris::snii::query { + +double decode_norm(uint8_t encoded) { + return encoded == 0 ? 1.0 : static_cast(encoded); +} + +uint8_t encode_norm(uint64_t doc_length) { + const uint64_t clamped = std::clamp(doc_length, 1, 255); + return static_cast(clamped); +} + +ScorerContext ScorerContext::make(uint64_t n, uint64_t df) { + ScorerContext ctx; + ctx.df_ = df; + const double nn = static_cast(n); + const double dff = static_cast(df); + // idf = log(1 + (N - df + 0.5) / (df + 0.5)); always positive for df <= N. + ctx.idf_ = std::log(1.0 + (nn - dff + 0.5) / (dff + 0.5)); + return ctx; +} + +ScorerContext ScorerContext::from_idf(double idf) { + ScorerContext ctx; + ctx.idf_ = idf; + return ctx; +} + +double ScorerContext::score(uint32_t tf, uint8_t encoded_norm, double avgdl, + const Bm25Params& params) const { + const double dl = decode_norm(encoded_norm); + const double tff = static_cast(tf); + const double denom = tff + params.k1 * (1.0 - params.b + params.b * dl / avgdl); + return idf_ * (tff * (params.k1 + 1.0)) / denom; +} + +double ScorerContext::max_score(uint32_t max_freq, uint8_t min_norm, double avgdl, + const Bm25Params& params) const { + // The score grows monotonically with tf and shrinks with dl, so the per-window + // upper bound uses the window's largest tf and smallest dl (min encoded norm). + return score(max_freq, min_norm, avgdl, params); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/bm25_scorer.h b/be/src/storage/index/snii/query/bm25_scorer.h new file mode 100644 index 00000000000000..5670af25174107 --- /dev/null +++ b/be/src/storage/index/snii/query/bm25_scorer.h @@ -0,0 +1,85 @@ +// 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. + +#pragma once + +#include + +// Bm25Scorer -- classic Okapi BM25 relevance scoring over SNII native stats. +// +// Per query term, idf is precomputed once from the collection statistics: +// idf = log(1 + (N - df + 0.5) / (df + 0.5)) +// where N = indexed doc count and df = the term's document frequency. The +// per-document contribution of a term then is: +// score = idf * (tf * (k1 + 1)) / (tf + k1 * (1 - b + b * dl / avgdl)) +// where tf is the in-doc term frequency, dl the document length decoded from the +// 1-byte encoded norm, and avgdl the average document length. +// +// Norm encode/decode (DOCUMENTED CONTRACT): the writer stores doc length as a +// byte-quantized value floor-clamped to [1, 255]; decode is the identity map +// back to a double length. encode_norm(len) = clamp(len, 1, 255); +// decode_norm(b) = (b == 0 ? 1.0 : (double)b). This keeps short docs (len <= 255) +// exact and saturates longer docs at 255, matching the reference oracle. +namespace doris::snii::query { + +// BM25 free parameters. Defaults are the classic Lucene/Elasticsearch values. +struct Bm25Params { + double k1 = 1.2; + double b = 0.75; +}; + +// Decodes a 1-byte encoded norm into a document length. byte 0 maps to 1.0 to +// avoid a zero-length divisor; otherwise it is the byte value itself. +double decode_norm(uint8_t encoded); + +// Encodes a document length into a 1-byte norm (clamped to [1, 255]). Provided +// so writers and test oracles share one quantization. +uint8_t encode_norm(uint64_t doc_length); + +// Per-term scoring context: the precomputed idf and the term's df. Built once per +// query term, then reused for every candidate document of that term. +class ScorerContext { +public: + // Builds the context from collection size n (indexed doc count) and the term's + // document frequency df. avgdl and params are supplied per score call. + static ScorerContext make(uint64_t n, uint64_t df); + + // Builds a context from a collection-scoped IDF that was computed outside + // the segment reader. This keeps segment-local TF/norm decoding separate + // from scanner-collection N/DF aggregation. + static ScorerContext from_idf(double idf); + + double idf() const { return idf_; } + uint64_t df() const { return df_; } + + // Scores one document occurrence: tf is the in-doc term frequency, encoded_norm + // the doc's 1-byte length norm, avgdl the collection average length. + double score(uint32_t tf, uint8_t encoded_norm, double avgdl, const Bm25Params& params) const; + + // Upper bound on score() over any document, given a window's maximum tf and the + // shortest doc length in the window (smallest dl maximizes the score). Used by + // the WAND-style block-max pruner. max_freq is the window's max tf; min_norm is + // the smallest encoded norm (=> smallest dl => largest score). + double max_score(uint32_t max_freq, uint8_t min_norm, double avgdl, + const Bm25Params& params) const; + +private: + double idf_ = 0.0; + uint64_t df_ = 0; +}; + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/boolean_query.cpp b/be/src/storage/index/snii/query/boolean_query.cpp new file mode 100644 index 00000000000000..ab26dee1005efb --- /dev/null +++ b/be/src/storage/index/snii/query/boolean_query.cpp @@ -0,0 +1,124 @@ +// 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 "storage/index/snii/query/boolean_query.h" + +#include +#include +#include +#include + +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/reader/dict_block_cache.h" + +namespace doris::snii::query { + +namespace { + +std::vector unique_terms(const std::vector& terms) { + std::vector out; + out.reserve(terms.size()); + for (const std::string& term : terms) out.emplace_back(term); + std::sort(out.begin(), out.end()); + out.erase(std::unique(out.begin(), out.end()), out.end()); + return out; +} + +Status resolve_or_postings(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* postings) { + postings->clear(); + // Request-scoped (stack-local, single-threaded) cache: OR terms that fall in the + // same on-demand DICT block read + zstd-decode + CRC-verify that block once + // instead of once per term. The resolved DictEntry is copied out, so the cache + // (and any pin it holds) can die when this returns. The shared reader stays const + // and lock-free -- no lock is ever held across the decode/IO. + reader::DictBlockCache dict_cache; + for (std::string_view term : unique_terms(terms)) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term, &found, &entry, &frq_base, &prx_base, &dict_cache)); + if (!found) continue; + + postings->push_back({std::move(entry), frq_base, prx_base}); + } + return Status::OK(); +} + +} // namespace + +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids) { + if (docids == nullptr) + return Status::Error("boolean_or: null out"); + docids->clear(); + if (terms.empty()) return Status::OK(); + + std::vector postings; + RETURN_IF_ERROR(resolve_or_postings(idx, terms, &postings)); + return internal::build_docid_union(idx, postings, docids); +} + +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids, QueryProfile* profile) { + QueryProfileScope profile_scope(idx.reader(), profile); + return boolean_or(idx, terms, docids); +} + +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + DocIdSink* sink) { + if (sink == nullptr) + return Status::Error("boolean_or: null sink"); + if (terms.empty()) return Status::OK(); + + std::vector postings; + RETURN_IF_ERROR(resolve_or_postings(idx, terms, &postings)); + return internal::emit_docid_union(idx, postings, sink); +} + +Status boolean_and(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids) { + if (docids == nullptr) + return Status::Error("boolean_and: null out"); + docids->clear(); + if (terms.empty()) return Status::OK(); + + io::BatchRangeFetcher round1(idx.reader()); + std::vector plans; + bool all_present = false; + RETURN_IF_ERROR(internal::plan_terms(idx, terms, &round1, &plans, &all_present, + /*need_positions=*/false)); + if (!all_present) return Status::OK(); + if (round1.pending() > 0) RETURN_IF_ERROR(round1.fetch()); + RETURN_IF_ERROR(internal::open_preludes(round1, &plans, + /*need_positions=*/false)); + return internal::build_docid_only_conjunction(idx, round1, plans, docids); +} + +Status boolean_and(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids, QueryProfile* profile) { + QueryProfileScope profile_scope(idx.reader(), profile); + return boolean_and(idx, terms, docids); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/boolean_query.h b/be/src/storage/index/snii/query/boolean_query.h new file mode 100644 index 00000000000000..40cbb25644fdc1 --- /dev/null +++ b/be/src/storage/index/snii/query/boolean_query.h @@ -0,0 +1,50 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// boolean_or -- MATCH_ANY semantics: return the sorted docid set containing at +// least one query term. Empty terms or all-absent terms produce an empty +// result. Duplicate input terms are ignored semantically and do not duplicate +// output docids. +namespace doris::snii::query { + +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids); +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids, QueryProfile* profile); +Status boolean_or(const reader::LogicalIndexReader& idx, const std::vector& terms, + DocIdSink* sink); + +// boolean_and (MATCH all-terms): sorted docid set of docs containing EVERY +// term, no positional constraint. Valid on docs-only indexes. Empty terms or +// any absent term -> empty result. +Status boolean_and(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids); +Status boolean_and(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids, QueryProfile* profile); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/count_query.cpp b/be/src/storage/index/snii/query/count_query.cpp new file mode 100644 index 00000000000000..3cc77ff9605e88 --- /dev/null +++ b/be/src/storage/index/snii/query/count_query.cpp @@ -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. + +#include "storage/index/snii/query/count_query.h" + +#include +#include + +#include "roaring/roaring.hh" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/query/internal/query_test_counters.h" + +namespace doris::snii::query { + +using format::DictEntry; +using reader::LogicalIndexReader; + +Status count_only_term_df(const LogicalIndexReader& idx, std::string_view term, uint64_t* count) { + if (count == nullptr) { + return Status::Error("count_only_term_df: null out"); + } + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term, &found, &entry, &frq_base, &prx_base)); + *count = found ? entry.df : 0; + SNII_QUERY_COUNT(count_fastpath_hits); + return Status::OK(); +} + +Status fabricate_null_disjoint_count_bitmap(uint64_t count, const roaring::Roaring& nulls, + roaring::Roaring* out) { + if (out == nullptr) { + return Status::Error( + "fabricate_null_disjoint_count_bitmap: null out"); + } + roaring::Roaring result; + if (count > 0) { + // [0, count + |nulls|) holds at least `count` non-null ids: at most + // |nulls| of its members are null. count counts only non-null docs, so + // the window end never exceeds the segment doc count (row space). + const uint64_t window_end = count + nulls.cardinality(); + if (window_end > uint64_t(std::numeric_limits::max()) + 1) { + return Status::Error( + "fabricate_null_disjoint_count_bitmap: count {} + null count {} exceeds the " + "uint32 docid domain (corrupt df or null bitmap)", + count, nulls.cardinality()); + } + result.addRange(0, window_end); + result -= nulls; + uint32_t last_kept = 0; + // Keep exactly the first `count` survivors (select ranks are 0-based). + if (!result.select(static_cast(count - 1), &last_kept)) { + return Status::Error( + "fabricate_null_disjoint_count_bitmap: window [0, {}) holds fewer than {} " + "non-null ids (corrupt df or null bitmap)", + window_end, count); + } + result.removeRange(uint64_t(last_kept) + 1, window_end); + } + *out = std::move(result); + return Status::OK(); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/count_query.h b/be/src/storage/index/snii/query/count_query.h new file mode 100644 index 00000000000000..bae76b4d91eed8 --- /dev/null +++ b/be/src/storage/index/snii/query/count_query.h @@ -0,0 +1,69 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// Forward-declare the CRoaring C++ bitmap so this header stays free of the +// (large) roaring include; the concrete type is only needed in the .cpp. +namespace roaring { +class Roaring; +} // namespace roaring + +// count_query -- G02 single-term count-only fast path primitives. They answer +// "how many docs match" from a dict entry alone, without reading .frq bytes. +// Multi-term queries, prefix/regexp/wildcard expansion, and phrases execute the +// normal query path. Deletes and extra predicates are a caller responsibility; +// see SniiIndexReader::_try_count_only_fastpath and the SegmentIterator guards +// in count_on_index_fastpath.h. +namespace doris::snii::query { + +// df of `term` in this segment without decoding postings. An absent term is a +// deterministic answer too: *count = 0 (mirrors term_query's empty result). +// Increments the count_fastpath_hits test seam. +Status count_only_term_df(const reader::LogicalIndexReader& idx, std::string_view term, + uint64_t* count); + +// Builds the fabricated count bitmap for a segment WITH a null bitmap: exactly +// `count` row ids DISJOINT from `nulls` (the first `count` non-null row ids, +// all < count + |nulls|). Why disjoint: the MATCH machinery unconditionally +// subtracts the segment null bitmap from every index result +// (FunctionMatchBase -> InvertedIndexResultBitmap::mask_out_null). Real +// postings never contain null docs -- the writer adds NO tokens for a null doc +// (scalar add_nulls) and a NULL array row is stored as an empty range (zero +// tokens) -- so that subtraction is a no-op on true results and df already IS +// the exact match count regardless of nulls. A naive [0, df) range however MAY +// collide with null row ids and be shrunk by mask_out_null; picking the ids +// from the non-null space makes the subtraction provably a no-op, preserving +// cardinality == df end to end. +// +// The window bound is doc-count-free: count counts only non-null docs, so +// count + |nulls| <= segment doc count and [0, count + |nulls|) always holds +// >= count non-null ids; every fabricated id therefore stays inside the +// segment's [0, num_rows) row space. Errors (id space would exceed the uint32 +// docid domain, or the window unexpectedly holds fewer than `count` survivors) +// only occur on a corrupt index; callers treat them as "fall through to the +// decode path", never as a fabricated answer. +Status fabricate_null_disjoint_count_bitmap(uint64_t count, const roaring::Roaring& nulls, + roaring::Roaring* out); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/docid_conjunction.cpp b/be/src/storage/index/snii/query/docid_conjunction.cpp new file mode 100644 index 00000000000000..1e3b892b4abfac --- /dev/null +++ b/be/src/storage/index/snii/query/docid_conjunction.cpp @@ -0,0 +1,929 @@ +// 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 "storage/index/snii/query/internal/docid_conjunction.h" + +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/reader/windowed_posting.h" + +namespace doris::snii::query::internal { + +using format::DictEntry; +using format::DictEntryEnc; +using format::DictEntryKind; +using format::FrqPreludeReader; +using format::WindowMeta; +using reader::LogicalIndexReader; + +namespace { + +using CandidateIt = std::vector::const_iterator; + +constexpr uint32_t kBoundedSpanBitsetDocs = 16 * 1024; +constexpr size_t kBoundedSpanBitsetWords = kBoundedSpanBitsetDocs / 64; +constexpr size_t kBoundedSpanBitsetMinInput = 32; + +struct CandidateRange { + size_t begin = 0; + size_t end = 0; +}; + +Status slim_frq_docs_len(const DictEntry& entry, uint64_t win_len, uint64_t* out) { + if (entry.frq_docs_len > win_len) { + return Status::Error( + "docid_conjunction: slim frq_docs_len exceeds frq window"); + } + *out = entry.frq_docs_len > 0 ? entry.frq_docs_len : win_len; + return Status::OK(); +} + +Status add_u64(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (rhs > std::numeric_limits::max() - lhs) { + return Status::Error(message); + } + *out = lhs + rhs; + return Status::OK(); +} + +Status posting_abs_offset(const LogicalIndexReader& idx, uint64_t base, uint64_t delta, + const char* message, uint64_t* out) { + uint64_t with_base = 0; + RETURN_IF_ERROR(add_u64(idx.section_refs().posting_region.offset, base, message, &with_base)); + return add_u64(with_base, delta, message, out); +} + +Status configure_term_plan(const LogicalIndexReader& idx, bool need_positions, + io::BatchRangeFetcher* fetcher, TermPlan* p) { + p->df = p->entry.df; + p->pod_ref = (p->entry.kind == DictEntryKind::kPodRef); + p->windowed = p->pod_ref && p->entry.enc == DictEntryEnc::kWindowed; + if (p->windowed) { + uint64_t prelude_abs = 0; + RETURN_IF_ERROR(posting_abs_offset(idx, p->frq_base, p->entry.frq_off_delta, + "docid_conjunction: prelude offset overflow", + &prelude_abs)); + p->prelude_handle = fetcher->add(prelude_abs, p->entry.prelude_len); + } else if (p->pod_ref) { + uint64_t foff = 0; + uint64_t flen = 0; + uint64_t poff = 0; + uint64_t plen = 0; + RETURN_IF_ERROR(idx.resolve_frq_window(p->entry, p->frq_base, &foff, &flen)); + uint64_t frq_fetch = flen; + RETURN_IF_ERROR(slim_frq_docs_len(p->entry, flen, &frq_fetch)); + p->frq_handle = fetcher->add(foff, frq_fetch); + if (need_positions) { + RETURN_IF_ERROR(idx.resolve_prx_window(p->entry, p->prx_base, &poff, &plen)); + p->prx_handle = fetcher->add(poff, plen); + } + } + return Status::OK(); +} + +std::vector all_windows(const FrqPreludeReader& prelude) { + std::vector ws(prelude.n_windows()); + for (uint32_t i = 0; i < prelude.n_windows(); ++i) ws[i] = i; + return ws; +} + +std::vector ascending_df_order(const std::vector& plans) { + std::vector order(plans.size()); + for (size_t i = 0; i < plans.size(); ++i) order[i] = i; + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return plans[a].df < plans[b].df; }); + return order; +} + +Status first_docid_in_window(const WindowMeta& meta, uint32_t window_ordinal, uint32_t* first) { + if (window_ordinal == 0) { + *first = 0; + return Status::OK(); + } + if (meta.win_base >= std::numeric_limits::max()) { + return Status::Error( + "docid_conjunction: window base exceeds docid range"); + } + *first = static_cast(meta.win_base + 1); + if (*first > meta.last_docid) { + return Status::Error( + "docid_conjunction: invalid window docid range"); + } + return Status::OK(); +} + +Status is_dense_full_window(const WindowMeta& meta, uint32_t window_ordinal, bool* full) { + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(meta, window_ordinal, &first)); + const uint64_t width = static_cast(meta.last_docid) - first + 1; + *full = meta.doc_count == width; + return Status::OK(); +} + +Status append_docid_range(uint32_t first, uint32_t last, std::vector* out) { + if (last < first) { + return Status::Error( + "docid_conjunction: invalid dense docid range"); + } + const uint64_t count64 = static_cast(last) - first + 1; + if (count64 > static_cast(std::numeric_limits::max() - out->size())) { + return Status::Error( + "docid_conjunction: dense docid range too large"); + } + out->reserve(out->size() + static_cast(count64)); + uint32_t docid = first; + while (true) { + out->push_back(docid); + if (docid == last) break; + ++docid; + } + return Status::OK(); +} + +CandidateRange find_candidate_range(const std::vector& candidates, size_t* search_begin, + uint32_t first, uint32_t last) { + const auto from = candidates.begin() + *search_begin; + const auto begin = std::lower_bound(from, candidates.end(), first); + const auto end = std::upper_bound(begin, candidates.end(), last); + *search_begin = static_cast(end - candidates.begin()); + return {.begin = static_cast(begin - candidates.begin()), + .end = static_cast(end - candidates.begin())}; +} + +void append_candidate_range(CandidateIt begin, CandidateIt end, std::vector* out) { + out->insert(out->end(), begin, end); +} + +void append_new_chunk_docids_to_out(const DocidChunk& chunk, size_t chunk_docids_begin, + std::vector* out) { + out->insert(out->end(), chunk.docids.begin() + chunk_docids_begin, chunk.docids.end()); +} + +void clear_ordinals_if_all_term_docs_selected(const std::vector& term_docids, + DocidChunk* chunk) { + if (chunk->docids.size() == term_docids.size() && !chunk->docids.empty() && + chunk->docids.front() == term_docids.front() && + chunk->docids.back() == term_docids.back()) { + chunk->prx_doc_ordinals.clear(); + } +} + +bool append_term_docs_if_candidates_cover_span(CandidateIt begin, CandidateIt end, + const std::vector& term_docids, + std::vector* out, DocidChunk* chunk) { + const uint32_t first = term_docids.front(); + const uint32_t last = term_docids.back(); + const uint64_t width = static_cast(last) - first + 1; + const size_t candidate_count = static_cast(end - begin); + if (width > candidate_count) { + return false; + } + + const auto span_begin = *begin == first ? begin : std::lower_bound(begin, end, first); + if (span_begin == end || *span_begin != first) { + return false; + } + if (static_cast(end - span_begin) < width) { + return false; + } + + const auto span_last = span_begin + static_cast(width) - 1; + if (*span_last != last) { + return false; + } + + const size_t chunk_docids_begin = chunk->docids.size(); + chunk->docids.insert(chunk->docids.end(), term_docids.begin(), term_docids.end()); + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + return true; +} + +Status append_candidate_range_with_ordinals(CandidateIt begin, CandidateIt end, uint32_t first, + uint32_t last, std::vector* out, + DocidChunk* chunk) { + const size_t candidate_count = static_cast(end - begin); + chunk->docids.reserve(candidate_count); + const uint64_t width = static_cast(last) - first + 1; + if (width > std::numeric_limits::max()) { + return Status::Error( + "docid_conjunction: dense window exceeds doc count range"); + } + chunk->prx_doc_count = static_cast(width); + const bool full_dense_range = + candidate_count == width && begin != end && *begin == first && *(end - 1) == last; + const size_t chunk_docids_begin = chunk->docids.size(); + if (full_dense_range) { + chunk->docids.insert(chunk->docids.end(), begin, end); + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + return Status::OK(); + } + chunk->prx_doc_ordinals.reserve(chunk->prx_doc_ordinals.size() + candidate_count); + for (auto it = begin; it != end; ++it) { + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back(*it - first); + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + return Status::OK(); +} + +bool intersect_dense_term_span_with_ordinals(CandidateIt begin, CandidateIt end, + const std::vector& term_docids, + size_t candidate_count, std::vector* out, + DocidChunk* chunk) { + const uint32_t first = term_docids.front(); + const uint32_t last = term_docids.back(); + const uint64_t width = static_cast(last) - first + 1; + if (term_docids.size() > width) { + return false; + } + const uint64_t missing_count = width - term_docids.size(); + if (missing_count != 0 && + (missing_count * 8 > width || missing_count >= candidate_count || + missing_count > static_cast(std::numeric_limits::max()))) { + return false; + } + + const size_t chunk_docids_begin = chunk->docids.size(); + if (missing_count == 0) { + for (auto it = begin; it != end; ++it) { + if (*it < first) { + continue; + } + if (*it > last) { + break; + } + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back(*it - first); + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return true; + } + + std::vector missing; + missing.reserve(static_cast(missing_count)); + uint32_t expect = first; + for (uint32_t docid : term_docids) { + while (expect < docid) { + missing.push_back(expect); + ++expect; + } + if (docid < std::numeric_limits::max()) { + expect = docid + 1; + } + } + while (expect <= last) { + missing.push_back(expect); + if (expect == std::numeric_limits::max()) { + break; + } + ++expect; + } + + size_t miss = 0; + for (auto it = begin; it != end; ++it) { + if (*it < first) { + continue; + } + if (*it > last) { + break; + } + while (miss < missing.size() && missing[miss] < *it) { + ++miss; + } + if (miss < missing.size() && missing[miss] == *it) { + continue; + } + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back(static_cast(*it - first - miss)); + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return true; +} + +bool intersect_bounded_span_with_ordinals(CandidateIt begin, CandidateIt end, + const std::vector& term_docids, + size_t candidate_count, std::vector* out, + DocidChunk* chunk) { + if (candidate_count < kBoundedSpanBitsetMinInput || + term_docids.size() < kBoundedSpanBitsetMinInput) { + return false; + } + + const uint32_t first = std::min(*begin, term_docids.front()); + const uint32_t last = std::max(*(end - 1), term_docids.back()); + const uint64_t width = static_cast(last) - first + 1; + if (width > kBoundedSpanBitsetDocs || term_docids.size() > width) { + return false; + } + + const auto word_count = static_cast((width + 63) >> 6); + std::array bits; + std::fill_n(bits.begin(), word_count, 0); + for (uint32_t docid : term_docids) { + const uint32_t off = docid - first; + bits[off >> 6] |= 1ULL << (off & 63); + } + + std::array ordinal_base; + uint32_t ordinal = 0; + for (size_t word = 0; word < word_count; ++word) { + ordinal_base[word] = ordinal; + ordinal += static_cast(__builtin_popcountll(bits[word])); + } + + const size_t chunk_docids_begin = chunk->docids.size(); + for (auto it = begin; it != end; ++it) { + const uint32_t off = *it - first; + const size_t word = off >> 6; + const uint64_t mask = 1ULL << (off & 63); + if ((bits[word] & mask) == 0) { + continue; + } + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back( + ordinal_base[word] + + static_cast(__builtin_popcountll(bits[word] & (mask - 1)))); + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return true; +} + +size_t log2_ceil(size_t n) { + if (n <= 1) return 1; + --n; + size_t bits = 0; + while (n != 0) { + ++bits; + n >>= 1; + } + return bits; +} + +void intersect_window_candidate_range(CandidateIt begin, CandidateIt end, + const std::vector& term_docids, uint32_t first, + uint32_t last, std::vector* out) { + const size_t candidate_count = static_cast(end - begin); + if (candidate_count == 0 || term_docids.empty()) return; + + const uint64_t width = static_cast(last) - first + 1; + const uint64_t missing_count = term_docids.size() <= width ? width - term_docids.size() : width; + if (term_docids.size() <= width && missing_count != 0 && missing_count * 8 <= width && + missing_count < candidate_count) { + std::vector missing; + missing.reserve(static_cast(missing_count)); + uint32_t expect = first; + for (uint32_t docid : term_docids) { + while (expect < docid) { + missing.push_back(expect); + ++expect; + } + if (docid < std::numeric_limits::max()) expect = docid + 1; + } + while (expect <= last) { + missing.push_back(expect); + if (expect == std::numeric_limits::max()) break; + ++expect; + } + size_t miss = 0; + for (auto it = begin; it != end; ++it) { + while (miss < missing.size() && missing[miss] < *it) ++miss; + if (miss == missing.size() || missing[miss] != *it) out->push_back(*it); + } + return; + } + + const size_t probes_per_candidate = log2_ceil(term_docids.size()) + 1; + if (candidate_count < term_docids.size() / probes_per_candidate) { + for (auto it = begin; it != end; ++it) { + if (std::binary_search(term_docids.begin(), term_docids.end(), *it)) { + out->push_back(*it); + } + } + return; + } + std::set_intersection(begin, end, term_docids.begin(), term_docids.end(), + std::back_inserter(*out)); +} + +Status intersect_window_candidate_range_with_ordinals(CandidateIt begin, CandidateIt end, + const std::vector& term_docids, + std::vector* out, + DocidChunk* chunk) { + if (term_docids.size() > std::numeric_limits::max()) { + return Status::Error( + "docid_conjunction: prx doc count exceeds u32"); + } + chunk->prx_doc_count = static_cast(term_docids.size()); + if (begin == end || term_docids.empty()) return Status::OK(); + + const size_t candidate_count = static_cast(end - begin); + const size_t max_matches = std::min(candidate_count, term_docids.size()); + out->reserve(out->size() + max_matches); + chunk->docids.reserve(chunk->docids.size() + max_matches); + if (candidate_count == term_docids.size() && *begin == term_docids.front() && + *(end - 1) == term_docids.back() && std::equal(begin, end, term_docids.begin())) { + const size_t chunk_docids_begin = chunk->docids.size(); + chunk->docids.insert(chunk->docids.end(), begin, end); + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + return Status::OK(); + } + if (append_term_docs_if_candidates_cover_span(begin, end, term_docids, out, chunk)) { + return Status::OK(); + } + + chunk->prx_doc_ordinals.reserve(chunk->prx_doc_ordinals.size() + max_matches); + if (intersect_dense_term_span_with_ordinals(begin, end, term_docids, candidate_count, out, + chunk)) { + return Status::OK(); + } + if (intersect_bounded_span_with_ordinals(begin, end, term_docids, candidate_count, out, + chunk)) { + return Status::OK(); + } + + const size_t probes_per_candidate = log2_ceil(term_docids.size()) + 1; + if (candidate_count < term_docids.size() / probes_per_candidate) { + const size_t chunk_docids_begin = chunk->docids.size(); + size_t doc_index = 0; + for (auto it = begin; it != end; ++it) { + const auto found = + std::lower_bound(term_docids.begin() + doc_index, term_docids.end(), *it); + if (found == term_docids.end()) break; + doc_index = static_cast(found - term_docids.begin()); + if (*found != *it) continue; + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back(static_cast(doc_index)); + ++doc_index; + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return Status::OK(); + } + + const size_t probes_per_term_doc = log2_ceil(candidate_count) + 1; + if (term_docids.size() < candidate_count / probes_per_term_doc) { + const size_t chunk_docids_begin = chunk->docids.size(); + auto candidate_it = begin; + for (size_t doc_index = 0; doc_index < term_docids.size(); ++doc_index) { + const uint32_t docid = term_docids[doc_index]; + candidate_it = std::lower_bound(candidate_it, end, docid); + if (candidate_it == end) break; + if (*candidate_it != docid) continue; + chunk->docids.push_back(docid); + chunk->prx_doc_ordinals.push_back(static_cast(doc_index)); + ++candidate_it; + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return Status::OK(); + } + + const size_t chunk_docids_begin = chunk->docids.size(); + size_t doc_index = 0; + for (auto it = begin; it != end; ++it) { + while (doc_index < term_docids.size() && term_docids[doc_index] < *it) { + ++doc_index; + } + if (doc_index == term_docids.size()) break; + if (term_docids[doc_index] != *it) continue; + chunk->docids.push_back(*it); + chunk->prx_doc_ordinals.push_back(static_cast(doc_index)); + ++doc_index; + } + append_new_chunk_docids_to_out(*chunk, chunk_docids_begin, out); + clear_ordinals_if_all_term_docs_selected(term_docids, chunk); + return Status::OK(); +} + +bool should_scan_all_windows(const LogicalIndexReader& idx, const TermPlan& p, + size_t candidate_count) { + const size_t window_count = p.prelude.n_windows(); + if (candidate_count > window_count * 64) return true; + + const uint64_t doc_count = idx.stats().doc_count; + const bool near_full = doc_count != 0 && static_cast(p.df) * 10 >= doc_count * 9; + return near_full && candidate_count > window_count * 4; +} + +Status decode_flat_docids_only(const io::BatchRangeFetcher& round1, const TermPlan& p, + std::vector* docids) { + Slice dd; + if (p.pod_ref) { + dd = round1.get(p.frq_handle); + } else { + RETURN_IF_ERROR(inline_dd_region(p.entry, &dd)); + } + return format::decode_dd_region(dd, p.entry.dd_meta, /*win_base=*/0, docids); +} + +struct WindowWork { + uint32_t ordinal = 0; + WindowMeta meta; + CandidateRange candidates; + size_t handle = 0; + bool dense_full = false; +}; + +Status emit_dense_full_window_docids(const WindowWork& f, const std::vector* candidates, + std::vector& out, DocidSource* source) { + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(f.meta, f.ordinal, &first)); + if (source != nullptr) { + DocidChunk chunk; + chunk.windowed = true; + chunk.window = f.ordinal; + chunk.prx_doc_count = f.meta.doc_count; + if (candidates == nullptr) { + RETURN_IF_ERROR(append_docid_range(first, f.meta.last_docid, &chunk.docids)); + } else { + const auto begin = candidates->begin() + f.candidates.begin; + const auto end = candidates->begin() + f.candidates.end; + RETURN_IF_ERROR(append_candidate_range_with_ordinals(begin, end, first, + f.meta.last_docid, &out, &chunk)); + } + source->chunks.push_back(std::move(chunk)); + } + if (candidates == nullptr) { + RETURN_IF_ERROR(append_docid_range(first, f.meta.last_docid, &out)); + } else if (source == nullptr) { + append_candidate_range(candidates->begin() + f.candidates.begin, + candidates->begin() + f.candidates.end, &out); + } + return Status::OK(); +} + +Status emit_decoded_window_docids(const WindowWork& f, const io::BatchRangeFetcher& fetcher, + const std::vector* candidates, + std::vector& out, DocidSource* source, + std::vector& docs, std::vector& freqs, + std::vector>& positions) { + docs.clear(); + freqs.clear(); + positions.clear(); + RETURN_IF_ERROR(reader::decode_window_slices(f.meta, fetcher.get(f.handle), Slice(), Slice(), + /*want_positions=*/false, /*want_freq=*/false, + &docs, &freqs, &positions)); + if (source != nullptr) { + DocidChunk chunk; + chunk.windowed = true; + chunk.window = f.ordinal; + if (candidates == nullptr) { + chunk.docids = docs; + if (docs.size() > std::numeric_limits::max()) { + return Status::Error( + "docid_conjunction: prx doc count exceeds u32"); + } + chunk.prx_doc_count = static_cast(docs.size()); + source->chunks.push_back(std::move(chunk)); + } else { + const auto begin = candidates->begin() + f.candidates.begin; + const auto end = candidates->begin() + f.candidates.end; + RETURN_IF_ERROR( + intersect_window_candidate_range_with_ordinals(begin, end, docs, &out, &chunk)); + if (!chunk.docids.empty()) { + source->chunks.push_back(std::move(chunk)); + } + } + } + if (candidates == nullptr) { + out.insert(out.end(), docs.begin(), docs.end()); + return Status::OK(); + } + if (source != nullptr) { + return Status::OK(); + } + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(f.meta, f.ordinal, &first)); + intersect_window_candidate_range(candidates->begin() + f.candidates.begin, + candidates->begin() + f.candidates.end, docs, first, + f.meta.last_docid, &out); + return Status::OK(); +} + +Status collect_windowed_docids_only(const LogicalIndexReader& idx, const TermPlan& p, + const std::vector& windows, + const std::vector* candidates, + std::vector* out, DocidSource* source) { + io::BatchRangeFetcher fetcher(idx.reader(), reader::kSameTermCoalesceGap); + std::vector work; + work.reserve(windows.size()); + out->reserve(candidates == nullptr ? p.entry.df : candidates->size()); + size_t candidate_search_begin = 0; + for (uint32_t w : windows) { + WindowMeta meta; + RETURN_IF_ERROR(p.prelude.window(w, &meta)); + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(meta, w, &first)); + CandidateRange candidate_range; + if (candidates != nullptr) { + candidate_range = find_candidate_range(*candidates, &candidate_search_begin, first, + meta.last_docid); + if (candidate_range.begin == candidate_range.end) { + continue; + } + } + bool dense_full = false; + RETURN_IF_ERROR(is_dense_full_window(meta, w, &dense_full)); + if (dense_full) { + work.push_back(WindowWork { + .ordinal = w, .meta = meta, .candidates = candidate_range, .dense_full = true}); + continue; + } + + reader::WindowAbsRange range; + RETURN_IF_ERROR(reader::windowed_window_range( + idx, p.entry, p.frq_base, p.prx_base, p.prelude, w, + /*want_positions=*/false, /*want_freq=*/false, &range)); + WindowWork f; + f.ordinal = w; + f.meta = meta; + f.candidates = candidate_range; + f.handle = fetcher.add(range.dd_off, range.dd_len); + work.push_back(f); + } + if (fetcher.pending() > 0) { + RETURN_IF_ERROR(fetcher.fetch()); + } + + std::vector docs; + std::vector freqs; + std::vector> positions; + for (const WindowWork& f : work) { + if (f.dense_full) { + RETURN_IF_ERROR(emit_dense_full_window_docids(f, candidates, *out, source)); + continue; + } + RETURN_IF_ERROR(emit_decoded_window_docids(f, fetcher, candidates, *out, source, docs, + freqs, positions)); + } + return Status::OK(); +} + +Status collect_docids_only(const LogicalIndexReader& idx, const io::BatchRangeFetcher& round1, + const TermPlan& p, const std::vector* candidates, + std::vector* out, DocidSource* source) { + if (p.windowed) { + std::vector windows; + if (candidates == nullptr) { + windows = all_windows(p.prelude); + } else if (should_scan_all_windows(idx, p, candidates->size())) { + // Dense candidate sets cover most windows; for near-full terms this also + // avoids a thousands-to-millions probe covering-window cursor pass with no + // byte win. + windows = all_windows(p.prelude); + } else { + p.prelude.select_covering_windows(*candidates, &windows); + } + return collect_windowed_docids_only(idx, p, windows, candidates, out, source); + } + + std::vector term_docids; + RETURN_IF_ERROR(decode_flat_docids_only(round1, p, &term_docids)); + if (source != nullptr) { + DocidChunk chunk; + if (term_docids.size() > std::numeric_limits::max()) { + return Status::Error( + "docid_conjunction: prx doc count exceeds u32"); + } + chunk.prx_doc_count = static_cast(term_docids.size()); + if (candidates == nullptr) { + chunk.docids = term_docids; + } else if (!term_docids.empty()) { + const auto begin = std::ranges::lower_bound(*candidates, term_docids.front()); + const auto end = std::upper_bound(begin, candidates->end(), term_docids.back()); + RETURN_IF_ERROR(intersect_window_candidate_range_with_ordinals(begin, end, term_docids, + out, &chunk)); + } + if (candidates == nullptr || !chunk.docids.empty()) { + source->chunks.push_back(std::move(chunk)); + } + } + if (candidates == nullptr) { + *out = std::move(term_docids); + return Status::OK(); + } + if (source != nullptr) { + return Status::OK(); + } + *out = intersect_sorted(*candidates, term_docids); + return Status::OK(); +} + +Status run_docid_only_conjunction_impl(const LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + const std::vector* initial_candidates, + std::vector* candidates, + std::vector* sources) { + if (sources != nullptr) { + sources->assign(plans.size(), DocidSource {}); + } + candidates->clear(); + if (plans.empty()) { + // No terms: the result is the initial candidate set verbatim (or empty). + if (initial_candidates != nullptr) { + *candidates = *initial_candidates; + } + return Status::OK(); + } + if (initial_candidates != nullptr && initial_candidates->empty()) { + return Status::OK(); + } + const std::vector order = ascending_df_order(plans); + for (size_t k = 0; k < order.size(); ++k) { + const size_t ti = order[k]; + std::vector next; + DocidSource* source = sources == nullptr ? nullptr : &(*sources)[ti]; + // k == 0 intersects against the (const) initial_candidates DIRECTLY, so + // the whole set is never copied once per plan -- the previous code seeded + // *candidates = *initial_candidates before the loop, which for a single + // plan (e.g. one phrase-prefix tail verified against the leading-term + // expected docids) was an O(|initial|) copy per call with no benefit. + // k > 0 chains on the previous term's already-whittled result. + const std::vector* input_candidates = k == 0 ? initial_candidates : candidates; + RETURN_IF_ERROR( + collect_docids_only(idx, round1, plans[ti], input_candidates, &next, source)); + if (source != nullptr && k + 1 == order.size()) { + source->docids_are_final_candidates = true; + } + *candidates = std::move(next); + if (candidates->empty()) { + return Status::OK(); + } + } + return Status::OK(); +} + +} // namespace + +Status resolve_query_term(const LogicalIndexReader& idx, std::string_view term, + ResolvedQueryTerm* resolved, bool* found) { + *found = false; + RETURN_IF_ERROR( + idx.lookup(term, found, &resolved->entry, &resolved->frq_base, &resolved->prx_base)); + return Status::OK(); +} + +Status resolve_query_terms_batch(const LogicalIndexReader& idx, + const std::vector& terms, + std::vector* resolved, + std::vector* found) { + DCHECK(std::ranges::is_sorted(terms)); + DCHECK(std::adjacent_find(terms.begin(), terms.end()) == terms.end()); + resolved->assign(terms.size(), ResolvedQueryTerm {}); + found->assign(terms.size(), 0); + std::vector lookup_results; + RETURN_IF_ERROR(idx.lookup_batch(terms, &lookup_results)); + for (size_t i = 0; i < terms.size(); ++i) { + (*found)[i] = lookup_results[i].found; + if (lookup_results[i].found) { + (*resolved)[i].entry = std::move(lookup_results[i].entry); + (*resolved)[i].frq_base = lookup_results[i].frq_base; + (*resolved)[i].prx_base = lookup_results[i].prx_base; + } + } + return Status::OK(); +} + +Status plan_terms(const LogicalIndexReader& idx, const std::vector& terms, + io::BatchRangeFetcher* fetcher, std::vector* plans, bool* all_present, + bool need_positions) { + *all_present = true; + plans->resize(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + ResolvedQueryTerm resolved; + bool found = false; + RETURN_IF_ERROR(resolve_query_term(idx, terms[i], &resolved, &found)); + if (!found) { + *all_present = false; + return Status::OK(); + } + TermPlan& p = (*plans)[i]; + p.order = i; + p.entry = std::move(resolved.entry); + p.frq_base = resolved.frq_base; + p.prx_base = resolved.prx_base; + RETURN_IF_ERROR(configure_term_plan(idx, need_positions, fetcher, &p)); + } + return Status::OK(); +} + +Status plan_resolved_terms(const LogicalIndexReader& idx, + const std::vector& terms, + io::BatchRangeFetcher* fetcher, std::vector* plans, + bool need_positions) { + plans->resize(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + TermPlan& p = (*plans)[i]; + p.order = i; + p.entry = terms[i].entry; + SNII_QUERY_COUNT(resolved_term_entry_copies); + p.frq_base = terms[i].frq_base; + p.prx_base = terms[i].prx_base; + RETURN_IF_ERROR(configure_term_plan(idx, need_positions, fetcher, &p)); + } + return Status::OK(); +} + +Status plan_resolved_terms(const LogicalIndexReader& idx, std::vector&& terms, + io::BatchRangeFetcher* fetcher, std::vector* plans, + bool need_positions) { + plans->resize(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + TermPlan& p = (*plans)[i]; + p.order = i; +#ifdef BE_TEST + const bool has_frq_payload = !terms[i].entry.frq_bytes.empty(); + const bool has_prx_payload = !terms[i].entry.prx_bytes.empty(); + const uint8_t* const frq_payload = terms[i].entry.frq_bytes.data(); + const uint8_t* const prx_payload = terms[i].entry.prx_bytes.data(); +#endif + p.entry = std::move(terms[i].entry); + SNII_QUERY_COUNT(resolved_term_entry_moves); +#ifdef BE_TEST + SNII_QUERY_ADD( + resolved_term_payload_pointer_reuses, + static_cast(has_frq_payload && p.entry.frq_bytes.data() == frq_payload) + + static_cast(has_prx_payload && + p.entry.prx_bytes.data() == prx_payload)); +#endif + p.frq_base = terms[i].frq_base; + p.prx_base = terms[i].prx_base; + RETURN_IF_ERROR(configure_term_plan(idx, need_positions, fetcher, &p)); + } + return Status::OK(); +} + +Status open_preludes(const io::BatchRangeFetcher& fetcher, std::vector* plans, + bool need_positions) { + for (TermPlan& p : *plans) { + if (!p.windowed) continue; + RETURN_IF_ERROR(FrqPreludeReader::open(fetcher.get(p.prelude_handle), &p.prelude)); + if (need_positions && !p.prelude.has_prx()) { + return Status::Error( + "docid_conjunction: windowed prelude has no positions"); + } + } + return Status::OK(); +} + +Status inline_dd_region(const DictEntry& entry, Slice* out) { + if (entry.dd_meta.disk_len > entry.frq_bytes.size()) { + return Status::Error( + "docid_conjunction: inline dd region exceeds frq bytes"); + } + *out = Slice(entry.frq_bytes.data(), static_cast(entry.dd_meta.disk_len)); + return Status::OK(); +} + +Status build_docid_only_conjunction(const LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + std::vector* candidates) { + return run_docid_only_conjunction_impl(idx, round1, plans, nullptr, candidates, nullptr); +} + +Status build_docid_only_conjunction(const LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + std::vector* candidates, + std::vector* sources) { + return run_docid_only_conjunction_impl(idx, round1, plans, nullptr, candidates, sources); +} + +Status filter_docids_by_conjunction(const LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + const std::vector& initial_candidates, + std::vector* candidates, + std::vector* sources) { + return run_docid_only_conjunction_impl(idx, round1, plans, &initial_candidates, candidates, + sources); +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/docid_posting_reader.cpp b/be/src/storage/index/snii/query/docid_posting_reader.cpp new file mode 100644 index 00000000000000..be58d24a88c869 --- /dev/null +++ b/be/src/storage/index/snii/query/docid_posting_reader.cpp @@ -0,0 +1,380 @@ +// 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 "storage/index/snii/query/internal/docid_posting_reader.h" + +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/reader/windowed_posting.h" + +namespace doris::snii::query::internal { + +using format::DictEntry; +using format::DictEntryEnc; +using format::DictEntryKind; +using format::FrqPreludeReader; +using format::WindowMeta; +using reader::LogicalIndexReader; + +namespace { + +Status decode_flat_docs(const DictEntry& entry, Slice dd_region, std::vector* docids) { + return format::decode_dd_region(dd_region, entry.dd_meta, + /*win_base=*/0, docids); +} + +Status decode_inline_docs(const DictEntry& entry, std::vector* docids) { + if (entry.dd_meta.disk_len > entry.frq_bytes.size()) { + return Status::Error( + "docid_posting_reader: inline dd region exceeds frq bytes"); + } + return decode_flat_docs( + entry, Slice(entry.frq_bytes.data(), static_cast(entry.dd_meta.disk_len)), + docids); +} + +Status slim_docs_fetch_len(const DictEntry& entry, uint64_t win_len, uint64_t* out) { + if (entry.frq_docs_len > win_len) { + return Status::Error( + "docid_posting_reader: slim frq_docs_len exceeds frq window"); + } + *out = entry.frq_docs_len > 0 ? entry.frq_docs_len : win_len; + return Status::OK(); +} + +Status add_u64(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (rhs > std::numeric_limits::max() - lhs) { + return Status::Error(message); + } + *out = lhs + rhs; + return Status::OK(); +} + +Status prelude_abs(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + uint64_t* out) { + uint64_t with_base = 0; + RETURN_IF_ERROR(add_u64(idx.section_refs().posting_region.offset, frq_base, + "docid_posting_reader: prelude offset overflow", &with_base)); + return add_u64(with_base, entry.frq_off_delta, "docid_posting_reader: prelude offset overflow", + out); +} + +Status validate_windowed_docs_prefix(const DictEntry& entry) { + if (entry.prelude_len == 0) { + return Status::Error( + "docid_posting_reader: windowed entry has no prelude"); + } + if (entry.prelude_len > entry.frq_docs_len) { + return Status::Error( + "docid_posting_reader: prelude_len exceeds docs prefix"); + } + if (entry.frq_docs_len > entry.frq_len) { + return Status::Error( + "docid_posting_reader: docs prefix exceeds frq_len"); + } + return Status::OK(); +} + +struct FlatPlan { + size_t out_index = 0; + const DictEntry* entry = nullptr; + size_t handle = 0; +}; + +struct WindowPlan { + size_t out_index = 0; + const ResolvedDocidPosting* posting = nullptr; + size_t prefix_handle = 0; +}; + +Status plan_flat_docs(const LogicalIndexReader& idx, const ResolvedDocidPosting& posting, + io::BatchRangeFetcher* fetcher, FlatPlan* plan) { + uint64_t win_abs = 0; + uint64_t win_len = 0; + RETURN_IF_ERROR(idx.resolve_frq_window(posting.entry, posting.frq_base, &win_abs, &win_len)); + uint64_t docs_len = 0; + RETURN_IF_ERROR(slim_docs_fetch_len(posting.entry, win_len, &docs_len)); + plan->handle = fetcher->add(win_abs, docs_len); + return Status::OK(); +} + +Status plan_window_prefix(const LogicalIndexReader& idx, WindowPlan* plan, + io::BatchRangeFetcher* fetcher) { + const ResolvedDocidPosting& posting = *plan->posting; + RETURN_IF_ERROR(validate_windowed_docs_prefix(posting.entry)); + uint64_t abs = 0; + RETURN_IF_ERROR(prelude_abs(idx, posting.entry, posting.frq_base, &abs)); + plan->prefix_handle = fetcher->add(abs, posting.entry.frq_docs_len); + return Status::OK(); +} + +// Records a non-inline (flat or windowed) posting into the per-encoding plan lists, +// adding the windowed prefix range to the shared fetcher. Flat ranges are added in a +// later pass (after all preludes are registered). Shared by the batched and streamed +// docid readers so both fetch the whole OR in one round. `posting` must outlive the +// plans (its address is captured); callers pass an element of a stable vector. +Status plan_noninline_posting(const LogicalIndexReader& idx, const ResolvedDocidPosting& posting, + size_t out_index, io::BatchRangeFetcher* fetcher, + std::vector* flat_plans, + std::vector* window_plans) { + if (posting.entry.enc == DictEntryEnc::kWindowed) { + WindowPlan plan; + plan.out_index = out_index; + plan.posting = &posting; + RETURN_IF_ERROR(plan_window_prefix(idx, &plan, fetcher)); + window_plans->push_back(std::move(plan)); + return Status::OK(); + } + FlatPlan plan; + plan.out_index = out_index; + plan.entry = &posting.entry; + flat_plans->push_back(plan); + return Status::OK(); +} + +Status window_dd_slice(Slice dd_block, const WindowMeta& meta, Slice* out) { + if (meta.dd_off > dd_block.size() || meta.dd_disk_len > dd_block.size() - meta.dd_off) { + return Status::Error( + "docid_posting_reader: window dd range out of prefix"); + } + *out = dd_block.subslice(static_cast(meta.dd_off), + static_cast(meta.dd_disk_len)); + return Status::OK(); +} + +Status first_docid_in_window(const WindowMeta& meta, uint32_t window_ordinal, uint32_t* first) { + if (window_ordinal == 0) { + *first = 0; + return Status::OK(); + } + if (meta.win_base >= std::numeric_limits::max()) { + return Status::Error( + "docid_posting_reader: window base exceeds docid range"); + } + *first = static_cast(meta.win_base + 1); + if (*first > meta.last_docid) { + return Status::Error( + "docid_posting_reader: invalid window docid range"); + } + return Status::OK(); +} + +Status is_dense_full_window(const WindowMeta& meta, uint32_t window_ordinal, bool* full) { + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(meta, window_ordinal, &first)); + const uint64_t width = static_cast(meta.last_docid) - first + 1; + *full = meta.doc_count == width; + return Status::OK(); +} + +Status decode_flat_plan(const io::BatchRangeFetcher& fetcher, const FlatPlan& plan, + std::vector* out) { + return decode_flat_docs(*plan.entry, fetcher.get(plan.handle), out); +} + +Status decode_window_prefix_plan(const io::BatchRangeFetcher& fetcher, const WindowPlan& plan, + DocIdSink* sink); + +Status decode_window_prefix_plan(const io::BatchRangeFetcher& fetcher, const WindowPlan& plan, + std::vector* out) { + VectorDocIdSink sink(*out); + return decode_window_prefix_plan(fetcher, plan, &sink); +} + +Status decode_window_prefix_plan(const io::BatchRangeFetcher& fetcher, const WindowPlan& plan, + DocIdSink* sink) { + const DictEntry& entry = plan.posting->entry; + const Slice prefix = fetcher.get(plan.prefix_handle); + if (entry.prelude_len > prefix.size()) { + return Status::Error( + "docid_posting_reader: short docs prefix"); + } + const size_t prelude_len = static_cast(entry.prelude_len); + FrqPreludeReader prelude; + RETURN_IF_ERROR(FrqPreludeReader::open(prefix.subslice(0, prelude_len), &prelude)); + const uint64_t dd_block_len = prelude.dd_block_len(); + if (dd_block_len > static_cast(std::numeric_limits::max()) - prelude_len) { + return Status::Error( + "docid_posting_reader: docs prefix length overflow"); + } + const size_t expected_prefix_len = prelude_len + static_cast(dd_block_len); + if (prefix.size() != expected_prefix_len) { + return Status::Error( + "docid_posting_reader: docs prefix length mismatch"); + } + const Slice dd_block = prefix.subslice(prelude_len, prefix.size() - prelude_len); + std::vector docs; + std::vector freqs; + std::vector> positions; + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta meta; + Slice dd_region; + RETURN_IF_ERROR(prelude.window(w, &meta)); + RETURN_IF_ERROR(window_dd_slice(dd_block, meta, &dd_region)); + bool dense_full = false; + RETURN_IF_ERROR(is_dense_full_window(meta, w, &dense_full)); + if (dense_full) { + uint32_t first = 0; + RETURN_IF_ERROR(first_docid_in_window(meta, w, &first)); + RETURN_IF_ERROR(sink->append_range(first, static_cast(meta.last_docid) + 1)); + continue; + } + docs.clear(); + freqs.clear(); + positions.clear(); + RETURN_IF_ERROR(reader::decode_window_slices( + meta, dd_region, Slice(), Slice(), /*want_positions=*/false, + /*want_freq=*/false, &docs, &freqs, &positions)); + RETURN_IF_ERROR(sink->append_sorted(docs)); + } + return Status::OK(); +} + +} // namespace + +Status read_docid_posting(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, std::vector* docids) { + if (docids == nullptr) { + return Status::Error("docid_posting_reader: null out"); + } + docids->clear(); + VectorDocIdSink sink(*docids); + return read_docid_posting(idx, entry, frq_base, prx_base, &sink); +} + +Status read_docid_posting(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, DocIdSink* sink) { + if (sink == nullptr) { + return Status::Error("docid_posting_reader: null sink"); + } + ResolvedDocidPosting posting {entry, frq_base, prx_base}; + if (posting.entry.kind == DictEntryKind::kInline) { + std::vector docs; + RETURN_IF_ERROR(decode_inline_docs(posting.entry, &docs)); + return sink->append_sorted(docs); + } + + io::BatchRangeFetcher docs_fetcher(idx.reader()); + if (posting.entry.enc == DictEntryEnc::kWindowed) { + WindowPlan plan; + plan.out_index = 0; + plan.posting = &posting; + RETURN_IF_ERROR(plan_window_prefix(idx, &plan, &docs_fetcher)); + if (docs_fetcher.pending() > 0) RETURN_IF_ERROR(docs_fetcher.fetch()); + return decode_window_prefix_plan(docs_fetcher, plan, sink); + } + + FlatPlan plan; + plan.out_index = 0; + plan.entry = &posting.entry; + RETURN_IF_ERROR(plan_flat_docs(idx, posting, &docs_fetcher, &plan)); + if (docs_fetcher.pending() > 0) RETURN_IF_ERROR(docs_fetcher.fetch()); + std::vector docs; + RETURN_IF_ERROR(decode_flat_plan(docs_fetcher, plan, &docs)); + return sink->append_sorted(docs); +} + +Status read_docid_postings_batched(const LogicalIndexReader& idx, + const std::vector& postings, + std::vector>* docids) { + if (docids == nullptr) { + return Status::Error( + "docid_posting_reader: null batched out"); + } + docids->clear(); + docids->resize(postings.size()); + + std::vector flat_plans; + std::vector window_plans; + io::BatchRangeFetcher docs_fetcher(idx.reader()); + + for (size_t i = 0; i < postings.size(); ++i) { + const ResolvedDocidPosting& posting = postings[i]; + if (posting.entry.kind == DictEntryKind::kInline) { + RETURN_IF_ERROR(decode_inline_docs(posting.entry, &(*docids)[i])); + continue; + } + RETURN_IF_ERROR( + plan_noninline_posting(idx, posting, i, &docs_fetcher, &flat_plans, &window_plans)); + } + + for (FlatPlan& plan : flat_plans) { + const ResolvedDocidPosting& posting = postings[plan.out_index]; + RETURN_IF_ERROR(plan_flat_docs(idx, posting, &docs_fetcher, &plan)); + } + if (docs_fetcher.pending() > 0) RETURN_IF_ERROR(docs_fetcher.fetch()); + + for (const FlatPlan& plan : flat_plans) { + RETURN_IF_ERROR(decode_flat_plan(docs_fetcher, plan, &(*docids)[plan.out_index])); + } + for (const WindowPlan& plan : window_plans) { + RETURN_IF_ERROR(decode_window_prefix_plan(docs_fetcher, plan, &(*docids)[plan.out_index])); + } + return Status::OK(); +} + +Status emit_docid_postings_streamed(const LogicalIndexReader& idx, + const std::vector& postings, + DocIdSink* sink) { + if (sink == nullptr) { + return Status::Error( + "docid_posting_reader: null streamed sink"); + } + + std::vector flat_plans; + std::vector window_plans; + io::BatchRangeFetcher docs_fetcher(idx.reader()); + // One scratch buffer reused across flat/inline postings: clear() keeps capacity, + // so at most one growth total instead of a fresh vector per posting. + std::vector scratch; + + for (size_t i = 0; i < postings.size(); ++i) { + const ResolvedDocidPosting& posting = postings[i]; + if (posting.entry.kind == DictEntryKind::kInline) { + scratch.clear(); + RETURN_IF_ERROR(decode_inline_docs(posting.entry, &scratch)); + RETURN_IF_ERROR(sink->append_sorted(scratch)); + continue; + } + RETURN_IF_ERROR( + plan_noninline_posting(idx, posting, i, &docs_fetcher, &flat_plans, &window_plans)); + } + + for (FlatPlan& plan : flat_plans) { + const ResolvedDocidPosting& posting = postings[plan.out_index]; + RETURN_IF_ERROR(plan_flat_docs(idx, posting, &docs_fetcher, &plan)); + } + if (docs_fetcher.pending() > 0) RETURN_IF_ERROR(docs_fetcher.fetch()); + + for (const FlatPlan& plan : flat_plans) { + scratch.clear(); + RETURN_IF_ERROR(decode_flat_plan(docs_fetcher, plan, &scratch)); + RETURN_IF_ERROR(sink->append_sorted(scratch)); + } + for (const WindowPlan& plan : window_plans) { + RETURN_IF_ERROR(decode_window_prefix_plan(docs_fetcher, plan, sink)); + } + return Status::OK(); +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/docid_set_ops.cpp b/be/src/storage/index/snii/query/docid_set_ops.cpp new file mode 100644 index 00000000000000..5f8931f4146e16 --- /dev/null +++ b/be/src/storage/index/snii/query/docid_set_ops.cpp @@ -0,0 +1,127 @@ +// 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 "storage/index/snii/query/internal/docid_set_ops.h" + +#include +#include +#include +#include + +namespace doris::snii::query::internal { + +std::vector intersect_sorted(const std::vector& a, + const std::vector& b) { + std::vector out; + out.reserve(std::min(a.size(), b.size())); + std::set_intersection(a.begin(), a.end(), b.begin(), b.end(), std::back_inserter(out)); + return out; +} + +void union_sorted_into(std::vector* acc, const std::vector& next) { + std::vector merged; + merged.reserve(acc->size() + next.size()); + std::set_union(acc->begin(), acc->end(), next.begin(), next.end(), std::back_inserter(merged)); + *acc = std::move(merged); +} + +std::vector union_sorted_many(const std::vector>& lists, + size_t reserve_cap) { + constexpr size_t kLinearFanInMax = 8; + struct Cursor { + uint32_t docid = 0; + size_t list = 0; + size_t offset = 0; + }; + struct GreaterDocId { + bool operator()(const Cursor& a, const Cursor& b) const { return a.docid > b.docid; } + }; + + size_t non_empty = 0; + size_t total = 0; + std::priority_queue, GreaterDocId> heap; + for (size_t i = 0; i < lists.size(); ++i) { + if (lists[i].empty()) continue; + ++non_empty; + total += lists[i].size(); + heap.push(Cursor {lists[i][0], i, 0}); + } + // The union is at most `total` (exactly that for disjoint inputs); reserve to it + // so the output grows in a single allocation rather than O(log) geometric + // reallocations. Cap guards against over-reserving for heavily-overlapping inputs. + const size_t reserve_hint = std::min(total, reserve_cap); + if (non_empty == 0) return {}; + if (non_empty == 1) { + for (const std::vector& docs : lists) { + if (!docs.empty()) return docs; + } + } + + if (non_empty <= kLinearFanInMax) { + std::vector offsets(lists.size(), 0); + std::vector out; + out.reserve(reserve_hint); + bool has_last = false; + uint32_t last = 0; + for (;;) { + bool found = false; + uint32_t next = 0; + for (size_t i = 0; i < lists.size(); ++i) { + if (offsets[i] >= lists[i].size()) continue; + const uint32_t docid = lists[i][offsets[i]]; + if (!found || docid < next) { + found = true; + next = docid; + } + } + if (!found) break; + if (!has_last || next != last) { + out.push_back(next); + last = next; + has_last = true; + } + for (size_t i = 0; i < lists.size(); ++i) { + while (offsets[i] < lists[i].size() && lists[i][offsets[i]] == next) { + ++offsets[i]; + } + } + } + return out; + } + + std::vector out; + out.reserve(reserve_hint); + bool has_last = false; + uint32_t last = 0; + while (!heap.empty()) { + const Cursor cur = heap.top(); + heap.pop(); + if (!has_last || cur.docid != last) { + out.push_back(cur.docid); + last = cur.docid; + has_last = true; + } + const size_t next_offset = cur.offset + 1; + const std::vector& docs = lists[cur.list]; + if (next_offset < docs.size()) { + heap.push(Cursor {docs[next_offset], cur.list, next_offset}); + } + } + return out; +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/docid_sink.h b/be/src/storage/index/snii/query/docid_sink.h new file mode 100644 index 00000000000000..604d3dde8bdbb8 --- /dev/null +++ b/be/src/storage/index/snii/query/docid_sink.h @@ -0,0 +1,90 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris::snii::query { + +// Bulk docid handoff for query operators. Each span is sorted ascending; callers +// that need a single vector can use VectorDocIdSink. +class DocIdSink { +public: + virtual ~DocIdSink() = default; + virtual Status append_sorted(std::span docids) = 0; + virtual Status append_range(uint32_t first, uint64_t last_exclusive) = 0; + + // True iff the sink deduplicates and globally orders on its own (e.g. a Roaring + // bitmap via addMany/addRange). For such sinks a multi-term OR can stream each + // posting straight in -- skipping the per-term vector materialization plus the + // K-way merge accumulator. Sinks that hand back a single globally-sorted, + // deduplicated vector (VectorDocIdSink) keep the default false, so callers + // materialize + merge before appending. The gate must stay conservative: + // streaming several postings into a non-dedup sink would break that contract. + virtual bool dedups() const { return false; } +}; + +class VectorDocIdSink final : public DocIdSink { +public: + explicit VectorDocIdSink(std::vector& docids) : docids_(docids) {} + + Status append_sorted(std::span docids) override { + docids_.insert(docids_.end(), docids.begin(), docids.end()); + return Status::OK(); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + if (last_exclusive <= first) { + return Status::OK(); + } + if (last_exclusive > static_cast(std::numeric_limits::max()) + 1) { + return Status::Error( + "docid_sink: range exceeds uint32 docid space"); + } + const uint64_t count = last_exclusive - first; + if (count > static_cast(docids_.max_size() - docids_.size())) { + return Status::Error("docid_sink: range too large"); + } + // GEOMETRIC BULK reserve -- never an exact one: append_range can be + // called once per docid run for a query, and an exact + // reserve(size()+count) caps capacity at "just enough" so the next + // append reallocates + memcpys the whole accumulated vector -- + // quadratic total memcpy across runs (same anti-pattern as the writer's + // add_nulls). Doubling on overflow keeps the O(count) amortization AND + // makes one large range pay at most one reallocation. + const size_t need = docids_.size() + static_cast(count); + if (need > docids_.capacity()) { + docids_.reserve(std::max(need, docids_.capacity() * 2)); + } + for (uint64_t docid = first; docid < last_exclusive; ++docid) { + docids_.push_back(static_cast(docid)); + } + return Status::OK(); + } + +private: + std::vector& docids_; +}; + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/docid_union.cpp b/be/src/storage/index/snii/query/docid_union.cpp new file mode 100644 index 00000000000000..f1296b18cd685e --- /dev/null +++ b/be/src/storage/index/snii/query/docid_union.cpp @@ -0,0 +1,58 @@ +// 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 "storage/index/snii/query/internal/docid_union.h" + +#include + +#include "storage/index/snii/query/internal/docid_set_ops.h" + +namespace doris::snii::query::internal { + +Status build_docid_union(const reader::LogicalIndexReader& idx, + const std::vector& postings, + std::vector* out) { + if (out == nullptr) + return Status::Error("docid_union: null out"); + out->clear(); + if (postings.empty()) return Status::OK(); + + std::vector> docs_by_posting; + RETURN_IF_ERROR(read_docid_postings_batched(idx, postings, &docs_by_posting)); + *out = union_sorted_many(docs_by_posting); + return Status::OK(); +} + +Status emit_docid_union(const reader::LogicalIndexReader& idx, + const std::vector& postings, DocIdSink* sink) { + if (sink == nullptr) + return Status::Error("docid_union: null sink"); + if (postings.empty()) return Status::OK(); + // A dedup-capable sink (Roaring) orders + dedups across postings itself, so stream + // each posting straight in over a single shared fetch round -- no per-term vector + // or K-way merge accumulator. A plain (non-dedup) sink keeps the materialize+merge + // path so its single-span contract (globally sorted, deduplicated) holds. + if (sink->dedups()) { + return emit_docid_postings_streamed(idx, postings, sink); + } + std::vector acc; + RETURN_IF_ERROR(build_docid_union(idx, postings, &acc)); + if (acc.empty()) return Status::OK(); + return sink->append_sorted(acc); +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/docid_conjunction.h b/be/src/storage/index/snii/query/internal/docid_conjunction.h new file mode 100644 index 00000000000000..ea63d9446a0071 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/docid_conjunction.h @@ -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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query::internal { + +struct ResolvedQueryTerm { + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; +}; + +struct TermPlan { + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + uint32_t df = 0; + size_t order = 0; + size_t frq_handle = 0; + size_t prx_handle = 0; + size_t prelude_handle = 0; + bool pod_ref = false; + bool windowed = false; + format::FrqPreludeReader prelude; +}; + +struct DocidChunk { + std::vector docids; + std::vector prx_doc_ordinals; + uint32_t prx_doc_count = 0; + bool windowed = false; + uint32_t window = 0; +}; + +struct DocidSource { + std::vector chunks; + bool docids_are_final_candidates = false; +}; + +Status resolve_query_term(const reader::LogicalIndexReader& idx, std::string_view term, + ResolvedQueryTerm* resolved, bool* found); + +// Resolves one sorted, duplicate-free term batch through bounded physical DICT +// reads. Results stay aligned with `terms`; absent terms have found[i]=0. +Status resolve_query_terms_batch(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* resolved, + std::vector* found); + +Status plan_terms(const reader::LogicalIndexReader& idx, const std::vector& terms, + io::BatchRangeFetcher* fetcher, std::vector* plans, bool* all_present, + bool need_positions); + +Status plan_resolved_terms(const reader::LogicalIndexReader& idx, + const std::vector& terms, + io::BatchRangeFetcher* fetcher, std::vector* plans, + bool need_positions); + +Status plan_resolved_terms(const reader::LogicalIndexReader& idx, + std::vector&& terms, io::BatchRangeFetcher* fetcher, + std::vector* plans, bool need_positions); + +Status open_preludes(const io::BatchRangeFetcher& fetcher, std::vector* plans, + bool need_positions); + +Status inline_dd_region(const format::DictEntry& entry, Slice* out); + +Status build_docid_only_conjunction(const reader::LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + std::vector* candidates); + +Status build_docid_only_conjunction(const reader::LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + std::vector* candidates, + std::vector* sources); + +Status filter_docids_by_conjunction(const reader::LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, + const std::vector& plans, + const std::vector& initial_candidates, + std::vector* candidates, + std::vector* sources); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/docid_posting_reader.h b/be/src/storage/index/snii/query/internal/docid_posting_reader.h new file mode 100644 index 00000000000000..0230a1f526039a --- /dev/null +++ b/be/src/storage/index/snii/query/internal/docid_posting_reader.h @@ -0,0 +1,62 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query::internal { + +struct ResolvedDocidPosting { + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; +}; + +// Decodes the docid-only posting for a resolved term. The caller owns term +// lookup and can batch/plan lookups independently; this module owns only the +// three posting encodings (inline, slim pod_ref, windowed pod_ref). +Status read_docid_posting(const reader::LogicalIndexReader& idx, const format::DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, std::vector* docids); + +Status read_docid_posting(const reader::LogicalIndexReader& idx, const format::DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, query::DocIdSink* sink); + +// Batch counterpart for multi-term docid-only operators. Windowed terms share one +// prelude fetch round and one docid fetch round, so OR-style operators pay by +// stage rather than by term. +Status read_docid_postings_batched(const reader::LogicalIndexReader& idx, + const std::vector& postings, + std::vector>* docids); + +// Streaming counterpart of read_docid_postings_batched for a dedup-capable sink +// (DocIdSink::dedups()==true, e.g. a Roaring bitmap). Shares the exact same single +// docid fetch round, but decodes each posting straight into the sink -- dense-full +// windows via append_range (run-preserving), the rest via append_sorted from one +// reused scratch buffer -- so no per-term vector or K-way merge accumulator is +// materialized. The sink dedups/orders across postings. One I/O round is preserved. +Status emit_docid_postings_streamed(const reader::LogicalIndexReader& idx, + const std::vector& postings, + query::DocIdSink* sink); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/docid_set_ops.h b/be/src/storage/index/snii/query/internal/docid_set_ops.h new file mode 100644 index 00000000000000..f651e930a1cdb3 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/docid_set_ops.h @@ -0,0 +1,39 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +namespace doris::snii::query::internal { + +std::vector intersect_sorted(const std::vector& a, + const std::vector& b); + +void union_sorted_into(std::vector* acc, const std::vector& next); + +// Sorted-deduplicated union of many sorted lists. The output is reserved by the +// summed input size (the union is at most the total of all inputs; for disjoint +// inputs it is exactly that), capped by `reserve_cap` so heavily-overlapping +// inputs (union << total) do not over-reserve. Default cap = no cap. +std::vector union_sorted_many(const std::vector>& lists, + size_t reserve_cap = std::numeric_limits::max()); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/docid_union.h b/be/src/storage/index/snii/query/internal/docid_union.h new file mode 100644 index 00000000000000..3243c082bbeec2 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/docid_union.h @@ -0,0 +1,38 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query::internal { + +// Reads already-resolved docid postings in planned batches, merges them as a +// sorted deduplicated union, then emits one bulk span to the sink. +Status build_docid_union(const reader::LogicalIndexReader& idx, + const std::vector& postings, + std::vector* out); + +Status emit_docid_union(const reader::LogicalIndexReader& idx, + const std::vector& postings, DocIdSink* sink); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/phrase_query_split.h b/be/src/storage/index/snii/query/internal/phrase_query_split.h new file mode 100644 index 00000000000000..08b746b55016b5 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/phrase_query_split.h @@ -0,0 +1,598 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +// phrase_query implements MATCH_PHRASE with WINDOW (sub-block) SKIPPING for +// high-df windowed terms (design spec section 6.2): +// 1. Resolve every term; reject if any is absent. +// 2. Batch-read each windowed term's prelude + each slim/inline term's full +// docid posting in one round; open the two-level prelude readers. +// 3. Pick the DRIVER = smallest-df term; materialize it fully -> the initial +// candidate docid set. +// 4. For every other term in ascending-df order, narrow the candidate set: +// - slim/inline: intersect with its (already decoded) full posting. +// - windowed: locate_window() the CURRENT candidates -> the SET of +// windows covering them; batch-fetch ONLY those windows' +// .frq docid regions; keep candidates present in some +// covering window. A high-df term thus reads +// O(candidates) windows instead of its whole O(df) +// posting. +// 5. Fetch PRX only for retained chunks and run the positional phrase check +// (term[0]@p, term[1]@p+1, ...) on the survivors. +// The result is identical to a full-read intersection; only the bytes read for +// high-df windowed terms shrink. +// +// Internal to the phrase-query implementation, which spans phrase_plan.cpp, +// phrase_position_source.cpp, phrase_emit.cpp, phrase_prefix_exec.cpp, +// phrase_planned_query.cpp and phrase_query.cpp. This header carries the types and +// functions those translation units share; nothing outside query/ may include it. +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; + +bool apply_common_grams_plan_debug_override(bool cost_prefers_gram, + CommonGramsPlanDebugOverride debug_override); + +class CommonGramsPlanningTimer { +public: + explicit CommonGramsPlanningTimer(format::PhraseQueryExecutionStats* stats) : stats_(stats) { + if (stats_ != nullptr) { + start_ = std::chrono::steady_clock::now(); + } + } + + ~CommonGramsPlanningTimer() { finish(); } + + CommonGramsPlanningTimer(const CommonGramsPlanningTimer&) = delete; + CommonGramsPlanningTimer& operator=(const CommonGramsPlanningTimer&) = delete; + + void finish() { + if (finished_) { + return; + } + finished_ = true; + if (stats_ == nullptr) { + return; + } + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_) + .count(); + stats_->common_grams_planning_ns += static_cast(std::max(1, elapsed)); + } + +private: + format::PhraseQueryExecutionStats* stats_ = nullptr; + std::chrono::steady_clock::time_point start_; + bool finished_ = false; +}; + +size_t position_span_size(std::pair span); + +bool should_use_monotonic_position_scan(std::pair anchor_span, + size_t checked_span_size, uint32_t anchor_offset, + uint32_t checked_offset); + +struct ExpectedTailPositions { + uint32_t docid = 0; + uint32_t phrase_frequency = 0; + size_t positions_begin = 0; + size_t positions_end = 0; +}; + +static_assert(sizeof(ExpectedTailPositions) == 3 * sizeof(uint64_t)); + +struct ExpectedTailPositionSet { + std::vector docs; + std::vector positions; + std::vector position_matched; + size_t matched_count = 0; + + void clear() { + docs.clear(); + positions.clear(); + position_matched.clear(); + matched_count = 0; + } + + void reserve_docs(size_t count) { + docs.reserve(count); + positions.reserve(count); + } +}; + +// One decoded chunk of a term's posting: a windowed term's covering window, or +// a slim/inline term's single posting. `docids` is decoded in the conjunction +// phase (and reused by the streaming cursor -- the dd region is decoded exactly +// once); `prx` is the on-disk positions bytes, decoded lazily by the cursor +// (once per chunk) during phrase verification. + +struct PosChunk { + std::vector docids; // ascending, absolute + // Empty means the chunk keeps every PRX doc in on-disk order. Non-empty means + // `docids[i]` corresponds to on-disk local document ordinal + // `prx_doc_ordinals[i]`, allowing PRX decode to skip positions for docs that + // were removed by the docid-only conjunction. + std::vector prx_doc_ordinals; + uint32_t prx_doc_count = 0; + Slice prx; // .prx window bytes (reference fetcher/round1/entry) + bool windowed = false; + uint32_t window = 0; +}; + +// A term's retained posting as an ordered list of chunks (windowed: covering +// windows in docid order; slim/inline: one). The referenced prx bytes live in +// `round1` / the per-term fetchers kept alive in phrase_query::owners for the +// whole query, so the cursor can decode positions during verification. + +struct PosSource { + std::vector chunks; + format::PrxDecodeContext* observer_context = nullptr; +}; + +struct PhraseExecutionState { + std::vector srcs; + std::vector> owners; + std::vector candidates; +}; + +struct PhraseTermMapping { + std::vector unique_terms; + std::vector phrase_plan_index; +}; + +struct PhysicalPhrasePlan { + std::vector unique_terms; + std::vector phrase_plan_index; + std::vector position_offsets; + std::vector common_gram_clauses; +}; + +bool has_common_grams_capability( + const LogicalIndexReader& idx, + const segment_v2::inverted_index::CommonGramsQueryIdentity* query_identity); + +bool entry_has_positions(const format::DictEntry& entry); + +Status build_physical_phrase_plan_prefix(const LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& query_info, + size_t clause_count, bool allow_common_grams, + PhysicalPhrasePlan* plan, bool* all_representable); + +Status build_physical_phrase_plan(const LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& query_info, + bool allow_common_grams, PhysicalPhrasePlan* plan, + bool* all_representable); + +size_t resolved_batch_index(const std::vector& batch_terms, std::string_view term); + +bool all_plan_terms_present(const PhysicalPhrasePlan& plan, + const std::vector& batch_terms, + const std::vector& found); + +uint64_t plan_visible_posting_bytes(const format::DictEntry& entry, bool need_positions); + +segment_v2::inverted_index::CommonGramsPlanRawCost phrase_plan_raw_cost( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + const std::vector& resolved, const std::vector& found, + bool need_positions); + +segment_v2::inverted_index::CommonGramsPlanRawCost alternative_clause_raw_cost( + const std::vector& terms, bool need_positions); + +void append_alternative_clause_cost( + const segment_v2::inverted_index::CommonGramsPlanRawCost& clause, + segment_v2::inverted_index::CommonGramsPlanRawCost* plan); + +segment_v2::inverted_index::CommonGramsPlanRawCost hybrid_verification_raw_cost( + const segment_v2::inverted_index::CommonGramsPlanRawCost& prefilter_cost, + const segment_v2::inverted_index::CommonGramsPlanRawCost& verification_cost); + +internal::ResolvedPhrasePlan materialize_resolved_phrase_plan( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + std::vector* resolved); + +internal::ResolvedPhrasePlan copy_resolved_phrase_plan( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + const std::vector& resolved); + +bool physical_phrase_plan_has_docs_only_term(const PhysicalPhrasePlan& plan, + const std::vector& batch_terms, + const std::vector& resolved); + +void append_physical_phrase_clause(const PhysicalPhrasePlan& source, size_t clause, + uint32_t position_offset, PhysicalPhrasePlan* target); + +struct HybridPositionedCover { + PhysicalPhrasePlan candidate_prefilter; + PhysicalPhrasePlan verification; +}; + +struct HybridExactPlanArtifact { + std::optional positioned_cover; +}; + +HybridExactPlanArtifact build_hybrid_exact_plan_artifact( + const PhysicalPhrasePlan& plain_plan, const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, + const std::vector& resolved); + +struct ResolvedMappedTail { + size_t batch_index = 0; + uint32_t expansion_ordinal = 0; +}; + +struct HybridPrefixMappedTails { + std::vector positioned_indices; + std::vector docs_only_indices; + std::vector docs_only_ordinals; +}; + +struct HybridPrefixPlanArtifact { + HybridPositionedCover plain_tail_cover; + HybridPrefixMappedTails mapped_tail_split; + std::optional positioned_tail_verification; + uint32_t plain_tail_position_offset = 0; + bool maps_tail_to_gram = false; +}; + +std::optional try_build_hybrid_prefix_plan_artifact( + const PhysicalPhrasePlan& plain_leading, const PhysicalPhrasePlan& gram_leading, + const std::vector& batch_terms, const std::vector& resolved, + const std::vector& mapped_tails, bool maps_tail_to_gram); + +struct HybridPrefixCandidateSet { + bool active = false; + std::vector docs; +}; + +Status build_hybrid_leading_candidates(const LogicalIndexReader& idx, + const PhysicalPhrasePlan& candidate_prefilter, + const std::vector& batch_terms, + const std::vector& resolved, + HybridPrefixCandidateSet* candidates); + +Status build_hybrid_docs_only_tail_candidates(const LogicalIndexReader& idx, + const std::vector& resolved, + const std::vector& gram_tail_indices, + const HybridPrefixCandidateSet& leading_candidates, + std::vector* candidates); + +Status execute_hybrid_exact_phrase_plan( + const LogicalIndexReader& idx, const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, const HybridExactPlanArtifact& artifact, + std::vector* resolved, std::vector* docids, + format::PrxDecodeContext* decode_context, bool* candidate_intersection_empty = nullptr); + +void append_resolved_phrase_clause(ResolvedQueryTerm term, uint32_t position_offset, + internal::ResolvedPhrasePlan* plan); + +internal::ResolvedPhrasePlan build_resolved_phrase_plan( + std::vector resolved_terms); + +Status resolve_and_execute_physical_phrase_plan(const LogicalIndexReader& idx, + const PhysicalPhrasePlan& plan, + std::vector* docids, + format::PrxDecodeContext* decode_context, + CommonGramsPlanningTimer& planning_timer); + +Status planned_exact_phrase_query_impl( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, format::PrxDecodeContext* decode_context, + ExactPhrasePlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + CommonGramsPlanDebugOverride debug_override); + +PhraseTermMapping build_phrase_term_mapping(const std::vector& terms); + +Status build_position_sources_for_candidates( + const LogicalIndexReader& idx, const io::BatchRangeFetcher& round1, + const std::vector& plans, std::vector* doc_sources, + const std::vector& candidates, + std::vector>* owners, std::vector* srcs, + format::PrxDecodeContext* observer_context); + +class PosChunkDecoder { +public: + explicit PosChunkDecoder(format::PrxDecodeContext* observer_context = nullptr) + : observer_context_(observer_context) {} + + void set_decode_state(format::PrxDecodeContext* observer_context) { + observer_context_ = observer_context; + } + + void reset() { + chunk_ = nullptr; + offsets_by_prx_ordinal_ = false; + } + + Status decode(const PosChunk& chunk) { + chunk_ = &chunk; + ByteSource ps(chunk.prx); + const bool selected_all = chunk.prx_doc_ordinals.empty(); + const bool decode_full = selected_all || should_decode_full_prx_window(chunk); + offsets_by_prx_ordinal_ = decode_full && !selected_all; + return internal::decode_and_validate_prx_frame( + &ps, chunk.prx_doc_ordinals, decode_full, selected_all, chunk.prx_doc_count, + chunk.docids.size(), &pflat_, &poff_, observer_context_); + } + + Status positions(size_t doc_index, std::pair* out) const { + if (chunk_ == nullptr || doc_index >= chunk_->docids.size()) { + return Status::Error( + "phrase_query: decoded chunk doc index out of range"); + } + const size_t pos_index = + offsets_by_prx_ordinal_ ? chunk_->prx_doc_ordinals[doc_index] : doc_index; + if (pos_index + 1 >= poff_.size()) { + return Status::Error( + "phrase_query: prx ordinal offset out of range"); + } + const uint32_t begin = poff_[pos_index]; + const uint32_t end = poff_[pos_index + 1]; + if (begin == end) { + *out = {nullptr, nullptr}; + return Status::OK(); + } + if (end > pflat_.size()) { + return Status::Error( + "phrase_query: prx offset out of range"); + } + *out = {pflat_.data() + begin, pflat_.data() + end}; + return Status::OK(); + } + + inline __attribute__((always_inline)) std::pair + positions_unchecked(size_t doc_index) const { + const size_t pos_index = + offsets_by_prx_ordinal_ ? chunk_->prx_doc_ordinals[doc_index] : doc_index; + const uint32_t begin = poff_[pos_index]; + const uint32_t end = poff_[pos_index + 1]; + if (begin == end) { + return {nullptr, nullptr}; + } + return {pflat_.data() + begin, pflat_.data() + end}; + } + +private: + static bool should_decode_full_prx_window(const PosChunk& chunk) { + return chunk.prx_doc_count != 0 && + static_cast(chunk.prx_doc_ordinals.size()) * 2 >= chunk.prx_doc_count; + } + + const PosChunk* chunk_ = nullptr; + bool offsets_by_prx_ordinal_ = false; + std::vector pflat_; + std::vector poff_; + format::PrxDecodeContext* observer_context_ = nullptr; +}; + +// Streaming position cursor over one term's retained chunks. It advances ONLY +// forward (callers seek ascending candidate docids), decodes each chunk's +// docids once (reused from the conjunction phase) and each chunk's positions at +// most once (lazily, into a flat CSR whose capacity is retained across chunks). +// No per-doc allocation, no per-candidate docid binary search: positions are +// addressed by the doc's local index within its chunk. This is the read-side +// dual of the windowed posting layout -- the S3-native batch fetch already +// pulled every needed chunk into memory; the cursor is pure in-memory column +// iteration. + +class PostingCursor { +public: + void init(const PosSource* src) { + src_ = src; + ci_ = 0; + li_ = 0; + decoded_pos_chunk_ = kNoChunk; + decoder_.set_decode_state(src->observer_context); + decoder_.reset(); + } + + // Positions the cursor at `target` (guaranteed present: candidates are the + // intersection of exactly these chunks' docids). Monotonic forward advance. + Status seek(uint32_t target) { + while (ci_ < src_->chunks.size() && + (src_->chunks[ci_].docids.empty() || src_->chunks[ci_].docids.back() < target)) { + ++ci_; + li_ = 0; + } + if (ci_ >= src_->chunks.size()) { + return Status::Error( + "phrase_query: cursor exhausted before target docid"); + } + const std::vector& d = src_->chunks[ci_].docids; + while (li_ < d.size() && d[li_] < target) { + ++li_; + } + if (li_ >= d.size() || d[li_] != target) { + return Status::Error( + "phrase_query: candidate missing from posting chunk"); + } + return Status::OK(); + } + + // [begin,end) of the current doc's positions, decoding the current chunk's + // .prx exactly once (cached). Must follow a seek that landed on a real doc. + Status positions(std::pair* out) { + if (ci_ >= src_->chunks.size() || li_ >= src_->chunks[ci_].docids.size()) { + return Status::Error( + "phrase_query: cursor positions out of range"); + } + if (decoded_pos_chunk_ != ci_) { + RETURN_IF_ERROR(decoder_.decode(src_->chunks[ci_])); + decoded_pos_chunk_ = ci_; + } + return decoder_.positions(li_, out); + } + + Status next(uint32_t* docid, std::pair* out) { + while (ci_ < src_->chunks.size() && + (src_->chunks[ci_].docids.empty() || li_ >= src_->chunks[ci_].docids.size())) { + ++ci_; + li_ = 0; + } + if (ci_ >= src_->chunks.size()) { + return Status::Error( + "phrase_query: cursor exhausted before next docid"); + } + *docid = src_->chunks[ci_].docids[li_]; + RETURN_IF_ERROR(positions(out)); + ++li_; + return Status::OK(); + } + +private: + static constexpr size_t kNoChunk = static_cast(-1); + + const PosSource* src_ = nullptr; + size_t ci_ = 0; // current chunk + size_t li_ = 0; // current local doc index within the chunk + size_t decoded_pos_chunk_ = kNoChunk; // which chunk decoder_ currently holds + PosChunkDecoder decoder_; +}; + +enum class PhraseCandidateMetric : uint8_t { + kExact, + kPrefixLeading, +}; + +Status build_phrase_execution_state(const LogicalIndexReader& idx, io::BatchRangeFetcher* round1, + std::vector* plans, PhraseExecutionState* state, + const std::vector* candidate_prefilter, + format::PrxDecodeContext* observer_context, + PhraseCandidateMetric candidate_metric); + +Status execute_phrase_plans(const LogicalIndexReader& idx, io::BatchRangeFetcher* round1, + std::vector* plans, + const std::vector& phrase_plan_index, + std::vector* docids, + format::PrxDecodeContext* observer_context, + std::vector* matches); + +Status execute_resolved_phrase_prefix_terms( + const LogicalIndexReader& idx, internal::ResolvedPhrasePlan exact_plan, + std::vector tail_terms, uint32_t tail_position_offset, + std::vector* docids, format::PrxDecodeContext* decode_context, + std::vector* matches = nullptr, + const std::vector* candidate_prefilter = nullptr); + +Status execute_hybrid_phrase_prefix_plan( + const LogicalIndexReader& idx, const HybridPrefixPlanArtifact& artifact, + const std::vector& batch_terms, const std::vector& resolved, + const std::vector& plain_tail_terms, std::vector* docids, + format::PrxDecodeContext* decode_context, CommonGramsPlanningTimer& planning_timer, + bool* candidate_intersection_empty); + +struct HybridPrefixCostEstimate { + segment_v2::inverted_index::CommonGramsPlanRawCost raw_cost; + uint64_t estimated_cost = 0; +}; + +HybridPrefixCostEstimate estimate_hybrid_prefix_plan_cost( + const HybridPrefixPlanArtifact& artifact, const std::vector& batch_terms, + const std::vector& resolved, const std::vector& found, + const std::vector& plain_tail_terms, uint32_t position_verify_factor); + +Status phrase_query_impl(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids, + format::PrxDecodeContext* decode_context, + std::vector* matches = nullptr); + +Status phrase_prefix_query_impl(const LogicalIndexReader& idx, + const std::vector& terms, + std::vector* const docids, int32_t max_expansions, + format::PrxDecodeContext* decode_context, + CommonGramsPlanningTimer* planning_timer, + std::vector* matches = nullptr); + +Status planned_phrase_prefix_query_impl( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, int32_t max_expansions, + format::PrxDecodeContext* decode_context, PhrasePrefixPlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + CommonGramsPlanDebugOverride debug_override); + +template +segment_v2::inverted_index::CommonGramsPlanRawCost alternative_clause_raw_cost( + const std::vector& terms, const std::vector& indices, + bool need_positions) { + segment_v2::inverted_index::CommonGramsPlanRawCost cost; + unsigned __int128 posting_bytes = 0; + unsigned __int128 candidate_df = 0; + for (size_t index : indices) { + DORIS_CHECK_LT(index, terms.size()); + posting_bytes += plan_visible_posting_bytes(terms[index].entry, need_positions); + candidate_df += terms[index].entry.df; + } + cost.posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + cost.estimated_candidate_df = candidate_df > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(candidate_df); + cost.clause_count = 1; + return cost; +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/internal/plain_term_routing.h b/be/src/storage/index/snii/query/internal/plain_term_routing.h new file mode 100644 index 00000000000000..787d237ddaf933 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/plain_term_routing.h @@ -0,0 +1,194 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/query/query_info.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query::internal { + +inline segment_v2::inverted_index::PlainTermKeyVersion plain_term_key_version( + const reader::LogicalIndexReader& idx) { + const auto* metadata = idx.common_grams_metadata(); + return metadata == nullptr ? segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw + : metadata->plain_term_key_version; +} + +inline Status route_plain_query_term_view(const reader::LogicalIndexReader& idx, + std::string_view logical_term, std::string* scratch, + std::string_view* physical_term, bool* representable) { + DORIS_CHECK(scratch != nullptr); + DORIS_CHECK(physical_term != nullptr); + DORIS_CHECK(representable != nullptr); + scratch->clear(); + *physical_term = {}; + *representable = false; + + const auto version = plain_term_key_version(idx); + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw && + segment_v2::inverted_index::legacy_raw_exact_requires_bypass(logical_term)) { + return Status::Error( + "SNII legacy raw term overlaps an internal term namespace"); + } + auto encoded = + segment_v2::inverted_index::try_encode_plain_term_view(logical_term, version, scratch); + if (!encoded.has_value()) { + return std::move(encoded.error()); + } + if (encoded->has_value()) { + *physical_term = **encoded; + *representable = true; + } + return Status::OK(); +} + +inline Status route_plain_query_term(const reader::LogicalIndexReader& idx, + std::string_view logical_term, std::string* physical_term, + bool* representable) { + DORIS_CHECK(physical_term != nullptr); + std::string scratch; + std::string_view physical_term_view; + RETURN_IF_ERROR(route_plain_query_term_view(idx, logical_term, &scratch, &physical_term_view, + representable)); + physical_term->assign(physical_term_view); + return Status::OK(); +} + +inline Status route_query_term_view(const reader::LogicalIndexReader& idx, + const segment_v2::TermInfo& term_info, std::string* scratch, + std::string_view* physical_term, bool* representable) { + DORIS_CHECK(term_info.is_single_term()); + if (term_info.key_kind == segment_v2::TermKeyKind::kCommonGram) { + return Status::Error( + "CommonGrams query plan requires segment capability validation"); + } + return route_plain_query_term_view(idx, term_info.get_single_term(), scratch, physical_term, + representable); +} + +inline Status route_query_term(const reader::LogicalIndexReader& idx, + const segment_v2::TermInfo& term_info, std::string* physical_term, + bool* representable) { + DORIS_CHECK(term_info.is_single_term()); + if (term_info.key_kind == segment_v2::TermKeyKind::kCommonGram) { + return Status::Error( + "CommonGrams query plan requires segment capability validation"); + } + return route_plain_query_term(idx, term_info.get_single_term(), physical_term, representable); +} + +inline Status route_query_terms(const reader::LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& query_info, + std::vector* routed_terms, bool* all_representable) { + DORIS_CHECK(routed_terms != nullptr); + DORIS_CHECK(all_representable != nullptr); + DORIS_CHECK(routed_terms->size() == query_info.term_infos.size()); + *all_representable = true; + const auto version = plain_term_key_version(idx); + size_t output_index = 0; + std::string scratch; + for (size_t i = 0; i < query_info.term_infos.size(); ++i) { + const auto& term_info = query_info.term_infos[i]; + DORIS_CHECK(term_info.is_single_term()); + if (term_info.key_kind == segment_v2::TermKeyKind::kCommonGram) { + return Status::Error( + "CommonGrams query plan requires segment capability validation"); + } + const auto logical_term = std::string_view((*routed_terms)[i]); + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw && + segment_v2::inverted_index::legacy_raw_exact_requires_bypass(logical_term)) { + return Status::Error( + "SNII legacy raw term overlaps an internal term namespace"); + } + auto physical_term = segment_v2::inverted_index::try_encode_plain_term_view( + logical_term, version, &scratch); + if (!physical_term.has_value()) { + return std::move(physical_term.error()); + } + if (!physical_term->has_value()) { + *all_representable = false; + continue; + } + if (output_index != i) { + if (scratch.empty()) { + (*routed_terms)[output_index] = std::move((*routed_terms)[i]); + } else { + (*routed_terms)[output_index] = std::move(scratch); + } + } else if (!scratch.empty()) { + (*routed_terms)[i] = std::move(scratch); + } + ++output_index; + } + routed_terms->resize(output_index); + return Status::OK(); +} + +inline Status route_plain_query_terms(const reader::LogicalIndexReader& idx, + const std::vector& logical_terms, + std::vector* physical_terms, + bool* all_representable) { + DORIS_CHECK(physical_terms != nullptr); + DORIS_CHECK(all_representable != nullptr); + physical_terms->clear(); + physical_terms->reserve(logical_terms.size()); + *all_representable = true; + for (const std::string& logical_term : logical_terms) { + std::string physical_term; + bool representable = false; + RETURN_IF_ERROR(route_plain_query_term(idx, logical_term, &physical_term, &representable)); + if (representable) { + physical_terms->push_back(std::move(physical_term)); + } else { + *all_representable = false; + } + } + return Status::OK(); +} + +inline Status route_plain_enumeration_prefix(const reader::LogicalIndexReader& idx, + std::string_view logical_prefix, + std::string* physical_prefix, bool* representable) { + DORIS_CHECK(physical_prefix != nullptr); + DORIS_CHECK(representable != nullptr); + physical_prefix->clear(); + *representable = false; + + const auto version = plain_term_key_version(idx); + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw && + segment_v2::inverted_index::legacy_raw_prefix_requires_bypass(logical_prefix)) { + return Status::Error( + "SNII legacy raw expansion overlaps an internal term namespace"); + } + auto encoded = segment_v2::inverted_index::try_encode_plain_term(logical_prefix, version, + physical_prefix); + if (!encoded.has_value()) { + return std::move(encoded.error()); + } + *representable = encoded.value(); + return Status::OK(); +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/position_math.h b/be/src/storage/index/snii/query/internal/position_math.h new file mode 100644 index 00000000000000..6db2a0ee4b599b --- /dev/null +++ b/be/src/storage/index/snii/query/internal/position_math.h @@ -0,0 +1,47 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +namespace doris::snii::query::internal { + +inline bool build_position_offsets(size_t count, std::vector* out) { + if (count >= std::numeric_limits::max()) { + return false; + } + out->clear(); + out->reserve(count); + uint32_t offset = 0; + while (out->size() < count) { + out->push_back(offset); + ++offset; + } + return true; +} + +inline bool add_position_offset(uint32_t start, uint32_t offset, uint32_t* out) { + if (start > std::numeric_limits::max() - offset) return false; + *out = start + offset; + return true; +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/query_test_counters.h b/be/src/storage/index/snii/query/internal/query_test_counters.h new file mode 100644 index 00000000000000..5e6b8acfd84e5d --- /dev/null +++ b/be/src/storage/index/snii/query/internal/query_test_counters.h @@ -0,0 +1,110 @@ +// 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. + +#pragma once + +#include + +// Deterministic op-count seam for the phrase-query hot path (T24/G01). The +// counters let the phrase UTs assert routing/complexity directly: +// - expected_docids_build : how many times the multi-tail phrase-prefix branch +// materializes the expected-docid vector for a query. +// Hoisted out of the per-tail loop, so == 1 per query +// (was == tail_hits when rebuilt inside the loop). +// - anchor_iterations : total sparsest-anchor outer-enumeration size summed +// over candidate docs (== docs x min_span). Anchoring +// on the shortest per-doc span instead of the hardcoded +// phrase-position-0 span shrinks this when the leading +// exact term is high-frequency. +// - monotonic_position_scans +// : number of non-anchor spans whose phrase-position +// lookup uses a forward-only scan instead of one +// binary search per anchor position. +// - prefix_expected_doc_visits +// : total expected-doc iterations performed by the +// multi-tail phrase-prefix verification and result +// materialization path. +// - count_fastpath_hits : count-only (G02) answers produced from dict-entry +// df alone for a single term, with NO posting decode +// (count_query.cpp). +// - resolved_term_entry_copies / moves +// : DictEntry ownership transfers performed by the two +// plan_resolved_terms overloads. +// - resolved_term_payload_pointer_reuses +// : non-empty inline FRQ/PRX vectors whose data pointer +// survives an entry move into TermPlan. +// - phrase_position_epoch_cache_hits / misses +// : same-document PhrasePositionLoader lookups served +// from the plan span cache vs loaded from its cursor. +// +// The seam is active only under SNII_QUERY_TEST_COUNTERS, which is auto-enabled by +// the library-wide BE_TEST define (be/CMakeLists.txt `if (MAKE_TEST)`) used to +// build doris_be_test. Because BE_TEST is applied to the whole BE tree, both the +// phrase_query.cpp increments AND the test translation unit that reads them observe +// the SAME process-wide singleton (the inline function below has one instance +// across every including TU). In a release build BE_TEST is undefined, the struct +// and singleton do not exist, and SNII_QUERY_COUNT/SNII_QUERY_ADD expand to +// ((void)0): zero overhead and NO global mutable state on the production query +// path. +// +// CONCURRENCY: the singleton is intentionally unsynchronized. It is a +// single-threaded, test-only seam -- one phrase query at a time -- and is never +// touched on the production path. Do NOT read or write it from concurrent tests. +// Reset it between test cases with `query_test_counters() = {}`. +#if defined(BE_TEST) && !defined(SNII_QUERY_TEST_COUNTERS) +#define SNII_QUERY_TEST_COUNTERS +#endif + +#ifdef SNII_QUERY_TEST_COUNTERS + +namespace doris::snii::query::internal { + +struct QueryTestCounters { + uint64_t expected_docids_build = 0; + uint64_t anchor_iterations = 0; + uint64_t monotonic_position_scans = 0; + uint64_t prefix_expected_doc_visits = 0; + uint64_t count_fastpath_hits = 0; + uint64_t resolved_term_entry_copies = 0; + uint64_t resolved_term_entry_moves = 0; + uint64_t resolved_term_payload_pointer_reuses = 0; + uint64_t phrase_position_epoch_cache_hits = 0; + uint64_t phrase_position_epoch_cache_misses = 0; +}; + +// `inline` gives a single shared instance across all TUs that include this header +// (phrase_query.cpp and the test), so counter increments made in the library are +// visible to the test that reads them. +inline QueryTestCounters& query_test_counters() { + static QueryTestCounters counters; + return counters; +} + +} // namespace doris::snii::query::internal + +// NOLINTBEGIN(clang-diagnostic-unused-macros): expanded by phrase_query.cpp, not by this header's TU +#define SNII_QUERY_COUNT(field) (++::doris::snii::query::internal::query_test_counters().field) +#define SNII_QUERY_ADD(field, n) \ + (::doris::snii::query::internal::query_test_counters().field += (n)) +// NOLINTEND(clang-diagnostic-unused-macros) + +#else + +#define SNII_QUERY_COUNT(field) ((void)0) +#define SNII_QUERY_ADD(field, n) ((void)0) + +#endif // SNII_QUERY_TEST_COUNTERS diff --git a/be/src/storage/index/snii/query/internal/regex_prefix.h b/be/src/storage/index/snii/query/internal/regex_prefix.h new file mode 100644 index 00000000000000..8fa872766b1e1d --- /dev/null +++ b/be/src/storage/index/snii/query/internal/regex_prefix.h @@ -0,0 +1,37 @@ +// 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. + +#pragma once + +#include + +#include +#include + +namespace doris::snii::query::internal { + +// Computes the dictionary-enumeration prefix used to narrow regexp term +// expansion. For left-anchored ("^") patterns it derives a tight common prefix +// from RE2::PossibleMatchRange (e.g. "^(order)" -> "order", where a naive literal +// scan stops at '(' and yields ""); otherwise it falls back to a conservative +// literal-prefix scan. The returned prefix only bounds how many dictionary terms +// visit_prefix_terms enumerates -- final term acceptance is always decided by +// RE2::FullMatch -- so an over-wide prefix can only enumerate extra terms and +// never drops a real match. `re` must be the compiled pattern (re.ok()). +std::string regex_enum_prefix(std::string_view pattern, const re2::RE2& re); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/resolved_phrase_plan.h b/be/src/storage/index/snii/query/internal/resolved_phrase_plan.h new file mode 100644 index 00000000000000..9bfd439753fb5b --- /dev/null +++ b/be/src/storage/index/snii/query/internal/resolved_phrase_plan.h @@ -0,0 +1,76 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/prx_decode_stats.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query { +struct PhraseMatch; +} + +namespace doris::snii::query::internal { + +struct ResolvedPhrasePlan { + std::vector unique_terms; + std::vector phrase_plan_index; + std::vector position_offsets; + + [[nodiscard]] bool is_valid() const { + if (phrase_plan_index.size() != position_offsets.size() || + unique_terms.empty() != phrase_plan_index.empty()) { + return false; + } + if (phrase_plan_index.empty()) { + return true; + } + + std::vector referenced(unique_terms.size(), 0); + for (size_t i = 0; i < phrase_plan_index.size(); ++i) { + if ((i == 0 && position_offsets[i] != 0) || + (i != 0 && position_offsets[i - 1] >= position_offsets[i])) { + return false; + } + const size_t plan_index = phrase_plan_index[i]; + if (plan_index >= unique_terms.size()) { + return false; + } + referenced[plan_index] = 1; + } + return std::ranges::all_of(referenced, [](uint8_t used) { return used != 0; }); + } +}; + +// Executes an already-resolved exact phrase plan. Planning policy and term-key +// semantics stay above this boundary; the executor treats terms as opaque and +// consumes the resolved entries so inline posting payloads can move into the +// execution plans without another allocation/copy. +Status execute_resolved_phrase_plan(const reader::LogicalIndexReader& idx, + ResolvedPhrasePlan&& plan, std::vector* docids, + format::PrxDecodeContext* observer_context = nullptr, + std::vector* matches = nullptr, + const std::vector* candidate_prefilter = nullptr); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/term_expansion.h b/be/src/storage/index/snii/query/internal/term_expansion.h new file mode 100644 index 00000000000000..6da0b4c3f8ed3f --- /dev/null +++ b/be/src/storage/index/snii/query/internal/term_expansion.h @@ -0,0 +1,47 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace doris::snii::query::internal { + +using TermMatcher = std::function; + +// Enumerates logical plain terms while retaining each matching physical +// DictEntry for direct posting resolution. PrefixHit::term is decoded logical +// text; PrefixHit::entry.term remains the physical dictionary key. +Status visit_expanded_plain_terms(const reader::LogicalIndexReader& idx, + std::string_view enum_prefix, const TermMatcher& matches, + const reader::LogicalIndexReader::PrefixHitVisitor& visitor, + int32_t max_expansions = 0); + +// Enumerates dictionary terms from `enum_prefix`, filters them with `matches`, +// and emits the sorted docid union for matching entries. PrefixHit carries the +// DictEntry and block bases, so callers avoid a second lookup per expanded term. +Status emit_expanded_docid_union(const reader::LogicalIndexReader& idx, + std::string_view enum_prefix, const TermMatcher& matches, + DocIdSink* const sink, int32_t max_expansions = 0); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/internal/wildcard_matcher.h b/be/src/storage/index/snii/query/internal/wildcard_matcher.h new file mode 100644 index 00000000000000..78eec25097a997 --- /dev/null +++ b/be/src/storage/index/snii/query/internal/wildcard_matcher.h @@ -0,0 +1,81 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace doris::snii::query::internal { + +// Glob matcher with reusable scratch. '*' matches >=0 bytes, '?' matches exactly +// one byte, every other byte is literal; matching is anchored at both ends (full +// match). The matching result is bit-for-bit identical to the former per-call DP +// in wildcard_query.cpp: the only change is that the two DP rows are constructed +// once and reused (assign(), never reallocated once capacity is large enough) +// across every term in a single expansion. A whole-dictionary scan therefore +// performs O(1) heap allocations for scratch instead of O(2N) -- two small +// std::vector constructions per visited term. +// +// The allocator is templated only so deterministic allocation-counting tests can +// inject a CountingAllocator; production constructs WildcardMatcher<> (default +// std::allocator). The matcher is request-scoped (a stack local of the calling +// wildcard_query frame), holds no shared mutable state, and is not thread-safe by +// design: each query owns its own instance. +template > +class WildcardMatcher { +public: + explicit WildcardMatcher(std::string_view pattern) : pattern_(pattern) {} + + bool operator()(std::string_view text) { + const size_t n = text.size() + 1; + prev_.assign(n, 0); // reuses the buffer; no realloc once capacity >= n + curr_.assign(n, 0); + prev_[0] = 1; + for (char p : pattern_) { + std::fill(curr_.begin(), curr_.end(), 0); + if (p == '*') { + curr_[0] = prev_[0]; + for (size_t i = 1; i < n; ++i) { + curr_[i] = prev_[i] || curr_[i - 1]; + } + } else { + for (size_t i = 1; i < n; ++i) { + curr_[i] = prev_[i - 1] && (p == '?' || p == text[i - 1]); + } + } + prev_.swap(curr_); + } + return prev_[text.size()] != 0; + } + + // Test-only debug accessor: the production path never depends on it. Reports + // the larger of the two scratch-row capacities so perf tests can assert the + // buffer stops reallocating after warmup. + size_t scratch_capacity() const { return std::max(prev_.capacity(), curr_.capacity()); } + +private: + std::string_view pattern_; + std::vector prev_; + std::vector curr_; +}; + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/phrase_cost.cpp b/be/src/storage/index/snii/query/phrase_cost.cpp new file mode 100644 index 00000000000000..fec24172cc5b2e --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_cost.cpp @@ -0,0 +1,276 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +uint64_t plan_visible_posting_bytes(const format::DictEntry& entry, bool need_positions) { + unsigned __int128 bytes = entry.kind == format::DictEntryKind::kInline + ? entry.inline_dd_disk_len + : entry.frq_docs_len; + if (need_positions) { + bytes += entry.kind == format::DictEntryKind::kInline ? entry.prx_bytes.size() + : entry.prx_len; + } + return bytes > std::numeric_limits::max() ? std::numeric_limits::max() + : static_cast(bytes); +} + +segment_v2::inverted_index::CommonGramsPlanRawCost phrase_plan_raw_cost( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + const std::vector& resolved, const std::vector& found, + bool need_positions) { + segment_v2::inverted_index::CommonGramsPlanRawCost cost; + unsigned __int128 posting_bytes = 0; + for (const std::string& term : plan.unique_terms) { + const size_t batch_index = resolved_batch_index(batch_terms, term); + if (found[batch_index] != 0) { + posting_bytes += + plan_visible_posting_bytes(resolved[batch_index].entry, need_positions); + } + } + cost.posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + cost.estimated_candidate_df = std::numeric_limits::max(); + for (size_t plan_index : plan.phrase_plan_index) { + const size_t batch_index = resolved_batch_index(batch_terms, plan.unique_terms[plan_index]); + if (found[batch_index] != 0) { + cost.estimated_candidate_df = + std::min(cost.estimated_candidate_df, resolved[batch_index].entry.df); + } + } + cost.clause_count = static_cast(plan.phrase_plan_index.size()); + if (cost.clause_count == 0) { + cost.estimated_candidate_df = 0; + } + return cost; +} + +segment_v2::inverted_index::CommonGramsPlanRawCost alternative_clause_raw_cost( + const std::vector& terms, bool need_positions) { + segment_v2::inverted_index::CommonGramsPlanRawCost cost; + unsigned __int128 posting_bytes = 0; + unsigned __int128 candidate_df = 0; + for (const auto& term : terms) { + posting_bytes += plan_visible_posting_bytes(term.entry, need_positions); + candidate_df += term.entry.df; + } + cost.posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + cost.estimated_candidate_df = candidate_df > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(candidate_df); + cost.clause_count = 1; + return cost; +} + +void append_alternative_clause_cost( + const segment_v2::inverted_index::CommonGramsPlanRawCost& clause, + segment_v2::inverted_index::CommonGramsPlanRawCost* plan) { + const unsigned __int128 posting_bytes = + static_cast(plan->posting_bytes_or_df_sum) + + clause.posting_bytes_or_df_sum; + plan->posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + plan->estimated_candidate_df = + plan->clause_count == 0 + ? clause.estimated_candidate_df + : std::min(plan->estimated_candidate_df, clause.estimated_candidate_df); + ++plan->clause_count; +} + +segment_v2::inverted_index::CommonGramsPlanRawCost hybrid_verification_raw_cost( + const segment_v2::inverted_index::CommonGramsPlanRawCost& prefilter_cost, + const segment_v2::inverted_index::CommonGramsPlanRawCost& verification_cost) { + auto cost = prefilter_cost; + const unsigned __int128 posting_bytes = + static_cast(cost.posting_bytes_or_df_sum) + + verification_cost.posting_bytes_or_df_sum; + cost.posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + cost.estimated_candidate_df = + std::min(cost.estimated_candidate_df, verification_cost.estimated_candidate_df); + cost.clause_count = verification_cost.clause_count; + return cost; +} + +bool physical_phrase_plan_has_docs_only_term(const PhysicalPhrasePlan& plan, + const std::vector& batch_terms, + const std::vector& resolved) { + for (const std::string& term : plan.unique_terms) { + if (!entry_has_positions(resolved[resolved_batch_index(batch_terms, term)].entry)) { + return true; + } + } + return false; +} + +void append_physical_phrase_clause(const PhysicalPhrasePlan& source, size_t clause, + uint32_t position_offset, PhysicalPhrasePlan* target) { + DORIS_CHECK_LT(clause, source.phrase_plan_index.size()); + DORIS_CHECK_EQ(source.phrase_plan_index.size(), source.position_offsets.size()); + DORIS_CHECK_EQ(source.phrase_plan_index.size(), source.common_gram_clauses.size()); + const size_t source_term = source.phrase_plan_index[clause]; + DORIS_CHECK_LT(source_term, source.unique_terms.size()); + const std::string& physical_term = source.unique_terms[source_term]; + + const auto unique = std::ranges::find(target->unique_terms, physical_term); + if (unique == target->unique_terms.end()) { + target->phrase_plan_index.push_back(target->unique_terms.size()); + target->unique_terms.push_back(physical_term); + } else { + target->phrase_plan_index.push_back( + static_cast(unique - target->unique_terms.begin())); + } + target->position_offsets.push_back(position_offset); + target->common_gram_clauses.push_back(source.common_gram_clauses[clause]); +} + +HybridPrefixCostEstimate estimate_hybrid_prefix_plan_cost( + const HybridPrefixPlanArtifact& artifact, const std::vector& batch_terms, + const std::vector& resolved, const std::vector& found, + const std::vector& plain_tail_terms, uint32_t position_verify_factor) { + const HybridPositionedCover& plain_tail_cover = artifact.plain_tail_cover; + const bool has_leading_prefilter = + !plain_tail_cover.candidate_prefilter.phrase_plan_index.empty(); + const auto leading_prefilter_cost = + phrase_plan_raw_cost(plain_tail_cover.candidate_prefilter, batch_terms, resolved, found, + /*need_positions=*/false); + unsigned __int128 posting_bytes = leading_prefilter_cost.posting_bytes_or_df_sum; + unsigned __int128 candidate_df_sum = 0; + unsigned __int128 position_verify_work = 0; + uint32_t max_clause_count = 0; + + const auto append_branch = [&](const PhysicalPhrasePlan& verification, + const auto& tail_verification_cost, + const auto* candidate_filter_cost) { + auto verification_cost = phrase_plan_raw_cost(verification, batch_terms, resolved, found, + /*need_positions=*/true); + append_alternative_clause_cost(tail_verification_cost, &verification_cost); + posting_bytes += verification_cost.posting_bytes_or_df_sum; + uint64_t candidate_df = verification_cost.estimated_candidate_df; + if (has_leading_prefilter) { + candidate_df = std::min(candidate_df, leading_prefilter_cost.estimated_candidate_df); + } + if (candidate_filter_cost != nullptr) { + posting_bytes += candidate_filter_cost->posting_bytes_or_df_sum; + candidate_df = std::min(candidate_df, candidate_filter_cost->estimated_candidate_df); + } + candidate_df_sum += candidate_df; + position_verify_work += + static_cast(candidate_df) * verification_cost.clause_count; + max_clause_count = std::max(max_clause_count, verification_cost.clause_count); + }; + + if (!artifact.maps_tail_to_gram) { + DORIS_CHECK(has_leading_prefilter); + DORIS_CHECK(artifact.mapped_tail_split.positioned_indices.empty()); + DORIS_CHECK(artifact.mapped_tail_split.docs_only_indices.empty()); + append_branch( + plain_tail_cover.verification, + alternative_clause_raw_cost(plain_tail_terms, /*need_positions=*/true), + static_cast(nullptr)); + } else { + const HybridPrefixMappedTails& split = artifact.mapped_tail_split; + DORIS_CHECK(!split.positioned_indices.empty() || !split.docs_only_indices.empty()); + DORIS_CHECK(has_leading_prefilter || !split.docs_only_indices.empty()); + if (!split.positioned_indices.empty()) { + DORIS_CHECK(artifact.positioned_tail_verification.has_value()); + append_branch(*artifact.positioned_tail_verification, + alternative_clause_raw_cost(resolved, split.positioned_indices, + /*need_positions=*/true), + static_cast( + nullptr)); + } + if (!split.docs_only_indices.empty()) { + const auto docs_only_filter_cost = + alternative_clause_raw_cost(resolved, split.docs_only_indices, + /*need_positions=*/false); + append_branch(plain_tail_cover.verification, + alternative_clause_raw_cost(plain_tail_terms, split.docs_only_ordinals, + /*need_positions=*/true), + &docs_only_filter_cost); + } + } + + HybridPrefixCostEstimate result; + result.raw_cost.posting_bytes_or_df_sum = posting_bytes > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(posting_bytes); + result.raw_cost.estimated_candidate_df = candidate_df_sum > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(candidate_df_sum); + result.raw_cost.clause_count = max_clause_count; + const unsigned __int128 estimated_cost = + posting_bytes + position_verify_work * position_verify_factor; + result.estimated_cost = estimated_cost > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(estimated_cost); + return result; +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/phrase_emit.cpp b/be/src/storage/index/snii/query/phrase_emit.cpp new file mode 100644 index 00000000000000..74824de445de64 --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_emit.cpp @@ -0,0 +1,682 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +namespace { +class PhrasePositionLoader { +public: + PhrasePositionLoader(size_t plan_count, std::vector& srcs) + : cursors_(plan_count), plan_spans_(plan_count), loaded_epoch_(plan_count, 0) { + for (size_t i = 0; i < plan_count; ++i) { + cursors_[i].init(&srcs[i]); + } + } + + void begin_doc(uint32_t docid) { + docid_ = docid; + ++epoch_; + if (epoch_ == 0) { + std::ranges::fill(loaded_epoch_, 0); + epoch_ = 1; + } + } + + Status positions_for_phrase_pos(const std::vector& phrase_plan_index, size_t phrase_pos, + std::pair* out) { + const size_t plan_index = phrase_plan_index[phrase_pos]; + if (loaded_epoch_[plan_index] != epoch_) { + RETURN_IF_ERROR(cursors_[plan_index].seek(docid_)); + RETURN_IF_ERROR(cursors_[plan_index].positions(&plan_spans_[plan_index])); + loaded_epoch_[plan_index] = epoch_; + SNII_QUERY_COUNT(phrase_position_epoch_cache_misses); + } else { + SNII_QUERY_COUNT(phrase_position_epoch_cache_hits); + } + *out = plan_spans_[plan_index]; + return Status::OK(); + } + +private: + std::vector cursors_; + std::vector> plan_spans_; + std::vector loaded_epoch_; + uint32_t docid_ = 0; + uint32_t epoch_ = 0; +}; + +class PhraseMatchCollector { +public: + PhraseMatchCollector(std::vector* docids, std::vector* matches) + : docids_(docids), matches_(matches) { + DCHECK(docids_ != nullptr || matches_ != nullptr); + } + + bool needs_frequency() const { return matches_ != nullptr; } + + void emit(uint32_t docid, uint32_t frequency) { + DCHECK_GT(frequency, 0); + if (docids_ != nullptr) { + docids_->push_back(docid); + } + if (matches_ != nullptr) { + matches_->push_back({.docid = docid, .frequency = frequency}); + } + } + +private: + std::vector* docids_; + std::vector* matches_; +}; + +bool contains_two_term_phrase(std::pair left_span, + std::pair right_span, + uint32_t right_delta) { + const uint32_t* left = left_span.first; + const uint32_t* right = right_span.first; + if (left == left_span.second || right == right_span.second) { + return false; + } + const uint32_t max_start = std::numeric_limits::max() - right_delta; + if (left + 1 == left_span.second && right + 1 == right_span.second) { + return *left <= max_start && *right == *left + right_delta; + } + while (left != left_span.second && right != right_span.second) { + if (*left > max_start) { + return false; + } + const uint32_t want = *left + right_delta; + while (right != right_span.second && *right < want) { + ++right; + } + if (right == right_span.second) { + return false; + } + if (*right == want) { + return true; + } + ++left; + } + return false; +} + +size_t select_phrase_verification_pair(const std::vector& plans, + const std::vector& phrase_plan_index) { + size_t best_left = 0; + uint64_t best_score = std::numeric_limits::max(); + for (size_t left = 0; left + 1 < phrase_plan_index.size(); ++left) { + const uint64_t score = static_cast(plans[phrase_plan_index[left]].df) + + plans[phrase_plan_index[left + 1]].df; + if (score < best_score) { + best_score = score; + best_left = left; + } + } + return best_left; +} + +class TwoTermPhraseStartCursor { +public: + TwoTermPhraseStartCursor(std::pair left_span, + std::pair right_span, + uint32_t right_delta, uint32_t left_offset) + : left_(left_span.first), + left_end_(left_span.second), + right_(right_span.first), + right_end_(right_span.second), + right_delta_(right_delta), + left_offset_(left_offset), + max_left_(std::numeric_limits::max() - right_delta) {} + + bool next(uint32_t* start) { + DCHECK(start != nullptr); + while (left_ != left_end_ && right_ != right_end_) { + if (*left_ > max_left_) { + return false; + } + const uint32_t want = *left_ + right_delta_; + while (right_ != right_end_ && *right_ < want) { + ++right_; + } + if (right_ == right_end_) { + return false; + } + const uint32_t left_position = *left_++; + if (*right_ == want && left_position >= left_offset_) { + *start = left_position - left_offset_; + return true; + } + } + return false; + } + +private: + const uint32_t* left_; + const uint32_t* left_end_; + const uint32_t* right_; + const uint32_t* right_end_; + uint32_t right_delta_; + uint32_t left_offset_; + uint32_t max_left_; +}; + +uint32_t count_two_term_phrase(std::pair left_span, + std::pair right_span, + uint32_t right_delta) { + TwoTermPhraseStartCursor starts(left_span, right_span, right_delta, /*left_offset=*/0); + uint32_t frequency = 0; + uint32_t start = 0; + while (starts.next(&start)) { + DCHECK_NE(frequency, std::numeric_limits::max()); + ++frequency; + } + return frequency; +} + +Status emit_two_term_phrase_streaming(const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector& srcs, + const std::vector& candidates, + PhraseMatchCollector* collector) { + const size_t left_plan = phrase_plan_index[0]; + const size_t right_plan = phrase_plan_index[1]; + const uint32_t right_delta = position_offsets[1] - position_offsets[0]; + + if (left_plan == right_plan) { + PostingCursor cursor; + cursor.init(&srcs[left_plan]); + for (uint32_t expected_docid : candidates) { + uint32_t docid = 0; + std::pair span; + RETURN_IF_ERROR(cursor.next(&docid, &span)); + if (docid != expected_docid) { + return Status::Error( + "phrase_query: repeated-term cursor/docid mismatch"); + } + const uint32_t frequency = collector->needs_frequency() + ? count_two_term_phrase(span, span, right_delta) + : contains_two_term_phrase(span, span, right_delta); + if (frequency != 0) { + collector->emit(docid, frequency); + } + } + return Status::OK(); + } + + PostingCursor left_cursor; + PostingCursor right_cursor; + left_cursor.init(&srcs[left_plan]); + right_cursor.init(&srcs[right_plan]); + for (uint32_t expected_docid : candidates) { + uint32_t left_docid = 0; + uint32_t right_docid = 0; + std::pair left_span; + std::pair right_span; + RETURN_IF_ERROR(left_cursor.next(&left_docid, &left_span)); + RETURN_IF_ERROR(right_cursor.next(&right_docid, &right_span)); + if (left_docid != expected_docid || right_docid != expected_docid) { + return Status::Error( + "phrase_query: two-term cursor/docid mismatch"); + } + const uint32_t frequency = + collector->needs_frequency() + ? count_two_term_phrase(left_span, right_span, right_delta) + : contains_two_term_phrase(left_span, right_span, right_delta); + if (frequency != 0) { + collector->emit(expected_docid, frequency); + } + } + return Status::OK(); +} + +void emit_two_term_phrase_chunk_pair(const PosChunk& left, const PosChunk& right, + const PosChunkDecoder& left_decoder, + const PosChunkDecoder& right_decoder, uint32_t right_delta, + PhraseMatchCollector* collector) { + size_t li = static_cast(std::ranges::lower_bound(left.docids, right.docids.front()) - + left.docids.begin()); + size_t ri = static_cast(std::ranges::lower_bound(right.docids, left.docids.front()) - + right.docids.begin()); + while (li < left.docids.size() && ri < right.docids.size()) { + const uint32_t left_docid = left.docids[li]; + const uint32_t right_docid = right.docids[ri]; + if (left_docid < right_docid) { + ++li; + continue; + } + if (right_docid < left_docid) { + ++ri; + continue; + } + + const std::pair left_span = + left_decoder.positions_unchecked(li); + const std::pair right_span = + right_decoder.positions_unchecked(ri); + const uint32_t frequency = + collector->needs_frequency() + ? count_two_term_phrase(left_span, right_span, right_delta) + : contains_two_term_phrase(left_span, right_span, right_delta); + if (frequency != 0) { + collector->emit(left_docid, frequency); + } + ++li; + ++ri; + } +} + +Status emit_two_term_phrase_chunk_merge(const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector& srcs, + PhraseMatchCollector* collector) { + const size_t left_plan = phrase_plan_index[0]; + const size_t right_plan = phrase_plan_index[1]; + const uint32_t right_delta = position_offsets[1] - position_offsets[0]; + const PosSource& left_src = srcs[left_plan]; + const PosSource& right_src = srcs[right_plan]; + + PosChunkDecoder left_decoder(left_src.observer_context); + PosChunkDecoder right_decoder(right_src.observer_context); + auto decoded_left_chunk = static_cast(-1); + auto decoded_right_chunk = static_cast(-1); + size_t left_chunk = 0; + size_t right_chunk = 0; + while (left_chunk < left_src.chunks.size() && right_chunk < right_src.chunks.size()) { + const PosChunk& left = left_src.chunks[left_chunk]; + const PosChunk& right = right_src.chunks[right_chunk]; + if (left.docids.empty()) { + ++left_chunk; + continue; + } + if (right.docids.empty()) { + ++right_chunk; + continue; + } + if (left.docids.back() < right.docids.front()) { + ++left_chunk; + continue; + } + if (right.docids.back() < left.docids.front()) { + ++right_chunk; + continue; + } + + if (decoded_left_chunk != left_chunk) { + RETURN_IF_ERROR(left_decoder.decode(left)); + decoded_left_chunk = left_chunk; + } + if (decoded_right_chunk != right_chunk) { + RETURN_IF_ERROR(right_decoder.decode(right)); + decoded_right_chunk = right_chunk; + } + + emit_two_term_phrase_chunk_pair(left, right, left_decoder, right_decoder, right_delta, + collector); + + const uint32_t left_last = left.docids.back(); + const uint32_t right_last = right.docids.back(); + if (left_last <= right_last) { + ++left_chunk; + } + if (right_last <= left_last) { + ++right_chunk; + } + } + return Status::OK(); +} + +bool phrase_start_matches_all_terms( + uint32_t start, size_t phrase_len, size_t pair_left, size_t pair_right, + const std::vector& position_offsets, + const std::vector>& span) { + for (size_t t = 0; t < phrase_len; ++t) { + if (t == pair_left || t == pair_right) { + continue; + } + uint32_t want = 0; + if (!internal::add_position_offset(start, position_offsets[t], &want)) { + return false; + } + if (!std::binary_search(span[t].first, span[t].second, want)) { + return false; + } + } + return true; +} + +Status emit_single_term_phrase_streaming(const std::vector& phrase_plan_index, + std::vector& srcs, + const std::vector& candidates, + PhraseMatchCollector* collector) { + PhrasePositionLoader loader(srcs.size(), srcs); + for (uint32_t d : candidates) { + loader.begin_doc(d); + std::pair single_span; + RETURN_IF_ERROR(loader.positions_for_phrase_pos(phrase_plan_index, 0, &single_span)); + if (single_span.first != single_span.second) { + const auto span_size = static_cast(single_span.second - single_span.first); + DCHECK_LE(span_size, std::numeric_limits::max()); + collector->emit(d, collector->needs_frequency() ? static_cast(span_size) : 1); + } + } + return Status::OK(); +} + +Status emit_multi_term_phrase_streaming(const std::vector& plans, + const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector& srcs, + const std::vector& candidates, + PhraseMatchCollector* collector) { + const size_t phrase_len = phrase_plan_index.size(); + PhrasePositionLoader loader(plans.size(), srcs); + std::vector> span(phrase_len); + const size_t pair_left = select_phrase_verification_pair(plans, phrase_plan_index); + const size_t pair_right = pair_left + 1; + for (uint32_t d : candidates) { + loader.begin_doc(d); + std::pair left_span; + std::pair right_span; + RETURN_IF_ERROR(loader.positions_for_phrase_pos(phrase_plan_index, pair_left, &left_span)); + RETURN_IF_ERROR( + loader.positions_for_phrase_pos(phrase_plan_index, pair_right, &right_span)); + + // `starts` retains raw pointers into the selected pair while the remaining + // clause spans are loaded below. Every unique plan owns an independent + // PostingCursor/PosChunkDecoder in PhrasePositionLoader; repeated phrase + // positions map back to one plan and reuse that plan's epoch-cached span. + TwoTermPhraseStartCursor starts(left_span, right_span, + position_offsets[pair_right] - position_offsets[pair_left], + position_offsets[pair_left]); + uint32_t start = 0; + if (!starts.next(&start)) { + continue; + } + + span[pair_left] = left_span; + span[pair_right] = right_span; + for (size_t pp = 0; pp < phrase_len; ++pp) { + if (pp == pair_left || pp == pair_right) { + continue; + } + RETURN_IF_ERROR(loader.positions_for_phrase_pos(phrase_plan_index, pp, &span[pp])); + } + + uint32_t frequency = 0; + bool has_previous_start = false; + uint32_t previous_start = 0; + const uint32_t* first_clause_position = span[0].first; + while (true) { + if (!collector->needs_frequency()) { + if (phrase_start_matches_all_terms(start, phrase_len, pair_left, pair_right, + position_offsets, span)) { + collector->emit(d, 1); + break; + } + } else if (!has_previous_start || start != previous_start) { + has_previous_start = true; + previous_start = start; + if (phrase_start_matches_all_terms(start, phrase_len, pair_left, pair_right, + position_offsets, span)) { + uint32_t first_clause_want = 0; + const bool representable = internal::add_position_offset( + start, position_offsets[0], &first_clause_want); + DCHECK(representable); + while (first_clause_position != span[0].second && + *first_clause_position < first_clause_want) { + ++first_clause_position; + } + const uint32_t* run_end = first_clause_position; + while (run_end != span[0].second && *run_end == first_clause_want) { + ++run_end; + } + const auto multiplicity = + static_cast(run_end - first_clause_position); + DCHECK_NE(multiplicity, 0); + DCHECK_LE(frequency, std::numeric_limits::max() - multiplicity); + frequency += multiplicity; + first_clause_position = run_end; + } + } + if (!starts.next(&start)) { + break; + } + } + if (frequency != 0) { + collector->emit(d, frequency); + } + } + return Status::OK(); +} + +// Single streaming pass over the candidates: for each (ascending) candidate, +// gather positions lazily, and test the consecutive-phrase predicate +// (term[0]@p, term[1]@p+1, ...). Multi-term phrases first test the cheapest +// adjacent pair by df before decoding the remaining terms for that document. +// Cursors decode each retained chunk at most once and address positions by +// local index -- no per-candidate docid binary search, no full-candidate +// position materialization. Candidates are ascending so the emitted docids are +// already sorted. + +Status emit_phrase_streaming(const std::vector& plans, + const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector& srcs, const std::vector& candidates, + PhraseMatchCollector* collector) { + const size_t phrase_len = phrase_plan_index.size(); + if (phrase_len == 1) { + return emit_single_term_phrase_streaming(phrase_plan_index, srcs, candidates, collector); + } + if (phrase_len == 2) { + if (phrase_plan_index[0] != phrase_plan_index[1]) { + return emit_two_term_phrase_chunk_merge(phrase_plan_index, position_offsets, srcs, + collector); + } + return emit_two_term_phrase_streaming(phrase_plan_index, position_offsets, srcs, candidates, + collector); + } + return emit_multi_term_phrase_streaming(plans, phrase_plan_index, position_offsets, srcs, + candidates, collector); +} + +// candidate_prefilter (optional): an ascending docid set the phrase must ALSO +// lie in. When provided, the leading-term conjunction is intersected with it so +// only docs in the prefilter get their positions read. Docs outside the +// prefilter cannot contribute (the caller guarantees the final answer is a +// subset), so this is result-preserving while cutting the position decode -- +// used by phrase-prefix to restrict the huge leading-phrase candidate set to +// the docs that also carry some tail expansion. + +} // namespace +Status build_phrase_execution_state(const LogicalIndexReader& idx, io::BatchRangeFetcher* round1, + std::vector* plans, PhraseExecutionState* state, + const std::vector* candidate_prefilter, + format::PrxDecodeContext* observer_context, + PhraseCandidateMetric candidate_metric) { + if (round1->pending() > 0) { + RETURN_IF_ERROR(round1->fetch()); + } + RETURN_IF_ERROR(internal::open_preludes(*round1, plans, + /*need_positions=*/true)); + + state->owners.clear(); + state->candidates.clear(); + std::vector doc_sources; + if (candidate_prefilter != nullptr) { + if (candidate_prefilter->empty()) { + return Status::OK(); + } + RETURN_IF_ERROR(internal::filter_docids_by_conjunction( + idx, *round1, *plans, *candidate_prefilter, &state->candidates, &doc_sources)); + } else { + RETURN_IF_ERROR(internal::build_docid_only_conjunction(idx, *round1, *plans, + &state->candidates, &doc_sources)); + } + if (observer_context != nullptr && observer_context->query_stats != nullptr) { + if (candidate_metric == PhraseCandidateMetric::kExact) { + observer_context->query_stats->exact_candidate_docs += state->candidates.size(); + observer_context->query_stats->exact_candidate_visits += state->candidates.size(); + } else { + observer_context->query_stats->prefix_leading_candidate_docs += + state->candidates.size(); + } + } + if (state->candidates.empty()) { + return Status::OK(); + } + RETURN_IF_ERROR(build_position_sources_for_candidates(idx, *round1, *plans, &doc_sources, + state->candidates, &state->owners, + &state->srcs, observer_context)); + return Status::OK(); +} + +namespace { +Status execute_phrase_plans_at_offsets(const LogicalIndexReader& idx, io::BatchRangeFetcher* round1, + std::vector* plans, + const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector* docids, + format::PrxDecodeContext* observer_context, + std::vector* matches, + const std::vector* candidate_prefilter = nullptr) { + DCHECK_EQ(phrase_plan_index.size(), position_offsets.size()); + PhraseExecutionState state; + RETURN_IF_ERROR(build_phrase_execution_state(idx, round1, plans, &state, candidate_prefilter, + observer_context, PhraseCandidateMetric::kExact)); + if (state.candidates.empty()) { + return Status::OK(); + } + + PhraseVerifyTimer verify_timer(observer_context); + PhraseMatchCollector collector(docids, matches); + RETURN_IF_ERROR(emit_phrase_streaming(*plans, phrase_plan_index, position_offsets, state.srcs, + state.candidates, &collector)); + verify_timer.commit_success(); + return Status::OK(); +} + +} // namespace +Status execute_phrase_plans(const LogicalIndexReader& idx, io::BatchRangeFetcher* round1, + std::vector* plans, + const std::vector& phrase_plan_index, + std::vector* docids, + format::PrxDecodeContext* observer_context, + std::vector* matches) { + std::vector position_offsets; + if (!internal::build_position_offsets(phrase_plan_index.size(), &position_offsets)) { + return Status::Error( + "phrase_query: phrase length exceeds doc position range"); + } + return execute_phrase_plans_at_offsets(idx, round1, plans, phrase_plan_index, position_offsets, + docids, observer_context, matches); +} + +} // namespace doris::snii::query::phrase_impl + +namespace doris::snii::query { + +using namespace phrase_impl; // NOLINT(google-build-using-namespace): module-internal impl namespace + +Status internal::execute_resolved_phrase_plan(const LogicalIndexReader& idx, + internal::ResolvedPhrasePlan&& plan, + std::vector* docids, + format::PrxDecodeContext* observer_context, + std::vector* matches, + const std::vector* candidate_prefilter) { + if (docids == nullptr && matches == nullptr) { + return Status::Error("resolved_phrase_plan: null out"); + } + if (docids != nullptr) { + docids->clear(); + } + if (matches != nullptr) { + matches->clear(); + } + DORIS_CHECK(plan.is_valid()); + if (plan.phrase_plan_index.empty()) { + return Status::OK(); + } + + if (plan.phrase_plan_index.size() == 1) { + DORIS_CHECK(matches == nullptr); + const internal::ResolvedQueryTerm& term = plan.unique_terms[plan.phrase_plan_index.front()]; + RETURN_IF_ERROR(internal::read_docid_posting(idx, term.entry, term.frq_base, term.prx_base, + docids)); + if (candidate_prefilter != nullptr) { + *docids = internal::intersect_sorted(*docids, *candidate_prefilter); + } + return Status::OK(); + } + + io::BatchRangeFetcher round1(idx.reader()); + std::vector plans; + RETURN_IF_ERROR(internal::plan_resolved_terms(idx, std::move(plan.unique_terms), &round1, + &plans, + /*need_positions=*/false)); + return execute_phrase_plans_at_offsets(idx, &round1, &plans, plan.phrase_plan_index, + plan.position_offsets, docids, observer_context, matches, + candidate_prefilter); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/phrase_plan.cpp b/be/src/storage/index/snii/query/phrase_plan.cpp new file mode 100644 index 00000000000000..d9facfddcf9b8d --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_plan.cpp @@ -0,0 +1,600 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +bool apply_common_grams_plan_debug_override(bool cost_prefers_gram, + CommonGramsPlanDebugOverride debug_override) { + switch (debug_override) { + case CommonGramsPlanDebugOverride::kNone: + return cost_prefers_gram; + case CommonGramsPlanDebugOverride::kForcePlain: + return false; + case CommonGramsPlanDebugOverride::kForceCommonGrams: + return true; + } + DORIS_CHECK(false); + return cost_prefers_gram; +} + +size_t position_span_size(std::pair span) { + if (span.first == span.second) { + return 0; + } + DCHECK(span.first != nullptr); + DCHECK(span.second != nullptr); + return static_cast(span.second - span.first); +} + +bool should_use_monotonic_position_scan(std::pair anchor_span, + size_t checked_span_size, uint32_t anchor_offset, + uint32_t checked_offset) { + const uint64_t anchor_count = position_span_size(anchor_span); + const uint64_t binary_search_upper_bound = + anchor_count * (static_cast(std::bit_width(checked_span_size)) + 1); + const uint64_t monotonic_scan_upper_bound = checked_span_size + 2 * anchor_count + 2; + + // Require a 2x comparison margin before paying even the O(1) validity + // checks and adding the scan-path branches. This keeps low-TF spans on the + // simpler binary-search path while retaining the high-TF dense case. + if (2 * monotonic_scan_upper_bound > binary_search_upper_bound) { + return false; + } + + // Scanning is considered only when every anchor yields a representable + // phrase start and checked-term position. Endpoint checks are sufficient + // because anchor positions are sorted; invalid boundary shapes stay on the + // existing per-anchor path without extra binary searches in this gate. + if (*anchor_span.first < anchor_offset) { + return false; + } + if (checked_offset <= anchor_offset) { + return true; + } + const uint32_t offset_delta = checked_offset - anchor_offset; + return anchor_span.second[-1] <= std::numeric_limits::max() - offset_delta; +} + +bool has_common_grams_capability( + const LogicalIndexReader& idx, + const segment_v2::inverted_index::CommonGramsQueryIdentity* query_identity) { + const auto* metadata = idx.common_grams_metadata(); + if (metadata == nullptr || query_identity == nullptr) { + return false; + } + if (idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1) { + return segment_v2::inverted_index::is_common_grams_query_compatible( + *metadata, *query_identity, + segment_v2::inverted_index::CommonGramsCoverage::kMixed); + } + return segment_v2::inverted_index::is_common_grams_query_compatible(*metadata, *query_identity); +} + +bool entry_has_positions(const format::DictEntry& entry) { + return entry.kind == format::DictEntryKind::kInline ? !entry.prx_bytes.empty() + : entry.prx_len != 0; +} + +Status build_physical_phrase_plan_prefix(const LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& query_info, + size_t clause_count, bool allow_common_grams, + PhysicalPhrasePlan* plan, bool* all_representable) { + plan->unique_terms.clear(); + plan->phrase_plan_index.clear(); + plan->position_offsets.clear(); + plan->common_gram_clauses.clear(); + *all_representable = true; + DORIS_CHECK_LE(clause_count, query_info.term_infos.size()); + if (clause_count == 0) { + return Status::OK(); + } + + const int32_t first_position = query_info.term_infos.front().position; + DORIS_CHECK_LE(clause_count, static_cast(std::numeric_limits::max())); + plan->phrase_plan_index.reserve(clause_count); + plan->position_offsets.reserve(clause_count); + plan->common_gram_clauses.reserve(clause_count); + for (size_t i = 0; i < clause_count; ++i) { + const segment_v2::TermInfo& term_info = query_info.term_infos[i]; + DORIS_CHECK(term_info.is_single_term()); + DORIS_CHECK_EQ(static_cast(term_info.position), + static_cast(first_position) + static_cast(i)); + + std::string physical_term; + if (term_info.key_kind == segment_v2::TermKeyKind::kCommonGram) { + DORIS_CHECK(allow_common_grams); + DORIS_CHECK(std::string_view(term_info.get_single_term()) + .starts_with(segment_v2::inverted_index::CG_V1_MARKER)); + physical_term = term_info.get_single_term(); + } else { + bool representable = false; + RETURN_IF_ERROR(internal::route_plain_query_term(idx, term_info.get_single_term(), + &physical_term, &representable)); + if (!representable) { + *all_representable = false; + return Status::OK(); + } + } + + auto unique = std::ranges::find(plan->unique_terms, physical_term); + if (unique == plan->unique_terms.end()) { + plan->phrase_plan_index.push_back(plan->unique_terms.size()); + plan->unique_terms.push_back(std::move(physical_term)); + } else { + plan->phrase_plan_index.push_back( + static_cast(unique - plan->unique_terms.begin())); + } + plan->position_offsets.push_back(static_cast(i)); + plan->common_gram_clauses.push_back( + static_cast(term_info.key_kind == segment_v2::TermKeyKind::kCommonGram)); + } + return Status::OK(); +} + +Status build_physical_phrase_plan(const LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& query_info, + bool allow_common_grams, PhysicalPhrasePlan* plan, + bool* all_representable) { + return build_physical_phrase_plan_prefix(idx, query_info, query_info.term_infos.size(), + allow_common_grams, plan, all_representable); +} + +size_t resolved_batch_index(const std::vector& batch_terms, std::string_view term) { + const auto it = std::ranges::lower_bound(batch_terms, term); + DORIS_CHECK(it != batch_terms.end()); + DORIS_CHECK_EQ(*it, term); + return static_cast(it - batch_terms.begin()); +} + +bool all_plan_terms_present(const PhysicalPhrasePlan& plan, + const std::vector& batch_terms, + const std::vector& found) { + for (const std::string& term : plan.unique_terms) { + if (found[resolved_batch_index(batch_terms, term)] == 0) { + return false; + } + } + return true; +} + +internal::ResolvedPhrasePlan materialize_resolved_phrase_plan( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + std::vector* resolved) { + internal::ResolvedPhrasePlan result; + result.phrase_plan_index = plan.phrase_plan_index; + result.position_offsets = plan.position_offsets; + result.unique_terms.reserve(plan.unique_terms.size()); + for (const std::string& term : plan.unique_terms) { + result.unique_terms.push_back( + std::move((*resolved)[resolved_batch_index(batch_terms, term)])); + } + return result; +} + +internal::ResolvedPhrasePlan copy_resolved_phrase_plan( + const PhysicalPhrasePlan& plan, const std::vector& batch_terms, + const std::vector& resolved) { + internal::ResolvedPhrasePlan result; + result.phrase_plan_index = plan.phrase_plan_index; + result.position_offsets = plan.position_offsets; + result.unique_terms.reserve(plan.unique_terms.size()); + for (const std::string& term : plan.unique_terms) { + result.unique_terms.push_back(resolved[resolved_batch_index(batch_terms, term)]); + } + return result; +} + +namespace { +PhysicalPhrasePlan build_hybrid_positioned_verification( + const PhysicalPhrasePlan& plain_plan, const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, const std::vector& resolved, + bool tail_covers_last_plain_clause, PhysicalPhrasePlan* candidate_prefilter) { + const size_t original_clause_count = plain_plan.phrase_plan_index.size(); + DORIS_CHECK_GT(original_clause_count, 0); + DORIS_CHECK_EQ(plain_plan.position_offsets.size(), original_clause_count); + DORIS_CHECK_EQ(plain_plan.common_gram_clauses.size(), original_clause_count); + DORIS_CHECK_EQ(gram_plan.position_offsets.size(), gram_plan.phrase_plan_index.size()); + DORIS_CHECK_EQ(gram_plan.common_gram_clauses.size(), gram_plan.phrase_plan_index.size()); + + PhysicalPhrasePlan verification; + std::vector positioned_gram_at(original_clause_count, + gram_plan.phrase_plan_index.size()); + for (size_t clause = 0; clause < gram_plan.phrase_plan_index.size(); ++clause) { + if (gram_plan.common_gram_clauses[clause] == 0) { + continue; + } + const size_t gram_term = gram_plan.phrase_plan_index[clause]; + DORIS_CHECK_LT(gram_term, gram_plan.unique_terms.size()); + const size_t batch_index = + resolved_batch_index(batch_terms, gram_plan.unique_terms[gram_term]); + if (!entry_has_positions(resolved[batch_index].entry)) { + if (candidate_prefilter != nullptr) { + append_physical_phrase_clause(gram_plan, clause, gram_plan.position_offsets[clause], + candidate_prefilter); + } + continue; + } + + const size_t original_offset = gram_plan.position_offsets[clause]; + DORIS_CHECK_LT(original_offset + 1, original_clause_count); + DORIS_CHECK_EQ(positioned_gram_at[original_offset], gram_plan.phrase_plan_index.size()); + positioned_gram_at[original_offset] = clause; + } + // Start from every positioned gram edge, then remove an edge only when both + // endpoint tokens remain covered by another positioned edge. The left-to-right + // pass produces a minimum-clause edge cover while retaining grams in preference + // to adding their two plain components during the pass below. + std::vector positioned_coverage(original_clause_count, 0); + if (tail_covers_last_plain_clause) { + ++positioned_coverage.back(); + } + for (size_t original_offset = 0; original_offset < original_clause_count; ++original_offset) { + if (positioned_gram_at[original_offset] == gram_plan.phrase_plan_index.size()) { + continue; + } + ++positioned_coverage[original_offset]; + ++positioned_coverage[original_offset + 1]; + } + for (size_t original_offset = 0; original_offset < original_clause_count; ++original_offset) { + if (positioned_gram_at[original_offset] == gram_plan.phrase_plan_index.size() || + positioned_coverage[original_offset] <= 1 || + positioned_coverage[original_offset + 1] <= 1) { + continue; + } + positioned_gram_at[original_offset] = gram_plan.phrase_plan_index.size(); + --positioned_coverage[original_offset]; + --positioned_coverage[original_offset + 1]; + } + + for (size_t original_offset = 0; original_offset < original_clause_count; ++original_offset) { + const size_t gram_clause = positioned_gram_at[original_offset]; + if (gram_clause != gram_plan.phrase_plan_index.size()) { + append_physical_phrase_clause(gram_plan, gram_clause, + static_cast(original_offset), &verification); + continue; + } + if (positioned_coverage[original_offset] == 0) { + DORIS_CHECK_EQ(plain_plan.position_offsets[original_offset], original_offset); + append_physical_phrase_clause(plain_plan, original_offset, + static_cast(original_offset), &verification); + } + } + DORIS_CHECK(!verification.phrase_plan_index.empty() || + (tail_covers_last_plain_clause && original_clause_count == 1)); + return verification; +} + +HybridPositionedCover build_hybrid_positioned_cover(const PhysicalPhrasePlan& plain_plan, + const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, + const std::vector& resolved, + bool tail_covers_last_plain_clause) { + HybridPositionedCover result; + result.verification = build_hybrid_positioned_verification( + plain_plan, gram_plan, batch_terms, resolved, tail_covers_last_plain_clause, + &result.candidate_prefilter); + return result; +} + +} // namespace +HybridExactPlanArtifact build_hybrid_exact_plan_artifact( + const PhysicalPhrasePlan& plain_plan, const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, + const std::vector& resolved) { + HybridExactPlanArtifact artifact; + if (gram_plan.phrase_plan_index.size() > 1 && + physical_phrase_plan_has_docs_only_term(gram_plan, batch_terms, resolved)) { + artifact.positioned_cover.emplace( + build_hybrid_positioned_cover(plain_plan, gram_plan, batch_terms, resolved, + /*tail_covers_last_plain_clause=*/false)); + DORIS_CHECK(!artifact.positioned_cover->candidate_prefilter.phrase_plan_index.empty()); + } + return artifact; +} + +namespace { +Status build_physical_phrase_plan_candidates(const LogicalIndexReader& idx, + const PhysicalPhrasePlan& plan, + const std::vector& batch_terms, + const std::vector& resolved, + std::vector* candidates) { + std::vector candidate_terms; + candidate_terms.reserve(plan.unique_terms.size()); + for (const std::string& term : plan.unique_terms) { + candidate_terms.push_back(resolved[resolved_batch_index(batch_terms, term)]); + } + + io::BatchRangeFetcher round1(idx.reader()); + std::vector plans; + RETURN_IF_ERROR(internal::plan_resolved_terms(idx, std::move(candidate_terms), &round1, &plans, + /*need_positions=*/false)); + if (round1.pending() > 0) { + RETURN_IF_ERROR(round1.fetch()); + } + RETURN_IF_ERROR(internal::open_preludes(round1, &plans, /*need_positions=*/false)); + return internal::build_docid_only_conjunction(idx, round1, plans, candidates); +} + +HybridPrefixMappedTails split_hybrid_prefix_mapped_tails( + const std::vector& resolved, + const std::vector& mapped_tails) { + HybridPrefixMappedTails result; + for (const ResolvedMappedTail& tail : mapped_tails) { + DORIS_CHECK_LT(tail.batch_index, resolved.size()); + if (entry_has_positions(resolved[tail.batch_index].entry)) { + result.positioned_indices.push_back(tail.batch_index); + } else { + result.docs_only_indices.push_back(tail.batch_index); + result.docs_only_ordinals.push_back(tail.expansion_ordinal); + } + } + return result; +} + +} // namespace +std::optional try_build_hybrid_prefix_plan_artifact( + const PhysicalPhrasePlan& plain_leading, const PhysicalPhrasePlan& gram_leading, + const std::vector& batch_terms, const std::vector& resolved, + const std::vector& mapped_tails, bool maps_tail_to_gram) { + const bool requires_plain_verification = + physical_phrase_plan_has_docs_only_term(gram_leading, batch_terms, resolved) || + std::ranges::any_of(mapped_tails, [&](const ResolvedMappedTail& tail) { + DORIS_CHECK_LT(tail.batch_index, resolved.size()); + return !entry_has_positions(resolved[tail.batch_index].entry); + }); + if (!requires_plain_verification) { + return std::nullopt; + } + + DORIS_CHECK(!plain_leading.phrase_plan_index.empty()); + DORIS_CHECK_LE(plain_leading.phrase_plan_index.size(), + static_cast(std::numeric_limits::max())); + HybridPrefixPlanArtifact artifact; + artifact.plain_tail_cover = + build_hybrid_positioned_cover(plain_leading, gram_leading, batch_terms, resolved, + /*tail_covers_last_plain_clause=*/false); + artifact.plain_tail_position_offset = + static_cast(plain_leading.phrase_plan_index.size()); + artifact.maps_tail_to_gram = maps_tail_to_gram; + if (!maps_tail_to_gram) { + DORIS_CHECK(mapped_tails.empty()); + return artifact; + } + + DORIS_CHECK(!mapped_tails.empty()); + artifact.mapped_tail_split = split_hybrid_prefix_mapped_tails(resolved, mapped_tails); + if (!artifact.mapped_tail_split.positioned_indices.empty()) { + artifact.positioned_tail_verification.emplace(build_hybrid_positioned_verification( + plain_leading, gram_leading, batch_terms, resolved, + /*tail_covers_last_plain_clause=*/true, + /*candidate_prefilter=*/nullptr)); + } + return artifact; +} + +Status build_hybrid_leading_candidates(const LogicalIndexReader& idx, + const PhysicalPhrasePlan& candidate_prefilter, + const std::vector& batch_terms, + const std::vector& resolved, + HybridPrefixCandidateSet* candidates) { + candidates->active = !candidate_prefilter.phrase_plan_index.empty(); + candidates->docs.clear(); + if (!candidates->active) { + return Status::OK(); + } + return build_physical_phrase_plan_candidates(idx, candidate_prefilter, batch_terms, resolved, + &candidates->docs); +} + +namespace { +Status build_tail_candidates_within_leading(const LogicalIndexReader& idx, + const std::vector& resolved, + const std::vector& tail_indices, + const std::vector& leading_candidates, + std::vector* candidates) { + std::vector tail_terms; + tail_terms.reserve(tail_indices.size()); + for (size_t index : tail_indices) { + DORIS_CHECK_LT(index, resolved.size()); + tail_terms.push_back(resolved[index]); + } + + io::BatchRangeFetcher round1(idx.reader()); + std::vector tail_plans; + RETURN_IF_ERROR(internal::plan_resolved_terms(idx, std::move(tail_terms), &round1, &tail_plans, + /*need_positions=*/false)); + if (round1.pending() > 0) { + RETURN_IF_ERROR(round1.fetch()); + } + RETURN_IF_ERROR(internal::open_preludes(round1, &tail_plans, /*need_positions=*/false)); + + candidates->clear(); + std::vector one_tail_plan; + one_tail_plan.reserve(1); + for (auto& tail_plan : tail_plans) { + one_tail_plan.clear(); + one_tail_plan.push_back(std::move(tail_plan)); + std::vector tail_matches; + RETURN_IF_ERROR(internal::filter_docids_by_conjunction( + idx, round1, one_tail_plan, leading_candidates, &tail_matches, nullptr)); + internal::union_sorted_into(candidates, tail_matches); + if (candidates->size() == leading_candidates.size()) { + break; + } + } + return Status::OK(); +} + +} // namespace +Status build_hybrid_docs_only_tail_candidates(const LogicalIndexReader& idx, + const std::vector& resolved, + const std::vector& gram_tail_indices, + const HybridPrefixCandidateSet& leading_candidates, + std::vector* candidates) { + DORIS_CHECK(!gram_tail_indices.empty()); + unsigned __int128 tail_df_sum = 0; + for (size_t index : gram_tail_indices) { + DORIS_CHECK_LT(index, resolved.size()); + tail_df_sum += resolved[index].entry.df; + } + if (leading_candidates.active && + static_cast(leading_candidates.docs.size()) <= tail_df_sum) { + return build_tail_candidates_within_leading(idx, resolved, gram_tail_indices, + leading_candidates.docs, candidates); + } + + std::vector tail_postings; + tail_postings.reserve(gram_tail_indices.size()); + for (size_t index : gram_tail_indices) { + const auto& tail = resolved[index]; + tail_postings.push_back({tail.entry, tail.frq_base, tail.prx_base}); + } + std::vector tail_candidates; + RETURN_IF_ERROR(internal::build_docid_union(idx, tail_postings, &tail_candidates)); + *candidates = leading_candidates.active + ? internal::intersect_sorted(leading_candidates.docs, tail_candidates) + : std::move(tail_candidates); + return Status::OK(); +} + +Status execute_hybrid_exact_phrase_plan( + const LogicalIndexReader& idx, const PhysicalPhrasePlan& gram_plan, + const std::vector& batch_terms, const HybridExactPlanArtifact& artifact, + std::vector* resolved, std::vector* docids, + format::PrxDecodeContext* decode_context, bool* candidate_intersection_empty) { + DORIS_CHECK(idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1); + if (candidate_intersection_empty != nullptr) { + *candidate_intersection_empty = false; + } + if (!artifact.positioned_cover.has_value()) { + auto resolved_gram = materialize_resolved_phrase_plan(gram_plan, batch_terms, resolved); + return internal::execute_resolved_phrase_plan(idx, std::move(resolved_gram), docids, + decode_context); + } + + const HybridPositionedCover& hybrid_plan = *artifact.positioned_cover; + std::vector gram_candidates; + RETURN_IF_ERROR(build_physical_phrase_plan_candidates( + idx, hybrid_plan.candidate_prefilter, batch_terms, *resolved, &gram_candidates)); + if (gram_candidates.empty()) { + if (candidate_intersection_empty != nullptr) { + *candidate_intersection_empty = true; + } + return Status::OK(); + } + auto resolved_verification = + materialize_resolved_phrase_plan(hybrid_plan.verification, batch_terms, resolved); + return internal::execute_resolved_phrase_plan(idx, std::move(resolved_verification), docids, + decode_context, nullptr, &gram_candidates); +} + +void append_resolved_phrase_clause(ResolvedQueryTerm term, uint32_t position_offset, + internal::ResolvedPhrasePlan* plan) { + const auto unique = std::ranges::find(plan->unique_terms, term.entry.term, + [](const ResolvedQueryTerm& candidate) { + return std::string_view(candidate.entry.term); + }); + if (unique == plan->unique_terms.end()) { + plan->phrase_plan_index.push_back(plan->unique_terms.size()); + plan->unique_terms.push_back(std::move(term)); + } else { + plan->phrase_plan_index.push_back(static_cast(unique - plan->unique_terms.begin())); + } + plan->position_offsets.push_back(position_offset); +} + +internal::ResolvedPhrasePlan build_resolved_phrase_plan( + std::vector resolved_terms) { + internal::ResolvedPhrasePlan plan; + plan.unique_terms.reserve(resolved_terms.size()); + plan.phrase_plan_index.reserve(resolved_terms.size()); + plan.position_offsets.reserve(resolved_terms.size()); + for (size_t i = 0; i < resolved_terms.size(); ++i) { + DORIS_CHECK_LE(i, static_cast(std::numeric_limits::max())); + append_resolved_phrase_clause(std::move(resolved_terms[i]), static_cast(i), + &plan); + } + return plan; +} + +Status resolve_and_execute_physical_phrase_plan(const LogicalIndexReader& idx, + const PhysicalPhrasePlan& plan, + std::vector* docids, + format::PrxDecodeContext* decode_context, + CommonGramsPlanningTimer& planning_timer) { + std::vector batch_terms = plan.unique_terms; + std::ranges::sort(batch_terms); + std::vector resolved; + std::vector found; + RETURN_IF_ERROR(internal::resolve_query_terms_batch(idx, batch_terms, &resolved, &found)); + if (!all_plan_terms_present(plan, batch_terms, found)) { + return Status::OK(); + } + auto resolved_plan = materialize_resolved_phrase_plan(plan, batch_terms, &resolved); + planning_timer.finish(); + return internal::execute_resolved_phrase_plan(idx, std::move(resolved_plan), docids, + decode_context); +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/phrase_planned_query.cpp b/be/src/storage/index/snii/query/phrase_planned_query.cpp new file mode 100644 index 00000000000000..a436f3cd8d2f0d --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_planned_query.cpp @@ -0,0 +1,743 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +Status planned_exact_phrase_query_impl( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, format::PrxDecodeContext* decode_context, + ExactPhrasePlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + CommonGramsPlanDebugOverride debug_override) { + if (docids == nullptr) { + return Status::Error( + "planned_exact_phrase_query: null out"); + } + docids->clear(); + DORIS_CHECK(!plain_query_info.has_common_gram()); + auto* query_stats = decode_context == nullptr ? nullptr : decode_context->query_stats; + CommonGramsPlanningTimer planning_timer(query_stats); + if (query_stats != nullptr) { + ++query_stats->common_grams_candidate_queries; + } + + PhysicalPhrasePlan plain_plan; + bool plain_representable = false; + RETURN_IF_ERROR(build_physical_phrase_plan(idx, plain_query_info, + /*allow_common_grams=*/false, &plain_plan, + &plain_representable)); + if (!plain_representable || plain_plan.phrase_plan_index.empty()) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + } + if (selected_plan != nullptr) { + *selected_plan = ExactPhrasePlanKind::kPlain; + } + return Status::OK(); + } + + const bool index_has_common_grams = has_common_grams_capability(idx, common_grams_identity); + const bool gram_capable = gram_query_info.has_common_gram() && index_has_common_grams; + if (!gram_capable) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + if (gram_query_info.has_common_gram()) { + ++query_stats->common_grams_fallback_incompatible; + } else { + ++query_stats->common_grams_fallback_no_gram; + } + } + if (selected_plan != nullptr) { + *selected_plan = ExactPhrasePlanKind::kPlain; + } + return resolve_and_execute_physical_phrase_plan(idx, plain_plan, docids, decode_context, + planning_timer); + } + + PhysicalPhrasePlan gram_plan; + bool gram_representable = false; + RETURN_IF_ERROR(build_physical_phrase_plan(idx, gram_query_info, + /*allow_common_grams=*/true, &gram_plan, + &gram_representable)); + DORIS_CHECK(gram_representable); + DORIS_CHECK(!gram_plan.phrase_plan_index.empty()); + + std::vector batch_terms = plain_plan.unique_terms; + batch_terms.insert(batch_terms.end(), gram_plan.unique_terms.begin(), + gram_plan.unique_terms.end()); + std::ranges::sort(batch_terms); + batch_terms.erase(std::unique(batch_terms.begin(), batch_terms.end()), batch_terms.end()); + + std::vector resolved; + std::vector found; + RETURN_IF_ERROR(internal::resolve_query_terms_batch(idx, batch_terms, &resolved, &found)); + const bool plain_present = all_plan_terms_present(plain_plan, batch_terms, found); + const bool gram_present = all_plan_terms_present(gram_plan, batch_terms, found); + if (!plain_present || !gram_present) { + if (query_stats != nullptr) { + if (plain_present) { + ++query_stats->common_grams_gram_plans; + ++query_stats->common_grams_authoritative_empty; + } else { + ++query_stats->common_grams_plain_plans; + ++query_stats->common_grams_authoritative_empty; + } + } + if (selected_plan != nullptr) { + *selected_plan = + plain_present ? ExactPhrasePlanKind::kCommonGrams : ExactPhrasePlanKind::kPlain; + } + return Status::OK(); + } + + std::optional hybrid_artifact; + if (idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1) { + hybrid_artifact.emplace( + build_hybrid_exact_plan_artifact(plain_plan, gram_plan, batch_terms, resolved)); + } + const bool plain_needs_positions = plain_plan.phrase_plan_index.size() > 1; + const bool hybrid_gram_verification = + hybrid_artifact.has_value() && hybrid_artifact->positioned_cover.has_value(); + const bool gram_needs_positions = + gram_plan.phrase_plan_index.size() > 1 && !hybrid_gram_verification; + const auto plain_raw_cost = + phrase_plan_raw_cost(plain_plan, batch_terms, resolved, found, plain_needs_positions); + segment_v2::inverted_index::CommonGramsPlanRawCost gram_raw_cost; + if (hybrid_gram_verification) { + const HybridPositionedCover& hybrid_plan = *hybrid_artifact->positioned_cover; + const auto gram_prefilter_raw_cost = + phrase_plan_raw_cost(hybrid_plan.candidate_prefilter, batch_terms, resolved, found, + /*need_positions=*/false); + const auto gram_verification_raw_cost = + phrase_plan_raw_cost(hybrid_plan.verification, batch_terms, resolved, found, + /*need_positions=*/true); + gram_raw_cost = + hybrid_verification_raw_cost(gram_prefilter_raw_cost, gram_verification_raw_cost); + } else { + gram_raw_cost = + phrase_plan_raw_cost(gram_plan, batch_terms, resolved, found, gram_needs_positions); + } + const uint64_t plain_cost = segment_v2::inverted_index::estimate_common_grams_plan_cost( + plain_raw_cost, plain_needs_positions ? cost_model.position_verify_factor : 0); + const uint64_t gram_cost = segment_v2::inverted_index::estimate_common_grams_plan_cost( + gram_raw_cost, (gram_needs_positions || hybrid_gram_verification) + ? cost_model.position_verify_factor + : 0); + const bool cost_prefers_gram = segment_v2::inverted_index::common_grams_plan_cost_wins( + plain_cost, gram_cost, cost_model.common_grams_cost_ratio_percent); + const bool use_gram = apply_common_grams_plan_debug_override(cost_prefers_gram, debug_override); + if (query_stats != nullptr) { + query_stats->common_grams_plain_posting_bytes += plain_raw_cost.posting_bytes_or_df_sum; + query_stats->common_grams_gram_posting_bytes += gram_raw_cost.posting_bytes_or_df_sum; + query_stats->common_grams_plain_estimated_candidate_df += + plain_raw_cost.estimated_candidate_df; + query_stats->common_grams_gram_estimated_candidate_df += + gram_raw_cost.estimated_candidate_df; + query_stats->common_grams_plain_estimated_cost += plain_cost; + query_stats->common_grams_gram_estimated_cost += gram_cost; + } + const ExactPhrasePlanKind chosen_kind = + use_gram ? ExactPhrasePlanKind::kCommonGrams : ExactPhrasePlanKind::kPlain; + if (query_stats != nullptr) { + if (use_gram) { + ++query_stats->common_grams_gram_plans; + } else { + ++query_stats->common_grams_plain_plans; + if (debug_override == CommonGramsPlanDebugOverride::kNone) { + ++query_stats->common_grams_fallback_cost; + } + } + } + if (selected_plan != nullptr) { + *selected_plan = chosen_kind; + } + planning_timer.finish(); + if (use_gram && + idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1) { + DORIS_CHECK(hybrid_artifact.has_value()); + bool candidate_intersection_empty = false; + RETURN_IF_ERROR(execute_hybrid_exact_phrase_plan( + idx, gram_plan, batch_terms, *hybrid_artifact, &resolved, docids, decode_context, + &candidate_intersection_empty)); + if (candidate_intersection_empty) { + if (query_stats != nullptr) { + ++query_stats->common_grams_authoritative_empty; + } + } + return Status::OK(); + } + auto resolved_plan = materialize_resolved_phrase_plan(use_gram ? gram_plan : plain_plan, + batch_terms, &resolved); + return internal::execute_resolved_phrase_plan(idx, std::move(resolved_plan), docids, + decode_context); +} + +Status phrase_query_impl(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids, + format::PrxDecodeContext* decode_context, + std::vector* matches) { + if (docids == nullptr && matches == nullptr) { + return Status::Error("phrase_query: null out"); + } + if (docids != nullptr) { + docids->clear(); + } + if (matches != nullptr) { + matches->clear(); + } + if (terms.empty()) { + return Status::OK(); + } + if (terms.size() == 1) { + DORIS_CHECK(matches == nullptr); + return term_query(idx, terms.front(), docids); + } + if (!idx.has_positions()) { + return Status::Error( + "phrase_query: index has no positions"); + } + io::BatchRangeFetcher round1(idx.reader()); + const PhraseTermMapping mapping = build_phrase_term_mapping(terms); + std::vector plans; + bool all_present = false; + RETURN_IF_ERROR(internal::plan_terms(idx, mapping.unique_terms, &round1, &plans, &all_present, + /*need_positions=*/false)); + if (!all_present) { + return Status::OK(); + } + return execute_phrase_plans(idx, &round1, &plans, mapping.phrase_plan_index, docids, + decode_context, matches); +} + +Status phrase_prefix_query_impl(const LogicalIndexReader& idx, + const std::vector& terms, + std::vector* const docids, int32_t max_expansions, + format::PrxDecodeContext* decode_context, + CommonGramsPlanningTimer* planning_timer, + std::vector* matches) { + if (docids == nullptr && matches == nullptr) { + return Status::Error("phrase_prefix_query: null out"); + } + if (docids != nullptr) { + docids->clear(); + } + if (matches != nullptr) { + matches->clear(); + } + if (terms.empty()) { + return Status::OK(); + } + if (terms.size() == 1) { + DORIS_CHECK(matches == nullptr); + if (planning_timer != nullptr) { + planning_timer->finish(); + } + return prefix_query(idx, terms.front(), docids, max_expansions); + } + std::vector exact_terms; + exact_terms.reserve(terms.size() - 1); + std::string physical_term_scratch; + for (size_t i = 0; i + 1 < terms.size(); ++i) { + std::string_view physical_term; + bool representable = false; + RETURN_IF_ERROR(internal::route_plain_query_term_view(idx, terms[i], &physical_term_scratch, + &physical_term, &representable)); + if (!representable) { + return Status::OK(); + } + ResolvedQueryTerm resolved; + bool found = false; + RETURN_IF_ERROR(internal::resolve_query_term(idx, physical_term, &resolved, &found)); + if (!found) { + return Status::OK(); + } + exact_terms.push_back(std::move(resolved)); + } + + // Expand the tail in the logical plain namespace. The visitor range-seeks + // past typed internal namespaces before counting max_expansions and decodes + // escaped physical keys before applying the logical prefix. + std::vector tail_hits; + RETURN_IF_ERROR(internal::visit_expanded_plain_terms( + idx, terms.back(), [](std::string_view) { return true; }, + [&](LogicalIndexReader::PrefixHit&& hit, bool*) { + tail_hits.push_back(std::move(hit)); + return Status::OK(); + }, + max_expansions)); + if (tail_hits.empty()) { + return Status::OK(); + } + std::vector tail_terms; + tail_terms.reserve(tail_hits.size()); + for (auto& hit : tail_hits) { + tail_terms.push_back(ResolvedQueryTerm { + .entry = std::move(hit.entry), .frq_base = hit.frq_base, .prx_base = hit.prx_base}); + } + auto exact_plan = build_resolved_phrase_plan(std::move(exact_terms)); + if (planning_timer != nullptr) { + planning_timer->finish(); + } + DORIS_CHECK_LE(terms.size() - 1, static_cast(std::numeric_limits::max())); + return execute_resolved_phrase_prefix_terms(idx, std::move(exact_plan), std::move(tail_terms), + static_cast(terms.size() - 1), docids, + decode_context, matches); +} + +Status planned_phrase_prefix_query_impl( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, int32_t max_expansions, + format::PrxDecodeContext* decode_context, PhrasePrefixPlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + CommonGramsPlanDebugOverride debug_override) { + if (docids == nullptr) { + return Status::Error( + "planned_phrase_prefix_query: null out"); + } + docids->clear(); + if (selected_plan != nullptr) { + *selected_plan = PhrasePrefixPlanKind::kPlain; + } + DORIS_CHECK(!plain_query_info.has_common_gram()); + auto* query_stats = decode_context == nullptr ? nullptr : decode_context->query_stats; + CommonGramsPlanningTimer planning_timer(query_stats); + if (query_stats != nullptr) { + ++query_stats->common_grams_candidate_queries; + } + if (plain_query_info.term_infos.empty()) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + } + return Status::OK(); + } + for (const auto& term_info : plain_query_info.term_infos) { + DORIS_CHECK(term_info.is_single_term()); + } + + PhysicalPhrasePlan plain_leading; + bool plain_representable = false; + RETURN_IF_ERROR(build_physical_phrase_plan_prefix( + idx, plain_query_info, plain_query_info.term_infos.size() - 1, + /*allow_common_grams=*/false, &plain_leading, &plain_representable)); + if (!plain_representable) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + } + return Status::OK(); + } + + const auto execute_plain = [&]() { + std::vector terms; + terms.reserve(plain_query_info.term_infos.size()); + for (const auto& term_info : plain_query_info.term_infos) { + terms.push_back(term_info.get_single_term()); + } + return phrase_prefix_query_impl(idx, terms, docids, max_expansions, decode_context, + &planning_timer); + }; + + const bool index_has_common_grams = has_common_grams_capability(idx, common_grams_identity); + + const bool can_plan_common_grams = gram_query_info.has_common_gram() && index_has_common_grams; + if (!can_plan_common_grams) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + if (gram_query_info.has_common_gram()) { + ++query_stats->common_grams_fallback_incompatible; + } else { + ++query_stats->common_grams_fallback_no_gram; + } + } + return execute_plain(); + } + DORIS_CHECK(!gram_query_info.term_infos.empty()); + DORIS_CHECK(gram_query_info.term_infos.back().is_single_term()); + + PhysicalPhrasePlan gram_leading; + bool gram_representable = false; + RETURN_IF_ERROR(build_physical_phrase_plan_prefix( + idx, gram_query_info, gram_query_info.term_infos.size() - 1, + /*allow_common_grams=*/true, &gram_leading, &gram_representable)); + if (!gram_representable) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + ++query_stats->common_grams_fallback_incompatible; + } + return execute_plain(); + } + + internal::ResolvedPhrasePlan selected_leading; + std::vector selected_tail_terms; + DORIS_CHECK_LE(plain_leading.phrase_plan_index.size(), + static_cast(std::numeric_limits::max())); + const uint32_t plain_tail_position_offset = + static_cast(plain_leading.phrase_plan_index.size()); + uint32_t selected_tail_position_offset = plain_tail_position_offset; + bool authoritative_empty = false; + bool execute_as_exact_phrase = false; + bool hybrid_executed = false; + const Status planning_status = [&]() -> Status { + const bool maps_tail_to_gram = + gram_query_info.term_infos.back().key_kind == segment_v2::TermKeyKind::kCommonGram; + std::vector logical_tail_terms; + std::vector plain_tail_terms; + std::vector tail_hits; + RETURN_IF_ERROR(internal::visit_expanded_plain_terms( + idx, plain_query_info.term_infos.back().get_single_term(), + [](std::string_view) { return true; }, + [&](LogicalIndexReader::PrefixHit&& hit, bool*) { + tail_hits.push_back(std::move(hit)); + return Status::OK(); + }, + max_expansions)); + if (tail_hits.empty()) { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + ++query_stats->common_grams_fallback_prefix_tail_empty; + ++query_stats->common_grams_authoritative_empty; + } + authoritative_empty = true; + return Status::OK(); + } + if (maps_tail_to_gram) { + logical_tail_terms.reserve(tail_hits.size()); + } else { + DORIS_CHECK(gram_query_info.term_infos.back().key_kind == + segment_v2::TermKeyKind::kPlain); + } + plain_tail_terms.reserve(tail_hits.size()); + for (auto& hit : tail_hits) { + if (maps_tail_to_gram) { + logical_tail_terms.push_back(std::move(hit.term)); + } + plain_tail_terms.push_back(ResolvedQueryTerm {.entry = std::move(hit.entry), + .frq_base = hit.frq_base, + .prx_base = hit.prx_base}); + } + + const auto select_plain_plan = [&]() -> Status { + if (query_stats != nullptr) { + ++query_stats->common_grams_plain_plans; + ++query_stats->common_grams_fallback_incompatible; + } + std::vector batch_terms = plain_leading.unique_terms; + std::ranges::sort(batch_terms); + std::vector resolved; + std::vector found; + RETURN_IF_ERROR( + internal::resolve_query_terms_batch(idx, batch_terms, &resolved, &found)); + if (!all_plan_terms_present(plain_leading, batch_terms, found)) { + if (query_stats != nullptr) { + ++query_stats->common_grams_authoritative_empty; + } + authoritative_empty = true; + return Status::OK(); + } + selected_leading = + materialize_resolved_phrase_plan(plain_leading, batch_terms, &resolved); + selected_tail_terms = std::move(plain_tail_terms); + selected_tail_position_offset = plain_tail_position_offset; + return Status::OK(); + }; + + std::vector mapped_tail_terms; + if (maps_tail_to_gram) { + DORIS_CHECK_GE(plain_query_info.term_infos.size(), 2U); + const std::string& left = + plain_query_info.term_infos[plain_query_info.term_infos.size() - 2] + .get_single_term(); + mapped_tail_terms.reserve(logical_tail_terms.size()); + for (const std::string& tail : logical_tail_terms) { + std::string gram; + auto encoded = + segment_v2::inverted_index::try_encode_common_gram(left, tail, &gram); + if (!encoded.has_value()) { + return std::move(encoded.error()); + } + if (!*encoded) { + return select_plain_plan(); + } + mapped_tail_terms.push_back(std::move(gram)); + } + } + + std::vector batch_terms = plain_leading.unique_terms; + batch_terms.insert(batch_terms.end(), gram_leading.unique_terms.begin(), + gram_leading.unique_terms.end()); + batch_terms.insert(batch_terms.end(), mapped_tail_terms.begin(), mapped_tail_terms.end()); + std::ranges::sort(batch_terms); + batch_terms.erase(std::unique(batch_terms.begin(), batch_terms.end()), batch_terms.end()); + + std::vector resolved; + std::vector found; + RETURN_IF_ERROR(internal::resolve_query_terms_batch(idx, batch_terms, &resolved, &found)); + const bool plain_present = all_plan_terms_present(plain_leading, batch_terms, found); + const bool gram_present = all_plan_terms_present(gram_leading, batch_terms, found); + if (!plain_present || !gram_present) { + if (query_stats != nullptr) { + if (plain_present) { + ++query_stats->common_grams_gram_plans; + ++query_stats->common_grams_authoritative_empty; + } else { + ++query_stats->common_grams_plain_plans; + ++query_stats->common_grams_authoritative_empty; + } + } + if (selected_plan != nullptr && plain_present) { + *selected_plan = PhrasePrefixPlanKind::kCommonGrams; + } + authoritative_empty = true; + return Status::OK(); + } + + std::vector present_mapped_tail_indices; + std::vector present_mapped_tails; + std::vector present_gram_tail_ordinals; + if (maps_tail_to_gram) { + present_mapped_tail_indices.reserve(mapped_tail_terms.size()); + present_mapped_tails.reserve(mapped_tail_terms.size()); + present_gram_tail_ordinals.reserve(mapped_tail_terms.size()); + for (size_t ordinal = 0; ordinal < mapped_tail_terms.size(); ++ordinal) { + const std::string& term = mapped_tail_terms[ordinal]; + const size_t batch_index = resolved_batch_index(batch_terms, term); + if (found[batch_index] != 0) { + present_mapped_tail_indices.push_back(batch_index); + DORIS_CHECK_LE(ordinal, + static_cast(std::numeric_limits::max())); + const uint32_t expansion_ordinal = static_cast(ordinal); + present_mapped_tails.push_back(ResolvedMappedTail { + .batch_index = batch_index, .expansion_ordinal = expansion_ordinal}); + present_gram_tail_ordinals.push_back(expansion_ordinal); + } + } + if (present_mapped_tail_indices.empty()) { + if (query_stats != nullptr) { + ++query_stats->common_grams_gram_plans; + ++query_stats->common_grams_authoritative_empty; + } + if (selected_plan != nullptr) { + *selected_plan = PhrasePrefixPlanKind::kCommonGrams; + } + authoritative_empty = true; + return Status::OK(); + } + } + + std::optional hybrid_artifact; + if (idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1) { + hybrid_artifact = try_build_hybrid_prefix_plan_artifact( + plain_leading, gram_leading, batch_terms, resolved, present_mapped_tails, + maps_tail_to_gram); + } + const bool hybrid_plan_requires_plain_verification = hybrid_artifact.has_value(); + const bool plain_needs_positions = !plain_leading.phrase_plan_index.empty(); + const bool gram_needs_positions = + !gram_leading.phrase_plan_index.empty() && !hybrid_plan_requires_plain_verification; + auto plain_raw_cost = phrase_plan_raw_cost(plain_leading, batch_terms, resolved, found, + plain_needs_positions); + const auto plain_tail_raw_cost = + alternative_clause_raw_cost(plain_tail_terms, plain_needs_positions); + append_alternative_clause_cost(plain_tail_raw_cost, &plain_raw_cost); + const uint64_t plain_cost = segment_v2::inverted_index::estimate_common_grams_plan_cost( + plain_raw_cost, plain_needs_positions ? cost_model.position_verify_factor : 0); + segment_v2::inverted_index::CommonGramsPlanRawCost gram_raw_cost; + uint64_t gram_cost = 0; + if (hybrid_plan_requires_plain_verification) { + const HybridPrefixCostEstimate hybrid_cost = estimate_hybrid_prefix_plan_cost( + *hybrid_artifact, batch_terms, resolved, found, plain_tail_terms, + cost_model.position_verify_factor); + gram_raw_cost = hybrid_cost.raw_cost; + gram_cost = hybrid_cost.estimated_cost; + } else { + gram_raw_cost = phrase_plan_raw_cost(gram_leading, batch_terms, resolved, found, + gram_needs_positions); + if (maps_tail_to_gram) { + append_alternative_clause_cost( + alternative_clause_raw_cost(resolved, present_mapped_tail_indices, + gram_needs_positions), + &gram_raw_cost); + } else { + append_alternative_clause_cost(plain_tail_raw_cost, &gram_raw_cost); + } + gram_cost = segment_v2::inverted_index::estimate_common_grams_plan_cost( + gram_raw_cost, gram_needs_positions ? cost_model.position_verify_factor : 0); + } + const bool cost_prefers_gram = segment_v2::inverted_index::common_grams_plan_cost_wins( + plain_cost, gram_cost, cost_model.common_grams_cost_ratio_percent); + const bool use_gram = + apply_common_grams_plan_debug_override(cost_prefers_gram, debug_override); + if (query_stats != nullptr) { + query_stats->common_grams_plain_posting_bytes += plain_raw_cost.posting_bytes_or_df_sum; + query_stats->common_grams_gram_posting_bytes += gram_raw_cost.posting_bytes_or_df_sum; + query_stats->common_grams_plain_estimated_candidate_df += + plain_raw_cost.estimated_candidate_df; + query_stats->common_grams_gram_estimated_candidate_df += + gram_raw_cost.estimated_candidate_df; + query_stats->common_grams_plain_estimated_cost += plain_cost; + query_stats->common_grams_gram_estimated_cost += gram_cost; + } + if (query_stats != nullptr) { + if (use_gram) { + ++query_stats->common_grams_gram_plans; + } else { + ++query_stats->common_grams_plain_plans; + if (debug_override == CommonGramsPlanDebugOverride::kNone) { + ++query_stats->common_grams_fallback_cost; + } + } + } + if (selected_plan != nullptr) { + *selected_plan = + use_gram ? PhrasePrefixPlanKind::kCommonGrams : PhrasePrefixPlanKind::kPlain; + } + + const bool hybrid_needs_plain_verification = + use_gram && hybrid_plan_requires_plain_verification; + if (hybrid_needs_plain_verification) { + DORIS_CHECK(hybrid_artifact.has_value()); + bool candidate_intersection_empty = false; + RETURN_IF_ERROR(execute_hybrid_phrase_prefix_plan( + idx, *hybrid_artifact, batch_terms, resolved, plain_tail_terms, docids, + decode_context, planning_timer, &candidate_intersection_empty)); + hybrid_executed = true; + if (candidate_intersection_empty) { + if (query_stats != nullptr) { + ++query_stats->common_grams_authoritative_empty; + } + authoritative_empty = true; + return Status::OK(); + } + return Status::OK(); + } + + const PhysicalPhrasePlan& selected_physical_leading = + use_gram ? gram_leading : plain_leading; + selected_leading = + materialize_resolved_phrase_plan(selected_physical_leading, batch_terms, &resolved); + if (!use_gram || !maps_tail_to_gram) { + selected_tail_terms = std::move(plain_tail_terms); + selected_tail_position_offset = plain_tail_position_offset; + return Status::OK(); + } + + if (present_mapped_tail_indices.size() == 1) { + DORIS_CHECK_GT(plain_tail_position_offset, 0U); + const uint32_t mapped_tail_position = plain_tail_position_offset - 1; + const size_t batch_index = present_mapped_tail_indices.front(); + const auto existing = + std::ranges::find(selected_leading.unique_terms, batch_terms[batch_index], + [](const ResolvedQueryTerm& term) { + return std::string_view(term.entry.term); + }); + if (existing == selected_leading.unique_terms.end()) { + append_resolved_phrase_clause(std::move(resolved[batch_index]), + mapped_tail_position, &selected_leading); + } else { + selected_leading.phrase_plan_index.push_back( + static_cast(existing - selected_leading.unique_terms.begin())); + selected_leading.position_offsets.push_back(mapped_tail_position); + } + execute_as_exact_phrase = true; + return Status::OK(); + } + + selected_tail_terms.reserve(present_mapped_tail_indices.size()); + for (size_t batch_index : present_mapped_tail_indices) { + const auto existing = + std::ranges::find(selected_leading.unique_terms, batch_terms[batch_index], + [](const ResolvedQueryTerm& term) { + return std::string_view(term.entry.term); + }); + if (existing == selected_leading.unique_terms.end()) { + selected_tail_terms.push_back(std::move(resolved[batch_index])); + } else { + selected_tail_terms.push_back(*existing); + } + } + DORIS_CHECK_GT(plain_tail_position_offset, 0U); + selected_tail_position_offset = plain_tail_position_offset - 1; + return Status::OK(); + }(); + RETURN_IF_ERROR(planning_status); + planning_timer.finish(); + if (authoritative_empty) { + return Status::OK(); + } + if (hybrid_executed) { + return Status::OK(); + } + if (execute_as_exact_phrase) { + return internal::execute_resolved_phrase_plan(idx, std::move(selected_leading), docids, + decode_context); + } + return execute_resolved_phrase_prefix_terms( + idx, std::move(selected_leading), std::move(selected_tail_terms), + selected_tail_position_offset, docids, decode_context); +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/phrase_position_source.cpp b/be/src/storage/index/snii/query/phrase_position_source.cpp new file mode 100644 index 00000000000000..d48aaf31a82256 --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_position_source.cpp @@ -0,0 +1,395 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +PhraseTermMapping build_phrase_term_mapping(const std::vector& terms) { + PhraseTermMapping mapping; + mapping.phrase_plan_index.reserve(terms.size()); + for (const std::string& term : terms) { + auto it = std::ranges::find(mapping.unique_terms, term); + if (it == mapping.unique_terms.end()) { + mapping.phrase_plan_index.push_back(mapping.unique_terms.size()); + mapping.unique_terms.push_back(term); + continue; + } + mapping.phrase_plan_index.push_back(static_cast(it - mapping.unique_terms.begin())); + } + return mapping; +} + +namespace { +Status append_prx_doc_ordinal(size_t ordinal, std::vector* out) { + if (ordinal > std::numeric_limits::max()) { + return Status::Error( + "phrase_query: prx doc ordinal exceeds u32"); + } + out->push_back(static_cast(ordinal)); + return Status::OK(); +} + +Status append_selected_ordinal(size_t doc_index, const std::vector& prx_doc_ordinals, + std::vector* selected_ordinals) { + if (!prx_doc_ordinals.empty()) { + selected_ordinals->push_back(prx_doc_ordinals[doc_index]); + return Status::OK(); + } + return append_prx_doc_ordinal(doc_index, selected_ordinals); +} + +Status append_selected_doc(size_t doc_index, uint32_t docid, + const std::vector& prx_doc_ordinals, + std::vector* selected_docids, + std::vector* selected_ordinals) { + selected_docids->push_back(docid); + return append_selected_ordinal(doc_index, prx_doc_ordinals, selected_ordinals); +} + +Status materialize_selected_prefix(size_t count, size_t capacity, + const std::vector& docids, + const std::vector& prx_doc_ordinals, + std::vector* selected_docids, + std::vector* selected_ordinals) { + selected_docids->reserve(capacity); + selected_ordinals->reserve(capacity); + selected_docids->insert(selected_docids->end(), docids.begin(), docids.begin() + count); + for (size_t i = 0; i < count; ++i) { + RETURN_IF_ERROR(append_selected_ordinal(i, prx_doc_ordinals, selected_ordinals)); + } + return Status::OK(); +} + +Status materialize_selected_prefix_if_needed(bool* selected_all, size_t count, size_t capacity, + const std::vector& docids, + const std::vector& prx_doc_ordinals, + std::vector* selected_docids, + std::vector* selected_ordinals) { + if (!*selected_all) { + return Status::OK(); + } + *selected_all = false; + return materialize_selected_prefix(count, capacity, docids, prx_doc_ordinals, selected_docids, + selected_ordinals); +} + +Status select_candidate_docs_for_prx(std::vector* docids, + std::vector* prx_doc_ordinals, + uint32_t prx_doc_count, + const std::vector& candidates, PosChunk* chunk) { + chunk->docids.clear(); + chunk->prx_doc_ordinals.clear(); + if (prx_doc_count == 0 && docids->size() > std::numeric_limits::max()) { + return Status::Error( + "phrase_query: prx doc count exceeds u32"); + } + chunk->prx_doc_count = + prx_doc_count == 0 ? static_cast(docids->size()) : prx_doc_count; + if (docids->empty() || candidates.empty()) { + return Status::OK(); + } + if (!prx_doc_ordinals->empty() && prx_doc_ordinals->size() != docids->size()) { + return Status::Error( + "phrase_query: prx ordinal/docid count mismatch"); + } + + std::vector selected_docids; + std::vector selected_ordinals; + bool selected_all = true; + const size_t selected_capacity = std::min(docids->size(), candidates.size()); + + auto candidate_it = std::ranges::lower_bound(candidates, docids->front()); + size_t candidate_index = static_cast(candidate_it - candidates.begin()); + for (size_t doc_index = 0; doc_index < docids->size(); ++doc_index) { + const uint32_t docid = (*docids)[doc_index]; + while (candidate_index < candidates.size() && candidates[candidate_index] < docid) { + ++candidate_index; + } + if (candidate_index == candidates.size()) { + RETURN_IF_ERROR(materialize_selected_prefix_if_needed( + &selected_all, doc_index, selected_capacity, *docids, *prx_doc_ordinals, + &selected_docids, &selected_ordinals)); + break; + } + if (candidates[candidate_index] != docid) { + RETURN_IF_ERROR(materialize_selected_prefix_if_needed( + &selected_all, doc_index, selected_capacity, *docids, *prx_doc_ordinals, + &selected_docids, &selected_ordinals)); + continue; + } + + if (!selected_all) { + RETURN_IF_ERROR(append_selected_doc(doc_index, docid, *prx_doc_ordinals, + &selected_docids, &selected_ordinals)); + } + ++candidate_index; + } + + if (selected_all) { + chunk->docids = std::move(*docids); + chunk->prx_doc_ordinals = std::move(*prx_doc_ordinals); + docids->clear(); + prx_doc_ordinals->clear(); + return Status::OK(); + } + if (selected_docids.empty()) { + return Status::OK(); + } + chunk->docids = std::move(selected_docids); + chunk->prx_doc_ordinals = std::move(selected_ordinals); + return Status::OK(); +} + +// PRX byte ranges for every candidate-bearing chunk across all phrase terms are +// added to one shared BatchRangeFetcher and fetched in a single batched round +// (T02). Pass 1 records, for each chunk that needs on-disk PRX bytes, where to +// write the fetched slice back: which plan's PosSource, which chunk within it, +// and the fetcher handle. + +struct PrxRangeAssignment { + size_t plan_index; + size_t chunk_index; + size_t handle; +}; + +void record_prx_assignment(std::vector* assignments, size_t plan_index, + size_t chunk_index, size_t handle) { + assignments->push_back(PrxRangeAssignment { + .plan_index = plan_index, .chunk_index = chunk_index, .handle = handle}); +} + +Status build_flat_position_source(const LogicalIndexReader& idx, + const io::BatchRangeFetcher& round1, DocidSource* doc_source, + const TermPlan& p, const std::vector& candidates, + size_t plan_index, io::BatchRangeFetcher* prx_fetcher, + std::vector* assignments, PosSource* src) { + PosChunk chunk; + std::vector docids; + std::vector prx_doc_ordinals; + const bool docids_are_final_candidates = + doc_source->docids_are_final_candidates && !doc_source->chunks.empty(); + if (!doc_source->chunks.empty()) { + DocidChunk& doc_chunk = doc_source->chunks.front(); + docids = std::move(doc_chunk.docids); + prx_doc_ordinals = std::move(doc_chunk.prx_doc_ordinals); + chunk.prx_doc_count = doc_chunk.prx_doc_count; + } + // pod_ref PRX bytes are read from the shared fetcher (one batched round for the + // whole phrase); inline PRX bytes already live in the dict entry. The pod_ref + // range is added unconditionally to keep the bytes read identical to the prior + // per-term fetch(); the handle is only recorded as an assignment when the chunk + // is kept (an empty chunk reads the same bytes but needs no backfill). + bool has_prx_handle = false; + size_t prx_handle = 0; + if (p.pod_ref) { + uint64_t poff = 0; + uint64_t plen = 0; + RETURN_IF_ERROR(idx.resolve_prx_window(p.entry, p.prx_base, &poff, &plen)); + prx_handle = prx_fetcher->add(poff, plen); + has_prx_handle = true; + } else { + chunk.prx = Slice(p.entry.prx_bytes); + } + if (docids.empty()) { + Slice dd; + if (p.pod_ref) { + dd = round1.get(p.frq_handle); + } else { + RETURN_IF_ERROR(internal::inline_dd_region(p.entry, &dd)); + } + RETURN_IF_ERROR(format::decode_dd_region(dd, p.entry.dd_meta, + /*win_base=*/0, &docids)); + if (docids.size() > std::numeric_limits::max()) { + return Status::Error( + "phrase_query: prx doc count exceeds u32"); + } + chunk.prx_doc_count = static_cast(docids.size()); + } + if (docids_are_final_candidates) { + chunk.docids = std::move(docids); + chunk.prx_doc_ordinals = std::move(prx_doc_ordinals); + if (!chunk.docids.empty()) { + if (has_prx_handle) { + record_prx_assignment(assignments, plan_index, src->chunks.size(), prx_handle); + } + src->chunks.push_back(std::move(chunk)); + } + return Status::OK(); + } + RETURN_IF_ERROR(select_candidate_docs_for_prx(&docids, &prx_doc_ordinals, chunk.prx_doc_count, + candidates, &chunk)); + if (!chunk.docids.empty()) { + if (has_prx_handle) { + record_prx_assignment(assignments, plan_index, src->chunks.size(), prx_handle); + } + src->chunks.push_back(std::move(chunk)); + } + return Status::OK(); +} + +bool chunk_may_contain_candidate(const DocidChunk& chunk, const std::vector& candidates) { + if (chunk.docids.empty() || candidates.empty()) { + return false; + } + const auto it = std::ranges::lower_bound(candidates, chunk.docids.front()); + return it != candidates.end() && *it <= chunk.docids.back(); +} + +Status decode_windowed_position_source(const LogicalIndexReader& idx, const TermPlan& p, + DocidSource* doc_source, + const std::vector& candidates, size_t plan_index, + io::BatchRangeFetcher* prx_fetcher, + std::vector* assignments, + PosSource* src) { + for (size_t i = 0; i < doc_source->chunks.size(); ++i) { + DocidChunk& doc_chunk = doc_source->chunks[i]; + if (!doc_source->docids_are_final_candidates && + !chunk_may_contain_candidate(doc_chunk, candidates)) { + continue; + } + if (!doc_chunk.windowed) { + return Status::Error( + "phrase_query: expected windowed doc chunk"); + } + PosChunk chunk; + if (doc_source->docids_are_final_candidates) { + chunk.docids = std::move(doc_chunk.docids); + chunk.prx_doc_ordinals = std::move(doc_chunk.prx_doc_ordinals); + chunk.prx_doc_count = doc_chunk.prx_doc_count; + } else { + RETURN_IF_ERROR( + select_candidate_docs_for_prx(&doc_chunk.docids, &doc_chunk.prx_doc_ordinals, + doc_chunk.prx_doc_count, candidates, &chunk)); + } + if (chunk.docids.empty()) { + continue; + } + + reader::WindowAbsRange range; + RETURN_IF_ERROR(reader::windowed_window_range( + idx, p.entry, p.frq_base, p.prx_base, p.prelude, doc_chunk.window, + /*want_positions=*/true, /*want_freq=*/false, &range)); + chunk.windowed = true; + chunk.window = doc_chunk.window; + const size_t prx_handle = prx_fetcher->add(range.prx_off, range.prx_len); + record_prx_assignment(assignments, plan_index, src->chunks.size(), prx_handle); + src->chunks.push_back(std::move(chunk)); + } + return Status::OK(); +} + +} // namespace +Status build_position_sources_for_candidates( + const LogicalIndexReader& idx, const io::BatchRangeFetcher& round1, + const std::vector& plans, std::vector* doc_sources, + const std::vector& candidates, + std::vector>* owners, std::vector* srcs, + format::PrxDecodeContext* observer_context) { + srcs->assign(plans.size(), PosSource {}); + for (PosSource& source : *srcs) { + source.observer_context = observer_context; + } + // All phrase terms share one PRX fetcher: pass 1 adds every candidate-bearing + // chunk's PRX range and records a backfill assignment; a single fetch() then + // issues one batched read (one serial round on a remote reader); pass 2 fills + // in each chunk's PRX slice. This collapses the prior per-term fetch() -- O(n) + // serial remote rounds for an n-term phrase -- into one. + auto prx_fetcher = + std::make_unique(idx.reader(), reader::kSameTermCoalesceGap); + std::vector assignments; + for (size_t i = 0; i < plans.size(); ++i) { + const TermPlan& p = plans[i]; + if (p.windowed) { + RETURN_IF_ERROR(decode_windowed_position_source(idx, p, &(*doc_sources)[i], candidates, + i, prx_fetcher.get(), &assignments, + &(*srcs)[i])); + continue; + } + RETURN_IF_ERROR(build_flat_position_source(idx, round1, &(*doc_sources)[i], p, candidates, + i, prx_fetcher.get(), &assignments, + &(*srcs)[i])); + } + if (prx_fetcher->pending() > 0) { + if (observer_context == nullptr || observer_context->stats == nullptr) { + RETURN_IF_ERROR(prx_fetcher->fetch()); + } else { + const auto fetch_start = std::chrono::steady_clock::now(); + RETURN_IF_ERROR(prx_fetcher->fetch()); + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - fetch_start) + .count(); + observer_context->stats->fetch_ns += + std::max(1, static_cast(elapsed)); + } + } + for (const PrxRangeAssignment& a : assignments) { + (*srcs)[a.plan_index].chunks[a.chunk_index].prx = prx_fetcher->get(a.handle); + } + // Keep the fetcher alive only when some chunk slice references its buffers. + if (!assignments.empty()) { + owners->push_back(std::move(prx_fetcher)); + } + return Status::OK(); +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/phrase_prefix_exec.cpp b/be/src/storage/index/snii/query/phrase_prefix_exec.cpp new file mode 100644 index 00000000000000..533962aedbe09c --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_prefix_exec.cpp @@ -0,0 +1,719 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query::phrase_impl { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; +using internal::PhraseVerifyTimer; + +namespace { +Status collect_expected_tail_positions(const std::vector& plans, + const std::vector& phrase_plan_index, + const std::vector& position_offsets, + std::vector& srcs, + const std::vector& candidates, + ExpectedTailPositionSet* out, + bool preserve_first_clause_multiplicity) { + const size_t n = phrase_plan_index.size(); + DCHECK(n > 1); + DCHECK_EQ(plans.size(), srcs.size()); + std::vector cur(plans.size()); + for (size_t i = 0; i < plans.size(); ++i) { + cur[i].init(&srcs[i]); + } + + std::vector> unique_span(plans.size()); + std::vector> span(n); + for (uint32_t d : candidates) { + for (size_t i = 0; i < plans.size(); ++i) { + RETURN_IF_ERROR(cur[i].seek(d)); + RETURN_IF_ERROR(cur[i].positions(&unique_span[plans[i].order])); + } + for (size_t i = 0; i < n; ++i) { + DCHECK_LT(phrase_plan_index[i], unique_span.size()); + span[i] = unique_span[phrase_plan_index[i]]; + } + + // Anchor the outer enumeration on the SPARSEST exact term (smallest + // per-doc position span), not the hardcoded phrase-position-0 term. The + // set of valid phrase starts is anchor-independent -- each valid start + // maps 1:1 to exactly one anchor position (anchor_pos = start + + // offset[anchor]) -- so enumerating the shortest span and binary-searching + // the others yields the identical result set with the fewest outer + // iterations. A leading high-frequency exact term no longer forces + // O(|span[0]|) work per candidate doc. + size_t anchor = 0; + auto best = position_span_size(span[0]); + if (!preserve_first_clause_multiplicity) { + for (size_t t = 1; t < n; ++t) { + const auto sz = position_span_size(span[t]); + if (sz < best) { + best = sz; + anchor = t; + } + } + } + const uint32_t anchor_off = position_offsets[anchor]; + SNII_QUERY_ADD(anchor_iterations, best); + + // Only the first non-anchor term is probed for every viable anchor. + // Later terms can be skipped after an earlier mismatch, so using the + // total anchor count to choose a forward scan for them could turn one + // binary lookup into a full-span walk. The first term uses a forward + // scan only when all anchors are valid and its conservative comparison + // upper bound is at most half that of repeated binary search. + const size_t first_checked_term = anchor == 0 ? 1 : 0; + size_t monotonic_position_scan_term = n; + if (should_use_monotonic_position_scan(span[anchor], + position_span_size(span[first_checked_term]), + anchor_off, position_offsets[first_checked_term])) { + monotonic_position_scan_term = first_checked_term; + SNII_QUERY_COUNT(monotonic_position_scans); + } + + const size_t expected_begin = out->positions.size(); + for (const uint32_t* p = span[anchor].first; p != span[anchor].second; ++p) { + const uint32_t anchor_pos = *p; + // Underflow guard: a general anchor (offset > 0) can sit at a position + // smaller than its offset, which would wrap `start`. Such a position + // admits no valid phrase start and is skipped. (The old span[0] anchor + // had offset 0 and could never underflow.) + if (anchor_pos < anchor_off) { + continue; + } + const uint32_t start = anchor_pos - anchor_off; + bool ok = true; + for (size_t t = 0; t < n; ++t) { + if (t == anchor) { + continue; // the anchor term's position is satisfied by construction + } + uint32_t want = 0; + if (!internal::add_position_offset(start, position_offsets[t], &want)) { + ok = false; + break; + } + if (t == monotonic_position_scan_term) { + while (span[t].first != span[t].second && *span[t].first < want) { + ++span[t].first; + } + if (span[t].first == span[t].second || *span[t].first != want) { + ok = false; + break; + } + } else if (!std::binary_search(span[t].first, span[t].second, want)) { + ok = false; + break; + } + } + uint32_t tail_pos = 0; + if (ok && internal::add_position_offset(start, position_offsets[n], &tail_pos)) { + out->positions.push_back(tail_pos); + } + } + const size_t expected_end = out->positions.size(); + if (expected_end != expected_begin) { + out->docs.push_back(ExpectedTailPositions { + .docid = d, .positions_begin = expected_begin, .positions_end = expected_end}); + } + } + return Status::OK(); +} + +Status collect_single_term_expected_tail_positions(std::vector& srcs, + const std::vector& candidates, + uint32_t tail_offset, + ExpectedTailPositionSet* out) { + PostingCursor cursor; + cursor.init(srcs.data()); + out->reserve_docs(out->docs.size() + candidates.size()); + + for (uint32_t d : candidates) { + RETURN_IF_ERROR(cursor.seek(d)); + std::pair span; + RETURN_IF_ERROR(cursor.positions(&span)); + + const size_t expected_begin = out->positions.size(); + for (const uint32_t* p = span.first; p != span.second; ++p) { + uint32_t tail_pos = 0; + if (internal::add_position_offset(*p, tail_offset, &tail_pos)) { + out->positions.push_back(tail_pos); + } + } + const size_t expected_end = out->positions.size(); + if (expected_end != expected_begin) { + out->docs.push_back(ExpectedTailPositions { + .docid = d, .positions_begin = expected_begin, .positions_end = expected_end}); + } + } + return Status::OK(); +} + +Status collect_expected_tail_positions(const LogicalIndexReader& idx, + internal::ResolvedPhrasePlan exact_plan, + uint32_t tail_position_offset, ExpectedTailPositionSet* out, + const std::vector* candidate_prefilter, + format::PrxDecodeContext* observer_context, + bool preserve_first_clause_multiplicity) { + out->clear(); + io::BatchRangeFetcher round1(idx.reader()); + std::vector plans; + RETURN_IF_ERROR(internal::plan_resolved_terms(idx, std::move(exact_plan.unique_terms), &round1, + &plans, + /*need_positions=*/false)); + + PhraseExecutionState state; + RETURN_IF_ERROR(build_phrase_execution_state(idx, &round1, &plans, &state, candidate_prefilter, + observer_context, + PhraseCandidateMetric::kPrefixLeading)); + if (state.candidates.empty()) { + return Status::OK(); + } + out->reserve_docs(state.candidates.size()); + DORIS_CHECK(!exact_plan.position_offsets.empty()); + DORIS_CHECK_GT(tail_position_offset, exact_plan.position_offsets.back()); + std::vector position_offsets = std::move(exact_plan.position_offsets); + position_offsets.push_back(tail_position_offset); + PhraseVerifyTimer verify_timer(observer_context); + if (exact_plan.phrase_plan_index.size() == 1) { + DORIS_CHECK_LT(position_offsets[0], position_offsets[1]); + RETURN_IF_ERROR(collect_single_term_expected_tail_positions( + state.srcs, state.candidates, position_offsets[1] - position_offsets[0], out)); + } else { + RETURN_IF_ERROR(collect_expected_tail_positions( + plans, exact_plan.phrase_plan_index, position_offsets, state.srcs, state.candidates, + out, preserve_first_clause_multiplicity)); + } + verify_timer.commit_success(); + return Status::OK(); +} + +bool contains_any_position(const ExpectedTailPositionSet& expected, + const ExpectedTailPositions& wanted, + std::pair actual) { + for (size_t i = wanted.positions_begin; i < wanted.positions_end; ++i) { + if (std::binary_search(actual.first, actual.second, expected.positions[i])) { + return true; + } + } + return false; +} + +uint32_t mark_matching_positions(ExpectedTailPositionSet* expected, + const ExpectedTailPositions& wanted, + std::pair actual) { + DCHECK_EQ(expected->position_matched.size(), expected->positions.size()); + size_t expected_index = wanted.positions_begin; + const uint32_t* actual_position = actual.first; + uint32_t added = 0; + while (expected_index < wanted.positions_end && actual_position != actual.second) { + const uint32_t expected_position = expected->positions[expected_index]; + if (expected_position < *actual_position) { + ++expected_index; + continue; + } + if (*actual_position < expected_position) { + ++actual_position; + continue; + } + const size_t expected_run_end = static_cast( + std::upper_bound(expected->positions.begin() + expected_index, + expected->positions.begin() + wanted.positions_end, + expected_position) - + expected->positions.begin()); + while (expected_index < expected_run_end) { + if (expected->position_matched[expected_index] == 0) { + expected->position_matched[expected_index] = 1; + DCHECK_NE(added, std::numeric_limits::max()); + ++added; + } + ++expected_index; + } + actual_position = std::upper_bound(actual_position, actual.second, expected_position); + } + return added; +} + +// Upper bound on prefix expansions whose position cursors are held resident at +// once. The old per-tail loop verified a single expansion at a time (one +// PosSource + one PRX buffer live); the merged sweep below holds up to this many +// tail PosSources + cursors + PRX buffers simultaneously so it can read every +// tail's docid/prx bytes in ONE batched round and verify them in a single +// forward pass. `max_expansions` may be unbounded (<= 0), so this hard cap keeps +// resident memory bounded independent of the query: expansions beyond the cap +// are processed as additional capped groups (each a fresh single fetch) whose +// matched-doc flags are accumulated. The cap is tightened because each cursor +// here also holds decoded PRX rather than plain docids. +constexpr size_t kMaxTailMergeBatch = 32; + +// Phrase-prefix only reads the residual tails' docid union to prefilter the +// leading-phrase candidate set when the smallest leading term's df reaches this +// -- i.e. when the leading candidate set is large enough that decoding all its +// positions dwarfs an extra docid-only union read. Scale the threshold with the +// segment so a fixed absolute gate does not disable the prefilter after rowset +// segmentation, while retaining the original 1<<16 cap for large segments. +constexpr uint32_t kMinPrefixLeadingPrefilterMinDf = 256; +constexpr uint32_t kMaxPrefixLeadingPrefilterMinDf = 1u << 16; +constexpr uint32_t kPrefixLeadingPrefilterDocFraction = 8; +// The tail union is decoded once for filtering and again for verification. +// Require enough leading PRX work to cover both reads and union overhead. +constexpr uint32_t kPrefixLeadingToTailDfRatio = 8; + +uint32_t prefix_leading_prefilter_min_df(const LogicalIndexReader& idx, + bool allow_segment_relative_gate) { + if (!allow_segment_relative_gate) { + return kMaxPrefixLeadingPrefilterMinDf; + } + const uint64_t segment_relative = + idx.stats().indexed_doc_count / kPrefixLeadingPrefilterDocFraction; + return static_cast(std::clamp( + segment_relative, kMinPrefixLeadingPrefilterMinDf, kMaxPrefixLeadingPrefilterMinDf)); +} + +// Merged multi-tail verification for ONE resident-capped group of prefix +// expansions (`tails`, already truncated by max_expansions upstream). This +// replaces the per-tail verify-then-union loop: instead of re-planning + TWO +// remote rounds (docid, then prx) + a separate doc-walk PER tail and unioning N +// result lists, it plans every tail into ONE shared round1 fetch, intersects +// each tail with `expected_docids` in memory (no I/O), builds every surviving +// tail's position source in ONE batched PRX round, then sweeps the group's tail +// cursors over the ascending `expected` docs a SINGLE time -- marking a doc as +// soon as ANY tail has a position adjacent to a leading match. +// +// The marked set is byte-identical to the per-tail path's +// UNION_{tail in group} { d : d in tail INTERSECT expected AND +// contains_any_position(expected, doc_d, pos_tail(d)) } +// because each tail's PosSource is still built from its OWN final-candidate +// docids (the shared-candidate argument is ignored for final-candidate sources), +// so pos_tail(d) and the per-doc position test are unchanged. Only the I/O rounds +// (2N -> 2) and the N separate unions (-> in-place flags) collapse. Bigram +// postings are NEVER consulted: every tail is verified against its unigram +// positions here. + +Status collect_merged_tail_matches(const LogicalIndexReader& idx, + std::vector tails, + ExpectedTailPositionSet* expected, + const std::vector& expected_docids, bool final_group, + std::vector* final_matches, + format::PrxDecodeContext* observer_context, + std::vector* frequency_matches) { + DCHECK(expected != nullptr); + DCHECK(final_matches != nullptr || frequency_matches != nullptr); + const size_t n = tails.size(); + if (n == 0 || expected->docs.empty()) { + return Status::OK(); + } + + // Plan every tail into ONE fetcher so their docid postings + windowed + // preludes are read in a single batched round (the per-tail path issued one + // round per tail). Each tail keeps its own single-term plan vector so the + // conjunction filter below consumes it directly, without slicing a shared + // plan vector (whose prelude readers own decoded directory buffers). + io::BatchRangeFetcher round1(idx.reader()); + std::vector> tail_plans(n); + for (size_t i = 0; i < n; ++i) { + std::vector one; + one.push_back(std::move(tails[i])); + RETURN_IF_ERROR(internal::plan_resolved_terms(idx, std::move(one), &round1, &tail_plans[i], + /*need_positions=*/false)); + } + if (round1.pending() > 0) { + RETURN_IF_ERROR(round1.fetch()); + } + for (size_t i = 0; i < n; ++i) { + RETURN_IF_ERROR(internal::open_preludes(round1, &tail_plans[i], + /*need_positions=*/true)); + } + + // Per-tail candidate docids (tail posting INTERSECT expected) and the aligned + // final-candidate doc sources feeding the batched position builder. The + // conjunction reads only already-fetched round1 bytes; a single-plan filter + // marks its one source docids_are_final_candidates, so the position builder + // materializes each tail's PosSource directly over its own candidate docs. + // Tails whose intersection is empty are dropped here (exactly as the old + // per-tail early return did), so no dead tail decodes its full posting. + std::vector> tail_candidates(n); + std::vector active_plans; + std::vector active_sources; + std::vector active_index; // active slot -> tail index (into tail_candidates) + for (size_t i = 0; i < n; ++i) { + std::vector tail_source; + RETURN_IF_ERROR(internal::filter_docids_by_conjunction( + idx, round1, tail_plans[i], expected_docids, &tail_candidates[i], &tail_source)); + if (tail_candidates[i].empty()) { + continue; // this expansion has no doc in the expected set: nothing to verify + } + active_plans.push_back(std::move(tail_plans[i].front())); + active_sources.push_back(tail_source.empty() ? DocidSource {} + : std::move(tail_source.front())); + active_index.push_back(i); + } + if (active_plans.empty() && (!final_group || expected->matched_count == 0)) { + return Status::OK(); + } + // An empty final group must still sweep expected->docs and emit matches that + // earlier resident groups recorded in-place. + + // ONE batched PRX round for every retained chunk across all surviving tails + // (vs one round per tail before). `candidates` is intentionally empty: every + // source is a final-candidate source, so the builder addresses positions by + // the source's own docids and never consults the shared candidate list. + std::vector> owners; + std::vector srcs; + const std::vector no_shared_candidates; + if (!active_plans.empty()) { + RETURN_IF_ERROR(build_position_sources_for_candidates(idx, round1, active_plans, + &active_sources, no_shared_candidates, + &owners, &srcs, observer_context)); + } + + // Single forward sweep over the ascending expected docs. For each doc probe + // only the tails that actually posted it (per-tail ascending cursor over + // tail_candidates), decode positions once, and emit the doc the instant one + // tail's positions land adjacent to a leading match. Cursors advance strictly + // forward because expected.docs is strictly ascending and each cursor is + // sought at most once per doc. + std::vector cursors(active_plans.size()); + for (size_t a = 0; a < active_plans.size(); ++a) { + cursors[a].init(&srcs[a]); + } + std::vector tail_pos(active_plans.size(), 0); + PhraseVerifyTimer verify_timer(observer_context); + DCHECK_EQ(expected->docs.size(), expected_docids.size()); + for (size_t expected_index = 0; expected_index < expected->docs.size(); ++expected_index) { + ExpectedTailPositions& doc = expected->docs[expected_index]; + DCHECK_EQ(doc.docid, expected_docids[expected_index]); + SNII_QUERY_COUNT(prefix_expected_doc_visits); + if (observer_context != nullptr && observer_context->query_stats != nullptr) { + ++observer_context->query_stats->prefix_tail_candidate_visits; + } + const uint32_t d = doc.docid; + if (frequency_matches != nullptr || doc.phrase_frequency == 0) { + bool matched = false; + uint32_t added_frequency = 0; + for (size_t a = 0; + a < active_plans.size() && (frequency_matches != nullptr || !matched); ++a) { + std::vector& cand = tail_candidates[active_index[a]]; + size_t& ti = tail_pos[a]; + while (ti < cand.size() && cand[ti] < d) { + ++ti; + } + if (ti >= cand.size() || cand[ti] != d) { + continue; // this expansion has no posting at d + } + RETURN_IF_ERROR(cursors[a].seek(d)); + std::pair actual; + RETURN_IF_ERROR(cursors[a].positions(&actual)); + if (frequency_matches == nullptr && contains_any_position(*expected, doc, actual)) { + matched = true; + } else if (frequency_matches != nullptr) { + const uint32_t added = mark_matching_positions(expected, doc, actual); + DCHECK_LE(added_frequency, std::numeric_limits::max() - added); + added_frequency += added; + } + } + if (matched || added_frequency != 0) { + if (doc.phrase_frequency == 0) { + ++expected->matched_count; + } + if (frequency_matches == nullptr) { + doc.phrase_frequency = 1; + } else { + DCHECK_LE(doc.phrase_frequency, + std::numeric_limits::max() - added_frequency); + doc.phrase_frequency += added_frequency; + } + } + } + if (final_group && doc.phrase_frequency != 0) { + if (final_matches != nullptr) { + final_matches->push_back(d); + } + if (frequency_matches != nullptr) { + frequency_matches->push_back( + PhraseMatch {.docid = d, .frequency = doc.phrase_frequency}); + } + } + } + verify_timer.commit_success(); + return Status::OK(); +} + +} // namespace +Status execute_resolved_phrase_prefix_terms( + const LogicalIndexReader& idx, internal::ResolvedPhrasePlan exact_plan, + std::vector tail_terms, uint32_t tail_position_offset, + std::vector* docids, format::PrxDecodeContext* decode_context, + std::vector* matches, const std::vector* candidate_prefilter) { + DORIS_CHECK(docids != nullptr || matches != nullptr); + if (tail_terms.empty()) { + return Status::OK(); + } + if (exact_plan.phrase_plan_index.empty()) { + DORIS_CHECK(matches == nullptr); + DORIS_CHECK_EQ(tail_position_offset, 0U); + if (tail_terms.size() == 1) { + const auto& tail = tail_terms.front(); + RETURN_IF_ERROR(internal::read_docid_posting(idx, tail.entry, tail.frq_base, + tail.prx_base, docids)); + if (candidate_prefilter != nullptr) { + *docids = internal::intersect_sorted(*docids, *candidate_prefilter); + } + return Status::OK(); + } + std::vector tail_postings; + tail_postings.reserve(tail_terms.size()); + for (const auto& tail : tail_terms) { + tail_postings.push_back({tail.entry, tail.frq_base, tail.prx_base}); + } + RETURN_IF_ERROR(internal::build_docid_union(idx, tail_postings, docids)); + if (candidate_prefilter != nullptr) { + *docids = internal::intersect_sorted(*docids, *candidate_prefilter); + } + return Status::OK(); + } + if (!idx.has_positions()) { + return Status::Error( + "phrase_prefix_query: index has no positions"); + } + DORIS_CHECK(!exact_plan.position_offsets.empty()); + DORIS_CHECK_GT(tail_position_offset, exact_plan.position_offsets.back()); + if (tail_terms.size() == 1) { + append_resolved_phrase_clause(std::move(tail_terms.front()), tail_position_offset, + &exact_plan); + return internal::execute_resolved_phrase_plan(idx, std::move(exact_plan), docids, + decode_context, matches, candidate_prefilter); + } + + uint32_t min_lead_df = std::numeric_limits::max(); + for (const ResolvedQueryTerm& term : exact_plan.unique_terms) { + min_lead_df = std::min(min_lead_df, term.entry.df); + } + uint64_t tail_df_sum = 0; + for (const ResolvedQueryTerm& tail : tail_terms) { + tail_df_sum += tail.entry.df; + } + const std::vector* prefilter = candidate_prefilter; + std::vector tail_union; + std::vector combined_prefilter; + const bool allow_segment_relative_prefilter = + candidate_prefilter == nullptr && exact_plan.phrase_plan_index.size() == 1; + const uint32_t leading_prefilter_min_df = + prefix_leading_prefilter_min_df(idx, allow_segment_relative_prefilter); + if (min_lead_df >= leading_prefilter_min_df && + tail_df_sum <= static_cast(min_lead_df) / kPrefixLeadingToTailDfRatio) { + std::vector tail_postings; + tail_postings.reserve(tail_terms.size()); + for (const ResolvedQueryTerm& tail : tail_terms) { + tail_postings.push_back({tail.entry, tail.frq_base, tail.prx_base}); + } + RETURN_IF_ERROR(internal::build_docid_union(idx, tail_postings, &tail_union)); + if (tail_union.empty()) { + return Status::OK(); + } + if (candidate_prefilter == nullptr) { + prefilter = &tail_union; + } else { + combined_prefilter = internal::intersect_sorted(*candidate_prefilter, tail_union); + if (combined_prefilter.empty()) { + return Status::OK(); + } + prefilter = &combined_prefilter; + } + } + + ExpectedTailPositionSet expected; + RETURN_IF_ERROR(collect_expected_tail_positions(idx, std::move(exact_plan), + tail_position_offset, &expected, prefilter, + decode_context, matches != nullptr)); + if (expected.docs.empty()) { + return Status::OK(); + } + if (matches != nullptr) { + expected.position_matched.assign(expected.positions.size(), 0); + } + + std::vector expected_docids; + expected_docids.reserve(expected.docs.size()); + for (const ExpectedTailPositions& doc : expected.docs) { + expected_docids.push_back(doc.docid); + } + SNII_QUERY_COUNT(expected_docids_build); + + std::vector final_matches; + // Keep the expected-doc set stable across groups. Compacting matched docs + // speculates that later tail groups intersect it; an empty intersection + // skips the sweep entirely, so the compaction work cannot be repaid safely. + for (size_t start = 0; start < tail_terms.size(); start += kMaxTailMergeBatch) { + const size_t end = std::min(start + kMaxTailMergeBatch, tail_terms.size()); + const bool final_group = end == tail_terms.size(); + std::vector group; + group.reserve(end - start); + for (size_t i = start; i < end; ++i) { + group.push_back(std::move(tail_terms[i])); + } + RETURN_IF_ERROR(collect_merged_tail_matches( + idx, std::move(group), &expected, expected_docids, final_group, + docids == nullptr ? nullptr : &final_matches, decode_context, matches)); + if (matches == nullptr && !final_group && expected.matched_count == expected.docs.size()) { + final_matches.reserve(expected.docs.size()); + for (const ExpectedTailPositions& doc : expected.docs) { + DCHECK_NE(doc.phrase_frequency, 0); + final_matches.push_back(doc.docid); + } + break; + } + } + if (docids != nullptr) { + *docids = std::move(final_matches); + } + return Status::OK(); +} + +namespace { +template +std::vector copy_resolved_terms(const std::vector& resolved, + const std::vector& indices) { + std::vector result; + result.reserve(indices.size()); + for (size_t index : indices) { + DORIS_CHECK_LT(index, resolved.size()); + result.push_back(resolved[index]); + } + return result; +} + +} // namespace +Status execute_hybrid_phrase_prefix_plan( + const LogicalIndexReader& idx, const HybridPrefixPlanArtifact& artifact, + const std::vector& batch_terms, const std::vector& resolved, + const std::vector& plain_tail_terms, std::vector* docids, + format::PrxDecodeContext* decode_context, CommonGramsPlanningTimer& planning_timer, + bool* candidate_intersection_empty) { + DORIS_CHECK(idx.common_grams_posting_policy() == format::CommonGramsPostingPolicy::kHybridV1); + DORIS_CHECK(docids != nullptr); + DORIS_CHECK(candidate_intersection_empty != nullptr); + docids->clear(); + *candidate_intersection_empty = false; + + const HybridPositionedCover& plain_tail_cover = artifact.plain_tail_cover; + const uint32_t plain_tail_position_offset = artifact.plain_tail_position_offset; + HybridPrefixCandidateSet leading_candidates; + RETURN_IF_ERROR(build_hybrid_leading_candidates(idx, plain_tail_cover.candidate_prefilter, + batch_terms, resolved, &leading_candidates)); + if (leading_candidates.active && leading_candidates.docs.empty()) { + planning_timer.finish(); + *candidate_intersection_empty = true; + return Status::OK(); + } + + if (!artifact.maps_tail_to_gram) { + DORIS_CHECK(leading_candidates.active); + DORIS_CHECK(artifact.mapped_tail_split.positioned_indices.empty()); + DORIS_CHECK(artifact.mapped_tail_split.docs_only_indices.empty()); + planning_timer.finish(); + RETURN_IF_ERROR(execute_resolved_phrase_prefix_terms( + idx, + copy_resolved_phrase_plan(plain_tail_cover.verification, batch_terms, resolved), + plain_tail_terms, plain_tail_position_offset, docids, decode_context, nullptr, + &leading_candidates.docs)); + *candidate_intersection_empty = docids->empty(); + return Status::OK(); + } + + const HybridPrefixMappedTails& split = artifact.mapped_tail_split; + DORIS_CHECK(!split.positioned_indices.empty() || !split.docs_only_indices.empty()); + DORIS_CHECK(leading_candidates.active || !split.docs_only_indices.empty()); + std::vector docs_only_tail_candidates; + if (!split.docs_only_indices.empty()) { + RETURN_IF_ERROR(build_hybrid_docs_only_tail_candidates( + idx, resolved, split.docs_only_indices, leading_candidates, + &docs_only_tail_candidates)); + } + planning_timer.finish(); + + if (!split.positioned_indices.empty()) { + DORIS_CHECK(artifact.positioned_tail_verification.has_value()); + std::vector positioned_docs; + RETURN_IF_ERROR(execute_resolved_phrase_prefix_terms( + idx, + copy_resolved_phrase_plan(*artifact.positioned_tail_verification, batch_terms, + resolved), + copy_resolved_terms(resolved, split.positioned_indices), + plain_tail_position_offset - 1, &positioned_docs, decode_context, nullptr, + leading_candidates.active ? &leading_candidates.docs : nullptr)); + internal::union_sorted_into(docids, positioned_docs); + } + + if (!split.docs_only_indices.empty() && !docs_only_tail_candidates.empty()) { + std::vector docs_only_docs; + RETURN_IF_ERROR(execute_resolved_phrase_prefix_terms( + idx, + copy_resolved_phrase_plan(plain_tail_cover.verification, batch_terms, resolved), + copy_resolved_terms(plain_tail_terms, split.docs_only_ordinals), + plain_tail_position_offset, &docs_only_docs, decode_context, nullptr, + &docs_only_tail_candidates)); + internal::union_sorted_into(docids, docs_only_docs); + } + *candidate_intersection_empty = docids->empty(); + return Status::OK(); +} + +} // namespace doris::snii::query::phrase_impl diff --git a/be/src/storage/index/snii/query/phrase_prx_validation.h b/be/src/storage/index/snii/query/phrase_prx_validation.h new file mode 100644 index 00000000000000..a4082360affba7 --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_prx_validation.h @@ -0,0 +1,55 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris::snii { +class ByteSource; +namespace format { +struct PrxDecodeContext; +} +} // namespace doris::snii + +namespace doris::snii::query::internal { + +// The production seam shared by PosChunkDecoder and focused shape-validation +// tests. A successful format decode commits its stats before caller-level CSR +// validation runs. +Status decode_and_validate_prx_frame(ByteSource* source, + std::span selected_doc_ordinals, + bool decode_full, bool all_docs_selected, + uint32_t expected_total_docs, size_t expected_selected_docs, + std::vector* pos_flat, + std::vector* pos_offsets, + format::PrxDecodeContext* decode_context); + +// Validates the CSR shape expected by phrase execution. Format-level decode +// statistics have already been committed when this function runs. +Status validate_prx_frame(std::span pos_flat, std::span pos_offsets, + uint32_t actual_total_docs, uint32_t expected_total_docs, + size_t expected_selected_docs, + std::span selected_doc_ordinals, + bool offsets_by_prx_ordinal, bool all_docs_selected); + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/phrase_query.cpp b/be/src/storage/index/snii/query/phrase_query.cpp new file mode 100644 index 00000000000000..98bf263952f1bb --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_query.cpp @@ -0,0 +1,319 @@ +// 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 "storage/index/snii/query/phrase_query.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/phrase_query_split.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/position_math.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "util/debug_points.h" + +namespace doris::snii::query { + +using query::internal::DocidChunk; +using query::internal::DocidSource; +using query::internal::ResolvedQueryTerm; +using query::internal::TermPlan; +using reader::LogicalIndexReader; + +CommonGramsPlanDebugOverride common_grams_plan_debug_override() { + CommonGramsPlanDebugOverride result = CommonGramsPlanDebugOverride::kNone; + DBUG_EXECUTE_IF(COMMON_GRAMS_FORCE_PLAIN_PLAN_DEBUG_POINT, + { result = CommonGramsPlanDebugOverride::kForcePlain; }); + DBUG_EXECUTE_IF(COMMON_GRAMS_FORCE_GRAM_PLAN_DEBUG_POINT, { + DORIS_CHECK(result != CommonGramsPlanDebugOverride::kForcePlain); + result = CommonGramsPlanDebugOverride::kForceCommonGrams; + }); + return result; +} + +using namespace phrase_impl; // NOLINT(google-build-using-namespace): module-internal impl namespace + +namespace internal { + +Status validate_prx_frame(std::span pos_flat, std::span pos_offsets, + uint32_t actual_total_docs, uint32_t expected_total_docs, + size_t expected_selected_docs, + std::span selected_doc_ordinals, + bool offsets_by_prx_ordinal, bool all_docs_selected) { + if (!all_docs_selected && expected_selected_docs != selected_doc_ordinals.size()) { + return Status::Error( + "phrase_query: selected prx ordinal-count mismatch"); + } + if (actual_total_docs != expected_total_docs) { + return Status::Error( + "phrase_query: prx total doc-count mismatch"); + } + const size_t expected_offsets = offsets_by_prx_ordinal + ? static_cast(expected_total_docs) + 1 + : expected_selected_docs + 1; + if (pos_offsets.size() != expected_offsets) { + return Status::Error( + offsets_by_prx_ordinal ? "phrase_query: full prx doc-count mismatch" + : "phrase_query: selected prx/doc-count mismatch"); + } + if (pos_offsets.back() != pos_flat.size()) { + return Status::Error( + "phrase_query: prx final offset mismatch"); + } + if (offsets_by_prx_ordinal && !selected_doc_ordinals.empty() && + static_cast(selected_doc_ordinals.back()) + 1 >= pos_offsets.size()) { + return Status::Error( + "phrase_query: prx ordinal offset out of range"); + } + return Status::OK(); +} + +Status decode_and_validate_prx_frame(ByteSource* source, + std::span selected_doc_ordinals, + bool decode_full, bool all_docs_selected, + uint32_t expected_total_docs, size_t expected_selected_docs, + std::vector* pos_flat, + std::vector* pos_offsets, + format::PrxDecodeContext* decode_context) { + DCHECK(decode_full || !all_docs_selected); + format::PrxDecodedShape decoded_shape; + format::PrxDecodeContext frame_context { + .stats = decode_context == nullptr ? nullptr : decode_context->stats, + .shape = &decoded_shape}; + if (decode_full) { + if (all_docs_selected) { + RETURN_IF_ERROR( + format::read_prx_window_csr(source, pos_flat, pos_offsets, &frame_context)); + } else { + RETURN_IF_ERROR(format::read_prx_window_csr_for_selection( + source, selected_doc_ordinals, pos_flat, pos_offsets, &frame_context)); + } + } else { + RETURN_IF_ERROR(format::read_prx_window_csr_selective( + source, selected_doc_ordinals, pos_flat, pos_offsets, &frame_context)); + } + return validate_prx_frame(*pos_flat, *pos_offsets, decoded_shape.total_docs, + expected_total_docs, expected_selected_docs, selected_doc_ordinals, + decode_full && !all_docs_selected, all_docs_selected); +} + +namespace { + +using PhraseVerifyClock = std::chrono::steady_clock; + +PhraseVerifyClock::time_point phrase_verify_clock_now() { +#ifdef BE_TEST + testing::note_phrase_verify_clock_read(); +#endif + return PhraseVerifyClock::now(); +} + +} // namespace + +uint64_t exclusive_phrase_verify_ns(uint64_t elapsed_ns, uint64_t decode_ns_before, + uint64_t decode_ns_after) { + DCHECK_GE(decode_ns_after, decode_ns_before); + const uint64_t decode_delta = decode_ns_after - decode_ns_before; + return elapsed_ns > decode_delta ? elapsed_ns - decode_delta : 0; +} + +PhraseVerifyTimer::PhraseVerifyTimer(format::PrxDecodeContext* decode_context) + : stats_(decode_context == nullptr ? nullptr : decode_context->stats) { + if (stats_ != nullptr) { + decode_ns_before_ = stats_->decode_ns; + start_ = phrase_verify_clock_now(); + } +} + +void PhraseVerifyTimer::commit_success() { + if (stats_ == nullptr) { + return; + } + const auto elapsed = + std::chrono::duration_cast(phrase_verify_clock_now() - start_) + .count(); + stats_->phrase_verify_ns += exclusive_phrase_verify_ns(static_cast(elapsed), + decode_ns_before_, stats_->decode_ns); +} + +} // namespace internal + +Status phrase_query(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids) { + return phrase_query_impl(idx, terms, docids, nullptr, nullptr); +} + +Status phrase_query(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids, QueryProfile* profile) { + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return phrase_query_impl(idx, terms, docids, profile == nullptr ? nullptr : &decode_context, + nullptr); +} + +Status phrase_query_with_frequencies(const LogicalIndexReader& idx, + const std::vector& terms, + std::vector* matches, QueryProfile* profile) { + if (matches == nullptr) { + return Status::Error( + "phrase_query_with_frequencies: null out"); + } + matches->clear(); + if (terms.size() < 2) { + return Status::Error( + "phrase_query_with_frequencies: at least two terms are required"); + } + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return phrase_query_impl(idx, terms, nullptr, profile == nullptr ? nullptr : &decode_context, + matches); +} + +Status planned_exact_phrase_query( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, QueryProfile* profile, ExactPhrasePlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + std::optional debug_override) { + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return planned_exact_phrase_query_impl( + idx, plain_query_info, gram_query_info, common_grams_identity, docids, + profile == nullptr ? nullptr : &decode_context, selected_plan, cost_model, + debug_override.has_value() ? *debug_override : common_grams_plan_debug_override()); +} + +Status planned_phrase_prefix_query( + const LogicalIndexReader& idx, const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, QueryProfile* profile, int32_t max_expansions, + PhrasePrefixPlanKind* selected_plan, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model, + std::optional debug_override) { + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return planned_phrase_prefix_query_impl( + idx, plain_query_info, gram_query_info, common_grams_identity, docids, max_expansions, + profile == nullptr ? nullptr : &decode_context, selected_plan, cost_model, + debug_override.has_value() ? *debug_override : common_grams_plan_debug_override()); +} + +Status phrase_prefix_query(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids, int32_t max_expansions) { + return phrase_prefix_query_impl(idx, terms, docids, max_expansions, nullptr, nullptr); +} + +Status phrase_prefix_query(const LogicalIndexReader& idx, const std::vector& terms, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions) { + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return phrase_prefix_query_impl(idx, terms, docids, max_expansions, + profile == nullptr ? nullptr : &decode_context, nullptr); +} + +Status phrase_prefix_query_with_frequencies(const LogicalIndexReader& idx, + const std::vector& terms, + std::vector* matches, + QueryProfile* profile, int32_t max_expansions) { + if (matches == nullptr) { + return Status::Error( + "phrase_prefix_query_with_frequencies: null out"); + } + matches->clear(); + if (terms.size() < 2) { + return Status::Error( + "phrase_prefix_query_with_frequencies: at least two terms are required"); + } + QueryProfileScope profile_scope(idx.reader(), profile); + format::PrxDecodeContext decode_context { + .stats = profile == nullptr ? nullptr : &profile->prx_decode_stats, + .query_stats = profile == nullptr ? nullptr : &profile->phrase_query_stats}; + return phrase_prefix_query_impl(idx, terms, nullptr, max_expansions, + profile == nullptr ? nullptr : &decode_context, nullptr, + matches); +} + +} // namespace doris::snii::query + +#ifdef BE_TEST +namespace doris::snii::query::internal::testing { +namespace { + +std::atomic& phrase_verify_clock_read_atomic() { + static std::atomic counter {0}; + return counter; +} + +} // namespace + +uint64_t phrase_verify_clock_read_count() { + return phrase_verify_clock_read_atomic().load(std::memory_order_relaxed); +} + +void reset_phrase_verify_clock_read_count() { + phrase_verify_clock_read_atomic().store(0, std::memory_order_relaxed); +} + +void note_phrase_verify_clock_read() { + phrase_verify_clock_read_atomic().fetch_add(1, std::memory_order_relaxed); +} + +} // namespace doris::snii::query::internal::testing +#endif diff --git a/be/src/storage/index/snii/query/phrase_query.h b/be/src/storage/index/snii/query/phrase_query.h new file mode 100644 index 00000000000000..879a2deb228de5 --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_query.h @@ -0,0 +1,131 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/query/query_info.h" +#include "storage/index/snii/format/prx_decode_stats.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// phrase_query -- MATCH_PHRASE: return the sorted docid set in which the terms +// occur consecutively (for some i, every term k appears at position pos+k in +// the same doc). It first builds the docid conjunction with docs-only posting +// reads, then fetches PRX only for chunks that can contain final candidates: +// 1. read preludes / docs-only posting ranges and intersect per-term docids; +// 2. fetch retained PRX chunks and stream positions for survivors; +// 3. for each surviving doc, check that some position p exists with +// term[0]@p, term[1]@p+1, ... term[n-1]@p+(n-1). +// An empty term list -> empty result. Any term absent -> empty result. +namespace doris::snii::query { + +enum class ExactPhrasePlanKind : uint8_t { + kPlain = 0, + kCommonGrams = 1, +}; + +enum class PhrasePrefixPlanKind : uint8_t { + kPlain = 0, + kCommonGrams = 1, +}; + +// Benchmark-only plan override. The debug points are inert unless the process-wide +// enable_debug_points switch is on, and planners apply them only after both complete plans resolve. +enum class CommonGramsPlanDebugOverride : uint8_t { + kNone = 0, + kForcePlain = 1, + kForceCommonGrams = 2, +}; + +inline constexpr char COMMON_GRAMS_FORCE_PLAIN_PLAN_DEBUG_POINT[] = + "snii.common_grams.force_plain_plan"; +inline constexpr char COMMON_GRAMS_FORCE_GRAM_PLAN_DEBUG_POINT[] = + "snii.common_grams.force_gram_plan"; + +CommonGramsPlanDebugOverride common_grams_plan_debug_override(); + +struct PhraseMatch { + uint32_t docid = 0; + uint32_t frequency = 0; + + bool operator==(const PhraseMatch&) const = default; +}; + +Status phrase_query(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids); +Status phrase_query(const reader::LogicalIndexReader& idx, const std::vector& terms, + std::vector* docids, QueryProfile* profile); + +// Scoring-only multi-term entry point. It runs the same opaque-term matcher as +// phrase_query(), but enumerates every valid phrase start in each matching doc. +Status phrase_query_with_frequencies(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* matches, + QueryProfile* profile = nullptr); + +// Selects between equivalent plain and CommonGrams exact-phrase plans above +// the existing opaque-term matcher. A missing or mismatched query identity +// forces the plain plan before any gram DICT lookup. +Status planned_exact_phrase_query( + const reader::LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, QueryProfile* profile = nullptr, + ExactPhrasePlanKind* selected_plan = nullptr, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model = {}, + std::optional debug_override = std::nullopt); + +Status planned_phrase_prefix_query( + const reader::LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& plain_query_info, + const segment_v2::InvertedIndexQueryInfo& gram_query_info, + const segment_v2::inverted_index::CommonGramsQueryIdentity* common_grams_identity, + std::vector* docids, QueryProfile* profile = nullptr, int32_t max_expansions = 0, + PhrasePrefixPlanKind* selected_plan = nullptr, + segment_v2::inverted_index::CommonGramsPlanCostModel cost_model = {}, + std::optional debug_override = std::nullopt); + +// phrase_prefix_query -- MATCH_PHRASE_PREFIX: the last item in `terms` is a +// term prefix and preceding items are exact terms. For example {"quick", "bro"} +// matches "quick brown" and "quick bronze". Empty terms -> empty result. +Status phrase_prefix_query(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* const docids, int32_t max_expansions = 0); +Status phrase_prefix_query(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions = 0); + +// Scoring-only multi-term entry point. Tail expansions are one logical phrase +// clause, so each phrase start contributes at most once even when several +// expanded terms share that position. +Status phrase_prefix_query_with_frequencies(const reader::LogicalIndexReader& idx, + const std::vector& terms, + std::vector* matches, + QueryProfile* profile = nullptr, + int32_t max_expansions = 0); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/phrase_verify_timer.h b/be/src/storage/index/snii/query/phrase_verify_timer.h new file mode 100644 index 00000000000000..bb04f33a64abed --- /dev/null +++ b/be/src/storage/index/snii/query/phrase_verify_timer.h @@ -0,0 +1,51 @@ +// 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. + +#pragma once + +#include +#include + +#include "storage/index/snii/format/prx_decode_stats.h" + +namespace doris::snii::query::internal { + +uint64_t exclusive_phrase_verify_ns(uint64_t elapsed_ns, uint64_t decode_ns_before, + uint64_t decode_ns_after); + +class PhraseVerifyTimer { +public: + explicit PhraseVerifyTimer(format::PrxDecodeContext* decode_context); + + void commit_success(); + +private: + format::PrxDecodeStats* stats_ = nullptr; + std::chrono::steady_clock::time_point start_; + uint64_t decode_ns_before_ = 0; +}; + +#ifdef BE_TEST +namespace testing { + +uint64_t phrase_verify_clock_read_count(); +void reset_phrase_verify_clock_read_count(); +void note_phrase_verify_clock_read(); + +} // namespace testing +#endif +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/prefix_query.cpp b/be/src/storage/index/snii/query/prefix_query.cpp new file mode 100644 index 00000000000000..632c80d8b099f8 --- /dev/null +++ b/be/src/storage/index/snii/query/prefix_query.cpp @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "storage/index/snii/query/prefix_query.h" + +#include +#include + +#include "storage/index/snii/query/internal/term_expansion.h" + +namespace doris::snii::query { + +using reader::LogicalIndexReader; + +Status prefix_query(const LogicalIndexReader& idx, std::string_view prefix, + std::vector* const docids, int32_t max_expansions) { + if (docids == nullptr) { + return Status::Error("prefix_query: null out"); + } + docids->clear(); + VectorDocIdSink sink(*docids); + return prefix_query(idx, prefix, &sink, max_expansions); +} + +Status prefix_query(const LogicalIndexReader& idx, std::string_view prefix, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions) { + QueryProfileScope profile_scope(idx.reader(), profile); + return prefix_query(idx, prefix, docids, max_expansions); +} + +Status prefix_query(const LogicalIndexReader& idx, std::string_view prefix, DocIdSink* const sink, + int32_t max_expansions) { + if (sink == nullptr) { + return Status::Error("prefix_query: null sink"); + } + + return internal::emit_expanded_docid_union( + idx, prefix, [](std::string_view) { return true; }, sink, max_expansions); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/prefix_query.h b/be/src/storage/index/snii/query/prefix_query.h new file mode 100644 index 00000000000000..6d17aa7b48b1c6 --- /dev/null +++ b/be/src/storage/index/snii/query/prefix_query.h @@ -0,0 +1,42 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// prefix_query -- MATCH_PREFIX semantics: enumerate dictionary terms with the +// requested prefix, then return the sorted docid set containing any enumerated +// term. Empty prefix enumerates all terms. No matching terms -> empty result. +namespace doris::snii::query { + +Status prefix_query(const reader::LogicalIndexReader& idx, std::string_view prefix, + std::vector* const docids, int32_t max_expansions = 0); +Status prefix_query(const reader::LogicalIndexReader& idx, std::string_view prefix, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions = 0); +Status prefix_query(const reader::LogicalIndexReader& idx, std::string_view prefix, + DocIdSink* const sink, int32_t max_expansions = 0); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/query_profile.cpp b/be/src/storage/index/snii/query/query_profile.cpp new file mode 100644 index 00000000000000..025272c421eadd --- /dev/null +++ b/be/src/storage/index/snii/query/query_profile.cpp @@ -0,0 +1,94 @@ +// 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 "storage/index/snii/query/query_profile.h" + +#include +#include +#include + +#include "storage/index/snii/io/file_reader.h" + +namespace doris::snii::query { + +namespace { + +#ifdef BE_TEST +std::atomic query_profile_clock_reads {0}; +#endif + +std::chrono::steady_clock::time_point query_profile_clock_now() { +#ifdef BE_TEST + query_profile_clock_reads.fetch_add(1, std::memory_order_relaxed); +#endif + return std::chrono::steady_clock::now(); +} + +} // namespace + +QueryProfileScope::QueryProfileScope(io::FileReader* reader, QueryProfile* profile) + : reader_(reader), profile_(profile) { + if (profile_ == nullptr) return; + + start_ = query_profile_clock_now(); + *profile_ = QueryProfile {}; + if (reader_ == nullptr) return; + + const io::IoMetrics* metrics = reader_->io_metrics(); + if (metrics == nullptr) return; + + profile_->has_io_metrics = true; + profile_->io_before = *metrics; +} + +QueryProfileScope::~QueryProfileScope() { + finish(); +} + +void QueryProfileScope::finish() { + if (profile_ == nullptr || finished_) return; + finished_ = true; + + const auto end = query_profile_clock_now(); + const auto elapsed = std::chrono::duration_cast(end - start_).count(); + profile_->elapsed_ns = std::max(1, static_cast(elapsed)); + + if (!profile_->has_io_metrics || reader_ == nullptr) return; + const io::IoMetrics* metrics = reader_->io_metrics(); + if (metrics == nullptr) { + profile_->has_io_metrics = false; + return; + } + profile_->io_after = *metrics; + profile_->io_delta = io::delta(profile_->io_after, profile_->io_before); +} + +#ifdef BE_TEST +namespace testing { + +uint64_t query_profile_clock_read_count() { + return query_profile_clock_reads.load(std::memory_order_relaxed); +} + +void reset_query_profile_clock_read_count() { + query_profile_clock_reads.store(0, std::memory_order_relaxed); +} + +} // namespace testing +#endif + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/query_profile.h b/be/src/storage/index/snii/query/query_profile.h new file mode 100644 index 00000000000000..2598deeef24f85 --- /dev/null +++ b/be/src/storage/index/snii/query/query_profile.h @@ -0,0 +1,67 @@ +// 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. + +#pragma once + +#include +#include + +#include "storage/index/snii/format/prx_decode_stats.h" +#include "storage/index/snii/io/io_metrics.h" + +namespace doris::snii::io { +class FileReader; +} + +namespace doris::snii::query { + +#ifdef BE_TEST +namespace testing { + +uint64_t query_profile_clock_read_count(); +void reset_query_profile_clock_read_count(); + +} // namespace testing +#endif + +struct QueryProfile { + uint64_t elapsed_ns = 0; + bool has_io_metrics = false; + io::IoMetrics io_before; + io::IoMetrics io_after; + io::IoMetrics io_delta; + format::PrxDecodeStats prx_decode_stats; + format::PhraseQueryExecutionStats phrase_query_stats; +}; + +class QueryProfileScope { +public: + QueryProfileScope(io::FileReader* reader, QueryProfile* profile); + ~QueryProfileScope(); + QueryProfileScope(const QueryProfileScope&) = delete; + QueryProfileScope& operator=(const QueryProfileScope&) = delete; + + void finish(); + +private: + io::FileReader* reader_ = nullptr; + QueryProfile* profile_ = nullptr; + std::chrono::steady_clock::time_point start_; + bool finished_ = false; +}; + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/regexp_query.cpp b/be/src/storage/index/snii/query/regexp_query.cpp new file mode 100644 index 00000000000000..d0115e1bc6907f --- /dev/null +++ b/be/src/storage/index/snii/query/regexp_query.cpp @@ -0,0 +1,144 @@ +// 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 "storage/index/snii/query/regexp_query.h" + +#include + +#include +#include +#include +#include +#include + +#include "storage/index/snii/query/internal/regex_prefix.h" +#include "storage/index/snii/query/internal/term_expansion.h" + +namespace doris::snii::query { + +namespace { + +bool is_regex_metachar(char c) { + switch (c) { + case '.': + case '^': + case '$': + case '|': + case '(': + case ')': + case '[': + case ']': + case '*': + case '+': + case '?': + case '{': + case '}': + case '\\': + return true; + default: + return false; + } +} + +std::string literal_prefix_for_regex(std::string_view pattern) { + std::string out; + size_t i = 0; + if (!pattern.empty() && pattern.front() == '^') { + i = 1; + } + for (; i < pattern.size(); ++i) { + const char c = pattern[i]; + if (is_regex_metachar(c)) { + break; + } + out.push_back(c); + } + return out; +} + +} // namespace + +namespace internal { + +std::string regex_enum_prefix(std::string_view pattern, const re2::RE2& re) { + // Left-anchored patterns can yield a tighter enumeration prefix via RE2's + // PossibleMatchRange than the conservative literal scan (which stops at the + // first metacharacter). The prefix only bounds how many dictionary terms are + // enumerated; final acceptance is still decided by RE2::FullMatch, so an + // over-wide prefix is always safe. Fall back to the literal scan otherwise. + if (!pattern.empty() && pattern.front() == '^' && re.ok()) { + std::string min_prefix; + std::string max_prefix; + if (re.PossibleMatchRange(&min_prefix, &max_prefix, 256) && !min_prefix.empty() && + !max_prefix.empty() && min_prefix.front() == max_prefix.front()) { + const auto mismatch_pair = std::ranges::mismatch(min_prefix, max_prefix); + const auto common_len = + static_cast(std::distance(min_prefix.begin(), mismatch_pair.in1)); + if (common_len > 0) { + return min_prefix.substr(0, common_len); + } + } + } + return literal_prefix_for_regex(pattern); +} + +} // namespace internal + +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, int32_t max_expansions) { + if (docids == nullptr) { + return Status::Error("regexp_query: null out"); + } + docids->clear(); + VectorDocIdSink sink(*docids); + return regexp_query(idx, pattern, &sink, max_expansions); +} + +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions) { + QueryProfileScope profile_scope(idx.reader(), profile); + return regexp_query(idx, pattern, docids, max_expansions); +} + +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + DocIdSink* const sink, int32_t max_expansions) { + if (sink == nullptr) { + return Status::Error("regexp_query: null sink"); + } + + re2::RE2::Options options; + options.set_log_errors(false); // Do not spam the BE log on user-supplied bad patterns. + const re2::RE2 re(re2::StringPiece(pattern.data(), pattern.size()), options); + if (!re.ok()) { + return Status::Error( + std::string("regexp_query: invalid regex: ") + re.error()); + } + + // RE2::FullMatch is anchored at both ends, matching std::regex_match whole-term + // semantics. Unsupported constructs (backreferences, lookaround) fail re.ok() + // above and return InvalidArgument rather than throwing. + const std::string enum_prefix = internal::regex_enum_prefix(pattern, re); + return internal::emit_expanded_docid_union( + idx, enum_prefix, + [&re](std::string_view term) { + return re2::RE2::FullMatch(re2::StringPiece(term.data(), term.size()), re); + }, + sink, max_expansions); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/regexp_query.h b/be/src/storage/index/snii/query/regexp_query.h new file mode 100644 index 00000000000000..b6b48ed57cfb6f --- /dev/null +++ b/be/src/storage/index/snii/query/regexp_query.h @@ -0,0 +1,44 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// regexp_query -- MATCH_REGEXP semantics over dictionary terms. The pattern is +// evaluated with RE2::FullMatch (anchored at both ends), so it must match the +// whole term. Invalid/unsupported patterns (e.g. backreferences, lookaround) +// return Status::InvalidArgument instead of throwing. Matching terms are executed +// as a sorted deduplicated docid union. +namespace doris::snii::query { + +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, int32_t max_expansions = 0); +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions = 0); +Status regexp_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + DocIdSink* const sink, int32_t max_expansions = 0); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/scoring_query.cpp b/be/src/storage/index/snii/query/scoring_query.cpp new file mode 100644 index 00000000000000..96afa8832b6237 --- /dev/null +++ b/be/src/storage/index/snii/query/scoring_query.cpp @@ -0,0 +1,913 @@ +// 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 "storage/index/snii/query/scoring_query.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/reader/dict_block_cache.h" +#include "storage/index/snii/reader/windowed_posting.h" + +namespace doris::snii::query { + +using format::DictEntry; +using format::DictEntryEnc; +using format::DictEntryKind; +using format::FrqPreludeReader; +using format::WindowMeta; +using reader::LogicalIndexReader; + +namespace { + +// One scored posting for one term in one doc. +struct TermPosting { + uint32_t docid = 0; + double score = 0.0; +}; + +// One window's block-max upper bound and the docid range it covers. block_max is +// true when max_score came from the frq_prelude columns (vs the exact-score +// fallback); both are valid upper bounds, so it is informational only. +struct WindowBound { + uint32_t first_docid = 0; // inclusive + uint32_t last_docid = 0; // inclusive + double max_score = 0.0; // block-max upper bound for any doc in this window + bool block_max = false; +}; + +// All scored postings of one query term plus its block-max metadata. +struct TermCursor { + std::vector postings; // ascending docid, exact per-doc scores + std::vector windows; // ascending, covering all postings + size_t pos = 0; // DAAT cursor into postings +}; + +uint32_t current_doc(const TermCursor& c) { + return c.pos < c.postings.size() ? c.postings[c.pos].docid + : std::numeric_limits::max(); +} + +// Reads one slim .frq window's bytes for a slim pod_ref/inline entry (prelude +// stripped). Windowed entries are handled separately via the prelude decode. +Status fetch_slim_window_bytes(const LogicalIndexReader& idx, const DictEntry& entry, + uint64_t frq_base, std::vector* window_owned, + Slice* window) { + if (entry.kind == DictEntryKind::kInline) { + *window = Slice(entry.frq_bytes); + return Status::OK(); + } + uint64_t win_abs = 0; + uint64_t win_len = 0; + RETURN_IF_ERROR(idx.resolve_frq_window(entry, frq_base, &win_abs, &win_len)); + io::BatchRangeFetcher fetcher(idx.reader()); + const size_t h = fetcher.add(win_abs, win_len); + RETURN_IF_ERROR(fetcher.fetch()); + Slice got = fetcher.get(h); + window_owned->assign(got.data(), got.data() + got.size()); + *window = Slice(*window_owned); + return Status::OK(); +} + +// Reads a windowed entry's frq_prelude (block-max columns live here). +Status fetch_prelude(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + FrqPreludeReader* out) { + const auto& region = idx.section_refs().posting_region; + const uint64_t prelude_abs = region.offset + frq_base + entry.frq_off_delta; + io::BatchRangeFetcher fetcher(idx.reader()); + const size_t h = fetcher.add(prelude_abs, entry.prelude_len); + RETURN_IF_ERROR(fetcher.fetch()); + return FrqPreludeReader::open(fetcher.get(h), out); +} + +// Builds per-window block-max bounds from a windowed entry's prelude. Each +// WindowMeta carries the window's max_freq / max_norm and its covered docid +// range (win_base+1 .. last_docid), so bounds come straight from the directory. +Status build_window_bounds(const FrqPreludeReader& prelude, const ScorerContext& ctx, double avgdl, + const Bm25Params& params, std::vector* windows) { + const uint32_t n = prelude.n_windows(); + for (uint32_t w = 0; w < n; ++w) { + WindowMeta m; + RETURN_IF_ERROR(prelude.window(w, &m)); + if (m.doc_count == 0) continue; + WindowBound wb; + wb.first_docid = static_cast(m.win_base) + (w == 0 ? 0u : 1u); + wb.last_docid = m.last_docid; + wb.max_score = ctx.max_score(m.max_freq, m.max_norm, avgdl, params); + wb.block_max = true; + windows->push_back(wb); + } + return Status::OK(); +} + +// Fallback single window covering all postings, bounded by the exact max score +// (always a valid upper bound, so pruning stays correct). +void single_window_fallback(const std::vector& postings, + std::vector* windows) { + if (postings.empty()) return; + WindowBound wb; + wb.first_docid = postings.front().docid; + wb.last_docid = postings.back().docid; + wb.block_max = false; + for (const auto& p : postings) wb.max_score = std::max(wb.max_score, p.score); + windows->push_back(wb); +} + +// Computes exact per-doc BM25 scores from decoded (docid, freq) vectors. +Status score_decoded(const stats::SniiStatsProvider& stats, const ScorerContext& ctx, double avgdl, + const Bm25Params& params, const std::vector& docids, + const std::vector& freqs, std::vector* out) { + out->reserve(docids.size()); + for (size_t i = 0; i < docids.size(); ++i) { + uint8_t norm = 0; + RETURN_IF_ERROR(stats.encoded_norm(docids[i], &norm)); + const uint32_t tf = i < freqs.size() ? freqs[i] : 1; + out->push_back({docids[i], ctx.score(tf, norm, avgdl, params)}); + } + return Status::OK(); +} + +// Decodes a slim/inline term's single .frq window ([dd_region][freq_region]) into +// docids/freqs using the entry's region metadata. +Status decode_slim(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + std::vector* docids, std::vector* freqs) { + std::vector owned; + Slice window; + RETURN_IF_ERROR(fetch_slim_window_bytes(idx, entry, frq_base, &owned, &window)); + const uint64_t dd_len = entry.dd_meta.disk_len; + if (dd_len > window.size()) { + return Status::Error( + "scoring_query: slim dd region exceeds window"); + } + Slice dd_region = window.subslice(0, static_cast(dd_len)); + RETURN_IF_ERROR(format::decode_dd_region(dd_region, entry.dd_meta, + /*win_base=*/0, docids)); + // G16-c freq-dropped segments (write_freq == false) carry a zero-length + // freq region on slim/inline entries. Fail with the SEMANTIC error and NOT + // with INVERTED_INDEX_FILE_CORRUPTED: the Doris segment iterator silently + // downgrades that code to a non-index evaluation, which would mask a + // by-design layout as data corruption once BM25 runs over mixed segments. + if (window.size() == static_cast(dd_len) && !docids->empty()) { + return Status::Error( + "scoring_query: freqs requested but the slim entry has no freq region " + "(freq-dropped positions index)"); + } + Slice freq_region = window.subslice(static_cast(dd_len), + window.size() - static_cast(dd_len)); + return format::decode_freq_region(freq_region, entry.freq_meta, docids->size(), freqs); +} + +// Builds the cursor for a windowed term: tiles all windows for exact scores and +// reads the prelude once for true per-window block-max bounds. +Status build_windowed_cursor(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const ScorerContext& ctx, const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, double avgdl, const Bm25Params& params, + TermCursor* cursor) { + reader::DecodedPosting posting; + // Scoring needs freqs for BM25: fetch the FULL windows (want_freq=true). + RETURN_IF_ERROR(reader::read_windowed_posting(idx, entry, frq_base, prx_base, + /*want_positions=*/false, + /*want_freq=*/true, &posting)); + RETURN_IF_ERROR(score_decoded(stats, ctx, avgdl, params, posting.docids, posting.freqs, + &cursor->postings)); + FrqPreludeReader prelude; + if (fetch_prelude(idx, entry, frq_base, &prelude).ok()) { + RETURN_IF_ERROR(build_window_bounds(prelude, ctx, avgdl, params, &cursor->windows)); + } + return Status::OK(); +} + +Status build_resolved_cursor(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const ScorerContext& ctx, const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, double avgdl, const Bm25Params& params, + TermCursor* cursor) { + const bool windowed = + entry.kind == DictEntryKind::kPodRef && entry.enc == DictEntryEnc::kWindowed; + if (windowed) { + RETURN_IF_ERROR(build_windowed_cursor(idx, stats, ctx, entry, frq_base, prx_base, avgdl, + params, cursor)); + } else { + std::vector docids; + std::vector freqs; + RETURN_IF_ERROR(decode_slim(idx, entry, frq_base, &docids, &freqs)); + RETURN_IF_ERROR(score_decoded(stats, ctx, avgdl, params, docids, freqs, &cursor->postings)); + } + if (cursor->windows.empty()) { + single_window_fallback(cursor->postings, &cursor->windows); + } + return Status::OK(); +} + +Status accumulate_decoded_candidate_scores(const stats::SniiStatsProvider& stats, + const ScorerContext& scorer, double avgdl, + const Bm25Params& params, + const std::vector& docids, + const std::vector& freqs, + std::span candidates, + std::span scores) { + DCHECK_EQ(candidates.size(), scores.size()); + DCHECK_EQ(docids.size(), freqs.size()); + size_t doc_index = 0; + size_t candidate_index = 0; + while (doc_index < docids.size() && candidate_index < candidates.size()) { + if (docids[doc_index] < candidates[candidate_index]) { + ++doc_index; + continue; + } + if (candidates[candidate_index] < docids[doc_index]) { + ++candidate_index; + continue; + } + uint8_t norm = 0; + RETURN_IF_ERROR(stats.encoded_norm(docids[doc_index], &norm)); + scores[candidate_index] += scorer.score(freqs[doc_index], norm, avgdl, params); + ++doc_index; + ++candidate_index; + } + return Status::OK(); +} + +struct CandidateWindowWork { + WindowMeta meta; + size_t candidate_begin = 0; + size_t candidate_end = 0; + size_t dd_handle = 0; + size_t freq_handle = 0; +}; + +Status accumulate_windowed_candidate_scores(const LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, + const std::vector& candidates, + const ScorerContext& scorer, double avgdl, + const Bm25Params& params, std::vector* scores) { + FrqPreludeReader prelude; + RETURN_IF_ERROR(reader::fetch_windowed_prelude(idx, entry, frq_base, &prelude)); + + std::vector windows; + const uint32_t window_count = prelude.n_windows(); + const bool candidates_cover_segment = + static_cast(candidates.size()) == stats.doc_count() && !candidates.empty() && + candidates.front() == 0 && + static_cast(candidates.back()) + 1 == stats.doc_count(); + bool scan_all = candidates_cover_segment; + if (scan_all) { + windows.resize(window_count); + std::iota(windows.begin(), windows.end(), 0); + } else { + prelude.select_covering_windows(candidates, &windows); + scan_all = windows.size() == window_count; + } + if (windows.empty()) { + return Status::OK(); + } + + // A sparse set must not be re-expanded into a near-full posting by merging + // across unselected windows. Dense/all-window reads retain the existing + // same-term coalescing policy and therefore the existing request shape. + const uint64_t coalesce_gap = scan_all ? reader::kSameTermCoalesceGap : 0; + io::BatchRangeFetcher fetcher(idx.reader(), coalesce_gap); + std::vector work; + work.reserve(windows.size()); + size_t candidate_cursor = 0; + for (uint32_t window : windows) { + CandidateWindowWork item; + RETURN_IF_ERROR(prelude.window(window, &item.meta)); + const uint32_t first_docid = + window == 0 ? 0 : static_cast(item.meta.win_base + 1); + while (candidate_cursor < candidates.size() && candidates[candidate_cursor] < first_docid) { + ++candidate_cursor; + } + item.candidate_begin = candidate_cursor; + while (candidate_cursor < candidates.size() && + candidates[candidate_cursor] <= item.meta.last_docid) { + ++candidate_cursor; + } + item.candidate_end = candidate_cursor; + + reader::WindowAbsRange range; + RETURN_IF_ERROR(reader::windowed_window_range( + idx, entry, frq_base, prx_base, prelude, window, + /*want_positions=*/false, /*want_freq=*/true, &range)); + item.dd_handle = fetcher.add(range.dd_off, range.dd_len); + item.freq_handle = fetcher.add(range.freq_off, range.freq_len); + work.push_back(std::move(item)); + } + RETURN_IF_ERROR(fetcher.fetch()); + + std::vector docids; + std::vector freqs; + std::vector> positions; + for (const auto& item : work) { + docids.clear(); + freqs.clear(); + RETURN_IF_ERROR(reader::decode_window_slices( + item.meta, fetcher.get(item.dd_handle), fetcher.get(item.freq_handle), Slice(), + /*want_positions=*/false, /*want_freq=*/true, &docids, &freqs, &positions)); + const size_t candidate_count = item.candidate_end - item.candidate_begin; + RETURN_IF_ERROR(accumulate_decoded_candidate_scores( + stats, scorer, avgdl, params, docids, freqs, + std::span(candidates) + .subspan(item.candidate_begin, candidate_count), + std::span(*scores).subspan(item.candidate_begin, candidate_count))); + } + return Status::OK(); +} + +Status accumulate_resolved_candidate_scores(const LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const DictEntry& entry, uint64_t frq_base, + uint64_t prx_base, + const std::vector& candidates, + const ScorerContext& scorer, double avgdl, + const Bm25Params& params, std::vector* scores) { + const bool windowed = + entry.kind == DictEntryKind::kPodRef && entry.enc == DictEntryEnc::kWindowed; + if (windowed) { + return accumulate_windowed_candidate_scores(idx, stats, entry, frq_base, prx_base, + candidates, scorer, avgdl, params, scores); + } + + std::vector docids; + std::vector freqs; + RETURN_IF_ERROR(decode_slim(idx, entry, frq_base, &docids, &freqs)); + return accumulate_decoded_candidate_scores(stats, scorer, avgdl, params, docids, freqs, + candidates, *scores); +} + +// Builds the cursor for one term: postings with exact scores + window bounds. +Status build_cursor(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const std::string& term, const Bm25Params& params, bool* found, + TermCursor* cursor) { + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term, found, &entry, &frq_base, &prx_base)); + if (!*found) return Status::OK(); + + const ScorerContext ctx = ScorerContext::make(stats.indexed_doc_count(), entry.df); + return build_resolved_cursor(idx, stats, ctx, entry, frq_base, prx_base, stats.avgdl(), params, + cursor); +} + +// Block-max upper bound for a term at a given docid: the max_score of the window +// covering docid (windows are ascending and contiguous). Beyond the last window +// the bound is 0 (the term cannot contribute). +double term_bound_at(const TermCursor& c, uint32_t docid) { + // Windows are ascending and contiguous; the first window whose last_docid is + // >= docid covers it. Its block-max is a valid upper bound for any contained + // doc, so it also bounds gaps between windows. + for (const auto& w : c.windows) { + if (docid <= w.last_docid) return w.max_score; + } + return 0.0; +} + +// Min-heap keyed on score (smallest at top) maintaining the top-K. +struct TopK { + explicit TopK(uint32_t k) : k_(k) {} + void offer(uint32_t docid, double score) { + if (heap_.size() < k_) { + heap_.push({score, docid}); + return; + } + if (heap_.empty()) return; + const Entry& worst = heap_.top(); // lowest score; ties: largest docid + const bool better = score > worst.first || (score == worst.first && docid < worst.second); + if (better) { + heap_.pop(); + heap_.push({score, docid}); + } + } + double threshold() const { return heap_.size() < k_ ? -1.0 : heap_.top().first; } + + using Entry = std::pair; + struct Cmp { + bool operator()(const Entry& a, const Entry& b) const { + if (a.first != b.first) return a.first > b.first; // min-score at top + return a.second < b.second; // for ties, largest docid at top (evictable) + } + }; + uint32_t k_; + std::priority_queue, Cmp> heap_; +}; + +void drain_sorted(TopK* topk, std::vector* out) { + std::vector all; + while (!topk->heap_.empty()) { + all.push_back({topk->heap_.top().second, topk->heap_.top().first}); + topk->heap_.pop(); + } + std::sort(all.begin(), all.end(), [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) return a.score > b.score; + return a.docid < b.docid; + }); + *out = std::move(all); +} + +Status build_cursors(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const std::vector& terms, const Bm25Params& params, + std::vector* cursors) { + for (const auto& term : terms) { + bool found = false; + TermCursor c; + RETURN_IF_ERROR(build_cursor(idx, stats, term, params, &found, &c)); + if (found && !c.postings.empty()) cursors->push_back(std::move(c)); + } + return Status::OK(); +} + +} // namespace + +Status scoring_query_candidates(const LogicalIndexReader& idx, + const stats::SniiStatsProvider& segment_stats, + const std::vector& terms, + const roaring::Roaring& final_candidates, double collection_avgdl, + const Bm25Params& params, std::vector* out) { + if (out == nullptr) { + return Status::Error( + "scoring_query_candidates: null out"); + } + out->clear(); + if (final_candidates.isEmpty()) { + return Status::OK(); + } + if (!(collection_avgdl > 0.0)) { + return Status::Error( + "scoring_query_candidates: collection avgdl must be positive"); + } + + std::vector candidate_docids; + candidate_docids.reserve(final_candidates.cardinality()); + for (uint32_t docid : final_candidates) { + candidate_docids.push_back(docid); + } + std::vector candidate_scores(candidate_docids.size(), 0.0); + reader::DictBlockCache dict_block_cache; + + for (const auto& term : terms) { + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term.physical_term, &found, &entry, &frq_base, &prx_base, + &dict_block_cache)); + if (!found) { + continue; + } + + const ScorerContext scorer = ScorerContext::from_idf(term.idf); + RETURN_IF_ERROR(accumulate_resolved_candidate_scores( + idx, segment_stats, entry, frq_base, prx_base, candidate_docids, scorer, + collection_avgdl, params, &candidate_scores)); + } + + std::vector scored_candidates; + scored_candidates.reserve(candidate_docids.size()); + for (size_t i = 0; i < candidate_docids.size(); ++i) { + scored_candidates.push_back({.docid = candidate_docids[i], .score = candidate_scores[i]}); + } + *out = std::move(scored_candidates); + return Status::OK(); +} + +Status scoring_query_exhaustive(const LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out) { + if (out == nullptr) + return Status::Error("scoring_query: null out"); + out->clear(); + if (k == 0) return Status::OK(); + + std::vector cursors; + RETURN_IF_ERROR(build_cursors(idx, stats, terms, params, &cursors)); + + std::unordered_map scores; + for (const auto& c : cursors) + for (const auto& p : c.postings) scores[p.docid] += p.score; + + std::vector all; + all.reserve(scores.size()); + for (const auto& [docid, score] : scores) all.push_back({docid, score}); + std::sort(all.begin(), all.end(), [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) return a.score > b.score; + return a.docid < b.docid; + }); + if (all.size() > k) all.resize(k); + *out = std::move(all); + return Status::OK(); +} + +namespace { + +// --- Phase C: selective-fetch (lazy window) WAND ----------------------------- +// +// A LazyTermCursor knows its per-window block-max bounds + docid ranges from the +// frq_prelude WITHOUT fetching any .frq window. Each window's exact (docid,score) +// postings are decoded on first access and cached, so a window is fetched at most +// once and ONLY when the WAND control flow touches a posting in it. Combined with +// window-level SkipTo (advance past whole windows whose last_docid < target via +// the prelude, never fetching them), the offer sequence is byte-identical to the +// eager scoring_query_wand path -- only the bytes read differ. +// +// Soundness: a window is fetched only when lazy_current_doc/lazy_skip_to land the +// cursor inside it, i.e. it covers a candidate the WAND pivot already proved can +// reach the running theta (bound >= theta). lazy_skip_to jumps the cursor to the +// SAME posting (first docid >= target) the eager per-doc walk would, so pivots, +// alignments and offers are identical to the eager path; only windows the eager +// path read-through-but-never-offered-from are skipped. Windows whose block-max +// bound never reaches theta are never the pivot, so never fetched. + +// One query term's lazily-fetched scoring state. +struct LazyTermCursor { + const LogicalIndexReader* idx = nullptr; + const stats::SniiStatsProvider* stats = nullptr; + ScorerContext ctx = ScorerContext::make(1, 1); + Bm25Params params; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + FrqPreludeReader prelude; + bool windowed = false; // false => slim/inline single block already materialized + + std::vector windows; // ascending; from prelude (or slim fallback) + std::vector postings; // sparse: only fetched windows are filled + std::vector win_start; // prefix offsets, size = windows.size()+1 + std::vector fetched; // size = windows.size() + size_t pos = 0; // virtual cursor over all windows' postings +}; + +// Total posting count across all windows (the virtual stream length). +uint32_t total_postings(const LazyTermCursor& c) { + return c.win_start.empty() ? 0 : c.win_start.back(); +} + +// Index of the window whose virtual range contains posting index p (p < total). +uint32_t window_of(const LazyTermCursor& c, uint32_t p) { + const auto it = std::upper_bound(c.win_start.begin(), c.win_start.end(), p); + return static_cast((it - c.win_start.begin()) - 1); +} + +// Fetches + decodes window w into the cursor's posting cache (idempotent). Only +// reached when the WAND proves window w can still contribute to the top-K. +Status materialize_window(LazyTermCursor* c, uint32_t w) { + if (c->fetched[w]) return Status::OK(); + WindowMeta meta; + RETURN_IF_ERROR(c->prelude.window(w, &meta)); + reader::WindowAbsRange r; + RETURN_IF_ERROR(reader::windowed_window_range( + *c->idx, c->entry, c->frq_base, c->prx_base, c->prelude, w, + /*want_positions=*/false, /*want_freq=*/true, &r)); + // Scoring needs docids + freqs: fetch the window's dd sub-range AND freq sub-range. + io::BatchRangeFetcher fetcher(c->idx->reader(), reader::kSameTermCoalesceGap); + const size_t dh = fetcher.add(r.dd_off, r.dd_len); + const size_t fh = fetcher.add(r.freq_off, r.freq_len); + RETURN_IF_ERROR(fetcher.fetch()); + std::vector docids; + std::vector freqs; + std::vector> pos; + RETURN_IF_ERROR(reader::decode_window_slices(meta, fetcher.get(dh), fetcher.get(fh), Slice(), + /*want_positions=*/false, + /*want_freq=*/true, &docids, &freqs, &pos)); + if (docids.size() != c->win_start[w + 1] - c->win_start[w]) { + return Status::Error( + "scoring_query: selective window doc-count drift"); + } + std::vector scored; + RETURN_IF_ERROR( + score_decoded(*c->stats, c->ctx, c->stats->avgdl(), c->params, docids, freqs, &scored)); + std::copy(scored.begin(), scored.end(), c->postings.begin() + c->win_start[w]); + c->fetched[w] = 1; + return Status::OK(); +} + +// Current docid at the cursor, fetching the covering window if needed. Exhausted +// cursor -> UINT32_MAX. +Status lazy_current_doc(LazyTermCursor* c, uint32_t* docid) { + if (c->pos >= total_postings(*c)) { + *docid = std::numeric_limits::max(); + return Status::OK(); + } + const uint32_t w = window_of(*c, static_cast(c->pos)); + RETURN_IF_ERROR(materialize_window(c, w)); + *docid = c->postings[c->pos].docid; + return Status::OK(); +} + +// Advances pos to the first posting with docid >= target, skipping ENTIRE windows +// whose last_docid < target WITHOUT fetching them (prelude-only), then fetching +// just the landing window. Lands on the same posting the eager per-doc walk would. +Status lazy_skip_to(LazyTermCursor* c, uint32_t target) { + const uint32_t total = total_postings(*c); + while (c->pos < total) { + const uint32_t w = window_of(*c, static_cast(c->pos)); + if (c->windows[w].last_docid >= target) break; + c->pos = c->win_start[w + 1]; // skip this window entirely (no fetch) + } + if (c->pos >= total) return Status::OK(); + const uint32_t w = window_of(*c, static_cast(c->pos)); + RETURN_IF_ERROR(materialize_window(c, w)); + while (c->pos < total && c->postings[c->pos].docid < target) ++c->pos; + return Status::OK(); +} + +// Initializes a lazy windowed cursor from the prelude alone: per-window block-max +// bounds + ranges + cache slots, with NO .frq window fetched. +Status build_lazy_windowed(LazyTermCursor* c) { + RETURN_IF_ERROR(reader::fetch_windowed_prelude(*c->idx, c->entry, c->frq_base, &c->prelude)); + RETURN_IF_ERROR( + build_window_bounds(c->prelude, c->ctx, c->stats->avgdl(), c->params, &c->windows)); + // build_window_bounds keeps only non-empty windows, in window order. Build the + // matching prefix-sum of doc_counts over those same non-empty windows so the + // bound list, win_start and fetched stay 1:1. + const uint32_t nb = static_cast(c->windows.size()); + c->win_start.assign(nb + 1, 0); + c->fetched.assign(nb, 0); + uint32_t bi = 0; + uint32_t acc = 0; + for (uint32_t w = 0; w < c->prelude.n_windows() && bi < nb; ++w) { + WindowMeta meta; + RETURN_IF_ERROR(c->prelude.window(w, &meta)); + if (meta.doc_count == 0) continue; + acc += meta.doc_count; + c->win_start[++bi] = acc; + } + c->postings.assign(acc, TermPosting {}); + return Status::OK(); +} + +// Initializes a slim/inline cursor: its single window is small, so fetch + score +// it eagerly (exactly as the existing path). One bound covers all its postings. +Status build_lazy_slim(LazyTermCursor* c) { + std::vector docids; + std::vector freqs; + RETURN_IF_ERROR(decode_slim(*c->idx, c->entry, c->frq_base, &docids, &freqs)); + RETURN_IF_ERROR(score_decoded(*c->stats, c->ctx, c->stats->avgdl(), c->params, docids, freqs, + &c->postings)); + single_window_fallback(c->postings, &c->windows); + c->win_start = {0, static_cast(c->postings.size())}; + c->fetched.assign(1, 1); // already materialized + return Status::OK(); +} + +// Builds a LazyTermCursor for one term: prelude-only for windowed terms (no .frq +// fetched), fully-materialized single window for slim/inline (small). +Status build_lazy_cursor(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const std::string& term, const Bm25Params& params, bool* found, + LazyTermCursor* c) { + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term, found, &c->entry, &c->frq_base, &prx_base)); + if (!*found) return Status::OK(); + c->idx = &idx; + c->stats = &stats; + c->params = params; + c->prx_base = prx_base; + c->ctx = ScorerContext::make(stats.indexed_doc_count(), c->entry.df); + c->windowed = + c->entry.kind == DictEntryKind::kPodRef && c->entry.enc == DictEntryEnc::kWindowed; + return c->windowed ? build_lazy_windowed(c) : build_lazy_slim(c); +} + +Status selective_build_cursors(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const std::vector& terms, const Bm25Params& params, + std::vector* cursors) { + for (const auto& term : terms) { + bool found = false; + LazyTermCursor c; + RETURN_IF_ERROR(build_lazy_cursor(idx, stats, term, params, &found, &c)); + if (found && total_postings(c) > 0) cursors->push_back(std::move(c)); + } + return Status::OK(); +} + +// Block-max upper bound for a lazy cursor at docid: block_max of the window +// covering docid (ascending, contiguous). Beyond the last window -> 0. Same +// semantics as term_bound_at over the eager cursor's window list. +double lazy_term_bound_at(const LazyTermCursor& c, uint32_t docid) { + for (const auto& w : c.windows) { + if (docid <= w.last_docid) return w.max_score; + } + return 0.0; +} + +// Sorts cursors ascending by current docid (materializing each cursor's current +// covering window), returning the smallest current docid via *front. +Status selective_sort_by_doc(std::vector* cursors, uint32_t* front) { + std::vector cur(cursors->size()); + for (size_t i = 0; i < cursors->size(); ++i) { + RETURN_IF_ERROR(lazy_current_doc(&(*cursors)[i], &cur[i])); + } + std::vector order(cursors->size()); + for (size_t i = 0; i < order.size(); ++i) order[i] = i; + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { return cur[a] < cur[b]; }); + std::vector sorted; + sorted.reserve(cursors->size()); + for (size_t i : order) sorted.push_back(std::move((*cursors)[i])); + *cursors = std::move(sorted); + *front = order.empty() ? std::numeric_limits::max() : cur[order.front()]; + return Status::OK(); +} + +// Finds the pivot term: the first cursor (current-docid order) at which the +// accumulated block-max bound reaches theta. >= keeps boundary ties (matching the +// exhaustive total order). *found=false when no remaining doc can beat theta. +Status selective_pivot(std::vector* cursors, double theta, size_t* pivot, + uint32_t* pivot_doc, bool* found) { + double bound = 0.0; + *found = false; + for (size_t i = 0; i < cursors->size(); ++i) { + uint32_t d = 0; + RETURN_IF_ERROR(lazy_current_doc(&(*cursors)[i], &d)); + if (d == std::numeric_limits::max()) break; + bound += lazy_term_bound_at((*cursors)[i], d); + if (bound >= theta) { + *pivot = i; + *pivot_doc = d; + *found = true; + return Status::OK(); + } + } + return Status::OK(); +} + +// Scores the aligned pivot doc exactly (summing all cursors AT pivot_doc) and +// advances those cursors by one posting. +Status selective_score_pivot(std::vector* cursors, uint32_t pivot_doc, TopK* topk) { + double doc_score = 0.0; + for (auto& c : *cursors) { + uint32_t d = 0; + RETURN_IF_ERROR(lazy_current_doc(&c, &d)); + if (d == pivot_doc) { + doc_score += c.postings[c.pos].score; // window already materialized + ++c.pos; + } + } + topk->offer(pivot_doc, doc_score); + return Status::OK(); +} + +// Advances the first lagging cursor (current doc < pivot_doc) up to pivot_doc. +Status selective_advance_lagging(std::vector* cursors, uint32_t pivot_doc) { + for (auto& c : *cursors) { + uint32_t d = 0; + RETURN_IF_ERROR(lazy_current_doc(&c, &d)); + if (d < pivot_doc) { + RETURN_IF_ERROR(lazy_skip_to(&c, pivot_doc)); + return Status::OK(); + } + } + return Status::OK(); +} + +// One WAND iteration body: sort, pick pivot, then either score (aligned) or skip +// a lagging cursor forward. *done=true ends the loop. +Status selective_step(std::vector* cursors, TopK* topk, bool* done) { + uint32_t front = 0; + RETURN_IF_ERROR(selective_sort_by_doc(cursors, &front)); + if (cursors->empty() || front == std::numeric_limits::max()) { + *done = true; + return Status::OK(); + } + size_t pivot = 0; + uint32_t pivot_doc = 0; + bool found_pivot = false; + RETURN_IF_ERROR(selective_pivot(cursors, topk->threshold(), &pivot, &pivot_doc, &found_pivot)); + if (!found_pivot) { + *done = true; + return Status::OK(); + } + if (front == pivot_doc) { + return selective_score_pivot(cursors, pivot_doc, topk); + } + return selective_advance_lagging(cursors, pivot_doc); +} + +Status selective_wand_loop(std::vector* cursors, TopK* topk) { + bool done = false; + while (!done) { + RETURN_IF_ERROR(selective_step(cursors, topk, &done)); + } + return Status::OK(); +} + +} // namespace + +Status scoring_query_wand_selective(const LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out) { + if (out == nullptr) + return Status::Error("scoring_query: null out"); + out->clear(); + if (k == 0) return Status::OK(); + + std::vector cursors; + RETURN_IF_ERROR(selective_build_cursors(idx, stats, terms, params, &cursors)); + + TopK topk(k); + RETURN_IF_ERROR(selective_wand_loop(&cursors, &topk)); + drain_sorted(&topk, out); + return Status::OK(); +} + +Status scoring_query_wand(const LogicalIndexReader& idx, const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out) { + if (out == nullptr) + return Status::Error("scoring_query: null out"); + out->clear(); + if (k == 0) return Status::OK(); + + std::vector cursors; + RETURN_IF_ERROR(build_cursors(idx, stats, terms, params, &cursors)); + + TopK topk(k); + // Document-at-a-time WAND with block-max bounds. + while (true) { + // Sort cursors by current docid (ascending; exhausted cursors sink). + std::sort(cursors.begin(), cursors.end(), [](const TermCursor& a, const TermCursor& b) { + return current_doc(a) < current_doc(b); + }); + if (cursors.empty() || + current_doc(cursors.front()) == std::numeric_limits::max()) { + break; + } + + const double theta = topk.threshold(); + // Accumulate block-max upper bounds in docid order to find the pivot term. + double bound = 0.0; + size_t pivot = 0; + bool found_pivot = false; + for (size_t i = 0; i < cursors.size(); ++i) { + const uint32_t d = current_doc(cursors[i]); + if (d == std::numeric_limits::max()) break; + bound += term_bound_at(cursors[i], d); + // Use >= (not >) so a doc whose upper bound only TIES the K-th threshold is + // still explored and exact-scored: under the (score desc, docid asc) total + // order a tie can still evict the current K-th entry (smaller docid wins), + // exactly as the exhaustive path would. Strict > would wrongly prune ties. + if (bound >= theta) { + pivot = i; + found_pivot = true; + break; + } + } + if (!found_pivot) break; // no doc can beat the threshold anymore. + + const uint32_t pivot_doc = current_doc(cursors[pivot]); + if (current_doc(cursors.front()) == pivot_doc) { + // All cursors at the pivot doc are aligned: score it exactly. + double doc_score = 0.0; + for (auto& c : cursors) { + if (current_doc(c) == pivot_doc) { + doc_score += c.postings[c.pos].score; + ++c.pos; + } + } + topk.offer(pivot_doc, doc_score); + } else { + // Advance a lagging cursor toward pivot_doc (skip docs it cannot win on). + for (auto& c : cursors) { + if (current_doc(c) < pivot_doc) { + while (c.pos < c.postings.size() && c.postings[c.pos].docid < pivot_doc) { + ++c.pos; + } + break; + } + } + } + } + drain_sorted(&topk, out); + return Status::OK(); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/scoring_query.h b/be/src/storage/index/snii/query/scoring_query.h new file mode 100644 index 00000000000000..e8e83a7881efdb --- /dev/null +++ b/be/src/storage/index/snii/query/scoring_query.h @@ -0,0 +1,98 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/stats/snii_stats_provider.h" + +// scoring_query -- top-K BM25 scored retrieval over one logical index for one or +// more query terms. Two entry points produce IDENTICAL rankings: +// - scoring_query_exhaustive(): scores every candidate document (the baseline +// correctness oracle). +// - scoring_query_wand(): a block-max / WAND-style optimization that uses the +// per-window max_freq / max_norm columns from the frq_prelude to bound each +// window's best possible score and SKIP windows that cannot enter the +// current top-K. A window without block-max stats (slim/inline entries or a +// missing prelude) is never pruned, so the result still equals the +// exhaustive ranking. +// +// Results are sorted by score descending; ties are broken by ascending docid so +// the ordering is deterministic and the two paths compare equal. +namespace doris::snii::query { + +// One scored hit. +struct ScoredDoc { + uint32_t docid = 0; + double score = 0.0; +}; + +// One logical scoring clause after its plain term has been routed to this +// segment's physical key. IDF remains collection-scoped and is never derived +// from the physical term's segment-local dictionary entry. +struct CollectionScoringTerm { + std::string physical_term; + double idf = 0.0; +}; + +// Scores every document in final_candidates using collection-scoped IDF and +// avgdl plus segment-local TF/norm. Results are returned in ascending docid +// order. Repeated terms are repeated scoring clauses. This path deliberately +// does not use the segment-local WAND bounds below. +Status scoring_query_candidates(const reader::LogicalIndexReader& idx, + const stats::SniiStatsProvider& segment_stats, + const std::vector& terms, + const roaring::Roaring& final_candidates, double collection_avgdl, + const Bm25Params& params, std::vector* out); + +// Exhaustive baseline: score every doc that contains any query term, return the +// top-k by score. params controls k1/b. Unknown terms are skipped. +Status scoring_query_exhaustive(const reader::LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out); + +// WAND-style block-max pruning. MUST return the same top-k as the exhaustive +// path. Windows whose block-max upper bound cannot beat the current k-th score +// are skipped; windows lacking block-max stats are scored fully. +Status scoring_query_wand(const reader::LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out); + +// SELECTIVE-FETCH block-max WAND (design spec section 5, "Phase C"). Same WAND / +// theta / >= tie machinery as scoring_query_wand, but it DEFERS the .frq window +// fetch: for each windowed term it first reads ONLY the frq_prelude (block-max +// columns), then fetches a term's .frq window lazily and at most once -- and ONLY +// when the running block-max bound proves a doc in that window can still reach the +// top-K (bound >= theta). A window the bound rules out is never fetched. The +// result (top-K docids AND scores, INCLUDING ties) is byte-identical to +// scoring_query_exhaustive / scoring_query_wand; only the bytes read differ. +// Slim/inline terms (no prelude) are fetched fully, exactly as today. +Status scoring_query_wand_selective(const reader::LogicalIndexReader& idx, + const stats::SniiStatsProvider& stats, + const std::vector& terms, uint32_t k, + const Bm25Params& params, std::vector* out); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/term_expansion.cpp b/be/src/storage/index/snii/query/term_expansion.cpp new file mode 100644 index 00000000000000..ea4b682df783bd --- /dev/null +++ b/be/src/storage/index/snii/query/term_expansion.cpp @@ -0,0 +1,168 @@ +// 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 "storage/index/snii/query/internal/term_expansion.h" + +#include +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/internal/docid_union.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/reader/dict_block_cache.h" + +namespace doris::snii::query::internal { +namespace { + +Status legacy_raw_prefix_exists(const reader::LogicalIndexReader& idx, std::string_view prefix, + bool* exists, reader::DictBlockCache* cache) { + DORIS_CHECK(exists != nullptr); + *exists = false; + return idx.visit_prefix_terms( + prefix, + [&](reader::LogicalIndexReader::PrefixHit&&, bool* stop) -> Status { + *exists = true; + *stop = true; + return Status::OK(); + }, + cache); +} + +Status prove_legacy_raw_has_no_reserved_terms(const reader::LogicalIndexReader& idx, + reader::DictBlockCache* cache) { + bool exists = false; + RETURN_IF_ERROR(legacy_raw_prefix_exists(idx, segment_v2::inverted_index::CG_V1_MARKER, &exists, + cache)); + if (!exists) { + RETURN_IF_ERROR( + legacy_raw_prefix_exists(idx, format::kPhraseBigramTermMarker, &exists, cache)); + } + if (exists) { + return Status::Error( + "SNII legacy raw expansion overlaps an existing internal term namespace"); + } + return Status::OK(); +} + +} // namespace + +Status visit_expanded_plain_terms(const reader::LogicalIndexReader& idx, + std::string_view enum_prefix, const TermMatcher& matches, + const reader::LogicalIndexReader::PrefixHitVisitor& visitor, + int32_t max_expansions) { + if (!matches || !visitor) { + return Status::Error( + "term_expansion: null matcher or visitor"); + } + + std::string physical_prefix; + bool representable = false; + reader::DictBlockCache dict_cache(/*max_entries=*/1); + const auto version = plain_term_key_version(idx); + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kLegacyRaw && + enum_prefix.empty()) { + RETURN_IF_ERROR(prove_legacy_raw_has_no_reserved_terms(idx, &dict_cache)); + representable = true; + } else { + RETURN_IF_ERROR( + route_plain_enumeration_prefix(idx, enum_prefix, &physical_prefix, &representable)); + } + if (!representable) { + return Status::OK(); + } + + int32_t count = 0; + bool stop_expansion = false; + std::string decoded_scratch; + const auto visit_hit = [&](reader::LogicalIndexReader::PrefixHit&& hit, bool* stop) -> Status { + std::string_view logical_term; + if (version != segment_v2::inverted_index::PlainTermKeyVersion::kEscapedV1 || + !hit.term.starts_with(segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX)) { + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kEscapedV1) { + DCHECK(!segment_v2::inverted_index::is_internal_term_key(hit.term)); + } + logical_term = hit.term; + decoded_scratch.clear(); + } else { + auto decoded = segment_v2::inverted_index::decode_plain_term_view(hit.term, version, + &decoded_scratch); + if (!decoded.has_value()) { + return std::move(decoded.error()); + } + logical_term = *decoded; + } + if (!matches(logical_term)) { + return Status::OK(); + } + if (!decoded_scratch.empty()) { + hit.term = decoded_scratch; + } + bool visitor_stop = false; + RETURN_IF_ERROR(visitor(std::move(hit), &visitor_stop)); + ++count; + *stop = visitor_stop || (max_expansions > 0 && count >= max_expansions); + stop_expansion = *stop; + return Status::OK(); + }; + + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kEscapedV1 && + physical_prefix.empty()) { + RETURN_IF_ERROR(idx.visit_term_range( + /*lower_inclusive=*/ {}, segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN, + visit_hit, &dict_cache)); + if (!stop_expansion && (max_expansions <= 0 || count < max_expansions)) { + RETURN_IF_ERROR( + idx.visit_term_range(segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_END, + /*upper_exclusive=*/std::nullopt, visit_hit, &dict_cache)); + } + } else { + RETURN_IF_ERROR(idx.visit_prefix_terms( + physical_prefix, + [&](reader::LogicalIndexReader::PrefixHit&& hit, bool* stop) -> Status { + if (version == segment_v2::inverted_index::PlainTermKeyVersion::kEscapedV1) { + DCHECK(!segment_v2::inverted_index::is_internal_term_key(hit.term)); + } + return visit_hit(std::move(hit), stop); + }, + &dict_cache)); + } + return Status::OK(); +} + +Status emit_expanded_docid_union(const reader::LogicalIndexReader& idx, + std::string_view enum_prefix, const TermMatcher& matches, + DocIdSink* const sink, int32_t max_expansions) { + if (sink == nullptr) { + return Status::Error("term_expansion: null sink"); + } + + std::vector postings; + RETURN_IF_ERROR(visit_expanded_plain_terms( + idx, enum_prefix, matches, + [&](reader::LogicalIndexReader::PrefixHit&& hit, bool*) { + postings.push_back({std::move(hit.entry), hit.frq_base, hit.prx_base}); + return Status::OK(); + }, + max_expansions)); + return emit_docid_union(idx, postings, sink); +} + +} // namespace doris::snii::query::internal diff --git a/be/src/storage/index/snii/query/term_query.cpp b/be/src/storage/index/snii/query/term_query.cpp new file mode 100644 index 00000000000000..633229eab67e97 --- /dev/null +++ b/be/src/storage/index/snii/query/term_query.cpp @@ -0,0 +1,58 @@ +// 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 "storage/index/snii/query/term_query.h" + +#include + +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" + +namespace doris::snii::query { + +using format::DictEntry; +using reader::LogicalIndexReader; + +Status term_query(const LogicalIndexReader& idx, std::string_view term, + std::vector* docids) { + if (docids == nullptr) + return Status::Error("term_query: null out"); + docids->clear(); + VectorDocIdSink sink(*docids); + return term_query(idx, term, &sink); +} + +Status term_query(const LogicalIndexReader& idx, std::string_view term, DocIdSink* sink) { + if (sink == nullptr) + return Status::Error("term_query: null sink"); + + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(idx.lookup(term, &found, &entry, &frq_base, &prx_base)); + if (!found) return Status::OK(); + return internal::read_docid_posting(idx, entry, frq_base, prx_base, sink); +} + +Status term_query(const LogicalIndexReader& idx, std::string_view term, + std::vector* docids, QueryProfile* profile) { + QueryProfileScope profile_scope(idx.reader(), profile); + return term_query(idx, term, docids); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/term_query.h b/be/src/storage/index/snii/query/term_query.h new file mode 100644 index 00000000000000..7296acd9c60702 --- /dev/null +++ b/be/src/storage/index/snii/query/term_query.h @@ -0,0 +1,41 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// term_query -- the simplest SNII query: return the sorted docid set that +// contains term. It runs the term lookup on the logical index, then issues a +// single batched .frq range read (one serial round) to decode the postings. +// Absent term -> empty result (OK status). +namespace doris::snii::query { + +Status term_query(const reader::LogicalIndexReader& idx, std::string_view term, + std::vector* docids); +Status term_query(const reader::LogicalIndexReader& idx, std::string_view term, DocIdSink* sink); +Status term_query(const reader::LogicalIndexReader& idx, std::string_view term, + std::vector* docids, QueryProfile* profile); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/wildcard_query.cpp b/be/src/storage/index/snii/query/wildcard_query.cpp new file mode 100644 index 00000000000000..e8313c499797aa --- /dev/null +++ b/be/src/storage/index/snii/query/wildcard_query.cpp @@ -0,0 +1,77 @@ +// 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 "storage/index/snii/query/wildcard_query.h" + +#include +#include +#include +#include + +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/internal/wildcard_matcher.h" + +namespace doris::snii::query { + +namespace { + +std::string literal_prefix_for_wildcard(std::string_view pattern) { + std::string out; + for (char c : pattern) { + if (c == '*' || c == '?') { + break; + } + out.push_back(c); + } + return out; +} + +} // namespace + +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, int32_t max_expansions) { + if (docids == nullptr) { + return Status::Error("wildcard_query: null out"); + } + docids->clear(); + VectorDocIdSink sink(*docids); + return wildcard_query(idx, pattern, &sink, max_expansions); +} + +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions) { + QueryProfileScope profile_scope(idx.reader(), profile); + return wildcard_query(idx, pattern, docids, max_expansions); +} + +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + DocIdSink* const sink, int32_t max_expansions) { + if (sink == nullptr) { + return Status::Error("wildcard_query: null sink"); + } + const std::string enum_prefix = literal_prefix_for_wildcard(pattern); + // Request-scoped matcher: its two DP scratch rows are reused across every + // visited dictionary term, so the whole-dictionary scan triggered by a + // leading wildcard performs O(1) scratch allocations instead of O(2N). + internal::WildcardMatcher<> matcher(pattern); + return internal::emit_expanded_docid_union( + idx, enum_prefix, [&matcher](std::string_view term) { return matcher(term); }, sink, + max_expansions); +} + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/query/wildcard_query.h b/be/src/storage/index/snii/query/wildcard_query.h new file mode 100644 index 00000000000000..66c08b18ae270e --- /dev/null +++ b/be/src/storage/index/snii/query/wildcard_query.h @@ -0,0 +1,42 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// wildcard_query -- MATCH_WILDCARD semantics over dictionary terms. `*` matches +// any byte sequence, `?` matches one byte, and all other bytes match literally. +// Matching terms are executed as a sorted deduplicated docid union. +namespace doris::snii::query { + +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, int32_t max_expansions = 0); +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + std::vector* const docids, QueryProfile* profile, + int32_t max_expansions = 0); +Status wildcard_query(const reader::LogicalIndexReader& idx, std::string_view pattern, + DocIdSink* const sink, int32_t max_expansions = 0); + +} // namespace doris::snii::query diff --git a/be/src/storage/index/snii/reader/dict_block_cache.h b/be/src/storage/index/snii/reader/dict_block_cache.h new file mode 100644 index 00000000000000..18c65978ee9263 --- /dev/null +++ b/be/src/storage/index/snii/reader/dict_block_cache.h @@ -0,0 +1,124 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_block.h" + +// DictBlockCache -- a REQUEST-SCOPED (per-query) MRU cache of decoded DICT +// blocks, keyed by block ordinal. +// +// Why request-scoped (and not a reader-level shared cache): the same DICT block +// is decoded once per LogicalIndexReader::lookup() today, so a multi-term query +// (phrase / boolean / conjunction) whose terms fall in the same block re-runs +// the zstd decompress + CRC verify + anchor parse for every term. Threading one +// of these caches through a single query's lookups collapses that to a single +// decode per unique block. +// +// CONCURRENCY: this object carries NO shared mutable state and is intentionally +// NOT thread-safe. It is meant to live on one query's stack/context and be used +// by a single thread; concurrent queries each own a separate cache. The shared +// LogicalIndexReader therefore stays const and lock-free -- no lock is ever held +// across a decode/IO. (The cross-query, lock-striped variant that would let +// queries share decoded blocks is deferred to the T26 concurrency work.) +namespace doris::snii::reader { + +// A decoded DICT block with stable backing storage. Heap-allocated and owned by +// a shared_ptr so the embedded DictBlockReader's Slice into `bytes` stays valid +// for the whole lifetime of any pin handed to a caller -- even after the block +// has been evicted from the cache. +struct DecodedDictBlock { + std::vector bytes; // decompressed (or raw) block bytes + format::DictBlockReader reader; // its Slice points into `bytes` +}; + +class DictBlockCache { +public: + // Loads (decodes) the block for an ordinal into a freshly heap-allocated + // DecodedDictBlock. Always invoked OUTSIDE any cache bookkeeping; it performs + // the file read + optional zstd decompress + CRC/anchor parse. + using Loader = std::function*)>; + + // A small fixed bound is enough for a single query: it only needs to keep the + // handful of distinct blocks touched while resolving one query's terms. + static constexpr size_t kDefaultMaxEntries = 8; + + DictBlockCache() = default; + explicit DictBlockCache(size_t max_entries) + : max_entries_(max_entries == 0 ? 1 : max_entries) {} + + // Returns the decoded block for `ordinal`, invoking `loader` only on a miss. + // The returned pin keeps the block alive for the caller's use regardless of + // any later eviction. On a hit, `loader` is not called (no re-decode). + Status get_or_load(uint32_t ordinal, const Loader& loader, + std::shared_ptr* out) { + if (auto it = index_.find(ordinal); it != index_.end()) { + order_.splice(order_.begin(), order_, it->second); // promote to MRU + *out = it->second->block; + return Status::OK(); + } + + std::shared_ptr loaded; + // decode happens here, never under a lock (explicit Status, header-safe: + // RETURN_IF_ERROR would need a bare `Status` in scope). + if (Status st = loader(&loaded); !st.ok()) { + return st; + } + order_.push_front(Entry {.ordinal = ordinal, .block = loaded}); + index_[ordinal] = order_.begin(); + evict_overflow(); + *out = std::move(loaded); + return Status::OK(); + } + + // Number of resident (non-evicted) entries -- bounded by max_entries(). + size_t size() const { return index_.size(); } + size_t max_entries() const { return max_entries_; } + +private: + struct Entry { + uint32_t ordinal = 0; + std::shared_ptr block; + }; + + // Drops least-recently-used entries until the bound holds. Evicting only + // releases the cache's reference; any pin a caller still holds keeps the + // block (and its reader's Slice) alive. + void evict_overflow() { + while (index_.size() > max_entries_) { + const Entry& victim = order_.back(); + index_.erase(victim.ordinal); + order_.pop_back(); + } + } + + size_t max_entries_ = kDefaultMaxEntries; + std::list order_; // front = most recently used + std::unordered_map::iterator> index_; +}; + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/logical_index_reader.cpp b/be/src/storage/index/snii/reader/logical_index_reader.cpp new file mode 100644 index 00000000000000..51a7d89da28b67 --- /dev/null +++ b/be/src/storage/index/snii/reader/logical_index_reader.cpp @@ -0,0 +1,1107 @@ +// 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 "storage/index/snii/reader/logical_index_reader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/metadata_blob.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/reader/dict_block_cache.h" + +namespace doris::snii::reader { + +struct LogicalIndexReader::NormsCacheState { + struct Data { + std::vector bytes; + format::NormsPodReader reader; + }; + + std::mutex mutex; + std::unique_ptr ready; + std::shared_future in_flight; +}; + +using format::BlockRef; +using format::bsbf_hash; +using format::DictBlockDirectoryReader; +using format::DictBlockReader; +using format::DictEntry; +using format::IndexTier; +using format::kBsbfBytesPerBlock; +using format::kBsbfHeaderSize; +using format::RegionRef; +using format::SampledTermIndexReader; + +namespace { +constexpr uint64_t kMaxDictBlockUncompBytes = 256ULL * 1024 * 1024; +constexpr uint64_t kDefaultDictResidentMaxBytes = 256ULL * 1024; +constexpr size_t kMaxDictLookupBatchRuns = 16; +constexpr uint64_t kMaxDictLookupBatchBytes = 4ULL * 1024 * 1024; +// Conservatively covers make_shared's control block and allocator bookkeeping +// for the state/data/vector allocations. The framed norms bytes are charged +// separately at their exact validated length. +constexpr size_t kNormsAllocationOverheadCharge = 128; + +// L0/L1 tiering threshold (bytes). Defaults to kBsbfResidentMaxBytes; the env +// SNII_BSBF_RESIDENT_MAX overrides it for tuning and for exercising the +// on-demand L1 path in tests without a 250K-term corpus. Read fresh each open. +uint64_t bsbf_resident_max_bytes() { + const char* s = std::getenv("SNII_BSBF_RESIDENT_MAX"); + if (s != nullptr) { + char* end = nullptr; + const unsigned long long v = std::strtoull(s, &end, 10); + if (end != s) { + return v; + } + } + return format::kBsbfResidentMaxBytes; +} + +uint64_t dict_resident_max_bytes() { + const char* s = std::getenv("SNII_DICT_RESIDENT_MAX"); + if (s != nullptr) { + char* end = nullptr; + const unsigned long long v = std::strtoull(s, &end, 10); + if (end != s) { + return v; + } + } + return kDefaultDictResidentMaxBytes; +} + +Status checked_size(uint64_t value, const char* error, size_t* out) { + if (value > std::numeric_limits::max()) { + return Status::Error(error); + } + *out = static_cast(value); + return Status::OK(); +} + +Status validate_norms_region(io::FileReader* reader, const RegionRef& norms, uint64_t doc_count, + size_t* length) { + *length = 0; + if (norms.length == 0) { + return Status::OK(); + } + const uint64_t file_size = reader->size(); + if (norms.offset > file_size || norms.length > file_size - norms.offset) { + return Status::Error( + "logical_index: norms region past end of file"); + } + if (doc_count > std::numeric_limits::max()) { + return Status::Error( + "logical_index: norms doc count exceeds uint32"); + } + const uint64_t payload_length = varint_len(doc_count) + doc_count; + const uint64_t expected_length = + 1 + varint_len(payload_length) + payload_length + sizeof(uint32_t); + if (norms.length != expected_length) { + return Status::Error( + "logical_index: norms region length mismatch"); + } + if (norms.length > std::numeric_limits::max()) { + return Status::Error( + "logical_index: norms region exceeds cache charge bounds"); + } + *length = static_cast(norms.length); + return Status::OK(); +} + +Status validate_null_bitmap_frame(Slice framed) { + ByteSource source(framed); + FramedSection section; + RETURN_IF_ERROR(SectionFramer::read(source, §ion)); + if (section.type != format::kNullBitmapSectionType) { + return Status::Error( + "logical_index: invalid null bitmap section type"); + } + if (source.remaining() != 0) { + return Status::Error( + "logical_index: trailing null bitmap frame bytes"); + } + + // NullBitmapReader owns the Roaring validation. Parse only the two prefix + // varints here to pin the region to exactly one complete payload; otherwise + // a valid frame could silently carry ignored bytes after roaring_bytes. + ByteSource payload(section.payload); + uint64_t doc_count = 0; + RETURN_IF_ERROR(payload.get_varint64(&doc_count)); + if (doc_count > std::numeric_limits::max()) { + return Status::Error( + "logical_index: null bitmap doc count exceeds uint32"); + } + uint64_t roaring_size = 0; + RETURN_IF_ERROR(payload.get_varint64(&roaring_size)); + if (roaring_size != payload.remaining()) { + return Status::Error( + "logical_index: null bitmap payload length mismatch"); + } + return Status::OK(); +} + +Status dict_block_memory_bytes(const BlockRef& ref, uint64_t* out) { + if ((ref.flags & format::block_ref_flags::kZstd) == 0) { + *out = ref.length; + return Status::OK(); + } + if (ref.uncomp_len == 0 || ref.uncomp_len > kMaxDictBlockUncompBytes) { + return Status::Error( + "dict block: zstd uncomp_len out of range"); + } + *out = ref.uncomp_len; + return Status::OK(); +} + +Status checked_memory_add(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (rhs > std::numeric_limits::max() - lhs) { + return Status::Error(message); + } + *out = lhs + rhs; + return Status::OK(); +} + +Status checked_memory_mul(uint64_t lhs, uint64_t rhs, const char* message, uint64_t* out) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return Status::Error(message); + } + *out = lhs * rhs; + return Status::OK(); +} + +// Decompresses a zstd dict block from its on-disk bytes into *out. The decode +// buffer in zstd_decompress is resize_uninitialized'd (T19) then fully written. +Status zstd_decompress_dict_block(Slice on_disk, const BlockRef& ref, std::vector* out) { + uint64_t memory_bytes = 0; + RETURN_IF_ERROR(dict_block_memory_bytes(ref, &memory_bytes)); + size_t uncomp_len = 0; + RETURN_IF_ERROR( + checked_size(memory_bytes, "dict block: zstd length out of range", &uncomp_len)); + return zstd_decompress(on_disk, uncomp_len, out); +} + +// Materializes the usable (uncompressed) bytes of a dict block from a view over +// its on-disk bytes -- a raw block is copied, a zstd block is decompressed. Used +// by the resident single-range path, where on_disk is a sub-slice of the shared +// region buffer (so a raw block must be copied, not aliased). +Status decompress_dict_block_payload(Slice on_disk, const BlockRef& ref, + std::vector* out) { + if ((ref.flags & format::block_ref_flags::kZstd) == 0) { + out->assign(on_disk.data(), on_disk.data() + on_disk.size()); + return Status::OK(); + } + return zstd_decompress_dict_block(on_disk, ref, out); +} + +Status read_dict_block_bytes(io::FileReader* reader, const BlockRef& ref, + std::vector* out) { + size_t read_len = 0; + RETURN_IF_ERROR(checked_size(ref.length, "dict block: on-disk length out of range", &read_len)); + + std::vector block_bytes; + RETURN_IF_ERROR(reader->read_at(ref.offset, read_len, &block_bytes)); + if (block_bytes.size() != read_len) { + return Status::Error( + "dict block: short read"); + } + + // Raw on-demand block: move the freshly read buffer in (no copy). + if ((ref.flags & format::block_ref_flags::kZstd) == 0) { + *out = std::move(block_bytes); + return Status::OK(); + } + return zstd_decompress_dict_block(Slice(block_bytes), ref, out); +} + +Status open_dict_block(io::FileReader* reader, const BlockRef& ref, IndexTier tier, + bool has_positions, std::vector* bytes, DictBlockReader* out) { + RETURN_IF_ERROR(read_dict_block_bytes(reader, ref, bytes)); + return DictBlockReader::open(Slice(*bytes), tier, has_positions, out); +} + +// Validates that block `ref` lies fully within dict_region and returns its byte +// range relative to the start of the region. Defends the single-range resident +// read against a corrupt directory ref (offset before the region, or a range +// that runs past it) before it is used to index the region buffer. +Status slice_dict_block_in_region(const BlockRef& ref, const RegionRef& dict_region, + size_t region_len, size_t* rel_off, size_t* len) { + if (ref.offset < dict_region.offset) { + return Status::Error( + "dict block: ref before dict region"); + } + size_t rel = 0; + size_t block_len = 0; + RETURN_IF_ERROR( + checked_size(ref.offset - dict_region.offset, "dict block: ref offset OOR", &rel)); + RETURN_IF_ERROR(checked_size(ref.length, "dict block: ref length OOR", &block_len)); + if (rel > region_len || block_len > region_len - rel) { + return Status::Error( + "dict block: ref past dict region"); + } + *rel_off = rel; + *len = block_len; + return Status::OK(); +} +} // namespace + +Status LogicalIndexReader::load_resident_dict_blocks() { + resident_dict_blocks_.clear(); + + const uint64_t max_bytes = dict_resident_max_bytes(); + if (max_bytes == 0 || dbd_.n_blocks() == 0) { + return Status::OK(); + } + + uint64_t total_bytes = 0; + for (uint32_t ord = 0; ord < dbd_.n_blocks(); ++ord) { + BlockRef ref {}; + RETURN_IF_ERROR(dbd_.get(ord, &ref)); + uint64_t block_bytes = 0; + RETURN_IF_ERROR(dict_block_memory_bytes(ref, &block_bytes)); + if (block_bytes > max_bytes - total_bytes) { + return Status::OK(); + } + total_bytes += block_bytes; + } + + // The resident blocks are physically contiguous within dict_region, so read + // the whole region in a SINGLE range read (was one read_at per block -> up to + // ~4 serial S3 rounds on a cold open) and decode each block from a sub-slice. + // The region buffer is <= the resident byte cap (<=256KB) and freed on return; + // each ResidentDictBlock keeps its own decoded copy. + const RegionRef& dict_region = section_refs().dict_region; + size_t region_len = 0; + RETURN_IF_ERROR( + checked_size(dict_region.length, "dict region: length out of range", ®ion_len)); + std::vector region; + RETURN_IF_ERROR(reader_->read_at(dict_region.offset, region_len, ®ion)); + if (region.size() != region_len) { + return Status::Error( + "dict region: short read"); + } + + resident_dict_blocks_.reserve(dbd_.n_blocks()); + for (uint32_t ord = 0; ord < dbd_.n_blocks(); ++ord) { + BlockRef ref {}; + RETURN_IF_ERROR(dbd_.get(ord, &ref)); + size_t rel_off = 0; + size_t block_len = 0; + RETURN_IF_ERROR( + slice_dict_block_in_region(ref, dict_region, region_len, &rel_off, &block_len)); + const Slice on_disk(region.data() + rel_off, block_len); + ResidentDictBlock block; + RETURN_IF_ERROR(decompress_dict_block_payload(on_disk, ref, &block.bytes)); + RETURN_IF_ERROR( + DictBlockReader::open(Slice(block.bytes), tier_, has_positions_, &block.reader)); + resident_dict_blocks_.push_back(std::move(block)); + } + return Status::OK(); +} + +Status LogicalIndexReader::dict_block_reader_for_ordinal( + uint32_t ordinal, DictBlockCache* cache, std::shared_ptr* pin, + const DictBlockReader** out) const { + pin->reset(); + if (!resident_dict_blocks_.empty()) { + if (resident_dict_blocks_.size() != dbd_.n_blocks() || + ordinal >= resident_dict_blocks_.size()) { + return Status::Error( + "logical_index: incomplete resident dict"); + } + // Resident blocks live for the reader lifetime: no pin needed. + *out = &resident_dict_blocks_[ordinal].reader; + return Status::OK(); + } + + // On-demand: decode into a heap-allocated DecodedDictBlock held by *pin so the + // reader's Slice never dangles. The loader (file read + optional zstd + CRC + + // anchor parse) runs OUTSIDE any cache bookkeeping; on a cache hit it is not + // called, so a block shared by several terms of one query decodes only once. + DictBlockCache::Loader loader = [&](std::shared_ptr* slot) -> Status { + BlockRef ref {}; + RETURN_IF_ERROR(dbd_.get(ordinal, &ref)); + auto block = std::make_shared(); + RETURN_IF_ERROR(open_dict_block(reader_, ref, tier_, has_positions_, &block->bytes, + &block->reader)); + *slot = std::move(block); + return Status::OK(); + }; + if (cache != nullptr) { + RETURN_IF_ERROR(cache->get_or_load(ordinal, loader, pin)); + } else { + RETURN_IF_ERROR(loader(pin)); + } + *out = &(*pin)->reader; + return Status::OK(); +} + +Status LogicalIndexReader::load_resident_bsbf() { + // Block-split bloom XFilter -- gated on RESIDENCY (P1 cold-read fix, see + // docs/perf/P1-cold-read-amplification.md). The bloom is set up and used ONLY + // when the whole (small) filter fits under the resident cap: it is read in + // full, verified, and kept in memory so probes are in-memory and enter the + // Doris searcher cache with the rest of the logical-index metadata. + // + // When NON-resident (the common case for a real text column, where the filter + // is many MB) the bloom is skipped ENTIRELY: not even the 28B header is read, + // and has_bsbf_ stays false. Every term then falls through to sti -> dict, + // which yields the true found/absent. At 1 MiB cache-block granularity a + // non-resident bloom never saves a physical block (an absent term still costs + // one dict block either way), so its 28B header + per-term 32B probes were pure + // cold read amplification. + const RegionRef& bsbf = core_.section_refs.bsbf; + if (open_mode_ != LogicalIndexOpenMode::kQuery || bsbf.length == 0 || + bsbf.length > bsbf_resident_max_bytes()) { + return Status::OK(); + } + if (bsbf.length <= kBsbfHeaderSize) { + return Status::Error( + "logical_index: bsbf section too small"); + } + const uint64_t num_bytes = bsbf.length - kBsbfHeaderSize; + std::vector head; + RETURN_IF_ERROR(reader_->read_at(bsbf.offset, bsbf.length, &head)); + if (head.size() < bsbf.length) { + return Status::Error( + "logical_index: short bsbf resident read"); + } + RETURN_IF_ERROR(format::BsbfHeader::parse(Slice(head.data(), kBsbfHeaderSize), bsbf.offset, + &bsbf_header_)); + // Cross-check the header geometry against the section ref. + if (bsbf_header_.num_bytes != num_bytes) { + return Status::Error( + "logical_index: bsbf header/section size mismatch"); + } + const Slice bitset(head.data() + kBsbfHeaderSize, bsbf_header_.num_bytes); + if (crc32c(bitset) != bsbf_header_.bitset_crc) { + return Status::Error( + "logical_index: bsbf bitset crc mismatch"); + } + bsbf_resident_bitset_.assign(bitset.data(), bitset.data() + bitset.size()); + has_bsbf_ = true; + bsbf_resident_ = true; + return Status::OK(); +} + +Status LogicalIndexReader::open(io::FileReader* file_reader, Slice core_frame, Slice sti_blob, + Slice dbd_blob, LogicalIndexReader* out, + LogicalIndexOpenMode open_mode) { + if (out == nullptr) { + return Status::Error("logical_index: null out"); + } + *out = LogicalIndexReader {}; + if (file_reader == nullptr) { + return Status::Error("logical_index: null file reader"); + } + if (core_frame.empty() || sti_blob.empty() || dbd_blob.empty()) { + return Status::Error( + "logical_index: empty mandatory metadata blob"); + } + + LogicalIndexReader candidate; + candidate.reader_ = file_reader; + candidate.open_mode_ = open_mode; + RETURN_IF_ERROR(format::decode_core_metadata(core_frame, &candidate.core_)); + candidate.tier_ = format::tier_of(candidate.core_.index_config); + candidate.has_positions_ = format::has_positions(candidate.core_.index_config); + size_t norms_length = 0; + RETURN_IF_ERROR(validate_norms_region(file_reader, candidate.core_.section_refs.norms, + candidate.core_.stats.doc_count, &norms_length)); + if (norms_length != 0) { + constexpr size_t fixed_charge = sizeof(NormsCacheState) + sizeof(NormsCacheState::Data) + + kNormsAllocationOverheadCharge; + if (norms_length > std::numeric_limits::max() - fixed_charge) { + return Status::Error( + "logical_index: norms cache charge overflow"); + } + candidate.norms_reserved_charge_ = fixed_charge + norms_length; + candidate.norms_cache_ = std::make_shared(); + } + // Raw frames alias the group read and compressed carriers materialize into + // transient scratch. Both readers own their decoded state. + { + std::vector scratch; + Slice frame; + RETURN_IF_ERROR(format::materialize_metadata_blob( + sti_blob, format::SectionType::kSampledTermIndex, + format::SectionType::kSampledTermIndexZstd, &scratch, &frame)); + RETURN_IF_ERROR(SampledTermIndexReader::open(frame, &candidate.sti_)); + } + { + std::vector scratch; + Slice frame; + RETURN_IF_ERROR(format::materialize_metadata_blob( + dbd_blob, format::SectionType::kDictBlockDirectory, + format::SectionType::kDictBlockDirectoryZstd, &scratch, &frame)); + RETURN_IF_ERROR(DictBlockDirectoryReader::open(frame, &candidate.dbd_)); + } + if (candidate.sti_.n_blocks() != candidate.dbd_.n_blocks()) { + return Status::Error( + "logical_index: sampled-term index and block directory count mismatch"); + } + if (open_mode == LogicalIndexOpenMode::kQuery) { + RETURN_IF_ERROR(candidate.load_resident_dict_blocks()); + } + RETURN_IF_ERROR(candidate.load_resident_bsbf()); + *out = std::move(candidate); + return Status::OK(); +} + +size_t LogicalIndexReader::memory_usage() const { + size_t bytes = sizeof(*this) + bsbf_resident_bitset_.capacity(); + if (core_.common_grams_metadata) { + const auto& common_grams = *core_.common_grams_metadata; + bytes += format::std_string_heap_bytes(common_grams.common_grams_dictionary_identity); + bytes += format::std_string_heap_bytes(common_grams.base_analyzer_fingerprint); + bytes += format::std_string_heap_bytes(common_grams.common_grams_fingerprint); + } + bytes += sti_.heap_bytes(); + bytes += dbd_.heap_bytes(); + for (const auto& block : resident_dict_blocks_) { + bytes += sizeof(block) + block.bytes.capacity() + block.reader.heap_bytes(); + } + // Norms are loaded lazily, but the searcher-cache charge is fixed when this + // reader is inserted. Reserve their complete framed size and heap state up + // front so later allocation is already accounted for. Saturation prevents + // an already-large resident metadata charge from wrapping around. + if (norms_reserved_charge_ > std::numeric_limits::max() - bytes) { + return std::numeric_limits::max(); + } + bytes += norms_reserved_charge_; + return bytes; +} + +Status LogicalIndexReader::open_norms(format::NormsPodReader* out) const { + if (out == nullptr) { + return Status::Error( + "logical_index: null norms reader"); + } + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + const RegionRef& norms = section_refs().norms; + if (norms.length == 0) { + return Status::Error( + "logical_index: index has no norms"); + } + DORIS_CHECK(norms_cache_ != nullptr); + + std::shared_future in_flight; + std::shared_ptr> completion; + { + std::lock_guard lock(norms_cache_->mutex); + if (norms_cache_->ready != nullptr) { + *out = norms_cache_->ready->reader; + return Status::OK(); + } + if (norms_cache_->in_flight.valid()) { + in_flight = norms_cache_->in_flight; + } else { + completion = std::make_shared>(); + in_flight = completion->get_future().share(); + norms_cache_->in_flight = in_flight; + } + } + + if (completion == nullptr) { + const Status status = in_flight.get(); + RETURN_IF_ERROR(status); + std::lock_guard lock(norms_cache_->mutex); + DORIS_CHECK(norms_cache_->ready != nullptr); + *out = norms_cache_->ready->reader; + return Status::OK(); + } + + auto data = std::make_unique(); + const size_t read_len = static_cast(norms.length); + Status status = reader_->read_at(norms.offset, read_len, &data->bytes); + if (status.ok() && data->bytes.size() != read_len) { + status = Status::Error( + "logical_index: short norms read"); + } + if (status.ok()) { + status = format::NormsPodReader::open(Slice(data->bytes), &data->reader); + } + { + std::lock_guard lock(norms_cache_->mutex); + if (status.ok()) { + norms_cache_->ready = std::move(data); + *out = norms_cache_->ready->reader; + } + norms_cache_->in_flight = std::shared_future {}; + } + completion->set_value(status); + return status; +} + +void LogicalIndexReader::release_compaction_norms() const { + DORIS_CHECK(open_mode_ == LogicalIndexOpenMode::kCompaction); + DORIS_CHECK(norms_cache_ != nullptr); + std::lock_guard lock(norms_cache_->mutex); + DORIS_CHECK(!norms_cache_->in_flight.valid()); + norms_cache_->ready.reset(); +} + +Status LogicalIndexReader::read_null_docids( + std::vector* out, const NullDocidsDecodeReservation& reserve_decode) const { + if (out == nullptr) { + return Status::Error( + "logical_index: null null-docids output"); + } + out->clear(); + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + + const RegionRef& ref = section_refs().null_bitmap; + if (ref.length == 0) { + if (stats().null_count != 0) { + return Status::Error( + "logical_index: null bitmap section missing"); + } + return Status::OK(); + } + + const uint64_t file_size = reader_->size(); + if (ref.offset > file_size || ref.length > file_size - ref.offset) { + return Status::Error( + "logical_index: null bitmap region past end of file"); + } + size_t read_len = 0; + RETURN_IF_ERROR( + checked_size(ref.length, "logical_index: null bitmap length out of range", &read_len)); + std::vector bytes; + RETURN_IF_ERROR(reader_->read_at(ref.offset, read_len, &bytes)); + if (bytes.size() != read_len) { + return Status::Error( + "logical_index: short null bitmap read"); + } + RETURN_IF_ERROR(validate_null_bitmap_frame(Slice(bytes))); + + uint64_t decoded_memory_bytes = 0; + RETURN_IF_ERROR( + format::NullBitmapReader::decoded_memory_bytes(Slice(bytes), &decoded_memory_bytes)); + if (reserve_decode) { + RETURN_IF_ERROR(reserve_decode(decoded_memory_bytes)); + } + + format::NullBitmapReader null_bitmap_reader; + RETURN_IF_ERROR(format::NullBitmapReader::open(Slice(bytes), &null_bitmap_reader)); + if (null_bitmap_reader.doc_count() != stats().doc_count) { + return Status::Error( + "logical_index: null bitmap doc count mismatch"); + } + if (null_bitmap_reader.null_count() != stats().null_count) { + return Status::Error( + "logical_index: null bitmap cardinality mismatch"); + } + + out->reserve(null_bitmap_reader.null_count()); + null_bitmap_reader.append_docids(*out); + for (uint32_t docid : *out) { + if (docid >= stats().doc_count) { + out->clear(); + return Status::Error( + "logical_index: null docid outside document domain"); + } + } + return Status::OK(); +} + +Status LogicalIndexReader::null_docids_scan_memory(NullDocidsScanMemory* out) const { + if (out == nullptr) { + return Status::Error( + "logical_index: null null-docids scan memory out"); + } + *out = NullDocidsScanMemory {}; + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + const RegionRef& ref = section_refs().null_bitmap; + if (ref.length == 0) { + if (stats().null_count != 0) { + return Status::Error( + "logical_index: null bitmap section missing"); + } + return Status::OK(); + } + const uint64_t file_size = reader_->size(); + if (ref.offset > file_size || ref.length > file_size - ref.offset) { + return Status::Error( + "logical_index: null bitmap region past end of file"); + } + size_t frame_bytes = 0; + RETURN_IF_ERROR(checked_size(ref.length, "logical_index: null bitmap length out of range", + &frame_bytes)); + RETURN_IF_ERROR(checked_memory_mul(stats().null_count, sizeof(uint32_t), + "logical_index: null docids output memory overflows", + &out->output_bytes)); + + out->frame_bytes = frame_bytes; + return Status::OK(); +} + +Status LogicalIndexReader::lookup(std::string_view term, bool* found, DictEntry* entry, + uint64_t* frq_base, uint64_t* prx_base, + DictBlockCache* cache) const { + *found = false; + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + + bool maybe = false; + uint32_t ordinal = 0; + RETURN_IF_ERROR(locate_candidate_dict_block(term, &maybe, &ordinal)); + if (!maybe) { + return Status::OK(); + } + + // Use a resident small-DICT block when present; otherwise read the DICT + // block on demand and parse it with the same validation path used at open. + // `pin` keeps an on-demand block alive through find_term (resident: null). + const DictBlockReader* br = nullptr; + std::shared_ptr pin; + RETURN_IF_ERROR(dict_block_reader_for_ordinal(ordinal, cache, &pin, &br)); + + bool hit = false; + RETURN_IF_ERROR(br->find_term(term, &hit, entry)); + if (!hit) { + return Status::OK(); + } + + *found = true; + *frq_base = br->frq_base(); + *prx_base = br->prx_base(); + return Status::OK(); +} + +Status LogicalIndexReader::locate_candidate_dict_block(std::string_view term, bool* maybe_present, + uint32_t* ordinal) const { + *maybe_present = false; + // A DEFINITELY-ABSENT term returns without a DICT read. The bloom is + // consulted only when resident; otherwise STI/DICT remains authoritative. + if (has_bsbf_) { + const uint64_t h = bsbf_hash(term); + bool maybe = true; + if (bsbf_resident_) { + const uint32_t blk = format::bsbf_block_index(h, bsbf_header_.num_blocks); + maybe = format::bsbf_block_contains( + h, + bsbf_resident_bitset_.data() + static_cast(blk) * kBsbfBytesPerBlock); + } + if (!maybe) { + return Status::OK(); + } + } + return sti_.locate(term, maybe_present, ordinal); +} + +Status LogicalIndexReader::collect_batch_lookup_groups( + const std::vector& terms, std::vector* candidates, + std::vector* groups) const { + candidates->clear(); + candidates->reserve(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + bool maybe = false; + uint32_t ordinal = 0; + RETURN_IF_ERROR(locate_candidate_dict_block(terms[i], &maybe, &ordinal)); + if (maybe) { + candidates->push_back({i, ordinal}); + } + } + + auto by_ordinal = [](const BatchLookupCandidate& lhs, const BatchLookupCandidate& rhs) { + return lhs.ordinal < rhs.ordinal; + }; + if (!std::ranges::is_sorted(*candidates, by_ordinal)) { + std::ranges::sort(*candidates, by_ordinal); + } + + groups->clear(); + for (size_t begin = 0; begin < candidates->size();) { + const uint32_t ordinal = (*candidates)[begin].ordinal; + size_t end = begin + 1; + while (end < candidates->size() && (*candidates)[end].ordinal == ordinal) { + ++end; + } + groups->push_back({ordinal, begin, end}); + begin = end; + } + return Status::OK(); +} + +Status LogicalIndexReader::resolve_batch_lookup_group( + const std::vector& terms, const std::vector& candidates, + const BatchLookupGroup& group, const DictBlockReader& block_reader, + std::vector* results) { + for (size_t i = group.begin; i < group.end; ++i) { + const size_t term_index = candidates[i].term_index; + BatchLookupResult& result = (*results)[term_index]; + RETURN_IF_ERROR(block_reader.find_term(terms[term_index], &result.found, &result.entry)); + if (result.found) { + result.frq_base = block_reader.frq_base(); + result.prx_base = block_reader.prx_base(); + } + } + return Status::OK(); +} + +Status LogicalIndexReader::lookup_batch_on_demand( + const std::vector& terms, const std::vector& candidates, + const std::vector& groups, + std::vector* results) const { + for (size_t wave_begin = 0; wave_begin < groups.size();) { + std::vector pending; + pending.reserve(kMaxDictLookupBatchRuns); + uint64_t pending_bytes = 0; + uint64_t pending_end = 0; + size_t pending_runs = 0; + size_t wave_end = wave_begin; + while (wave_end < groups.size()) { + BlockRef ref {}; + RETURN_IF_ERROR(dbd_.get(groups[wave_end].ordinal, &ref)); + const uint64_t ref_end = ref.offset + ref.length; + const bool starts_new_run = pending.empty() || ref.offset > pending_end; + if (!pending.empty() && ((starts_new_run && pending_runs == kMaxDictLookupBatchRuns) || + ref.length > kMaxDictLookupBatchBytes || + pending_bytes > kMaxDictLookupBatchBytes - ref.length)) { + break; + } + pending.push_back({wave_end, ref, 0}); + pending_bytes += ref.length; + if (starts_new_run) { + ++pending_runs; + } + pending_end = std::max(pending_end, ref_end); + ++wave_end; + } + DORIS_CHECK(!pending.empty()); + io::BatchRangeFetcher fetcher(reader_, /*coalesce_gap=*/0); + for (PendingBatchLookupBlock& block : pending) { + block.handle = fetcher.add(block.ref.offset, block.ref.length); + } + RETURN_IF_ERROR(fetcher.fetch()); + + for (const PendingBatchLookupBlock& block : pending) { + const Slice on_disk = fetcher.get(block.handle); + std::vector decoded; + Slice payload = on_disk; + if ((block.ref.flags & format::block_ref_flags::kZstd) != 0) { + RETURN_IF_ERROR(zstd_decompress_dict_block(on_disk, block.ref, &decoded)); + payload = Slice(decoded); + } + DictBlockReader block_reader; + RETURN_IF_ERROR(DictBlockReader::open(payload, tier_, has_positions_, &block_reader)); + RETURN_IF_ERROR(resolve_batch_lookup_group(terms, candidates, groups[block.group_index], + block_reader, results)); + } + wave_begin = wave_end; + } + return Status::OK(); +} + +Status LogicalIndexReader::lookup_batch(const std::vector& terms, + std::vector* results) const { + DORIS_CHECK(results != nullptr); + DCHECK(std::ranges::is_sorted(terms)); + DCHECK(std::adjacent_find(terms.begin(), terms.end()) == terms.end()); + results->assign(terms.size(), BatchLookupResult {}); + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + + std::vector candidates; + std::vector groups; + RETURN_IF_ERROR(collect_batch_lookup_groups(terms, &candidates, &groups)); + if (groups.empty()) { + return Status::OK(); + } + + // Resident dictionaries always stay zero-I/O. One-block batches keep the + // existing synchronous read-through path. + if (!resident_dict_blocks_.empty() || groups.size() == 1) { + for (const BatchLookupGroup& group : groups) { + const DictBlockReader* block_reader = nullptr; + std::shared_ptr pin; + RETURN_IF_ERROR(dict_block_reader_for_ordinal(group.ordinal, /*cache=*/nullptr, &pin, + &block_reader)); + RETURN_IF_ERROR( + resolve_batch_lookup_group(terms, candidates, group, *block_reader, results)); + } + return Status::OK(); + } + + return lookup_batch_on_demand(terms, candidates, groups, results); +} + +Status LogicalIndexReader::decode_dict_block(uint32_t ordinal, std::vector* entries, + uint64_t* frq_base, uint64_t* prx_base) const { + if (entries == nullptr || frq_base == nullptr || prx_base == nullptr) { + return Status::Error( + "logical_index: null decode_dict_block out"); + } + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + entries->clear(); + // Same resolution path lookup() uses (resident block or one on-demand range + // read + zstd + CRC); `pin` keeps an on-demand block alive through + // decode_all. No request-scoped cache: a sequential full scan touches each + // block exactly once. + const DictBlockReader* br = nullptr; + std::shared_ptr pin; + RETURN_IF_ERROR(dict_block_reader_for_ordinal(ordinal, /*cache=*/nullptr, &pin, &br)); + RETURN_IF_ERROR(br->decode_all(entries)); + *frq_base = br->frq_base(); + *prx_base = br->prx_base(); + return Status::OK(); +} + +Status LogicalIndexReader::dict_block_scan_memory(uint32_t ordinal, + DictBlockScanMemory* out) const { + if (out == nullptr) { + return Status::Error( + "logical_index: null dict block scan memory out"); + } + *out = DictBlockScanMemory {}; + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + + BlockRef ref {}; + RETURN_IF_ERROR(dbd_.get(ordinal, &ref)); + uint64_t plain_bytes = 0; + RETURN_IF_ERROR(dict_block_memory_bytes(ref, &plain_bytes)); + + // During a compressed decode the fetched bytes coexist with the plain block. + // The reader then owns the plain bytes plus anchor vectors/strings. Counting + // every entry as an anchor is conservative for all valid anchor intervals. + uint64_t anchor_slots = 0; + RETURN_IF_ERROR(checked_memory_mul(ref.n_entries, sizeof(uint32_t) + sizeof(std::string), + "logical_index: dict decode anchor memory overflows", + &anchor_slots)); + uint64_t decode_bytes = 0; + RETURN_IF_ERROR(checked_memory_add(ref.length, plain_bytes, + "logical_index: dict decode bytes overflow", &decode_bytes)); + RETURN_IF_ERROR(checked_memory_add(decode_bytes, anchor_slots, + "logical_index: dict decode slots overflow", &decode_bytes)); + RETURN_IF_ERROR(checked_memory_add(decode_bytes, plain_bytes, + "logical_index: dict decode terms overflow", &decode_bytes)); + + uint64_t entry_slots = 0; + RETURN_IF_ERROR(checked_memory_mul(ref.n_entries, sizeof(DictEntry), + "logical_index: dict entry slots overflow", &entry_slots)); + uint64_t entry_payload = 0; + RETURN_IF_ERROR(checked_memory_mul(plain_bytes, 2, "logical_index: dict entry payload overflow", + &entry_payload)); + RETURN_IF_ERROR(checked_memory_add(entry_slots, entry_payload, + "logical_index: dict entries memory overflow", + &out->entries_bytes)); + out->decode_bytes = decode_bytes; + return Status::OK(); +} + +Status LogicalIndexReader::visit_prefix_terms(std::string_view prefix, + const PrefixHitVisitor& visitor, + DictBlockCache* cache) const { + if (!visitor) { + return Status::Error( + "logical_index: null prefix visitor"); + } + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + + // Seek the start block: the SampledTermIndex block whose first term <= prefix + // (terms with `prefix` are >= prefix, so they begin in that block or later). + // If the prefix sorts before every sample (or is empty), start at block 0. + uint32_t start = 0; + if (!prefix.empty()) { + bool maybe = false; + uint32_t ordinal = 0; + RETURN_IF_ERROR(sti_.locate(prefix, &maybe, &ordinal)); + if (maybe) { + start = ordinal; + } + } + + for (uint32_t ord = start; ord < dbd_.n_blocks(); ++ord) { + const DictBlockReader* br = nullptr; + std::shared_ptr pin; + RETURN_IF_ERROR(dict_block_reader_for_ordinal(ord, cache, &pin, &br)); + + // Stream this block's prefix range: anchor-jump past pre-prefix segments, + // decode only the bodies we keep, and stop at the first term past the + // range (decode_all materialized every entry of every scanned block). + // The visitor still owns final term acceptance, so results are identical; + // `br`/`pin` stay alive across this synchronous call. + bool prefix_exhausted = false; + bool visitor_stopped = false; + RETURN_IF_ERROR(br->visit_prefix_range( + prefix, /*accept_key=*/ {}, + [&](DictEntry&& e, bool* stop) -> Status { + PrefixHit hit; + hit.term = e.term; + hit.entry = std::move(e); + hit.frq_base = br->frq_base(); + hit.prx_base = br->prx_base(); + RETURN_IF_ERROR(visitor(std::move(hit), stop)); + visitor_stopped = *stop; + return Status::OK(); + }, + &prefix_exhausted)); + if (visitor_stopped || prefix_exhausted) { + return Status::OK(); + } + } + return Status::OK(); +} + +Status LogicalIndexReader::visit_term_range(std::string_view lower_inclusive, + std::optional upper_exclusive, + const PrefixHitVisitor& visitor, + DictBlockCache* cache) const { + if (!visitor) { + return Status::Error( + "logical_index: null range visitor"); + } + if (reader_ == nullptr) { + return Status::Error("logical_index: not opened"); + } + if (upper_exclusive.has_value() && *upper_exclusive <= lower_inclusive) { + return Status::OK(); + } + if (upper_exclusive.has_value()) { + bool upper_reaches_dictionary = false; + uint32_t upper_ordinal = 0; + RETURN_IF_ERROR(sti_.locate(*upper_exclusive, &upper_reaches_dictionary, &upper_ordinal)); + if (!upper_reaches_dictionary) { + return Status::OK(); + } + } + + uint32_t start = 0; + if (!lower_inclusive.empty()) { + bool maybe = false; + uint32_t ordinal = 0; + RETURN_IF_ERROR(sti_.locate(lower_inclusive, &maybe, &ordinal)); + if (maybe) { + start = ordinal; + } + } + + for (uint32_t ord = start; ord < dbd_.n_blocks(); ++ord) { + const DictBlockReader* block_reader = nullptr; + std::shared_ptr pin; + RETURN_IF_ERROR(dict_block_reader_for_ordinal(ord, cache, &pin, &block_reader)); + + bool range_exhausted = false; + bool visitor_stopped = false; + RETURN_IF_ERROR(block_reader->visit_term_range( + lower_inclusive, upper_exclusive, /*accept_key=*/ {}, + [&](DictEntry&& entry, bool* stop) -> Status { + PrefixHit hit; + hit.term = entry.term; + hit.entry = std::move(entry); + hit.frq_base = block_reader->frq_base(); + hit.prx_base = block_reader->prx_base(); + RETURN_IF_ERROR(visitor(std::move(hit), stop)); + visitor_stopped = *stop; + return Status::OK(); + }, + &range_exhausted)); + if (visitor_stopped || range_exhausted) { + return Status::OK(); + } + } + return Status::OK(); +} + +Status LogicalIndexReader::prefix_terms(std::string_view prefix, std::vector* const out, + int32_t max_terms, DictBlockCache* cache) const { + if (out == nullptr) { + return Status::Error("logical_index: null out"); + } + out->clear(); + return visit_prefix_terms( + prefix, + [&](PrefixHit&& hit, bool* stop) { + out->push_back(std::move(hit)); + *stop = max_terms > 0 && out->size() >= static_cast(max_terms); + return Status::OK(); + }, + cache); +} + +namespace { + +// Validates a pod_ref window locator against the posting region and returns the +// absolute window range (after the prelude). Rejects corrupt locators rather +// than letting size_t underflow / uint64 overflow reach read_at. +Status resolve_window(const format::RegionRef& section, uint64_t base, uint64_t off_delta, + uint64_t total_len, uint64_t prelude_len, uint64_t* abs_off, uint64_t* len) { + if (prelude_len > total_len) { + return Status::Error( + "logical_index: prelude_len exceeds window len"); + } + const uint64_t in_region = base + off_delta; + if (in_region < base) { + return Status::Error( + "logical_index: locator overflow"); + } + if (in_region > section.length || total_len > section.length - in_region) { + return Status::Error( + "logical_index: window past posting region"); + } + *abs_off = section.offset + in_region + prelude_len; + *len = total_len - prelude_len; + return Status::OK(); +} + +} // namespace + +Status LogicalIndexReader::resolve_frq_window(const format::DictEntry& entry, uint64_t frq_base, + uint64_t* abs_off, uint64_t* len) const { + return resolve_window(section_refs().posting_region, frq_base, entry.frq_off_delta, + entry.frq_len, entry.prelude_len, abs_off, len); +} + +Status LogicalIndexReader::resolve_prx_window(const format::DictEntry& entry, uint64_t prx_base, + uint64_t* abs_off, uint64_t* len) const { + // .prx windows carry no prelude (prelude_len = 0); both spans live in the + // same posting region (prx span precedes frq span for the same term). + return resolve_window(section_refs().posting_region, prx_base, entry.prx_off_delta, + entry.prx_len, 0, abs_off, len); +} + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/logical_index_reader.h b/be/src/storage/index/snii/reader/logical_index_reader.h new file mode 100644 index 00000000000000..f7fdee0a902f18 --- /dev/null +++ b/be/src/storage/index/snii/reader/logical_index_reader.h @@ -0,0 +1,285 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/bsbf.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/metadata_blob.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/format/stats_block.h" +#include "storage/index/snii/io/file_reader.h" + +// LogicalIndexReader -- read-side counterpart of LogicalIndexWriter for one +// logical index. It owns decoded Core, SampledTermIndex, and DICT block-directory +// state from one adjacent metadata group, and resolves a query term to its +// DictEntry through the documented lookup flow: +// XFilter (reject absent) -> SampledTermIndex (candidate block ordinal) -> +// DICT block directory (block range) -> resident small-DICT block or one +// range read of the DICT block -> DictBlockReader::find_term. +// +// lookup() also returns the block's frq_base/prx_base (captured by the +// DictBlockReader) so callers can resolve a pod_ref entry's absolute .frq/.prx +// offsets via the writer's contract. Both deltas index into the SAME +// interleaved posting region (prx_base == frq_base; the prx span precedes the +// frq span): +// abs_frq = posting_region.offset + frq_base + entry.frq_off_delta +// abs_prx = posting_region.offset + prx_base + entry.prx_off_delta +// +// The reader retains no raw metadata-group bytes after open. +namespace doris::snii::format { +class NormsPodReader; +} + +namespace doris::snii::reader { + +// Forward-declared: this widely-included header only names DictBlockCache* and +// shared_ptr*; the full definitions are pulled into the +// .cpp and into tests that construct a cache. Keeps the request-scoped cache +// header out of the ~500 TUs that transitively include this one. +struct DecodedDictBlock; +class DictBlockCache; + +enum class LogicalIndexOpenMode : uint8_t { + kQuery, + kCompaction, +}; + +struct DictBlockScanMemory { + uint64_t decode_bytes = 0; + uint64_t entries_bytes = 0; +}; + +struct NullDocidsScanMemory { + uint64_t frame_bytes = 0; + uint64_t output_bytes = 0; +}; + +class LogicalIndexReader { +public: + LogicalIndexReader() = default; + + // Parses one mandatory Core/STI/DBD metadata group and binds the reader to + // file_reader. The reader retains decoded state, not the input byte slices. + static Status open(io::FileReader* file_reader, Slice core_frame, Slice sti_blob, + Slice dbd_blob, LogicalIndexReader* out, + LogicalIndexOpenMode open_mode = LogicalIndexOpenMode::kQuery); + + // Resolves term to a DictEntry. *found=false when the term is absent (XFilter + // rejection, out-of-range sample, or DICT-block miss). On a hit, *entry is + // filled and *frq_base / *prx_base carry the candidate block's bases. + // + // `cache` is an optional REQUEST-SCOPED DictBlockCache: when a single query + // threads one cache through its per-term lookups, an on-demand DICT block hit + // by several terms is decoded once instead of once per term. nullptr keeps the + // pre-existing behavior (each lookup materializes its own block). The cache is + // caller-owned, single-threaded, and never mutates this (const) reader. + Status lookup(std::string_view term, bool* found, format::DictEntry* entry, uint64_t* frq_base, + uint64_t* prx_base, DictBlockCache* cache = nullptr) const; + + struct BatchLookupResult { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + }; + + // Resolves one sorted, duplicate-free term batch. Terms are first mapped to + // candidate DICT ordinals through the same XFilter/STI path as lookup(), then + // distinct on-demand blocks are fetched concurrently in bounded waves. + // Results stay aligned with `terms`; absent terms have found=false. + Status lookup_batch(const std::vector& terms, + std::vector* results) const; + + // One enumerated term whose key has the requested prefix, with its DictEntry + // and the owning DICT block's frq/prx bases (for posting resolution). + struct PrefixHit { + std::string term; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + }; + + using PrefixHitVisitor = std::function; + + // Ordered term enumeration: every term with `prefix`, in lexicographic order, + // by seeking the start DICT block via the SampledTermIndex and scanning + // forward across contiguous blocks until the terms pass the prefix range. + // Empty prefix enumerates all terms. This is the contiguous-DICT-block design + // the term-anchor layout was built for (MATCH_PHRASE_PREFIX / prefix / range + // queries). The visitor form avoids materializing all hits when callers only + // need a bounded expansion. + Status visit_prefix_terms(std::string_view prefix, const PrefixHitVisitor& visitor, + DictBlockCache* cache = nullptr) const; + Status visit_term_range(std::string_view lower_inclusive, + std::optional upper_exclusive, + const PrefixHitVisitor& visitor, DictBlockCache* cache = nullptr) const; + Status prefix_terms(std::string_view prefix, std::vector* const out, + int32_t max_terms = 0, DictBlockCache* cache = nullptr) const; + + // ---- Sequential whole-dictionary access (T2.3, compaction index merge) ---- + // Number of DICT blocks in this index (0 for an empty dictionary). + uint32_t n_dict_blocks() const { return dbd_.n_blocks(); } + // Decodes EVERY entry of DICT block `ordinal` in lexicographic order into + // *entries (each self-contained, owning its term and any inline posting + // bytes) and returns the block's frq/prx bases. One block is materialized + // at a time so a full-dictionary scan (SniiSegmentTermCursor) holds a single + // block's entries, never the whole vocabulary. ordinal must be + // < n_dict_blocks(). + Status decode_dict_block(uint32_t ordinal, std::vector* entries, + uint64_t* frq_base, uint64_t* prx_base) const; + // Returns conservative pre-allocation charges for the on-demand block decode + // and its fully materialized DictEntry vector. Compaction cursors reserve both + // before decoding so MEM_LIMIT_EXCEEDED is returned before the normal block + // allocations whenever the shared merge cap cannot admit them. + Status dict_block_scan_memory(uint32_t ordinal, DictBlockScanMemory* out) const; + + // Resolves a pod_ref entry's absolute .frq / .prx window byte range, + // validating the locator against the posting_region length (defends against + // corrupt entries: prelude_len > frq_len underflow, or off_delta+len past the + // region). Both windows resolve against the single posting_region. *abs_off + // is the absolute file offset of the window (after prelude); *len its byte + // length. + Status resolve_frq_window(const format::DictEntry& entry, uint64_t frq_base, uint64_t* abs_off, + uint64_t* len) const; + Status resolve_prx_window(const format::DictEntry& entry, uint64_t prx_base, uint64_t* abs_off, + uint64_t* len) const; + + const format::SectionRefs& section_refs() const { return core_.section_refs; } + const format::StatsBlock& stats() const { return core_.stats; } + format::IndexTier tier() const { return tier_; } + bool has_positions() const { return has_positions_; } + LogicalIndexOpenMode open_mode() const { return open_mode_; } + const segment_v2::inverted_index::CommonGramsSegmentMetadata* common_grams_metadata() const { + return core_.common_grams_metadata ? &*core_.common_grams_metadata : nullptr; + } + format::CommonGramsPostingPolicy common_grams_posting_policy() const { + return core_.common_grams_posting_policy; + } + io::FileReader* reader() const { return reader_; } + + // Returns a reader over the validated norms section. The first call reads + // and validates the section; later calls share the immutable reader-owned + // bytes. The full on-disk section is reserved in memory_usage() before this + // LogicalIndexReader enters the searcher cache, so lazy loading cannot make + // the cache under-report its eventual resident size. + Status open_norms(format::NormsPodReader* out) const; + // Compaction scans one source norm vector at a time. This charge matches the + // reader's full cache accounting; release_compaction_norms() drops the loaded + // frame after its values have been scattered into destination vectors. + size_t compaction_norms_cache_charge() const { return norms_reserved_charge_; } + void release_compaction_norms() const; + + // Reads and validates the sparse null-bitmap side POD, then returns its + // docids in ascending order. Work is O(null_count), not O(doc_count), which + // lets compaction remap NULL rows without scanning the complete document + // domain. A missing section is valid only when StatsBlock::null_count is 0. + using NullDocidsDecodeReservation = std::function; + Status read_null_docids(std::vector* out, + const NullDocidsDecodeReservation& reserve_decode = + NullDocidsDecodeReservation()) const; + Status null_docids_scan_memory(NullDocidsScanMemory* out) const; + size_t memory_usage() const; + +private: + struct NormsCacheState; + struct BatchLookupCandidate { + size_t term_index = 0; + uint32_t ordinal = 0; + }; + struct BatchLookupGroup { + uint32_t ordinal = 0; + size_t begin = 0; + size_t end = 0; + }; + struct PendingBatchLookupBlock { + size_t group_index = 0; + format::BlockRef ref; + size_t handle = 0; + }; + io::FileReader* reader_ = nullptr; + format::IndexTier tier_ = format::IndexTier::kT1; + bool has_positions_ = false; + LogicalIndexOpenMode open_mode_ = LogicalIndexOpenMode::kQuery; + format::CoreMetadata core_; + format::SampledTermIndexReader sti_; + format::DictBlockDirectoryReader dbd_; + format::BsbfHeader bsbf_header_; // resident header (from section ref) + bool has_bsbf_ = false; + // L0 tiering: when the bsbf section is small (<= kBsbfResidentMaxBytes) its + // whole bitset is loaded here at open -> in-memory probe, no per-lookup + // round. Larger filters keep only the parsed header here, so the small + // header enters Doris searcher cache and lookup reads just one 32-byte body + // block for an L1 probe. + bool bsbf_resident_ = false; + std::vector bsbf_resident_bitset_; + + // Small DICT blocks are opened once with the index so exact lookups avoid an + // otherwise serial S3 round for the term dictionary. Empty means the + // dictionary exceeded the resident threshold and lookup/prefix enumeration + // read blocks on demand. Each DictBlockReader holds a Slice into the owning + // bytes. + struct ResidentDictBlock { + std::vector bytes; + format::DictBlockReader reader; + }; + Status load_resident_dict_blocks(); + Status load_resident_bsbf(); + // Resolves the DictBlockReader for `ordinal`. Resident blocks return a pointer + // into the reader-owned resident set with *pin left null (stable for the reader + // lifetime). On-demand blocks are decoded (optionally via the request-scoped + // `cache`) into a heap-allocated DecodedDictBlock; *pin holds it alive so *out + // never dangles under a later cache eviction. Callers must keep *pin alive for + // as long as they use *out. + Status dict_block_reader_for_ordinal(uint32_t ordinal, DictBlockCache* cache, + std::shared_ptr* pin, + const format::DictBlockReader** out) const; + Status locate_candidate_dict_block(std::string_view term, bool* maybe_present, + uint32_t* ordinal) const; + Status collect_batch_lookup_groups(const std::vector& terms, + std::vector* candidates, + std::vector* groups) const; + static Status resolve_batch_lookup_group(const std::vector& terms, + const std::vector& candidates, + const BatchLookupGroup& group, + const format::DictBlockReader& block_reader, + std::vector* results); + Status lookup_batch_on_demand(const std::vector& terms, + const std::vector& candidates, + const std::vector& groups, + std::vector* results) const; + std::vector resident_dict_blocks_; + std::shared_ptr norms_cache_; + size_t norms_reserved_charge_ = 0; +}; + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/snii_segment_reader.cpp b/be/src/storage/index/snii/reader/snii_segment_reader.cpp new file mode 100644 index 00000000000000..e769946b9eecfa --- /dev/null +++ b/be/src/storage/index/snii/reader/snii_segment_reader.cpp @@ -0,0 +1,193 @@ +// 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 "storage/index/snii/reader/snii_segment_reader.h" + +#include +#include +#include +#include + +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/tail_pointer.h" + +namespace doris::snii::reader { +namespace { + +Status corrupted(std::string_view message) { + return Status::Error(message); +} + +Status read_tail_pointer(io::FileReader* reader, format::TailPointer* tail, + uint64_t* footer_offset) { + const size_t footer_size = format::tail_pointer_size(); + const uint64_t total = reader->size(); + if (total < footer_size) { + return corrupted("segment: file smaller than tail pointer"); + } + *footer_offset = total - footer_size; + std::vector bytes; + RETURN_IF_ERROR(reader->read_at(*footer_offset, footer_size, &bytes)); + return format::decode_tail_pointer(Slice(bytes), tail); +} + +// Proves Core -> STI -> DBD are adjacent and end at or before directory_offset, so open_index() can +// cover the whole group with one range read. It also bounds core.length + sti.length + dbd.length by +// directory_offset (itself bounded by the file size), which is why open_index() may sum and narrow +// those on-disk 64-bit lengths to size_t without a further overflow check. +Status validate_metadata_group(const format::LogicalIndexMetadataRef& entry, + uint64_t directory_offset) { + const auto& core = entry.core_metadata; + const auto& sti = entry.sampled_term_index; + const auto& dbd = entry.dict_block_directory; + if (core.offset > directory_offset || core.length > directory_offset - core.offset) { + return corrupted("segment: Core metadata reference is outside metadata area"); + } + const uint64_t sti_offset = core.offset + core.length; + if (sti.offset != sti_offset) { + return corrupted("segment: STI metadata is not adjacent to Core metadata"); + } + if (sti.length > directory_offset - sti.offset) { + return corrupted("segment: STI metadata reference is outside metadata area"); + } + const uint64_t dbd_offset = sti.offset + sti.length; + if (dbd.offset != dbd_offset) { + return corrupted("segment: DBD metadata is not adjacent to STI metadata"); + } + if (dbd.length > directory_offset - dbd.offset) { + return corrupted("segment: DBD metadata reference is outside metadata area"); + } + return Status::OK(); +} + +Status find_metadata_ref(const format::MetadataDirectory& directory, uint64_t index_id, + std::string_view suffix, const format::LogicalIndexMetadataRef** out) { + *out = directory.find(index_id, suffix); + if (*out == nullptr) { + return Status::Error( + "segment: logical index not found"); + } + return Status::OK(); +} + +} // namespace + +Status SniiSegmentReader::open(io::FileReader* const reader, SniiSegmentReader* const out) { + if (reader == nullptr) { + return Status::Error("segment: null reader"); + } + if (out == nullptr) { + return Status::Error("segment: null out"); + } + *out = {}; + + // The per-segment bootstrap header at offset zero remains an inspect-tool record and is + // intentionally not read here. The footer validates the exact format version and its own CRC, + // and it locates the metadata directory. Reading only the file tail avoids an otherwise + // redundant offset-zero cache block or remote round trip on cold queries. Future incompatible + // evolution must bump the footer format version rather than rely on a min-reader-version change + // under a stable format version. + format::TailPointer tail; + uint64_t footer_offset = 0; + RETURN_IF_ERROR(read_tail_pointer(reader, &tail, &footer_offset)); + if (tail.directory_offset > footer_offset || + tail.directory_length > footer_offset - tail.directory_offset) { + return corrupted("segment: metadata directory reference overlaps footer or EOF"); + } + if (tail.directory_length > static_cast(std::numeric_limits::max())) { + return corrupted("segment: metadata directory exceeds protobuf parse limit"); + } + // The protobuf parse limit above already bounds the length, so narrowing it is safe. + const auto directory_length = static_cast(tail.directory_length); + + std::vector directory_bytes; + RETURN_IF_ERROR(reader->read_at(tail.directory_offset, directory_length, &directory_bytes)); + if (crc32c(Slice(directory_bytes)) != tail.directory_crc32c) { + return corrupted("segment: metadata directory crc32c mismatch"); + } + + format::MetadataDirectory directory; + RETURN_IF_ERROR(format::MetadataDirectory::decode(Slice(directory_bytes), &directory)); + for (const auto& entry : directory.entries()) { + RETURN_IF_ERROR(validate_metadata_group(entry, tail.directory_offset)); + } + + out->reader_ = reader; + out->directory_ = std::move(directory); + return Status::OK(); +} + +Status SniiSegmentReader::index_exists(uint64_t index_id, std::string_view suffix, + bool* const exists) const { + if (exists == nullptr) { + return Status::Error("segment: null exists out"); + } + if (reader_ == nullptr) { + return Status::Error("segment: not opened"); + } + *exists = directory_.find(index_id, suffix) != nullptr; + return Status::OK(); +} + +Status SniiSegmentReader::open_index(uint64_t index_id, std::string_view suffix, + LogicalIndexReader* const out, + LogicalIndexOpenMode open_mode) const { + if (out == nullptr) { + return Status::Error("segment: null index out"); + } + *out = {}; + if (reader_ == nullptr) { + return Status::Error("segment: not opened"); + } + const format::LogicalIndexMetadataRef* entry = nullptr; + RETURN_IF_ERROR(find_metadata_ref(directory_, index_id, suffix, &entry)); + + // Safe to sum and narrow: open() ran validate_metadata_group on every directory entry. + const auto core_length = static_cast(entry->core_metadata.length); + const auto sti_length = static_cast(entry->sampled_term_index.length); + const auto dbd_length = static_cast(entry->dict_block_directory.length); + const size_t group_length = core_length + sti_length + dbd_length; + std::vector group; + RETURN_IF_ERROR(reader_->read_at(entry->core_metadata.offset, group_length, &group)); + DORIS_CHECK_EQ(group.size(), group_length); + const Slice bytes(group); + return LogicalIndexReader::open( + reader_, bytes.subslice(0, core_length), bytes.subslice(core_length, sti_length), + bytes.subslice(core_length + sti_length, dbd_length), out, open_mode); +} + +Status SniiSegmentReader::section_refs_for_index(uint64_t index_id, std::string_view suffix, + format::SectionRefs* const out) const { + if (out == nullptr) { + return Status::Error("segment: null section refs out"); + } + if (reader_ == nullptr) { + return Status::Error("segment: not opened"); + } + const format::LogicalIndexMetadataRef* entry = nullptr; + RETURN_IF_ERROR(find_metadata_ref(directory_, index_id, suffix, &entry)); + std::vector core_bytes; + RETURN_IF_ERROR(reader_->read_at(entry->core_metadata.offset, entry->core_metadata.length, + &core_bytes)); + format::CoreMetadata core; + RETURN_IF_ERROR(format::decode_core_metadata(Slice(core_bytes), &core)); + *out = core.section_refs; + return Status::OK(); +} + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/snii_segment_reader.h b/be/src/storage/index/snii/reader/snii_segment_reader.h new file mode 100644 index 00000000000000..42c29f867c6e3c --- /dev/null +++ b/be/src/storage/index/snii/reader/snii_segment_reader.h @@ -0,0 +1,79 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/metadata_directory.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// SniiSegmentReader -- entry point for the SNII segment read path. It opens a +// single .idx container through a (possibly metered) io::FileReader and exposes +// its logical indexes. open() reads only the file tail: +// 1. the fixed tail pointer (last tail_pointer_size() bytes), which also gates +// the container format_version ('TAIL' magic + format_version exact-match + +// tail crc), and +// 2. the raw protobuf logical-index metadata directory. +// The bootstrap header at offset 0 is still WRITTEN on disk (for inspect tooling) +// but is intentionally NOT read at open: its only runtime role (the container +// version gate) is already covered, more strictly, by the tail pointer, so +// skipping it avoids a redundant offset-0 cache block / remote round-trip per +// segment on cold queries. +// Per-index metadata groups are read lazily by open_index() so opening one logical +// index does not read every other logical index's metadata. +// +// open_index() then materializes one LogicalIndexReader from the metadata group +// of a given (index_id, suffix); query functions operate on that reader. +namespace doris::snii::reader { + +class SniiSegmentReader { +public: + SniiSegmentReader() = default; + + // Reads the tail pointer + raw metadata directory from reader (the offset-0 + // bootstrap header is not read; the tail pointer gates the container version). + // reader must outlive the returned SniiSegmentReader and every + // LogicalIndexReader opened from it. reader == nullptr / out == nullptr -> + // InvalidArgument; structural problems -> Corruption / Unsupported. + static Status open(io::FileReader* const reader, SniiSegmentReader* const out); + + uint32_t n_logical_indexes() const { return static_cast(directory_.size()); } + + Status index_exists(uint64_t index_id, std::string_view suffix, bool* const exists) const; + + // Loads the adjacent Core/STI/DBD group for (index_id, suffix) and builds a + // LogicalIndexReader bound to the same FileReader. Absent index -> NotFound. + Status open_index(uint64_t index_id, std::string_view suffix, LogicalIndexReader* const out, + LogicalIndexOpenMode open_mode = LogicalIndexOpenMode::kQuery) const; + Status section_refs_for_index(uint64_t index_id, std::string_view suffix, + format::SectionRefs* const out) const; + + io::FileReader* reader() const { return reader_; } + +private: + io::FileReader* reader_ = nullptr; + format::MetadataDirectory directory_; +}; + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/windowed_posting.cpp b/be/src/storage/index/snii/reader/windowed_posting.cpp new file mode 100644 index 00000000000000..17d931640fb428 --- /dev/null +++ b/be/src/storage/index/snii/reader/windowed_posting.cpp @@ -0,0 +1,308 @@ +// 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 "storage/index/snii/reader/windowed_posting.h" + +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/batch_range_fetcher.h" + +namespace doris::snii::reader { + +using format::DictEntry; +using format::FrqPreludeReader; +using format::FrqRegionMeta; +using format::WindowMeta; + +namespace { + +// Resolves the absolute file offset of the prelude bytes for a windowed entry. +// The frq span lives in the interleaved posting region (after the term's prx span). +uint64_t prelude_abs(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base) { + const auto& region = idx.section_refs().posting_region; + return region.offset + frq_base + entry.frq_off_delta; +} + +// Validates that [off, off+len) fits within [0, total). +Status in_bounds(uint64_t off, uint64_t len, uint64_t total) { + if (off > total || len > total - off) { + return Status::Error( + "windowed_posting: range out of section"); + } + return Status::OK(); +} + +// Block geometry of a windowed entry's grouped .frq payload (all offsets absolute). +struct BlockGeometry { + uint64_t dd_block_off = 0; // absolute start of the dd-block + uint64_t dd_block_len = 0; + uint64_t freq_block_off = 0; // absolute start of the freq-block + uint64_t freq_block_len = 0; + uint64_t frq_region_len = 0; // entry.frq_len - prelude_len (dd-block + freq-block) +}; + +// Derives the dd-block / freq-block absolute ranges from the entry + prelude, +// validating they tile the post-prelude .frq region exactly. +Status resolve_blocks(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t frq_base, + const FrqPreludeReader& prelude, BlockGeometry* g) { + if (entry.prelude_len > entry.frq_len) { + return Status::Error( + "windowed_posting: prelude_len exceeds frq_len"); + } + const uint64_t frq_window_start = prelude_abs(idx, entry, frq_base) + entry.prelude_len; + g->frq_region_len = entry.frq_len - entry.prelude_len; + g->dd_block_len = prelude.dd_block_len(); + g->freq_block_len = prelude.freq_block_len(); + // dd-block + freq-block must fit exactly within the post-prelude region. + if (g->dd_block_len > g->frq_region_len || + g->freq_block_len > g->frq_region_len - g->dd_block_len) { + return Status::Error( + "windowed_posting: blocks exceed frq region"); + } + g->dd_block_off = frq_window_start; + g->freq_block_off = frq_window_start + g->dd_block_len; + return Status::OK(); +} + +// Per-window decode state for the full-posting path. +struct WindowSlices { + WindowMeta meta; + Slice dd_region; + Slice freq_region; + Slice prx_window; +}; + +// Carves window w's dd (and freq when want_freq) sub-slices out of the fetched +// blocks, validating each locator against its block length. +Status carve_region_slices(const WindowMeta& m, Slice dd_block, Slice freq_block, bool want_freq, + WindowSlices* out) { + RETURN_IF_ERROR(in_bounds(m.dd_off, m.dd_disk_len, dd_block.size())); + out->dd_region = + dd_block.subslice(static_cast(m.dd_off), static_cast(m.dd_disk_len)); + if (!want_freq) { + return Status::OK(); + } + RETURN_IF_ERROR(in_bounds(m.freq_off, m.freq_disk_len, freq_block.size())); + out->freq_region = freq_block.subslice(static_cast(m.freq_off), + static_cast(m.freq_disk_len)); + return Status::OK(); +} + +// Decodes window w from the fetched blocks (+ optional prx slice) and appends to out. +Status append_window(const WindowSlices& ws, bool want_positions, bool want_freq, + DecodedPosting* out) { + std::vector docids, freqs; + std::vector> pos; + RETURN_IF_ERROR(decode_window_slices(ws.meta, ws.dd_region, ws.freq_region, ws.prx_window, + want_positions, want_freq, &docids, &freqs, &pos)); + out->docids.insert(out->docids.end(), docids.begin(), docids.end()); + out->freqs.insert(out->freqs.end(), freqs.begin(), freqs.end()); + if (want_positions) { + for (auto& v : pos) { + out->positions.push_back(std::move(v)); + } + } + return Status::OK(); +} + +} // namespace + +Status fetch_windowed_prelude(const LogicalIndexReader& idx, const DictEntry& entry, + uint64_t frq_base, FrqPreludeReader* prelude) { + if (entry.prelude_len == 0) { + return Status::Error( + "windowed_posting: windowed entry has no prelude"); + } + if (entry.prelude_len > entry.frq_len) { + return Status::Error( + "windowed_posting: prelude_len exceeds frq_len"); + } + const uint64_t prelude_offset = prelude_abs(idx, entry, frq_base); + io::BatchRangeFetcher fetcher(idx.reader()); + const size_t h = fetcher.add(prelude_offset, entry.prelude_len); + RETURN_IF_ERROR(fetcher.fetch()); + return FrqPreludeReader::open(fetcher.get(h), prelude); +} + +Status windowed_window_range(const LogicalIndexReader& idx, const DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, const FrqPreludeReader& prelude, + uint32_t w, bool want_positions, bool want_freq, WindowAbsRange* out) { + if (out == nullptr) { + return Status::Error("windowed_posting: null range"); + } + *out = WindowAbsRange {}; + BlockGeometry g; + RETURN_IF_ERROR(resolve_blocks(idx, entry, frq_base, prelude, &g)); + WindowMeta meta; + RETURN_IF_ERROR(prelude.window(w, &meta)); + + // dd sub-range within the dd-block. + RETURN_IF_ERROR(in_bounds(meta.dd_off, meta.dd_disk_len, g.dd_block_len)); + out->dd_off = g.dd_block_off + meta.dd_off; + out->dd_len = meta.dd_disk_len; + + if (want_freq) { + // Symmetric to the positions guard below: a G16 freq-elided posting + // (freq-dropped index or prune-mode bigram) declares has_freq=false in + // its prelude flags. INVALID_ARGUMENT and not FILE_CORRUPTED: the + // Doris segment iterator silently downgrades the corruption code to a + // non-index evaluation, which would mask this by-design layout. + if (!prelude.has_freq()) { + return Status::Error( + "windowed_posting: freqs requested but prelude has none"); + } + RETURN_IF_ERROR(in_bounds(meta.freq_off, meta.freq_disk_len, g.freq_block_len)); + out->freq_off = g.freq_block_off + meta.freq_off; + out->freq_len = meta.freq_disk_len; + } + + if (!want_positions) { + return Status::OK(); + } + if (!prelude.has_prx()) { + return Status::Error( + "windowed_posting: positions requested but prelude has none"); + } + const uint64_t prx_region_start = + idx.section_refs().posting_region.offset + prx_base + entry.prx_off_delta; + RETURN_IF_ERROR(in_bounds(meta.prx_off, meta.prx_len, entry.prx_len)); + out->prx_off = prx_region_start + meta.prx_off; + out->prx_len = meta.prx_len; + return Status::OK(); +} + +Status decode_window_slices(const WindowMeta& meta, Slice dd_region, Slice freq_region, + Slice prx_window, bool want_positions, bool want_freq, + std::vector* docids, std::vector* freqs, + std::vector>* positions) { + FrqRegionMeta dd_meta; + dd_meta.zstd = meta.dd_zstd; + dd_meta.uncomp_len = meta.dd_uncomp_len; + dd_meta.disk_len = meta.dd_disk_len; + dd_meta.crc = meta.crc_dd; + dd_meta.verify_crc = meta.verify_crc; + RETURN_IF_ERROR(format::decode_dd_region(dd_region, dd_meta, meta.win_base, docids)); + if (docids->size() != meta.doc_count) { + return Status::Error( + "windowed_posting: frq doc_count mismatch"); + } + if (want_freq) { + FrqRegionMeta freq_meta; + freq_meta.zstd = meta.freq_zstd; + freq_meta.uncomp_len = meta.freq_uncomp_len; + freq_meta.disk_len = meta.freq_disk_len; + freq_meta.crc = meta.crc_freq; + freq_meta.verify_crc = meta.verify_crc; + RETURN_IF_ERROR(format::decode_freq_region(freq_region, freq_meta, meta.doc_count, freqs)); + } else { + freqs->clear(); + } + if (!want_positions) { + return Status::OK(); + } + + ByteSource psrc(prx_window); + RETURN_IF_ERROR(format::read_prx_window(&psrc, positions)); + if (!psrc.eof()) { + return Status::Error( + "windowed_posting: trailing bytes after prx frame"); + } + if (positions->size() != docids->size()) { + return Status::Error( + "windowed_posting: prx/frq doc-count mismatch"); + } + return Status::OK(); +} + +namespace { + +// Fetches the dd-block (always), the freq-block (when want_freq) and the whole .prx +// region (when want_positions) of a windowed entry in ONE batch and returns the +// in-memory block slices. The dd-block is a single contiguous range -> the +// docid-only / phrase path reads it as one Range GET (the byte-saving core). +Status fetch_blocks(const LogicalIndexReader& idx, const DictEntry& entry, uint64_t prx_base, + const BlockGeometry& g, bool want_positions, bool want_freq, + io::BatchRangeFetcher* fetcher, size_t* dd_h, size_t* freq_h, size_t* prx_h) { + *dd_h = fetcher->add(g.dd_block_off, g.dd_block_len); + if (want_freq) { + *freq_h = fetcher->add(g.freq_block_off, g.freq_block_len); + } + if (want_positions) { + const uint64_t prx_region_start = + idx.section_refs().posting_region.offset + prx_base + entry.prx_off_delta; + *prx_h = fetcher->add(prx_region_start, entry.prx_len); + } + return fetcher->fetch(); +} + +} // namespace + +Status read_windowed_posting(const LogicalIndexReader& idx, const DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, bool want_positions, + bool want_freq, DecodedPosting* out) { + if (out == nullptr) { + return Status::Error("windowed_posting: null out"); + } + *out = DecodedPosting {}; + + FrqPreludeReader prelude; + RETURN_IF_ERROR(fetch_windowed_prelude(idx, entry, frq_base, &prelude)); + if (want_positions && !prelude.has_prx()) { + return Status::Error( + "windowed_posting: positions requested but prelude has none"); + } + // G16 freq-elided postings (freq-dropped index or prune-mode bigram) + // declare has_freq=false; fail with the semantic error instead of a deep + // region-decode corruption. INVALID_ARGUMENT so the Doris segment + // iterator's corruption downgrade never masks this by-design layout. + if (want_freq && !prelude.has_freq()) { + return Status::Error( + "windowed_posting: freqs requested but prelude has none"); + } + BlockGeometry g; + RETURN_IF_ERROR(resolve_blocks(idx, entry, frq_base, prelude, &g)); + + io::BatchRangeFetcher fetcher(idx.reader()); + size_t dd_h = 0, freq_h = 0, prx_h = 0; + RETURN_IF_ERROR(fetch_blocks(idx, entry, prx_base, g, want_positions, want_freq, &fetcher, + &dd_h, &freq_h, &prx_h)); + const Slice dd_block = fetcher.get(dd_h); + const Slice freq_block = want_freq ? fetcher.get(freq_h) : Slice(); + const Slice prx_region = want_positions ? fetcher.get(prx_h) : Slice(); + + const uint32_t n = prelude.n_windows(); + for (uint32_t w = 0; w < n; ++w) { + WindowSlices ws; + RETURN_IF_ERROR(prelude.window(w, &ws.meta)); + RETURN_IF_ERROR(carve_region_slices(ws.meta, dd_block, freq_block, want_freq, &ws)); + if (want_positions) { + RETURN_IF_ERROR(in_bounds(ws.meta.prx_off, ws.meta.prx_len, prx_region.size())); + ws.prx_window = prx_region.subslice(static_cast(ws.meta.prx_off), + static_cast(ws.meta.prx_len)); + } + RETURN_IF_ERROR(append_window(ws, want_positions, want_freq, out)); + } + return Status::OK(); +} + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/reader/windowed_posting.h b/be/src/storage/index/snii/reader/windowed_posting.h new file mode 100644 index 00000000000000..dee527430aff08 --- /dev/null +++ b/be/src/storage/index/snii/reader/windowed_posting.h @@ -0,0 +1,121 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// WindowedPostingReader -- shared read-side decode of a windowed term's posting +// from its two-level frq_prelude + GROUPED dd-block / freq-block (design 1.6). +// +// A windowed pod_ref entry's .frq payload is laid out +// [prelude][dd-block][freq-block] +// where the dd-block concatenates every window's dd_region and the freq-block +// every window's freq_region. The docs-only prefix [prelude][dd-block] is ONE +// contiguous run. This helper: +// 1. range-fetches the prelude (prelude_len bytes) and parses the directory, +// 2. range-fetches the WHOLE dd-block in ONE contiguous range (and, for +// scoring, +// the whole freq-block in one more range), +// 3. decodes each window's dd region (and freq region) from the in-memory +// blocks +// via the prelude metadata (dd_off/dd_disk_len, freq_off/freq_disk_len), +// and concatenates the per-window docids / freqs / positions. +// +// The slim/inline single-window path is handled by the term/phrase/scoring +// callers directly; this helper is for enc=windowed entries only. +namespace doris::snii::reader { + +// Coalesce gap (bytes) used when batch-fetching MULTIPLE dd sub-ranges of the +// SAME term (the phrase window-skip path): dd regions of one term are +// contiguous in the dd-block, so merging reads separated by <= this gap into +// one physical Range GET trades a little over-read for fewer remote GETs (the +// design's higher-priority metric). Only applied to same-term multi-window +// batches, never to cross-term. +inline constexpr uint64_t kSameTermCoalesceGap = 16 * 1024; + +// Full decoded posting for one windowed term (docids ascending across windows). +struct DecodedPosting { + std::vector docids; + std::vector freqs; // aligned with docids + std::vector> positions; // aligned; empty when no prx +}; + +// Decodes the entire windowed posting. want_positions requires the index to +// have positions (and the entry to carry prx). want_freq selects whether the +// freq-block is fetched + decoded: when false ONLY the contiguous +// [prelude][dd-block] prefix is fetched (docid-only / phrase callers) and +// DecodedPosting.freqs stays empty; when true the freq-block is additionally +// fetched (scoring). Returns Corruption on any prelude/block inconsistency +// (doc-count mismatch, out-of-range offsets). +Status read_windowed_posting(const LogicalIndexReader& idx, const format::DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, bool want_positions, + bool want_freq, DecodedPosting* out); + +// --- Sub-block (window) skipping helpers (shared with phrase / selective WAND) +// -- +// +// These expose the per-window dd/freq/prx addressing within the grouped blocks +// so the skip path can fetch ONLY the windows covering candidate docids (their +// dd sub-ranges within the dd-block, near-contiguous and coalesce-friendly) +// instead of the whole posting, without duplicating the offset arithmetic. + +// Absolute file byte ranges of one window's regions. dd is always valid; freq +// is valid only when want_freq; prx is valid only when want_positions (and +// has_prx). +struct WindowAbsRange { + uint64_t dd_off = 0; + uint64_t dd_len = 0; + uint64_t freq_off = 0; + uint64_t freq_len = 0; + uint64_t prx_off = 0; + uint64_t prx_len = 0; +}; + +// Fetches + parses the two-level prelude of a windowed entry (one batched +// read). +Status fetch_windowed_prelude(const LogicalIndexReader& idx, const format::DictEntry& entry, + uint64_t frq_base, format::FrqPreludeReader* prelude); + +// Computes the absolute file ranges of window w's dd region (and freq region +// when want_freq, and .prx window when want_positions), fully validated against +// the POD sections (anti-DoS: rejects out-of-range offsets and overflowing +// locators). +Status windowed_window_range(const LogicalIndexReader& idx, const format::DictEntry& entry, + uint64_t frq_base, uint64_t prx_base, + const format::FrqPreludeReader& prelude, uint32_t w, + bool want_positions, bool want_freq, WindowAbsRange* out); + +// Decodes one window's docids (and per-doc positions when want_positions, and +// per-doc freqs when want_freq) from already-fetched byte slices: dd_region is +// the window's dd sub-slice; freq_region its freq sub-slice (ignored when +// !want_freq); prx_window its .prx bytes. The decoded docids are absolute +// (win_base applied). Returns Corruption on any doc-count mismatch between the +// prelude, dd/freq and prx. +Status decode_window_slices(const format::WindowMeta& meta, Slice dd_region, Slice freq_region, + Slice prx_window, bool want_positions, bool want_freq, + std::vector* docids, std::vector* freqs, + std::vector>* positions); + +} // namespace doris::snii::reader diff --git a/be/src/storage/index/snii/snii_doris_adapter.cpp b/be/src/storage/index/snii/snii_doris_adapter.cpp new file mode 100644 index 00000000000000..e4a5875184535f --- /dev/null +++ b/be/src/storage/index/snii/snii_doris_adapter.cpp @@ -0,0 +1,405 @@ +// 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 "storage/index/snii/snii_doris_adapter.h" + +#include + +#include +#include +#include +#include + +#include "common/cast_set.h" +#include "cpp/sync_point.h" +#include "runtime/exec_env.h" +#include "runtime/thread_context.h" +#include "storage/index/snii/common/uninitialized_buffer.h" +#include "util/countdown_latch.h" +#include "util/threadpool.h" + +namespace doris::segment_v2::snii_doris { +namespace { +// Per-call cap on concurrently dispatched physical segment reads. A coalesced +// batch of at most this many segments is served as a single concurrent round +// (mirrors the "at most one serial round" contract that the test-side +// MeteredFileReader measures, and the S3 standalone reader's 16-way fan-out). +constexpr size_t kMaxConcurrentReads = 16; +} // namespace + +thread_local const io::IOContext* DorisSniiFileReader::_scoped_io_ctx = nullptr; +ThreadPool* DorisSniiFileReader::_io_pool_for_test = nullptr; + +Status DorisSniiFileWriter::append(::doris::snii::Slice data) { + if (_writer == nullptr) { + return Status::Error("doris writer is null"); + } + return _writer->append(Slice(reinterpret_cast(data.data()), data.size())); +} + +Status DorisSniiFileWriter::finalize() { + if (_writer == nullptr) { + return Status::Error("doris writer is null"); + } + return Status::OK(); +} + +uint64_t DorisSniiFileWriter::bytes_written() const { + return _writer == nullptr ? 0 : _writer->bytes_appended(); +} + +DorisSniiFileReader::DorisSniiFileReader(io::FileReaderSPtr reader, const io::IOContext* io_ctx, + bool direct_remote_io) + : _reader(std::move(reader)), + _direct_remote_io(direct_remote_io), + _default_io_ctx(_make_index_io_context(io_ctx)) {} + +io::IOContext DorisSniiFileReader::_make_index_io_context(const io::IOContext* io_ctx) { + io::IOContext index_io_ctx; + if (io_ctx != nullptr) { + index_io_ctx = *io_ctx; + } + index_io_ctx.is_inverted_index = true; + // is_index_data is inherited from io_ctx: META scopes set it true at the source + // (index_file_reader), non-meta reads default to false. + return index_io_ctx; +} + +DorisSniiFileReader::ScopedIOContext::ScopedIOContext(const io::IOContext* io_ctx) + : _previous(_scoped_io_ctx), _io_ctx(DorisSniiFileReader::_make_index_io_context(io_ctx)) { + _scoped_io_ctx = &_io_ctx; +} + +DorisSniiFileReader::ScopedIOContext::~ScopedIOContext() { + _scoped_io_ctx = _previous; +} + +Status DorisSniiFileReader::read_at(uint64_t offset, size_t len, std::vector* out) { + if (out == nullptr) { + return Status::Error("output buffer is null"); + } + RETURN_IF_ERROR(_check_read_range(offset, len)); + RETURN_IF_ERROR(_read_at(offset, len, out, _current_io_ctx())); + if (len > 0) { + _record_read_stats(cast_set(len), cast_set(len), 1, 1); + } + return Status::OK(); +} + +// NOLINTNEXTLINE(readability-non-const-parameter): out is the SNII read output buffer. +Status DorisSniiFileReader::_read_at(uint64_t offset, size_t len, std::vector* out, + const io::IOContext* io_ctx) const { + TEST_SYNC_POINT_RETURN_WITH_VALUE("DorisSniiFileReader::_read_at", + Status::IOError("injected SNII read failure"), offset, len); + DCHECK(_reader != nullptr); + DCHECK(out != nullptr); + DCHECK(_check_read_range(offset, len).ok()); + if (len == 0) { + out->clear(); + return Status::OK(); + } + out->resize(len); + size_t bytes_read = 0; + auto status = _reader->read_at(offset, Slice(out->data(), len), &bytes_read, io_ctx); + if (!status.ok()) { + return status; + } + if (bytes_read != len) { + return Status::Error( + fmt::format("short read at offset {}, expect {}, got {}", offset, len, bytes_read)); + } + return Status::OK(); +} + +// NOLINTBEGIN(readability-non-const-parameter): outs is the SNII batch read output buffer. +Status DorisSniiFileReader::read_batch(const std::vector<::doris::snii::io::Range>& ranges, + std::vector>* outs) { + if (outs == nullptr) { + return Status::Error("output buffers is null"); + } + outs->clear(); + outs->resize(ranges.size()); + if (ranges.empty()) { + return Status::OK(); + } + + // ----- Phase 1: plan (serial, lock-free) ----- + // No section-classification lock exists on this reader, so the whole plan is + // a plain in-memory scan; the NO-IO-UNDER-LOCK red line holds trivially (no + // lock is taken anywhere in read_batch). + struct IndexedRange { + uint64_t offset = 0; + size_t len = 0; + size_t index = 0; + }; + int64_t request_bytes = 0; + std::vector sorted; + sorted.reserve(ranges.size()); + for (size_t i = 0; i < ranges.size(); ++i) { + RETURN_IF_ERROR(_check_read_range(ranges[i].offset, ranges[i].len)); + request_bytes += cast_set(ranges[i].len); + if (ranges[i].len == 0) { + continue; + } + sorted.push_back({ranges[i].offset, ranges[i].len, i}); + } + if (sorted.empty()) { + return Status::OK(); + } + // F27: callers (BatchRangeFetcher::fetch) already pass offset-sorted ranges; + // only pay for a sort when the input is actually out of order. + auto by_offset = [](const IndexedRange& lhs, const IndexedRange& rhs) { + return lhs.offset < rhs.offset; + }; + if (!std::ranges::is_sorted(sorted, by_offset)) { + std::ranges::sort(sorted, by_offset); + } + + // Coalesce the sorted ranges into disjoint physical segments. + struct Seg { + uint64_t offset = 0; + size_t len = 0; + size_t begin = 0; // first index into `sorted` covered by this segment + size_t end = 0; // one-past-last index into `sorted` + bool single = false; + }; + constexpr uint64_t max_coalesced_gap = 4096; + constexpr uint64_t max_coalesced_read = 1ULL << 20; + std::vector segs; + for (size_t begin = 0; begin < sorted.size();) { + uint64_t read_offset = sorted[begin].offset; + uint64_t read_end = sorted[begin].offset + sorted[begin].len; + size_t end = begin + 1; + while (end < sorted.size()) { + const uint64_t next_end = sorted[end].offset + sorted[end].len; + if ((sorted[end].offset > read_end && + sorted[end].offset - read_end > max_coalesced_gap) || + next_end - read_offset > max_coalesced_read) { + break; + } + read_end = std::max(read_end, next_end); + ++end; + } + Seg seg; + seg.offset = read_offset; + seg.len = cast_set(read_end - read_offset); + seg.begin = begin; + seg.end = end; + // A single-range group exactly covers its segment, so it can be read + // straight into the caller's output slot with no temp + no second copy. + seg.single = (end == begin + 1); + segs.push_back(seg); + begin = end; + } + + // Resolve per-segment target buffers, io contexts and the shared sink on the + // calling thread: workers (which run on tracker-less pool threads) must not + // allocate, and per-segment private cache-stat slots keep disjoint physical + // reads from racing on the shared FileCacheStatistics. + const size_t num_segs = segs.size(); + const io::IOContext* base_io_ctx = _current_io_ctx(); + std::vector> tmp_bufs(num_segs); + std::vector*> targets(num_segs); + std::vector seg_stats(num_segs); + std::vector seg_io_ctx(num_segs); + std::vector seg_status(num_segs); + int64_t read_bytes = 0; + for (size_t s = 0; s < num_segs; ++s) { + const Seg& seg = segs[s]; + std::vector* target = + seg.single ? &(*outs)[sorted[seg.begin].index] : &tmp_bufs[s]; + ::doris::snii::resize_uninitialized(*target, seg.len); + targets[s] = target; + seg_io_ctx[s] = *base_io_ctx; + seg_io_ctx[s].file_cache_stats = + base_io_ctx->file_cache_stats != nullptr ? &seg_stats[s] : nullptr; + read_bytes += cast_set(seg.len); + } + + // ----- Phase 2: physical reads (lock-free; concurrent when a pool exists) ----- + auto run_segment = [&](size_t s) { + seg_status[s] = _read_at(segs[s].offset, segs[s].len, targets[s], &seg_io_ctx[s]); + }; + ThreadPool* pool = _select_io_pool(); + if (pool != nullptr && num_segs > 1) { + // Carried onto the pool threads below. They are "Orphan" threads with no MemTrackerLimiter + // of their own (buffered_reader.cpp:426), and the read does allocate down there: in cloud + // mode _read_at reaches CachedRemoteFileReader::read_at_impl -> + // _read_from_indirect_cache -> _read_remote_blocks_into_cache -> _execute_remote_read -> + // _execute_s3_fallback, which does `new char[span_size]`. Without this the span is charged + // to no tracker at all, which memory_orphan_check() (on by default) treats as a bug. + // Same shape as the hedged-read path in cached_remote_file_reader.cpp:500-505. + const std::shared_ptr parent_resource_ctx = + thread_context()->resource_ctx(); + for (size_t base = 0; base < num_segs; base += kMaxConcurrentReads) { + const size_t wave_end = std::min(base + kMaxConcurrentReads, num_segs); + ::doris::CountDownLatch latch(cast_set(wave_end - base)); + for (size_t s = base; s < wave_end; ++s) { + Status submit_st = + pool->submit_func([&run_segment, &latch, s, parent_resource_ctx]() { + std::unique_ptr attach_task; + if (parent_resource_ctx != nullptr) { + attach_task = std::make_unique(parent_resource_ctx); + } + run_segment(s); + latch.count_down(); + }); + if (!submit_st.ok()) { + // Pool full/shut down: read this segment inline; never skip + // the count_down or the latch would not drain. + run_segment(s); + latch.count_down(); + } + } + latch.wait(); + } + } else { + // Serial fallback: no executor (e.g. tools without ExecEnv) or a single + // segment (avoids micro-batch scheduling overhead). + for (size_t s = 0; s < num_segs; ++s) { + run_segment(s); + } + } + + // ----- Phase 3: merge stats, first-error, scatter, account (serial) ----- + // Fold every segment's private stats back FIRST: physical IO that already + // happened (including partial work inside a segment that then failed) must + // reach the query profile even when another segment of this batch errors. + if (base_io_ctx->file_cache_stats != nullptr) { + for (size_t s = 0; s < num_segs; ++s) { + _merge_file_cache_statistics(base_io_ctx->file_cache_stats, seg_stats[s]); + } + } + Status first_error = Status::OK(); + int64_t completed_request_bytes = 0; + int64_t completed_read_bytes = 0; + for (size_t s = 0; s < num_segs; ++s) { + if (!seg_status[s].ok()) { + if (first_error.ok()) { + first_error = seg_status[s]; + } + continue; + } + completed_read_bytes += cast_set(segs[s].len); + for (size_t i = segs[s].begin; i < segs[s].end; ++i) { + completed_request_bytes += cast_set(sorted[i].len); + } + } + if (!first_error.ok()) { + // Record only what actually completed so the logical counters stay + // truthful for the failed batch; ranges/rounds reflect what was issued. + _record_read_stats(completed_request_bytes, completed_read_bytes, + cast_set(num_segs), + cast_set(_compute_num_waves(num_segs))); + return first_error; + } + for (size_t s = 0; s < num_segs; ++s) { + const Seg& seg = segs[s]; + if (seg.single) { + continue; // already read in place + } + const std::vector& bytes = tmp_bufs[s]; + for (size_t i = seg.begin; i < seg.end; ++i) { + const uint64_t pos = sorted[i].offset - seg.offset; + auto& out = (*outs)[sorted[i].index]; + out.assign(bytes.begin() + cast_set(pos), + bytes.begin() + cast_set(pos + sorted[i].len)); + } + } + _record_read_stats(request_bytes, read_bytes, cast_set(num_segs), + cast_set(_compute_num_waves(num_segs))); + return Status::OK(); +} +// NOLINTEND(readability-non-const-parameter) + +uint64_t DorisSniiFileReader::size() const { + return _reader == nullptr ? 0 : _reader->size(); +} + +const io::IOContext* DorisSniiFileReader::_current_io_ctx() const { + return _scoped_io_ctx != nullptr ? _scoped_io_ctx : &_default_io_ctx; +} + +void DorisSniiFileReader::_record_read_stats(int64_t request_bytes, int64_t read_bytes, + int64_t range_read_count, + int64_t serial_read_rounds) const { + const auto* io_ctx = _current_io_ctx(); + if (io_ctx->file_cache_stats == nullptr) { + return; + } + auto* stats = io_ctx->file_cache_stats; + stats->inverted_index_request_bytes += request_bytes; + stats->inverted_index_read_bytes += read_bytes; + stats->inverted_index_range_read_count += range_read_count; + stats->inverted_index_serial_read_rounds += serial_read_rounds; + if (_direct_remote_io) { + // No CachedRemoteFileReader below us: every byte this reader fetched was a + // direct remote GET, so account it as physical remote IO here. + stats->inverted_index_remote_physical_read_bytes += read_bytes; + } +} + +void DorisSniiFileReader::set_io_thread_pool_for_test(ThreadPool* pool) { + _io_pool_for_test = pool; +} + +ThreadPool* DorisSniiFileReader::_select_io_pool() { + if (_io_pool_for_test != nullptr) { + return _io_pool_for_test; + } + if (ExecEnv::ready()) { + return ExecEnv::GetInstance()->buffered_reader_prefetch_thread_pool(); + } + return nullptr; +} + +size_t DorisSniiFileReader::_compute_num_waves(size_t seg_count) { + if (seg_count == 0) { + return 0; + } + return (seg_count + kMaxConcurrentReads - 1) / kMaxConcurrentReads; +} + +void DorisSniiFileReader::_merge_file_cache_statistics(io::FileCacheStatistics* dst, + const io::FileCacheStatistics& src) { + if (dst == nullptr) { + return; + } + // Delegate to the canonical field list so per-wave private stats can never + // silently drop fields the general layer adds (e.g. write_cache_io_timer, + // remote_only_on_miss_*), which a hand-rolled copy here used to do. + dst->merge_from(src); +} + +Status DorisSniiFileReader::_check_read_range(uint64_t offset, size_t len) const { + if (_reader == nullptr) { + return Status::Error("doris reader is null"); + } + if (offset > std::numeric_limits::max() - len) { + return Status::Error( + fmt::format("read range overflows: offset {}, len {}", offset, len)); + } + const uint64_t end = offset + len; + if (end > _reader->size()) { + return Status::Error( + fmt::format("read range exceeds file size: offset {}, len {}, file size {}", offset, + len, _reader->size())); + } + return Status::OK(); +} + +} // namespace doris::segment_v2::snii_doris diff --git a/be/src/storage/index/snii/snii_doris_adapter.h b/be/src/storage/index/snii/snii_doris_adapter.h new file mode 100644 index 00000000000000..8d5966e0f7a90d --- /dev/null +++ b/be/src/storage/index/snii/snii_doris_adapter.h @@ -0,0 +1,114 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "io/fs/file_reader.h" +#include "io/fs/file_writer.h" +#include "io/io_common.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/file_writer.h" +#include "util/slice.h" + +namespace doris { +class ThreadPool; +} // namespace doris + +namespace doris::segment_v2::snii_doris { + +class DorisSniiFileWriter final : public ::doris::snii::io::FileWriter { +public: + explicit DorisSniiFileWriter(io::FileWriter* writer) : _writer(writer) {} + + Status append(::doris::snii::Slice data) override; + Status finalize() override; + uint64_t bytes_written() const override; + +private: + io::FileWriter* _writer = nullptr; +}; + +class DorisSniiFileReader final : public ::doris::snii::io::FileReader { +public: + class ScopedIOContext { + public: + explicit ScopedIOContext(const io::IOContext* io_ctx); + ~ScopedIOContext(); + + ScopedIOContext(const ScopedIOContext&) = delete; + ScopedIOContext& operator=(const ScopedIOContext&) = delete; + + private: + const io::IOContext* _previous = nullptr; + io::IOContext _io_ctx; + }; + + // `direct_remote_io` marks a reader whose byte ranges are served straight from + // remote storage with no CachedRemoteFileReader in between (the NO_CACHE + // diagnostic bypass on a non-local filesystem). Only that layer would normally + // account physical remote bytes, so this reader then counts its own reads as + // physical remote IO. Never set it for cached or local readers: the former + // double-counts, the latter reports local disk as remote fetch volume. + explicit DorisSniiFileReader(io::FileReaderSPtr reader, const io::IOContext* io_ctx = nullptr, + bool direct_remote_io = false); + + Status read_at(uint64_t offset, size_t len, std::vector* out) override; + Status read_batch(const std::vector<::doris::snii::io::Range>& ranges, + std::vector>* outs) override; + uint64_t size() const override; + + // Test-only: inject (or clear with nullptr) the thread pool used to fan out + // batch segment reads. nullptr (default) routes to the BE buffered-reader + // prefetch pool when ExecEnv is ready, otherwise a serial fallback. + static void set_io_thread_pool_for_test(ThreadPool* pool); + +private: + static io::IOContext _make_index_io_context(const io::IOContext* io_ctx); + Status _check_read_range(uint64_t offset, size_t len) const; + Status _read_at(uint64_t offset, size_t len, std::vector* out, + const io::IOContext* io_ctx) const; + const io::IOContext* _current_io_ctx() const; + void _record_read_stats(int64_t request_bytes, int64_t read_bytes, int64_t range_read_count, + int64_t serial_read_rounds) const; + + // Selects the executor for parallel batch segment reads: the test seam if + // set, else the BE prefetch pool when ExecEnv is ready, else nullptr (the + // caller then reads segments serially). + static ThreadPool* _select_io_pool(); + // Number of dependent serial dispatch waves for `seg_count` physical + // segments given the per-call concurrency cap. This is the F19 metric: + // a coalesced batch of <= cap segments is a single concurrent round. + static size_t _compute_num_waves(size_t seg_count); + // Folds a per-segment private stats slot back into the shared sink (disjoint + // physical reads never race on it). Delegates to FileCacheStatistics::merge_from + // so the field list can never drift from io_common.h. + static void _merge_file_cache_statistics(io::FileCacheStatistics* dst, + const io::FileCacheStatistics& src); + + io::FileReaderSPtr _reader; + const bool _direct_remote_io = false; + io::IOContext _default_io_ctx; + static thread_local const io::IOContext* _scoped_io_ctx; + // Test seam for the batch-read executor; see set_io_thread_pool_for_test. + static ThreadPool* _io_pool_for_test; +}; + +} // namespace doris::segment_v2::snii_doris diff --git a/be/src/storage/index/snii/snii_index_reader.cpp b/be/src/storage/index/snii/snii_index_reader.cpp new file mode 100644 index 00000000000000..58a1679fb34de3 --- /dev/null +++ b/be/src/storage/index/snii/snii_index_reader.cpp @@ -0,0 +1,1198 @@ +// 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 "storage/index/snii/snii_index_reader.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "runtime/exec_env.h" +#include "runtime/query_context.h" +#include "runtime/runtime_profile.h" +#include "runtime/runtime_state.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/index_reader_helper.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/analyzer/segment_analyzer_context.h" +#include "storage/index/inverted/common/single_flight.h" +#include "storage/index/inverted/inverted_index_cache.h" +#include "storage/index/inverted/inverted_index_iterator.h" +#include "storage/index/inverted/token_filter/common_grams_filter.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/count_query.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/scoring_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/index/snii/snii_prx_profile.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "util/defer_op.h" +#include "util/time.h" + +#ifdef BE_TEST +namespace doris::snii::testing { +namespace { + +std::atomic prx_execution_profile_scope_constructions {0}; +std::atomic prx_execution_profile_scope_flushes {0}; + +} // namespace + +void record_prx_execution_profile_scope_construction() { + prx_execution_profile_scope_constructions.fetch_add(1, std::memory_order_relaxed); +} + +void record_prx_execution_profile_scope_flush() { + prx_execution_profile_scope_flushes.fetch_add(1, std::memory_order_relaxed); +} + +void reset_prx_execution_profile_scope_counters() { + prx_execution_profile_scope_constructions.store(0, std::memory_order_relaxed); + prx_execution_profile_scope_flushes.store(0, std::memory_order_relaxed); +} + +uint64_t prx_execution_profile_scope_construction_count() { + return prx_execution_profile_scope_constructions.load(std::memory_order_relaxed); +} + +uint64_t prx_execution_profile_scope_flush_count() { + return prx_execution_profile_scope_flushes.load(std::memory_order_relaxed); +} + +} // namespace doris::snii::testing +#endif + +namespace doris::segment_v2 { + +namespace { + +class RoaringDocIdSink final : public ::doris::snii::query::DocIdSink { +public: + explicit RoaringDocIdSink(roaring::Roaring* bitmap) : _bitmap(bitmap) { + DCHECK(_bitmap != nullptr); + } + + Status append_sorted(std::span docids) override { + if (!docids.empty()) { + _bitmap->addMany(docids.size(), docids.data()); + } + return Status::OK(); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + if (last_exclusive > first) { + _bitmap->addRange(first, last_exclusive); + } + return Status::OK(); + } + + // Roaring addMany/addRange deduplicate and order natively, so multi-term OR + // can stream each posting straight into the bitmap (no per-term vector + merge). + bool dedups() const override { return true; } + +private: + roaring::Roaring* _bitmap; +}; + +struct SniiQueryExecutionResult { + std::shared_ptr bitmap; + std::vector<::doris::snii::query::PhraseMatch> phrase_matches; +}; + +std::vector to_terms(const InvertedIndexQueryInfo& query_info) { + std::vector terms; + terms.reserve(query_info.term_infos.size()); + for (const auto& term_info : query_info.term_infos) { + DCHECK(term_info.is_single_term()); + terms.push_back(term_info.get_single_term()); + } + return terms; +} + +bool uses_plain_term_frequency_scoring(InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info) { + return query_type == InvertedIndexQueryType::MATCH_ANY_QUERY || + query_type == InvertedIndexQueryType::MATCH_ALL_QUERY || + (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && + query_info.term_infos.size() == 1); +} + +bool uses_phrase_frequency_scoring(InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info) { + return query_info.term_infos.size() > 1 && + (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY || + query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY); +} + +Status score_plain_term_candidates(const IndexQueryContextPtr& context, + std::string_view column_name, + const InvertedIndexQueryInfo& query_info, + const ::doris::snii::reader::LogicalIndexReader& logical_reader, + const ::doris::snii::stats::SniiStatsProvider& segment_stats, + const roaring::Roaring& final_candidates) { + DORIS_CHECK(context->collection_statistics != nullptr); + DORIS_CHECK(context->collection_similarity != nullptr); + + const std::wstring field_name = StringUtil::string_to_wstring(std::string(column_name)); + const double collection_avgdl = + context->collection_statistics->get_or_calculate_avg_dl(field_name); + std::vector<::doris::snii::query::CollectionScoringTerm> scoring_terms; + scoring_terms.reserve(query_info.term_infos.size()); + for (const auto& term_info : query_info.term_infos) { + DORIS_CHECK(term_info.is_single_term()); + std::string physical_term; + bool representable = false; + RETURN_IF_ERROR(::doris::snii::query::internal::route_query_term( + logical_reader, term_info, &physical_term, &representable)); + if (!representable) { + continue; + } + const std::string& logical_term = term_info.get_single_term(); + const double idf = context->collection_statistics->get_or_calculate_idf( + field_name, StringUtil::string_to_wstring(logical_term)); + scoring_terms.push_back({.physical_term = std::move(physical_term), .idf = idf}); + } + DORIS_CHECK(final_candidates.isEmpty() || !scoring_terms.empty()); + + std::vector<::doris::snii::query::ScoredDoc> scored_docs; + RETURN_IF_ERROR(::doris::snii::query::scoring_query_candidates( + logical_reader, segment_stats, scoring_terms, final_candidates, collection_avgdl, + ::doris::snii::query::Bm25Params {}, &scored_docs)); + for (const auto& scored_doc : scored_docs) { + context->collection_similarity->collect(scored_doc.docid, + static_cast(scored_doc.score)); + } + return Status::OK(); +} + +Status score_phrase_matches(const IndexQueryContextPtr& context, std::string_view column_name, + InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, + const ::doris::snii::reader::LogicalIndexReader& logical_reader, + const ::doris::snii::stats::SniiStatsProvider& segment_stats, + const roaring::Roaring& final_candidates, + const std::vector<::doris::snii::query::PhraseMatch>& matches) { + DORIS_CHECK(context->collection_statistics != nullptr); + DORIS_CHECK(context->collection_similarity != nullptr); + DORIS_CHECK(uses_phrase_frequency_scoring(query_type, query_info)); + DORIS_CHECK_EQ(final_candidates.cardinality(), matches.size()); + + const std::wstring field_name = StringUtil::string_to_wstring(std::string(column_name)); + const double collection_avgdl = + context->collection_statistics->get_or_calculate_avg_dl(field_name); + const size_t idf_term_count = query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY + ? query_info.term_infos.size() - 1 + : query_info.term_infos.size(); + double idf_sum = 0.0; + for (size_t i = 0; i < idf_term_count; ++i) { + const auto& term_info = query_info.term_infos[i]; + DORIS_CHECK(term_info.is_single_term()); + idf_sum += context->collection_statistics->get_or_calculate_idf( + field_name, StringUtil::string_to_wstring(term_info.get_single_term())); + } + + const auto scorer = ::doris::snii::query::ScorerContext::from_idf(idf_sum); + std::vector<::doris::snii::query::ScoredDoc> scored_docs; + scored_docs.reserve(matches.size()); + for (const auto& match : matches) { + DCHECK(final_candidates.contains(match.docid)); + DCHECK_NE(match.frequency, 0); + uint8_t norm = 0; + RETURN_IF_ERROR(segment_stats.encoded_norm(match.docid, &norm)); + scored_docs.push_back({.docid = match.docid, + .score = scorer.score(match.frequency, norm, collection_avgdl, + ::doris::snii::query::Bm25Params {})}); + } + for (const auto& scored_doc : scored_docs) { + context->collection_similarity->collect(scored_doc.docid, + static_cast(scored_doc.score)); + } + return Status::OK(); +} + +void parse_phrase_slop(std::string* query, InvertedIndexQueryInfo* query_info) { + DCHECK(query != nullptr); + DCHECK(query_info != nullptr); + const auto is_digits = [](std::string_view str) { + return std::all_of(str.begin(), str.end(), [](unsigned char c) { return std::isdigit(c); }); + }; + + const size_t last_space_pos = query->find_last_of(' '); + if (last_space_pos == std::string::npos) { + return; + } + const size_t tilde_pos = last_space_pos + 1; + if (tilde_pos >= query->size() - 1 || (*query)[tilde_pos] != '~') { + return; + } + + const size_t slop_pos = tilde_pos + 1; + std::string_view slop_str(query->data() + slop_pos, query->size() - slop_pos); + if (slop_str.empty()) { + return; + } + + bool ordered = false; + if (slop_str.size() == 1) { + if (!std::isdigit(static_cast(slop_str[0]))) { + return; + } + } else if (slop_str.back() == '+') { + ordered = true; + slop_str.remove_suffix(1); + } + + if (!is_digits(slop_str)) { + return; + } + auto result = std::from_chars(slop_str.begin(), slop_str.end(), query_info->slop); + if (result.ec != std::errc()) { + return; + } + query_info->ordered = ordered; + *query = query->substr(0, last_space_pos); +} + +std::shared_ptr docids_to_bitmap(const std::vector& docids) { + auto result = std::make_shared(); + if (!docids.empty()) { + result->addMany(docids.size(), docids.data()); + } + result->runOptimize(); + return result; +} + +// Runs `compute` under single-flight keyed by `key`: concurrent identical queries collapse to a +// single execution and the followers reuse the leader's bitmap. `compute(out)` fills *out and +// returns its Status; on overall success *result receives the bitmap. See SingleFlight for why +// this matters under a cold cache with parallel scanners hitting the same segment. +template +Status run_query_single_flight( + ::doris::segment_v2::inverted_index::SingleFlight< + std::pair>>& flight, + const std::string& key, std::shared_ptr* result, +#ifdef BE_TEST + SniiIndexReader::SingleFlightFollowerJoinedObserver follower_joined_observer, + void* follower_joined_opaque, + SniiIndexReader::SingleFlightLeaderBeforeComputeObserver leader_before_compute_observer, + void* leader_before_compute_opaque, +#endif + Compute&& compute) { + auto follower = flight.join_or_lead(key); + if (follower.has_value()) { +#ifdef BE_TEST + if (follower_joined_observer != nullptr) { + follower_joined_observer(follower_joined_opaque); + } +#endif + auto [leader_status, leader_bitmap] = follower->get(); + if (leader_status.ok() && leader_bitmap != nullptr) { + *result = std::move(leader_bitmap); + return Status::OK(); + } + // Leader failed; fall through and compute independently (rare error path). + } + const bool is_leader = !follower.has_value(); +#ifdef BE_TEST + if (is_leader && leader_before_compute_observer != nullptr) { + leader_before_compute_observer(leader_before_compute_opaque); + } +#endif + + Status status = Status::OK(); + std::shared_ptr bitmap; + { + // Publish to any waiting followers on every exit path (including errors). + DEFER(if (is_leader) { flight.publish(key, std::make_pair(status, bitmap)); }); + status = compute(&bitmap); + } + RETURN_IF_ERROR(status); + *result = std::move(bitmap); + return Status::OK(); +} + +Status execute_snii_query(const ::doris::snii::reader::LogicalIndexReader& logical_reader, + InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, std::string_view search_str, + const std::vector& terms, int32_t max_expansions, + bool collect_phrase_frequency, SniiQueryExecutionResult* result, + ::doris::snii::query::QueryProfile* profile) { + result->bitmap = std::make_shared(); + result->phrase_matches.clear(); + DORIS_CHECK(!collect_phrase_frequency || uses_phrase_frequency_scoring(query_type, query_info)); + RoaringDocIdSink sink(result->bitmap.get()); + std::vector docids; + bool emitted_to_sink = false; + Status status; + switch (query_type) { + case InvertedIndexQueryType::EQUAL_QUERY: + case InvertedIndexQueryType::MATCH_ANY_QUERY: + status = terms.size() == 1 + ? ::doris::snii::query::term_query(logical_reader, terms.front(), &sink) + : ::doris::snii::query::boolean_or(logical_reader, terms, &sink); + emitted_to_sink = true; + break; + case InvertedIndexQueryType::MATCH_ALL_QUERY: + if (terms.size() == 1) { + status = ::doris::snii::query::term_query(logical_reader, terms.front(), &sink); + emitted_to_sink = true; + } else { + status = ::doris::snii::query::boolean_and(logical_reader, terms, &docids); + } + break; + case InvertedIndexQueryType::MATCH_PHRASE_QUERY: + if (query_info.slop != 0) { + return Status::Error( + "SNII does not support sloppy phrase query yet"); + } + if (terms.size() == 1) { + status = ::doris::snii::query::term_query(logical_reader, terms.front(), &sink); + emitted_to_sink = true; + } else { + status = collect_phrase_frequency + ? ::doris::snii::query::phrase_query_with_frequencies( + logical_reader, terms, &result->phrase_matches, profile) + : ::doris::snii::query::phrase_query(logical_reader, terms, &docids, + profile); + } + break; + case InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY: + if (terms.size() == 1) { + status = ::doris::snii::query::prefix_query(logical_reader, terms.front(), &sink, + max_expansions); + emitted_to_sink = true; + } else { + status = collect_phrase_frequency + ? ::doris::snii::query::phrase_prefix_query_with_frequencies( + logical_reader, terms, &result->phrase_matches, profile, + max_expansions) + : ::doris::snii::query::phrase_prefix_query( + logical_reader, terms, &docids, profile, max_expansions); + } + break; + case InvertedIndexQueryType::MATCH_REGEXP_QUERY: + status = ::doris::snii::query::regexp_query(logical_reader, search_str, &sink, + max_expansions); + emitted_to_sink = true; + break; + case InvertedIndexQueryType::WILDCARD_QUERY: + status = ::doris::snii::query::wildcard_query(logical_reader, search_str, &sink, + max_expansions); + emitted_to_sink = true; + break; + case InvertedIndexQueryType::LESS_THAN_QUERY: + case InvertedIndexQueryType::LESS_EQUAL_QUERY: + case InvertedIndexQueryType::GREATER_THAN_QUERY: + case InvertedIndexQueryType::GREATER_EQUAL_QUERY: + case InvertedIndexQueryType::RANGE_QUERY: + return Status::Error( + "SNII inverted index storage format does not support BKD/range query"); + default: + return Status::Error( + "SNII unsupported inverted index query type {}", query_type_to_string(query_type)); + } + RETURN_IF_ERROR(status); + if (collect_phrase_frequency) { + for (const auto& match : result->phrase_matches) { + result->bitmap->add(match.docid); + } + result->bitmap->runOptimize(); + } else if (emitted_to_sink) { + result->bitmap->runOptimize(); + } else { + result->bitmap = docids_to_bitmap(docids); + } + return Status::OK(); +} + +} // namespace + +Status SniiIndexReader::new_iterator(std::unique_ptr* iterator) { + if (*iterator == nullptr) { + *iterator = InvertedIndexIterator::create_unique(); + } + dynamic_cast(iterator->get()) + ->add_reader(_reader_type, + dynamic_pointer_cast(shared_from_this())); + return Status::OK(); +} + +Status SniiIndexReader::_parse_query_terms( + const IndexQueryContextPtr& context, std::string search_str, + InvertedIndexQueryType query_type, const InvertedIndexAnalyzerCtx* analyzer_ctx, + InvertedIndexQueryInfo* query_info, + std::optional purpose_override) { + DCHECK(query_info != nullptr); + if (query_type == InvertedIndexQueryType::MATCH_REGEXP_QUERY || + query_type == InvertedIndexQueryType::WILDCARD_QUERY) { + query_info->term_infos.emplace_back(search_str, 0); + return Status::OK(); + } + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY) { + parse_phrase_slop(&search_str, query_info); + } + + const bool actual_similarity = + context->collection_similarity && + IndexReaderHelper::is_need_similarity_score(query_type, &_index_meta); + const auto purpose = purpose_override.value_or(inverted_index::select_analysis_purpose( + query_type, query_info->slop, actual_similarity)); + SCOPED_RAW_TIMER(&context->stats->inverted_index_analyzer_timer); + try { + if (analyzer_ctx != nullptr && !analyzer_ctx->should_tokenize()) { + query_info->term_infos.emplace_back(search_str); + } else { + auto analyzer = analyzer_ctx == nullptr ? nullptr : analyzer_ctx->get_analyzer(purpose); + if (analyzer != nullptr) { + auto reader = inverted_index::InvertedIndexAnalyzer::create_reader( + analyzer_ctx->char_filter_map); + reader->init(search_str.data(), static_cast(search_str.size()), true); + query_info->term_infos = inverted_index::InvertedIndexAnalyzer::get_analyse_result( + reader, analyzer.get()); + } else { + query_info->term_infos = inverted_index::InvertedIndexAnalyzer::get_analyse_result( + search_str, _index_meta.properties(), purpose); + } + } + } catch (const CLuceneError& e) { + return Status::Error( + "SNII analyze query failed: {}", e.what()); + } catch (const Exception& e) { + return Status::Error( + "SNII analyze query failed: {}", e.what()); + } + return Status::OK(); +} + +Status SniiIndexReader::_get_logical_reader( + const IndexQueryContextPtr& context, InvertedIndexCacheHandle* searcher_cache_handle, + std::unique_ptr<::doris::snii::reader::LogicalIndexReader>* uncached_reader, + const ::doris::snii::reader::LogicalIndexReader** logical_reader) { + DCHECK(searcher_cache_handle != nullptr); + DCHECK(uncached_reader != nullptr); + DCHECK(logical_reader != nullptr); + + const bool enable_searcher_cache = + context->runtime_state != nullptr && + context->runtime_state->query_options().enable_inverted_index_searcher_cache; + const auto index_file_key = _index_file_reader->get_index_file_cache_key(&_index_meta); + InvertedIndexSearcherCache::CacheKey searcher_cache_key(index_file_key); + + bool cache_hit = false; + if (enable_searcher_cache) { + SCOPED_RAW_TIMER(&context->stats->inverted_index_lookup_timer); + cache_hit = InvertedIndexSearcherCache::instance()->lookup(searcher_cache_key, + searcher_cache_handle); + } + + if (cache_hit) { + context->stats->inverted_index_searcher_cache_hit++; + *logical_reader = searcher_cache_handle->get_snii_logical_reader(); + if (*logical_reader == nullptr) { + return Status::InternalError("SNII searcher cache entry has no logical reader"); + } + return Status::OK(); + } + + SCOPED_RAW_TIMER(&context->stats->inverted_index_searcher_open_timer); + context->stats->inverted_index_searcher_cache_miss++; +#ifdef BE_TEST + if (_searcher_open_observer != nullptr) { + _searcher_open_observer(_searcher_open_opaque); + } +#endif + RETURN_IF_ERROR( + _index_file_reader->init(config::inverted_index_read_buffer_size, context->io_ctx)); + auto opened_reader = + DORIS_TRY(_index_file_reader->open_snii_index(&_index_meta, context->io_ctx)); + + if (!enable_searcher_cache) { + *logical_reader = opened_reader.get(); + *uncached_reader = std::move(opened_reader); + return Status::OK(); + } + + const size_t reader_size = std::max(opened_reader->memory_usage(), 1); + auto* cache_value = new InvertedIndexSearcherCache::CacheValue( + std::move(opened_reader), reader_size, UnixMillis(), _index_file_reader); + InvertedIndexSearcherCache::instance()->insert(searcher_cache_key, cache_value, + searcher_cache_handle); + *logical_reader = searcher_cache_handle->get_snii_logical_reader(); + if (*logical_reader == nullptr) { + return Status::InternalError("SNII searcher cache insert produced empty logical reader"); + } + return Status::OK(); +} + +Status SniiIndexReader::query(const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + const InvertedIndexAnalyzerCtx* analyzer_ctx) { + return _query(context, column_name, query_value, query_type, bit_map, nullptr, analyzer_ctx); +} + +Status SniiIndexReader::query_with_null_bitmap( + const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx) { + DORIS_CHECK(null_bitmap_cache_handle != nullptr); + return _query(context, column_name, query_value, query_type, bit_map, null_bitmap_cache_handle, + analyzer_ctx); +} + +Status SniiIndexReader::_query(const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx) { + const bool track_requested_null_time = null_bitmap_cache_handle != nullptr; + const int64_t query_ns_before = + track_requested_null_time ? context->stats->inverted_index_query_timer : 0; + int64_t requested_null_ns = 0; + DEFER({ + if (!track_requested_null_time) { + return; + } + const int64_t inclusive_query_ns = + context->stats->inverted_index_query_timer - query_ns_before; + DORIS_CHECK_GE(inclusive_query_ns, 0); + const int64_t exclusive_query_ns = + inclusive_query_ns > requested_null_ns ? inclusive_query_ns - requested_null_ns : 0; + context->stats->inverted_index_query_timer = query_ns_before + exclusive_query_ns; + }); + SCOPED_RAW_TIMER(&context->stats->inverted_index_query_timer); + const std::string search_str = query_value.get(); + const auto finish_query = + [&](const ::doris::snii::reader::LogicalIndexReader* reader) -> Status { + if (null_bitmap_cache_handle == nullptr) { + return Status::OK(); + } + const int64_t null_ns_before = context->stats->inverted_index_query_null_bitmap_timer; + Status status = _read_null_bitmap(context, null_bitmap_cache_handle, reader); + const int64_t null_ns_after = context->stats->inverted_index_query_null_bitmap_timer; + DORIS_CHECK_GE(null_ns_after, null_ns_before); + requested_null_ns += null_ns_after - null_ns_before; + return status; + }; + + if (int ignore_above = + std::stoi(get_parser_ignore_above_value_from_properties(_index_meta.properties())); + _reader_type == InvertedIndexReaderType::STRING_TYPE && search_str.size() > ignore_above) { + return Status::Error( + "query value is too long, evaluate skipped."); + } + + const bool actual_similarity = + context->collection_similarity && + IndexReaderHelper::is_need_similarity_score(query_type, &_index_meta); + const int32_t max_expansions = + context->runtime_state == nullptr + ? 50 + : context->runtime_state->query_options().inverted_index_max_expansions; + InvertedIndexQueryInfo query_info; + std::string plain_analysis_str = search_str; + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY) { + parse_phrase_slop(&plain_analysis_str, &query_info); + } + const bool common_grams_phrase_shape = + (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && query_info.slop == 0) || + query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY; + const bool common_grams_query_eligible = common_grams_phrase_shape && !actual_similarity; + const bool raw_pattern_query = query_type == InvertedIndexQueryType::MATCH_REGEXP_QUERY || + query_type == InvertedIndexQueryType::WILDCARD_QUERY; + // Lucene-style CommonGrams: the plan decision is purely local -- the segment's own analyzer + // identity says whether gram terms can exist, and the mutable BE config + // enable_common_grams_query_plan is the only runtime switch (its cache generation bumps on + // every flip, keeping result-cache entries from crossing plan modes). + const auto common_grams_safety = config::common_grams_query_plan_config_snapshot(); + const inverted_index::CommonGramsPlanCostModel common_grams_cost_model { + .position_verify_factor = common_grams_safety.position_verify_factor, + .common_grams_cost_ratio_percent = common_grams_safety.plan_cost_ratio_percent, + .generation = common_grams_safety.cost_model_generation}; + const auto has_common_grams_analyzer = [](const InvertedIndexAnalyzerCtx* ctx) { + return ctx != nullptr && ctx->analyzer_provider != nullptr && + ctx->analyzer_provider->uses_common_grams() && + ctx->has_complete_common_grams_identity(); + }; + const bool safety_requires_plain = !common_grams_safety.enabled; + // The raw cache key cannot prove whether the immutable segment analyzer has CommonGrams until + // its metadata is open. Delay every eligible forced-plain lookup, then restore ordinary cache + // access below only for a segment that cannot contain gram terms. + const bool initial_force_plain = common_grams_query_eligible && safety_requires_plain; + const bool initial_allow_result_cache = !actual_similarity && !initial_force_plain; + const bool defer_result_cache_lookup = !actual_similarity && !initial_allow_result_cache; + const InvertedIndexRawQuerySemantic raw_semantic { + .raw_query_bytes = search_str, + .query_type = query_type, + .slop = query_info.slop, + .ordered = query_info.ordered, + .max_expansions = max_expansions, + .common_grams_cache_generation = common_grams_safety.cache_generation}; + const auto index_file_key = _index_file_reader->get_index_file_cache_key(&_index_meta); + InvertedIndexQueryCache::CacheKey cache_key {index_file_key, column_name, query_type, + raw_semantic.encode()}; + std::string single_flight_key = cache_key.encode(); + auto* cache = InvertedIndexQueryCache::instance(); + InvertedIndexQueryCacheHandle cache_handler; + bool allow_result_cache = initial_allow_result_cache; + if (handle_query_cache(context, cache, cache_key, &cache_handler, bit_map, + allow_result_cache)) { + return finish_query(nullptr); + } + + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && query_info.slop != 0) { + return Status::Error( + "SNII does not support sloppy phrase query yet"); + } + + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(context->io_ctx); + InvertedIndexCacheHandle searcher_cache_handle; + std::unique_ptr<::doris::snii::reader::LogicalIndexReader> uncached_reader; + const ::doris::snii::reader::LogicalIndexReader* logical_reader = nullptr; + RETURN_IF_ERROR(_get_logical_reader(context, &searcher_cache_handle, &uncached_reader, + &logical_reader)); + + std::optional rebuilt_analyzer_context; + const auto* common_grams_metadata = logical_reader->common_grams_metadata(); + if (!raw_pattern_query) { + auto rebuilt_result = inverted_index::maybe_rebuild_segment_analyzer_context( + analyzer_ctx, common_grams_metadata, _index_meta.properties(), + ExecEnv::GetInstance()->index_policy_mgr()); + if (!rebuilt_result.has_value()) { + if (common_grams_query_eligible) { + ++context->stats->snii_stats.common_grams_fallback_base_analyzer_mismatch; + } + return std::move(rebuilt_result.error()); + } + rebuilt_analyzer_context = std::move(*rebuilt_result); + } + const InvertedIndexAnalyzerCtx* effective_analyzer_context = + rebuilt_analyzer_context ? &*rebuilt_analyzer_context : analyzer_ctx; + const bool effective_common_grams_configured = + common_grams_query_eligible && has_common_grams_analyzer(effective_analyzer_context); + const bool segment_may_contain_common_grams = + common_grams_metadata != nullptr && common_grams_metadata->common_grams_coverage != + inverted_index::CommonGramsCoverage::kNone; + const bool force_plain = + common_grams_query_eligible && safety_requires_plain && + (effective_common_grams_configured || segment_may_contain_common_grams); + allow_result_cache = !actual_similarity && !force_plain; + if (defer_result_cache_lookup && allow_result_cache && + handle_query_cache(context, cache, cache_key, &cache_handler, bit_map, + allow_result_cache)) { + return finish_query(logical_reader); + } + InvertedIndexQueryInfo execution_query_info = query_info; + const auto plain_purpose = common_grams_query_eligible + ? std::optional(inverted_index::AnalysisPurpose::kPlainQuery) + : std::nullopt; + RETURN_IF_ERROR(_parse_query_terms(context, plain_analysis_str, query_type, + effective_analyzer_context, &execution_query_info, + plain_purpose)); + if (execution_query_info.term_infos.empty()) { + auto msg = fmt::format("token parser result is empty for SNII query '{}'", search_str); + if (is_match_query(query_type)) { + LOG(WARNING) << msg; + bit_map = std::make_shared(); + insert_query_cache(context, cache, cache_key, bit_map, &cache_handler, + allow_result_cache); + return finish_query(logical_reader); + } + return Status::Error(msg); + } + if (execution_query_info.has_common_gram()) { + return Status::Error( + "CommonGrams term escaped the plain query analyzer"); + } + if (actual_similarity && query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY && + execution_query_info.term_infos.size() == 1) { + return Status::Error( + "SNII scoring does not support a single-token phrase-prefix query"); + } + std::vector terms = to_terms(execution_query_info); + + // G02 count-only fast path: the SegmentIterator asserted (via the context + // flag) that only the match COUNT of this predicate matters, so eligible + // shapes are answered from dict-entry df without decoding postings. Placed + // AFTER the query-cache lookup (a cached row-accurate bitmap is free and + // counts correctly) and BEFORE single-flight; the fabricated [0, df) bitmap + // is returned early and NEVER inserted into the query cache or published to + // single-flight followers -- both are keyed identically to row-accurate + // queries and must only ever serve real row ids. + if (context->count_on_index_fastpath) { + bool count_handled = false; + std::shared_ptr count_bitmap; + RETURN_IF_ERROR(_try_count_only_fastpath(context, query_type, execution_query_info, terms, + &count_handled, &count_bitmap, logical_reader)); + if (count_handled) { + bit_map = std::move(count_bitmap); + RETURN_IF_ERROR(finish_query(logical_reader)); + // G03 reply: tell the SegmentIterator the bitmap is count-shaped + // (cardinality exact, row ids fabricated) so it may short-circuit + // row emission. Deliberately NOT set on the cache-hit return above + // or on the decode path below -- those bitmaps are row-accurate + // and keep today's emission. + context->count_on_index_fastpath_hit = true; + return Status::OK(); + } + } + + // Under a cold cache, parallel scanners _lazy_init the same segment concurrently and each + // would otherwise miss the searcher/query caches and redundantly open + decode this segment's + // index. Collapse identical concurrent queries into one shared execution (see SingleFlight). + static ::doris::segment_v2::inverted_index::SingleFlight< + std::pair>> + query_single_flight; + std::shared_ptr result_bitmap; + std::vector<::doris::snii::query::PhraseMatch> phrase_matches; + auto* phrase_matches_out = + actual_similarity && uses_phrase_frequency_scoring(query_type, execution_query_info) + ? &phrase_matches + : nullptr; + Status single_flight_status; + if (!allow_result_cache) { + single_flight_status = + _compute_query_bitmap(context, + {.query_type = query_type, + .query_info = execution_query_info, + .search_str = search_str, + .max_expansions = max_expansions, + .common_grams_query_shape = common_grams_query_eligible, + .force_plain = force_plain, + .common_grams_cost_model = common_grams_cost_model, + .analyzer_ctx = effective_analyzer_context, + .physical_raw_query_key = single_flight_key, + .logical_reader = logical_reader}, + &terms, &result_bitmap, phrase_matches_out); + } else { + DORIS_CHECK(phrase_matches_out == nullptr); + single_flight_status = run_query_single_flight( + query_single_flight, single_flight_key, &result_bitmap, +#ifdef BE_TEST + _single_flight_follower_joined_observer, _single_flight_follower_joined_opaque, + _single_flight_leader_before_compute_observer, + _single_flight_leader_before_compute_opaque, +#endif + [&](std::shared_ptr* out) { + auto status = _compute_query_bitmap( + context, + {.query_type = query_type, + .query_info = execution_query_info, + .search_str = search_str, + .max_expansions = max_expansions, + .common_grams_query_shape = common_grams_query_eligible, + .force_plain = force_plain, + .common_grams_cost_model = common_grams_cost_model, + .analyzer_ctx = effective_analyzer_context, + .physical_raw_query_key = single_flight_key, + .logical_reader = logical_reader}, + &terms, out, nullptr); + if (status.ok()) { + insert_query_cache(context, cache, cache_key, *out, &cache_handler, + allow_result_cache); + } + return status; + }); + } + RETURN_IF_ERROR(single_flight_status); + DORIS_CHECK(result_bitmap != nullptr); + if (actual_similarity && !result_bitmap->isEmpty()) { + ::doris::snii::stats::SniiStatsProvider segment_stats; + RETURN_IF_ERROR( + ::doris::snii::stats::SniiStatsProvider::open(logical_reader, &segment_stats)); + if (phrase_matches_out != nullptr) { + RETURN_IF_ERROR(score_phrase_matches(context, column_name, query_type, + execution_query_info, *logical_reader, + segment_stats, *result_bitmap, phrase_matches)); + } else if (uses_plain_term_frequency_scoring(query_type, execution_query_info)) { + RETURN_IF_ERROR(score_plain_term_candidates(context, column_name, execution_query_info, + *logical_reader, segment_stats, + *result_bitmap)); + } + } + bit_map = result_bitmap; + return finish_query(logical_reader); +} + +Status SniiIndexReader::_compute_query_bitmap( + const IndexQueryContextPtr& context, const SniiQueryBitmapRequest& request, + std::vector* preanalyzed_terms, std::shared_ptr* out, + std::vector<::doris::snii::query::PhraseMatch>* phrase_matches) { + // Bound once so the body below reads the same as before the request object was introduced; + // renaming 71 uses would have buried the actual change. + const InvertedIndexQueryType query_type = request.query_type; + const InvertedIndexQueryInfo& request_query_info = request.query_info; + const bool common_grams_query_shape = request.common_grams_query_shape; + const bool force_plain = request.force_plain; + const inverted_index::CommonGramsPlanCostModel common_grams_cost_model = + request.common_grams_cost_model; + const InvertedIndexAnalyzerCtx* analyzer_ctx = request.analyzer_ctx; + const std::string_view physical_raw_query_key = request.physical_raw_query_key; + const std::string_view search_str = request.search_str; + const int32_t max_expansions = request.max_expansions; + const ::doris::snii::reader::LogicalIndexReader* logical_reader = request.logical_reader; + + DORIS_CHECK(preanalyzed_terms != nullptr); + DORIS_CHECK(logical_reader != nullptr); + DORIS_CHECK(request_query_info.term_infos.size() == preanalyzed_terms->size()); + if (phrase_matches != nullptr) { + phrase_matches->clear(); + } + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && request_query_info.slop != 0) { + return Status::Error( + "SNII does not support sloppy phrase query yet"); + } + const auto* common_grams_metadata = logical_reader->common_grams_metadata(); + InvertedIndexQueryInfo query_info = request_query_info; + std::vector routed_terms = *preanalyzed_terms; + auto* terms = &routed_terms; + + const auto* common_grams_identity = + analyzer_ctx == nullptr ? nullptr : analyzer_ctx->get_common_grams_identity(); + const bool common_grams_configured = common_grams_query_shape && analyzer_ctx != nullptr && + analyzer_ctx->analyzer_provider != nullptr && + analyzer_ctx->analyzer_provider->uses_common_grams() && + analyzer_ctx->has_complete_common_grams_identity(); + const bool common_grams_candidate = + common_grams_configured && query_info.term_infos.size() >= 2 && !force_plain; + const bool common_grams_forced_plain = + common_grams_configured && query_info.term_infos.size() >= 2 && force_plain; + enum class CommonGramsPlainFallback : uint8_t { kNone, kNoGram, kIncompatible, kKillSwitch }; + CommonGramsPlainFallback common_grams_plain_fallback = + common_grams_forced_plain ? CommonGramsPlainFallback::kKillSwitch + : CommonGramsPlainFallback::kNone; + const bool common_grams_compatible = + common_grams_metadata != nullptr && common_grams_identity != nullptr && + (logical_reader->common_grams_posting_policy() == + ::doris::snii::format::CommonGramsPostingPolicy::kHybridV1 + ? inverted_index::is_common_grams_query_compatible( + *common_grams_metadata, *common_grams_identity, + inverted_index::CommonGramsCoverage::kMixed) + : inverted_index::is_common_grams_query_compatible(*common_grams_metadata, + *common_grams_identity)); + const auto* common_grams_word_set = + common_grams_configured ? analyzer_ctx->analyzer_provider->common_grams_word_set() + : nullptr; + const auto common_grams_query_mode = + query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY + ? inverted_index::CommonGramsQueryMode::kExact + : inverted_index::CommonGramsQueryMode::kPhrasePrefix; + const bool proven_no_common_gram = + common_grams_candidate && common_grams_compatible && common_grams_word_set != nullptr && + !inverted_index::common_grams_query_may_use_gram( + *preanalyzed_terms, common_grams_query_mode, *common_grams_word_set); + if (proven_no_common_gram && query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY) { + DORIS_CHECK(phrase_matches == nullptr); + ::doris::snii::SniiPrxExecutionProfileScope execution_profile(*context->stats); + InvertedIndexQueryInfo empty_gram_query_info; + std::vector docids; + RETURN_IF_ERROR(::doris::snii::query::planned_exact_phrase_query( + *logical_reader, query_info, empty_gram_query_info, common_grams_identity, &docids, + execution_profile.profile(), nullptr, common_grams_cost_model, + ::doris::snii::query::CommonGramsPlanDebugOverride::kNone)); + *out = docids_to_bitmap(docids); + return Status::OK(); + } + if (proven_no_common_gram) { + common_grams_plain_fallback = CommonGramsPlainFallback::kNoGram; + } else if (common_grams_candidate && common_grams_compatible) { + DORIS_CHECK(phrase_matches == nullptr); + DORIS_CHECK(!physical_raw_query_key.empty()); + const auto debug_override = ::doris::snii::query::common_grams_plan_debug_override(); + ::doris::snii::SniiPrxExecutionProfileScope execution_profile(*context->stats); + + // The gram-side analysis always runs. Without a memoized plan choice nothing can tell us + // in advance that the plain plan wins, and one analyzer pass over a phrase string is noise + // next to the posting decode that follows it. + InvertedIndexQueryInfo gram_query_info; + const auto gram_purpose = query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY + ? inverted_index::AnalysisPurpose::kExactPhraseQuery + : inverted_index::AnalysisPurpose::kPhrasePrefixQuery; + RETURN_IF_ERROR(_parse_query_terms(context, std::string(search_str), query_type, + analyzer_ctx, &gram_query_info, gram_purpose)); + + std::vector docids; + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY) { + RETURN_IF_ERROR(::doris::snii::query::planned_exact_phrase_query( + *logical_reader, query_info, gram_query_info, common_grams_identity, &docids, + execution_profile.profile(), nullptr, common_grams_cost_model, debug_override)); + } else { + DORIS_CHECK(query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY); + RETURN_IF_ERROR(::doris::snii::query::planned_phrase_prefix_query( + *logical_reader, query_info, gram_query_info, common_grams_identity, &docids, + execution_profile.profile(), max_expansions, nullptr, common_grams_cost_model, + debug_override)); + } + *out = docids_to_bitmap(docids); + return Status::OK(); + } else if (common_grams_candidate) { + common_grams_plain_fallback = CommonGramsPlainFallback::kIncompatible; + } + switch (query_type) { + case InvertedIndexQueryType::EQUAL_QUERY: + case InvertedIndexQueryType::MATCH_ANY_QUERY: + case InvertedIndexQueryType::MATCH_ALL_QUERY: + case InvertedIndexQueryType::MATCH_PHRASE_QUERY: { + bool all_representable = false; + RETURN_IF_ERROR(::doris::snii::query::internal::route_query_terms( + *logical_reader, query_info, terms, &all_representable)); + if (terms->empty() && (query_type == InvertedIndexQueryType::EQUAL_QUERY || + query_type == InvertedIndexQueryType::MATCH_ANY_QUERY)) { + *out = std::make_shared(); + return Status::OK(); + } + if (!all_representable && (query_type == InvertedIndexQueryType::MATCH_ALL_QUERY || + query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY)) { + *out = std::make_shared(); + return Status::OK(); + } + break; + } + default: + break; + } + SniiQueryExecutionResult query_result; + const bool phrase_can_decode_prx = + query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && query_info.slop == 0; + const bool needs_prx_profile = + terms->size() > 1 && (phrase_can_decode_prx || + query_type == InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY); + if (needs_prx_profile) { + ::doris::snii::SniiPrxExecutionProfileScope execution_profile(*context->stats); + const Status execution_status = execute_snii_query( + *logical_reader, query_type, query_info, search_str, *terms, max_expansions, + phrase_matches != nullptr, &query_result, execution_profile.profile()); + if (common_grams_plain_fallback != CommonGramsPlainFallback::kNone) { + auto& plan_stats = execution_profile.profile()->phrase_query_stats; + ++plan_stats.common_grams_candidate_queries; + ++plan_stats.common_grams_plain_plans; + switch (common_grams_plain_fallback) { + case CommonGramsPlainFallback::kNoGram: + ++plan_stats.common_grams_fallback_no_gram; + break; + case CommonGramsPlainFallback::kIncompatible: + ++plan_stats.common_grams_fallback_incompatible; + break; + case CommonGramsPlainFallback::kKillSwitch: + ++plan_stats.common_grams_fallback_kill_switch; + break; + case CommonGramsPlainFallback::kNone: + break; + } + } + RETURN_IF_ERROR(execution_status); + } else { + RETURN_IF_ERROR(execute_snii_query(*logical_reader, query_type, query_info, search_str, + *terms, max_expansions, phrase_matches != nullptr, + &query_result, nullptr)); + } + *out = std::move(query_result.bitmap); + if (phrase_matches != nullptr) { + *phrase_matches = std::move(query_result.phrase_matches); + } + return Status::OK(); +} + +#ifdef BE_TEST +Status SniiIndexReader::_compute_query_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, + std::string_view search_str, + std::vector* terms, + int32_t max_expansions, + std::shared_ptr* out) { + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(context->io_ctx); + InvertedIndexCacheHandle searcher_cache_handle; + std::unique_ptr<::doris::snii::reader::LogicalIndexReader> uncached_reader; + const ::doris::snii::reader::LogicalIndexReader* logical_reader = nullptr; + RETURN_IF_ERROR(_get_logical_reader(context, &searcher_cache_handle, &uncached_reader, + &logical_reader)); + return _compute_query_bitmap(context, + {.query_type = query_type, + .query_info = query_info, + .search_str = search_str, + .max_expansions = max_expansions, + .logical_reader = logical_reader}, + terms, out, nullptr); +} +#endif + +Status SniiIndexReader::_try_count_only_fastpath( + const IndexQueryContextPtr& context, InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, const std::vector& terms, + bool* handled, std::shared_ptr* out, + const ::doris::snii::reader::LogicalIndexReader* preopened_reader) { + *handled = false; + // Shape guard: only exact-term query types. Prefix/regexp/wildcard/ + // phrase-prefix expand the term set, so no single dict entry carries the + // count; range types never reach SNII anyway. + switch (query_type) { + case InvertedIndexQueryType::EQUAL_QUERY: + case InvertedIndexQueryType::MATCH_ANY_QUERY: + case InvertedIndexQueryType::MATCH_ALL_QUERY: + case InvertedIndexQueryType::MATCH_PHRASE_QUERY: + break; + default: + return Status::OK(); + } + // The normal path rejects sloppy phrases with a BYPASS (downgrade to raw + // evaluation); the count path must fall through so the downgrade happens. + if (query_type == InvertedIndexQueryType::MATCH_PHRASE_QUERY && query_info.slop != 0) { + return Status::OK(); + } + if (terms.size() != 1) { + // Multi-term MATCH_ANY (OR) / MATCH_ALL (AND) counts are not derivable + // from per-term dfs (overlap unknown), and phrases require positional + // verification, so execute the normal query path. + return Status::OK(); + } + + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(context->io_ctx); + InvertedIndexCacheHandle searcher_cache_handle; + std::unique_ptr<::doris::snii::reader::LogicalIndexReader> uncached_reader; + const ::doris::snii::reader::LogicalIndexReader* logical_reader = preopened_reader; + if (logical_reader == nullptr) { + RETURN_IF_ERROR(_get_logical_reader(context, &searcher_cache_handle, &uncached_reader, + &logical_reader)); + } + + std::string physical_term_scratch; + std::string_view physical_term; + bool representable = false; + DORIS_CHECK(query_info.term_infos.size() == 1); + RETURN_IF_ERROR(::doris::snii::query::internal::route_query_term_view( + *logical_reader, query_info.term_infos.front(), &physical_term_scratch, &physical_term, + &representable)); + uint64_t count = 0; + if (representable) { + RETURN_IF_ERROR( + ::doris::snii::query::count_only_term_df(*logical_reader, physical_term, &count)); + } + + // Null handling. df is the exact match count REGARDLESS of nulls: the + // writer adds no tokens for a null doc (scalar add_nulls; a NULL array row + // is an empty range), so postings -- and therefore df -- never include + // null rows, exactly matching MATCH's "null never matches" semantics. The + // fabricated bitmap however flows through FunctionMatchBase -> + // InvertedIndexResultBitmap::mask_out_null, which subtracts the segment's + // REAL null bitmap from it; a dense [0, df) range colliding with null row + // ids would be shrunk below df. So on a segment WITH a null bitmap, load + // it (query-cache backed, the same read the normal MATCH path performs) + // and fabricate df ids DISJOINT from it, making that subtraction a + // provable no-op. Segments without a null section (the writer omits it + // when no row is null) keep the trivial [0, df) range. + auto result = std::make_shared(); + if (count > 0 && logical_reader->section_refs().null_bitmap.length > 0) { + InvertedIndexQueryCacheHandle null_bitmap_cache_handle; + RETURN_IF_ERROR(_read_null_bitmap(context, &null_bitmap_cache_handle, logical_reader)); + std::shared_ptr nulls = null_bitmap_cache_handle.get_bitmap(); + // Fall through on a missing bitmap behind the cache handle or a + // fabrication failure (df + null count breaching the docid domain): + // both mean a corrupt index, and the row-accurate decode -- which + // intersects real ids -- must own the answer rather than a blind + // fabrication. The count_fastpath_hits test seam already counted the + // dict lookup above; production correctness is unaffected. + if (nulls == nullptr) { + return Status::OK(); + } + if (!nulls->isEmpty()) { + if (!::doris::snii::query::fabricate_null_disjoint_count_bitmap(count, *nulls, + result.get()) + .ok()) { + return Status::OK(); + } + } else { + result->addRange(0, count); + } + } else if (count > 0) { + result->addRange(0, count); + } + *out = std::move(result); + *handled = true; + return Status::OK(); +} + +Status SniiIndexReader::try_query(const IndexQueryContextPtr& /*context*/, + const std::string& /*column_name*/, const Field& /*query_value*/, + InvertedIndexQueryType /*query_type*/, size_t* /*count*/) { + return Status::Error("SNII does not support try_query"); +} + +Status SniiIndexReader::read_null_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryCacheHandle* cache_handle, + lucene::store::Directory* /*dir*/) { + return _read_null_bitmap(context, cache_handle, nullptr); +} + +Status SniiIndexReader::_read_null_bitmap( + const IndexQueryContextPtr& context, InvertedIndexQueryCacheHandle* cache_handle, + const ::doris::snii::reader::LogicalIndexReader* preopened_reader) { + SCOPED_RAW_TIMER(&context->stats->inverted_index_query_null_bitmap_timer); + auto index_file_key = _index_file_reader->get_index_file_cache_key(&_index_meta); + InvertedIndexQueryCache::CacheKey cache_key { + index_file_key, "", InvertedIndexQueryType::UNKNOWN_QUERY, "null_bitmap"}; + auto* cache = InvertedIndexQueryCache::instance(); + if (cache->lookup(cache_key, cache_handle)) { + return Status::OK(); + } + + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(context->io_ctx); + InvertedIndexCacheHandle searcher_cache_handle; + std::unique_ptr<::doris::snii::reader::LogicalIndexReader> uncached_reader; + const ::doris::snii::reader::LogicalIndexReader* logical_reader = preopened_reader; + if (logical_reader == nullptr) { + RETURN_IF_ERROR(_get_logical_reader(context, &searcher_cache_handle, &uncached_reader, + &logical_reader)); + } + auto null_bitmap = std::make_shared(); + const auto& ref = logical_reader->section_refs().null_bitmap; + if (ref.length > 0) { + std::vector bytes; + RETURN_IF_ERROR(logical_reader->reader()->read_at(ref.offset, ref.length, &bytes)); + ::doris::snii::format::NullBitmapReader reader; + RETURN_IF_ERROR(::doris::snii::format::NullBitmapReader::open(::doris::snii::Slice(bytes), + &reader)); + reader.copy_to(null_bitmap.get()); + null_bitmap->runOptimize(); + } + cache->insert(cache_key, null_bitmap, cache_handle); + return Status::OK(); +} + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/snii/snii_index_reader.h b/be/src/storage/index/snii/snii_index_reader.h new file mode 100644 index 00000000000000..da8a8dea33ee0a --- /dev/null +++ b/be/src/storage/index/snii/snii_index_reader.h @@ -0,0 +1,179 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_grams_query_cost.h" +#include "storage/index/inverted/inverted_index_query_type.h" +#include "storage/index/inverted/inverted_index_reader.h" + +namespace doris::snii::reader { +class LogicalIndexReader; +} // namespace doris::snii::reader + +namespace doris::snii::query { +struct PhraseMatch; +} // namespace doris::snii::query + +namespace doris::segment_v2 { + +// One query plus the plan the caller chose for it. This is a parameter object rather than a +// parameter list because _compute_query_bitmap() took fourteen positional arguments, two of them +// adjacent bools (common_grams_query_shape, force_plain) that no call site could tell apart +// without counting commas. +struct SniiQueryBitmapRequest { + InvertedIndexQueryType query_type; + const InvertedIndexQueryInfo& query_info; + std::string_view search_str; + int32_t max_expansions = 0; + + // Plan decisions the caller has already made. Both bools are false on the plain path. + bool common_grams_query_shape = false; + bool force_plain = false; + inverted_index::CommonGramsPlanCostModel common_grams_cost_model {}; + const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr; + // Identifies the physical query for the single-flight key; empty when unused. + std::string_view physical_raw_query_key {}; + + const ::doris::snii::reader::LogicalIndexReader* logical_reader = nullptr; +}; + +class SniiIndexReader final : public InvertedIndexReader { + ENABLE_FACTORY_CREATOR(SniiIndexReader); + +public: +#ifdef BE_TEST + using SingleFlightFollowerJoinedObserver = void (*)(void*) noexcept; + using SingleFlightLeaderBeforeComputeObserver = void (*)(void*) noexcept; + using SearcherOpenObserver = void (*)(void*) noexcept; +#endif + + SniiIndexReader(const TabletIndex* index_meta, + const std::shared_ptr& index_file_reader, + InvertedIndexReaderType reader_type) + : InvertedIndexReader(index_meta, index_file_reader), _reader_type(reader_type) {} + + Status new_iterator(std::unique_ptr* iterator) override; + Status query(const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr) override; + Status query_with_null_bitmap(const IndexQueryContextPtr& context, + const std::string& column_name, const Field& query_value, + InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr) override; + Status try_query(const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + size_t* count) override; + Status read_null_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryCacheHandle* cache_handle, + lucene::store::Directory* dir = nullptr) override; + InvertedIndexReaderType type() override { return _reader_type; } + +#ifdef BE_TEST + void set_single_flight_follower_joined_observer_for_test( + SingleFlightFollowerJoinedObserver observer, void* opaque) { + _single_flight_follower_joined_observer = observer; + _single_flight_follower_joined_opaque = opaque; + } + void set_single_flight_leader_before_compute_observer_for_test( + SingleFlightLeaderBeforeComputeObserver observer, void* opaque) { + _single_flight_leader_before_compute_observer = observer; + _single_flight_leader_before_compute_opaque = opaque; + } + void set_searcher_open_observer_for_test(SearcherOpenObserver observer, void* opaque) { + _searcher_open_observer = observer; + _searcher_open_opaque = opaque; + } +#endif + +private: + Status _query(const IndexQueryContextPtr& context, const std::string& column_name, + const Field& query_value, InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + InvertedIndexQueryCacheHandle* null_bitmap_cache_handle, + const InvertedIndexAnalyzerCtx* analyzer_ctx); + Status _parse_query_terms( + const IndexQueryContextPtr& context, std::string search_str, + InvertedIndexQueryType query_type, const InvertedIndexAnalyzerCtx* analyzer_ctx, + InvertedIndexQueryInfo* query_info, + std::optional purpose_override = std::nullopt); + Status _get_logical_reader( + const IndexQueryContextPtr& context, InvertedIndexCacheHandle* searcher_cache_handle, + std::unique_ptr<::doris::snii::reader::LogicalIndexReader>* uncached_reader, + const ::doris::snii::reader::LogicalIndexReader** logical_reader); + Status _read_null_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryCacheHandle* cache_handle, + const ::doris::snii::reader::LogicalIndexReader* preopened_reader); + // Opens the segment index and runs the query, producing the result bitmap. Invoked as the + // single-flight "compute" step by query(); see SingleFlight for the concurrency rationale. + Status _compute_query_bitmap(const IndexQueryContextPtr& context, + const SniiQueryBitmapRequest& request, + std::vector* terms, + std::shared_ptr* out, + std::vector<::doris::snii::query::PhraseMatch>* phrase_matches); +#ifdef BE_TEST + Status _compute_query_bitmap(const IndexQueryContextPtr& context, + InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, + std::string_view search_str, std::vector* terms, + int32_t max_expansions, std::shared_ptr* out); +#endif + // G02 count-only fast path. Only called when the caller (SegmentIterator) + // set context->count_on_index_fastpath, i.e. the match count alone decides + // the scan result. On *handled = true, *out is a bitmap of cardinality df + // (row ids NOT real) built from a single exact term's dict-entry df WITHOUT + // decoding postings. On a segment without a null bitmap the fabricated ids + // are the dense range [0, df); on a segment WITH one they are the first df + // NON-NULL row ids (see + // fabricate_null_disjoint_count_bitmap) so that the unconditional + // FunctionMatchBase -> mask_out_null subtraction of the real null bitmap + // is a no-op and the cardinality stays df -- which is already the exact + // match count, because postings never contain null docs. Falls through + // (*handled = false) for every other shape: every multi-term query + // (including phrase and OR/AND), prefix/regexp/wildcard/phrase-prefix + // expansion, and sloppy phrase. + // On *handled = true, query() also raises + // context->count_on_index_fastpath_hit (G03) so the SegmentIterator may + // short-circuit row emission for the count-shaped bitmap. + Status _try_count_only_fastpath( + const IndexQueryContextPtr& context, InvertedIndexQueryType query_type, + const InvertedIndexQueryInfo& query_info, const std::vector& terms, + bool* handled, std::shared_ptr* out, + const ::doris::snii::reader::LogicalIndexReader* preopened_reader = nullptr); + + InvertedIndexReaderType _reader_type; +#ifdef BE_TEST + SingleFlightFollowerJoinedObserver _single_flight_follower_joined_observer = nullptr; + void* _single_flight_follower_joined_opaque = nullptr; + SingleFlightLeaderBeforeComputeObserver _single_flight_leader_before_compute_observer = nullptr; + void* _single_flight_leader_before_compute_opaque = nullptr; + SearcherOpenObserver _searcher_open_observer = nullptr; + void* _searcher_open_opaque = nullptr; +#endif +}; + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/snii/snii_index_writer.cpp b/be/src/storage/index/snii/snii_index_writer.cpp new file mode 100644 index 00000000000000..698f7eb9f34327 --- /dev/null +++ b/be/src/storage/index/snii/snii_index_writer.cpp @@ -0,0 +1,502 @@ +// 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 "storage/index/snii/snii_index_writer.h" + +#include + +#include +#include +#include + +#include "common/cast_set.h" +#include "common/config.h" +#include "common/logging.h" +#include "storage/index/index_file_writer.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/query/query_info.h" +#include "storage/index/inverted/token_filter/common_grams_filter.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/writer/global_memory_limiter.h" +#include "storage/tablet/tablet_schema.h" + +namespace doris::segment_v2 { +namespace { + +Status validate_common_grams_metadata_seed( + const inverted_index::CommonGramsSegmentMetadata& metadata) { + using namespace inverted_index; + auto status = validate_common_grams_segment_metadata(metadata); + if (!status.ok() || metadata.plain_term_key_version != PlainTermKeyVersion::kEscapedV1 || + metadata.common_grams_coverage != CommonGramsCoverage::kComplete || + metadata.common_grams_semantics_version != COMMON_GRAMS_SEMANTICS_VERSION_V1 || + metadata.common_grams_key_version != COMMON_GRAMS_KEY_VERSION_V1 || + metadata.scoring_coverage != ScoringCoverage::kComplete || + metadata.scoring_stats_version != COMMON_GRAMS_SCORING_STATS_VERSION_V1 || + metadata.norm_semantics_version != COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1) { + return Status::Error( + "SNII CommonGrams metadata identity seed is incomplete or incompatible"); + } + return Status::OK(); +} + +} // namespace + +SniiIndexColumnWriter::SniiIndexColumnWriter( + IndexFileWriter* index_file_writer, const TabletIndex* index_meta, bool, + std::optional common_grams_metadata_seed) + : _index_file_writer(index_file_writer), + _index_meta(index_meta), + _common_grams_build_enabled(config::enable_common_grams_index_build), + _common_grams_metadata_seed(std::move(common_grams_metadata_seed)) {} + +Status SniiIndexColumnWriter::init() { + _should_analyzer = + inverted_index::InvertedIndexAnalyzer::should_analyzer(_index_meta->properties()); + _has_positions = get_parser_phrase_support_string_from_properties(_index_meta->properties()) == + INVERTED_INDEX_PARSER_PHRASE_SUPPORT_YES; + _config = _has_positions ? ::doris::snii::format::IndexConfig::kDocsPositions + : ::doris::snii::format::IndexConfig::kDocsOnly; + auto ignore_above_value = + get_parser_ignore_above_value_from_properties(_index_meta->properties()); + _ignore_above = cast_set(std::stoul(ignore_above_value)); + const auto spill_threshold = + static_cast(config::inverted_index_ram_buffer_size * 1024 * 1024); + _memory_reporter = std::make_unique<::doris::snii::writer::MemoryReporter>( + nullptr, spill_threshold, + ::doris::snii::writer::MemoryReporter::CapPolicy::kSpillThreshold); + _term_buffer = std::make_unique<::doris::snii::writer::SpimiTermBuffer>( + _has_positions, spill_threshold, _memory_reporter.get()); + // G09: join the PROCESS-WIDE build-RAM limiter. The per-writer spill threshold above + // bounds one writer; a load keeps (tablets x concurrency) writers alive at + // once, none of which may ever reach it -- the global registry bounds their + // SUM by asking the largest buffers to spill early (advisory flags honored + // on each writer's own thread; byte-identical output). Budget refreshed + // from the mutable config at every writer init; 0 disables (no + // registration, zero per-token overhead beyond the G08 path). + // G09 anti-storm knobs (see the config comments): the forced-spill floor + // gates both the owner-side honor (a request is a pending no-op until the + // reclaimable arena regrows past it) and the limiter's victim eligibility, + // and the run-file cap merge-compacts a writer's spill runs so the final + // k-way merge's fd fan-in stays bounded. Applied unconditionally -- the + // floor also protects test-seam requests, and the cap also bounds + // per-writer gate-2 runs when the global limiter is off. + _term_buffer->set_forced_spill_min_arena_bytes( + static_cast(std::max(config::snii_forced_spill_min_arena_bytes, 0))); + _term_buffer->set_max_run_files( + static_cast(std::max(config::snii_spill_max_run_files_per_buffer, 0))); + const int64_t global_budget = config::snii_index_writer_global_memory_bytes; + if (global_budget > 0) { + auto* global_limiter = ::doris::snii::writer::GlobalMemoryLimiter::instance(); + global_limiter->set_budget_bytes(global_budget); + global_limiter->set_min_victim_arena_bytes(config::snii_forced_spill_min_arena_bytes); + _term_buffer->attach_global_limiter(global_limiter); + } + _analyzer_config.analyzer_name = get_analyzer_name_from_properties(_index_meta->properties()); + _analyzer_config.parser_type = get_inverted_index_parser_type_from_string( + get_parser_string_from_properties(_index_meta->properties())); + _analyzer_config.parser_mode = + get_parser_mode_string_from_properties(_index_meta->properties()); + _analyzer_config.char_filter_map = + get_parser_char_filter_map_from_properties(_index_meta->properties()); + _analyzer_config.lower_case = + get_parser_lowercase_from_properties(_index_meta->properties()); + _analyzer_config.stop_words = get_parser_stopwords_from_properties(_index_meta->properties()); + try { + _char_string_reader = inverted_index::InvertedIndexAnalyzer::create_reader( + _analyzer_config.char_filter_map); + if (_should_analyzer) { + auto analyzer_provider = + inverted_index::InvertedIndexAnalyzer::create_analyzer_provider( + &_analyzer_config); + const bool policy_uses_common_grams = analyzer_provider->uses_common_grams(); + _uses_common_grams = _common_grams_build_enabled && policy_uses_common_grams; + if (policy_uses_common_grams && !_uses_common_grams) { + _common_grams_metadata_seed.reset(); + } + if (_uses_common_grams) { + if (!_has_positions) { + close_on_error(); + return Status::Error( + "SNII CommonGrams requires phrase positions"); + } + const auto base_analyzer_fingerprint = + analyzer_provider->base_analyzer_fingerprint(); + if (base_analyzer_fingerprint.empty()) { + close_on_error(); + return Status::Error( + "SNII CommonGrams analyzer has no immutable base fingerprint"); + } + const auto* identity = analyzer_provider->common_grams_identity(); + if (identity == nullptr) { + close_on_error(); + return Status::Error( + "SNII CommonGrams analyzer has no immutable dictionary identity"); + } + if (_common_grams_metadata_seed.has_value()) { + if (!inverted_index::common_grams_identity_matches(*_common_grams_metadata_seed, + *identity)) { + close_on_error(); + return Status::Error( + "SNII CommonGrams metadata seed does not match the analyzer " + "identity"); + } + } else { + _common_grams_metadata_seed = + inverted_index::make_common_grams_segment_metadata(*identity); + } + auto status = validate_common_grams_metadata_seed(*_common_grams_metadata_seed); + if (!status.ok()) { + close_on_error(); + return status; + } + _config = ::doris::snii::format::IndexConfig::kDocsPositionsScoring; + } else if (_common_grams_metadata_seed.has_value()) { + close_on_error(); + return Status::Error( + "SNII CommonGrams metadata cannot be attached to a plain analyzer"); + } + _analyzer = analyzer_provider->get_analyzer( + _uses_common_grams ? inverted_index::AnalysisPurpose::kSniiTransientIndex + : inverted_index::AnalysisPurpose::kPlainQuery); + } else if (_common_grams_metadata_seed.has_value()) { + close_on_error(); + return Status::Error( + "SNII CommonGrams metadata cannot be attached to a keyword analyzer"); + } + } catch (const CLuceneError& e) { + return Status::Error( + "SNII create analyzer failed: {}", e.what()); + } catch (const Exception& e) { + return Status::Error( + "SNII create analyzer failed: {}", e.what()); + } + if (_uses_common_grams) { + _term_buffer->enable_common_gram_pair_keys(); + } + return Status::OK(); +} + +void SniiIndexColumnWriter::set_direct_load(bool is_direct_load) { + // The PRX compression-tier hint must be stable for one index. The first + // pre-row call wins; repeat or late calls are ignored and logged. + DCHECK(!_direct_load_marked && _rid == 0); + if (_direct_load_marked || _rid != 0) { + LOG_EVERY_N(WARNING, 100) << "SNII set_direct_load(" << is_direct_load + << ") ignored (already_marked=" << _direct_load_marked + << ", rows_fed=" << _rid << ") for index " + << (_index_meta != nullptr ? _index_meta->index_id() : -1) + << "; keeping the first-captured PRX tier decision"; + return; + } + _direct_load_marked = true; + _is_direct_load = is_direct_load; +} + +Status SniiIndexColumnWriter::_add_value_tokens(const Slice& value, uint32_t docid, + uint32_t position_base, uint32_t* max_position, + uint32_t* semantic_length) { + DCHECK(max_position != nullptr); + DCHECK(semantic_length != nullptr); + *max_position = position_base; + *semantic_length = 0; + if ((!_should_analyzer && value.size > _ignore_above) || (_should_analyzer && value.empty())) { + return Status::OK(); + } + + // T1a: tokens STREAM from the analyzer straight into the SPIMI buffer as + // string_views (the buffer interns the bytes into its own storage) -- no + // per-row vector and no per-token std::string materialization + // (the old get_analyse_result lane; profile: 3.4-4.7% of import CPU burned + // in token realloc). Golden-byte pins (snii_writer_golden_bytes_test.cpp) + // hold this path byte-identical to the materializing one it replaced. + auto consume_token = [&](std::string_view term, int32_t token_position, bool retain_positions) { + const uint32_t position = + _has_positions ? position_base + cast_set(token_position) : 0; + _term_buffer->add_token(term, docid, position, retain_positions); + *max_position = std::max(*max_position, position); + }; + + if (!_should_analyzer) { + // Keyword lane: the whole value is one exact-match token at position 0 + // (an EMPTY value is a valid keyword token, mirrored from the old lane). + consume_token(std::string_view(value.data, value.size), 0, _has_positions); + } else { + try { + _char_string_reader->init(value.data, cast_set(value.size), false); + if (_uses_common_grams) { + auto* token_stream = _analyzer->reusableTokenStream(L"", _char_string_reader); + if (_common_grams_filter == nullptr) { + _common_grams_filter = + dynamic_cast(token_stream); + DORIS_CHECK(_common_grams_filter != nullptr); + } else { + DCHECK_EQ(static_cast(_common_grams_filter), + token_stream); + } + _common_grams_filter->reset(); + + int32_t position = 0; + std::optional<::doris::snii::writer::ClassifiedPlainTerm> previous; + size_t previous_logical_term_size = 0; + inverted_index::SniiCommonGramsIndexEvent event; + while (_common_grams_filter->next_snii_index_event(&event)) { + const auto current = _term_buffer->intern_classified_plain_term( + event.plain_term, event.logical_term, + _common_grams_filter->common_words()); + const bool has_preceding_gram = + previous.has_value() && (previous->is_common || current.is_common) && + inverted_index::common_gram_component_sizes_encodable( + previous_logical_term_size, event.logical_term.size()); + const uint32_t gram_position = position_base + cast_set(position); + const uint32_t physical_position = + position_base + cast_set(position + 1); + if (has_preceding_gram) { + DCHECK(previous.has_value()); + _term_buffer->add_common_gram_and_plain( + previous->id, current.id, docid, gram_position, physical_position, + previous->is_common && current.is_common); + } else { + _term_buffer->add_plain_token(current.id, docid, physical_position); + } + ++position; + ++*semantic_length; + *max_position = std::max(*max_position, physical_position); + previous = current; + previous_logical_term_size = event.logical_term.size(); + } + } else { + std::unique_ptr owned_token_stream( + _analyzer->tokenStream(L"", _char_string_reader)); + auto* token_stream = owned_token_stream.get(); + // EXACT InvertedIndexAnalyzer::get_analyse_result semantics, + // including the subtle one: an empty token's position increment is + // dropped WITH the token (not accumulated into the next). + lucene::analysis::Token token; + int32_t position = 0; + while (token_stream->next(&token)) { + if (token.termLength() != 0) { + const std::string_view term(token.termBuffer(), + token.termLength()); + position += token.getPositionIncrement(); + consume_token(term, position, _has_positions); + } + } + token_stream->close(); + } + } catch (const CLuceneError& e) { + return _latch_analysis_failure(Status::Error( + "SNII analyze value failed: {}", e.what())); + } catch (const Exception& e) { + return _latch_analysis_failure(Status::Error( + "SNII analyze value failed: {}", e.what())); + } + } + return Status::OK(); +} + +Status SniiIndexColumnWriter::add_values(const std::string /*name*/, const void* values, + size_t count) { + if (!_failure_status.ok()) { + return _failure_status; + } + const auto* v = reinterpret_cast(values); + for (size_t i = 0; i < count; ++i) { + uint32_t max_position = 0; + uint32_t semantic_length = 0; + RETURN_IF_ERROR(_add_value_tokens(*v, _rid, 0, &max_position, &semantic_length)); + if (_uses_common_grams) { + _encoded_norms.push_back(::doris::snii::query::encode_norm(semantic_length)); + _report_encoded_norms_capacity(); + _scoring_token_count += semantic_length; + } + ++v; + ++_rid; + } + return Status::OK(); +} + +Status SniiIndexColumnWriter::add_array_values(size_t field_size, const void* value_ptr, + const uint8_t* nested_null_map, + const uint8_t* offsets_ptr, size_t count) { + if (!_failure_status.ok()) { + return _failure_status; + } + if (_uses_common_grams) { + return _latch_analysis_failure(Status::Error( + "SNII CommonGrams does not support ARRAY fields")); + } + if (count == 0) { + return Status::OK(); + } + const auto* offsets = reinterpret_cast(offsets_ptr); + size_t start_off = 0; + for (size_t i = 0; i < count; ++i) { + auto array_elem_size = offsets[i + 1] - offsets[i]; + uint32_t position_base = 0; + for (auto j = start_off; j < start_off + array_elem_size; ++j) { + if (nested_null_map != nullptr && nested_null_map[j] == 1) { + continue; + } + const auto* value = reinterpret_cast( + reinterpret_cast(value_ptr) + j * field_size); + uint32_t max_position = position_base; + uint32_t semantic_length = 0; + RETURN_IF_ERROR(_add_value_tokens(*value, _rid, position_base, &max_position, + &semantic_length)); + position_base = max_position + 1; + } + start_off += array_elem_size; + ++_rid; + } + return Status::OK(); +} + +void SniiIndexColumnWriter::_report_null_docids_capacity(bool release_all) { + if (_memory_reporter == nullptr) { + return; + } + const int64_t now = + release_all ? 0 : static_cast(_null_docids.capacity() * sizeof(uint32_t)); + if (now != _null_docids_charged_bytes) { + _memory_reporter->report(now - _null_docids_charged_bytes); + _null_docids_charged_bytes = now; + } +} + +void SniiIndexColumnWriter::_report_encoded_norms_capacity(bool release_all) { + if (_memory_reporter == nullptr) { + return; + } + const int64_t now = release_all ? 0 : static_cast(_encoded_norms.capacity()); + if (now != _encoded_norms_charged_bytes) { + _memory_reporter->report(now - _encoded_norms_charged_bytes); + _encoded_norms_charged_bytes = now; + } +} + +Status SniiIndexColumnWriter::add_nulls(uint32_t count) { + if (!_failure_status.ok()) { + return _failure_status; + } + // GEOMETRIC BULK reserve -- never an exact one: append_nullable calls + // add_nulls once per NULL RUN (thousands to millions of calls on a large + // interleaved-null segment), and an exact reserve(size()+count) caps + // capacity at "just enough" -- the NEXT call then reallocates and memcpys + // the WHOLE array, defeating geometric growth and turning total memcpy + // quadratic: O(runs x array_bytes). On an agentlogs full-compaction segment + // (12.4M rows, 22% interleaved nulls) that was TBs of memcpy per tablet -- + // the compaction ran 8+x slower than V3 (whose add_nulls is a roaring + // addRange). Doubling on overflow keeps the O(count) amortization AND makes + // one large run pay at most one reallocation. + const size_t need = _null_docids.size() + count; + if (need > _null_docids.capacity()) { + _null_docids.reserve(std::max(need, _null_docids.capacity() * 2)); + } + for (uint32_t i = 0; i < count; ++i) { + _null_docids.push_back(_rid + i); + } + _rid += count; + if (_uses_common_grams) { + _encoded_norms.insert(_encoded_norms.end(), count, ::doris::snii::query::encode_norm(0)); + _report_encoded_norms_capacity(); + } + _report_null_docids_capacity(); + return Status::OK(); +} + +Status SniiIndexColumnWriter::add_array_nulls(const uint8_t* null_map, size_t num_rows) { + if (!_failure_status.ok()) { + return _failure_status; + } + DCHECK(_rid >= num_rows); + if (num_rows == 0 || null_map == nullptr) { + return Status::OK(); + } + const auto first_row = _rid - num_rows; + for (size_t i = 0; i < num_rows; ++i) { + if (null_map[i] == 1) { + _null_docids.push_back(cast_set(first_row + i)); + } + } + _report_null_docids_capacity(); + return Status::OK(); +} + +Status SniiIndexColumnWriter::finish() { + if (!_failure_status.ok()) { + return _failure_status; + } + DCHECK(_term_buffer != nullptr); + auto status = _term_buffer->status(); + if (!status.ok()) { + return Status::InternalError("SNII term buffer error: {}", status.to_string()); + } + // Ownership of _null_docids hands off to the flush below (transient, + // flush-scoped); release the accumulation-phase charge so the retained + // reporter (and the LOAD MemTracker behind it) balances to zero. + _report_null_docids_capacity(/*release_all=*/true); + IndexFileWriter::SniiAddIndexOptions options {}; + options.is_direct_load = _is_direct_load; + if (_uses_common_grams) { + options.encoded_norms = std::move(_encoded_norms); + options.common_grams_metadata = _build_common_grams_metadata(); + options.common_grams_posting_policy = + ::doris::snii::format::CommonGramsPostingPolicy::kHybridV1; + } + status = _index_file_writer->add_snii_index( + _index_meta, cast_set(_rid), std::move(_null_docids), _term_buffer.get(), + _config, std::move(options), _memory_reporter.get()); + _report_encoded_norms_capacity(/*release_all=*/true); + RETURN_IF_ERROR(status); + _index_file_writer->retain_snii_memory_reporter(std::move(_memory_reporter)); + _term_buffer.reset(); + return Status::OK(); +} + +inverted_index::CommonGramsSegmentMetadata SniiIndexColumnWriter::_build_common_grams_metadata() + const { + DORIS_CHECK(_uses_common_grams); + DORIS_CHECK(_common_grams_metadata_seed.has_value()); + auto metadata = *_common_grams_metadata_seed; + metadata.common_grams_coverage = inverted_index::CommonGramsCoverage::kMixed; + metadata.scoring_doc_count = _rid; + metadata.scoring_token_count = _scoring_token_count; + return metadata; +} + +Status SniiIndexColumnWriter::_latch_analysis_failure(Status status) { + DORIS_CHECK(!status.ok()); + DORIS_CHECK(_failure_status.ok()); + _failure_status = std::move(status); + close_on_error(); + return _failure_status; +} + +void SniiIndexColumnWriter::close_on_error() { + _term_buffer.reset(); + // Balance the LOAD MemTracker mirror before dropping the reporter. + _report_null_docids_capacity(/*release_all=*/true); + _report_encoded_norms_capacity(/*release_all=*/true); + _memory_reporter.reset(); + _null_docids.clear(); + std::vector().swap(_encoded_norms); + _scoring_token_count = 0; +} + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/snii/snii_index_writer.h b/be/src/storage/index/snii/snii_index_writer.h new file mode 100644 index 00000000000000..c81e4aa2c0baaf --- /dev/null +++ b/be/src/storage/index/snii/snii_index_writer.h @@ -0,0 +1,142 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "storage/index/index_writer.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/inverted_index_parser.h" +#include "storage/index/inverted/query/query_info.h" +#include "storage/index/inverted/util/reader.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "util/slice.h" + +namespace lucene::analysis { +class Analyzer; +} + +namespace doris::segment_v2::inverted_index { +class CommonGramsFilter; +} + +namespace doris::segment_v2 { + +class SniiIndexColumnWriter final : public IndexColumnWriter { +public: + SniiIndexColumnWriter(IndexFileWriter* index_file_writer, const TabletIndex* index_meta, bool, + std::optional + common_grams_metadata_seed = std::nullopt); + ~SniiIndexColumnWriter() override = default; + + Status init() override; + void set_direct_load(bool is_direct_load) override; + Status add_values(const std::string name, const void* values, size_t count) override; + Status add_array_values(size_t field_size, const void* value_ptr, + const uint8_t* nested_null_map, const uint8_t* offsets_ptr, + size_t count) override; + Status add_nulls(uint32_t count) override; + Status add_array_nulls(const uint8_t* null_map, size_t num_rows) override; + Status finish() override; + int64_t size() const override { return 0; } + void close_on_error() override; + +#ifdef BE_TEST + // TEST-ONLY view of the accumulated null docids: the growth-policy + // regression pin asserts add_nulls keeps geometric growth (an exact + // reserve(size+count) per null RUN made total memcpy quadratic -- the + // agentlogs full-compaction pathology). + const std::vector& null_docids_for_test() const { return _null_docids; } + ::doris::snii::writer::SpimiTermBuffer* term_buffer_for_test() const { + return _term_buffer.get(); + } + ::doris::snii::writer::MemoryReporter* memory_reporter_for_test() const { + return _memory_reporter.get(); + } + const std::vector& encoded_norms_for_test() const { return _encoded_norms; } + uint64_t scoring_token_count_for_test() const { return _scoring_token_count; } + ::doris::snii::format::IndexConfig config_for_test() const { return _config; } + bool has_common_grams_metadata_seed_for_test() const { + return _common_grams_metadata_seed.has_value(); + } + inverted_index::CommonGramsSegmentMetadata common_grams_metadata_for_test() const { + return _build_common_grams_metadata(); + } + void set_analysis_for_test(inverted_index::ReaderPtr reader, + std::shared_ptr analyzer) { + _should_analyzer = true; + _char_string_reader = std::move(reader); + _analyzer = std::move(analyzer); + } +#endif + +private: + Status _add_value_tokens(const Slice& value, uint32_t docid, uint32_t position_base, + uint32_t* max_position, uint32_t* semantic_length); + inverted_index::CommonGramsSegmentMetadata _build_common_grams_metadata() const; + // Mirrors _null_docids' capacity into _memory_reporter (delta-charged); + // release_all zeroes the charge (finish() handoff / close_on_error()). + void _report_null_docids_capacity(bool release_all = false); + void _report_encoded_norms_capacity(bool release_all = false); + Status _latch_analysis_failure(Status status); + + IndexFileWriter* _index_file_writer = nullptr; + const TabletIndex* _index_meta = nullptr; + bool _should_analyzer = false; + bool _has_positions = false; + const bool _common_grams_build_enabled; + bool _uses_common_grams = false; + // Latch: set_direct_load() ran. The first call wins; a repeat or late call + // is ignored (and logged) so one index keeps one stable compression-tier + // decision. + bool _direct_load_marked = false; + // Captured by set_direct_load() under the same latch: this writer serves a + // stream/broker load (DataWriteType::TYPE_DIRECT). Consumed at finish() to + // route the prx region to the load-tier zstd level (patch C, + // config::snii_prx_zstd_level_direct_load). + bool _is_direct_load = false; + uint32_t _ignore_above = 0; + uint32_t _rid = 0; + ::doris::snii::format::IndexConfig _config = ::doris::snii::format::IndexConfig::kDocsOnly; + InvertedIndexAnalyzerConfig _analyzer_config; + inverted_index::ReaderPtr _char_string_reader; + std::shared_ptr _analyzer; + inverted_index::CommonGramsFilter* _common_grams_filter = nullptr; + std::unique_ptr<::doris::snii::writer::MemoryReporter> _memory_reporter; + std::unique_ptr<::doris::snii::writer::SpimiTermBuffer> _term_buffer; + std::vector _null_docids; + std::vector _encoded_norms; + uint64_t _scoring_token_count = 0; + std::optional _common_grams_metadata_seed; + // Bytes of _null_docids capacity currently mirrored into _memory_reporter + // (and through it Doris's LOAD MemTracker). Re-charged on growth in + // add_nulls / add_array_nulls, released in finish() / close_on_error() -- + // without it a large interleaved-null segment accumulates untracked RSS the + // G09 limiter cannot see. + int64_t _null_docids_charged_bytes = 0; + int64_t _encoded_norms_charged_bytes = 0; + Status _failure_status = Status::OK(); +}; + +} // namespace doris::segment_v2 diff --git a/be/src/storage/index/snii/snii_prx_profile.h b/be/src/storage/index/snii/snii_prx_profile.h new file mode 100644 index 00000000000000..c343c5b54fa076 --- /dev/null +++ b/be/src/storage/index/snii/snii_prx_profile.h @@ -0,0 +1,308 @@ +// 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. + +#pragma once + +#include +#include + +#include "runtime/runtime_profile.h" +#include "storage/index/snii/format/prx_decode_stats.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/olap_common.h" + +namespace doris::snii { + +#ifdef BE_TEST +namespace testing { + +void record_prx_execution_profile_scope_construction(); +void record_prx_execution_profile_scope_flush(); +void reset_prx_execution_profile_scope_counters(); +uint64_t prx_execution_profile_scope_construction_count(); +uint64_t prx_execution_profile_scope_flush_count(); + +} // namespace testing +#endif + +inline void add_prx_decode_stats(OlapReaderStatistics* target, + const format::PrxDecodeStats& delta) { + SniiQueryStats& stats = target->snii_stats; + stats.prx_raw_frames += static_cast(delta.raw_frames); + stats.prx_zstd_frames += static_cast(delta.zstd_frames); + stats.prx_pfor_frames += static_cast(delta.pfor_frames); + stats.prx_plaintext_bytes += static_cast(delta.plaintext_bytes); + stats.prx_total_docs += static_cast(delta.total_docs); + stats.prx_selected_docs += static_cast(delta.selected_docs); + stats.prx_total_positions += static_cast(delta.total_positions); + stats.prx_selected_positions += static_cast(delta.selected_positions); + stats.prx_fetch_ns += static_cast(delta.fetch_ns); + stats.prx_decode_ns += static_cast(delta.decode_ns); + stats.prx_phrase_verify_ns += static_cast(delta.phrase_verify_ns); +} + +inline void add_phrase_query_stats(OlapReaderStatistics* target, + const format::PhraseQueryExecutionStats& delta) { + SniiQueryStats& stats = target->snii_stats; + stats.phrase_candidate_docs += static_cast(delta.exact_candidate_docs); + stats.phrase_candidate_visits += static_cast(delta.exact_candidate_visits); + stats.phrase_prefix_leading_candidate_docs += + static_cast(delta.prefix_leading_candidate_docs); + stats.phrase_prefix_tail_candidate_visits += + static_cast(delta.prefix_tail_candidate_visits); + stats.common_grams_candidate_queries += + static_cast(delta.common_grams_candidate_queries); + stats.common_grams_plain_plans += static_cast(delta.common_grams_plain_plans); + stats.common_grams_gram_plans += static_cast(delta.common_grams_gram_plans); + stats.common_grams_fallback_no_gram += + static_cast(delta.common_grams_fallback_no_gram); + stats.common_grams_fallback_incompatible += + static_cast(delta.common_grams_fallback_incompatible); + stats.common_grams_fallback_kill_switch += + static_cast(delta.common_grams_fallback_kill_switch); + stats.common_grams_fallback_cost += static_cast(delta.common_grams_fallback_cost); + stats.common_grams_fallback_base_analyzer_mismatch += + static_cast(delta.common_grams_fallback_base_analyzer_mismatch); + stats.common_grams_fallback_prefix_tail_empty += + static_cast(delta.common_grams_fallback_prefix_tail_empty); + stats.common_grams_authoritative_empty += + static_cast(delta.common_grams_authoritative_empty); + stats.common_grams_plain_posting_bytes += + static_cast(delta.common_grams_plain_posting_bytes); + stats.common_grams_gram_posting_bytes += + static_cast(delta.common_grams_gram_posting_bytes); + stats.common_grams_plain_estimated_candidate_df += + static_cast(delta.common_grams_plain_estimated_candidate_df); + stats.common_grams_gram_estimated_candidate_df += + static_cast(delta.common_grams_gram_estimated_candidate_df); + stats.common_grams_plain_estimated_cost += + static_cast(delta.common_grams_plain_estimated_cost); + stats.common_grams_gram_estimated_cost += + static_cast(delta.common_grams_gram_estimated_cost); + stats.common_grams_planning_ns += static_cast(delta.common_grams_planning_ns); +} + +// Exists only around an actual SNII compute execution. Query-cache hits, +// count-only fast paths, and single-flight followers never construct this +// scope, so they cannot contribute another execution's PRX work. The destructor +// flushes already-committed frames on both success and early error returns. +class SniiPrxExecutionProfileScope { +public: + explicit SniiPrxExecutionProfileScope(OlapReaderStatistics& target) : target_(target) { +#ifdef BE_TEST + testing::record_prx_execution_profile_scope_construction(); +#endif + } + ~SniiPrxExecutionProfileScope() { + add_prx_decode_stats(&target_, profile_.prx_decode_stats); + add_phrase_query_stats(&target_, profile_.phrase_query_stats); +#ifdef BE_TEST + testing::record_prx_execution_profile_scope_flush(); +#endif + } + + SniiPrxExecutionProfileScope(const SniiPrxExecutionProfileScope&) = delete; + SniiPrxExecutionProfileScope& operator=(const SniiPrxExecutionProfileScope&) = delete; + SniiPrxExecutionProfileScope(SniiPrxExecutionProfileScope&&) = delete; + SniiPrxExecutionProfileScope& operator=(SniiPrxExecutionProfileScope&&) = delete; + + query::QueryProfile* profile() { return &profile_; } + +private: + OlapReaderStatistics& target_; + query::QueryProfile profile_; +}; + +class SniiPrxRuntimeProfileCounters { +public: + static constexpr std::array counter_names() { + return {"SniiPrxRawFrames", + "SniiPrxZstdFrames", + "SniiPrxPforFrames", + "SniiPrxPlaintextBytes", + "SniiPrxTotalDocs", + "SniiPrxSelectedDocs", + "SniiPrxTotalPositions", + "SniiPrxSelectedPositions", + "SniiPrxFetchTime", + "SniiPrxInclusiveDecodeTime", + "SniiPrxExclusivePhraseVerifyTime"}; + } + + void initialize(RuntimeProfile* profile) { + raw_frames_ = profile->add_nonzero_counter("SniiPrxRawFrames", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + zstd_frames_ = profile->add_nonzero_counter("SniiPrxZstdFrames", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + pfor_frames_ = profile->add_nonzero_counter("SniiPrxPforFrames", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + plaintext_bytes_ = profile->add_nonzero_counter("SniiPrxPlaintextBytes", TUnit::BYTES, + RuntimeProfile::ROOT_COUNTER, 1); + total_docs_ = profile->add_nonzero_counter("SniiPrxTotalDocs", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + selected_docs_ = profile->add_nonzero_counter("SniiPrxSelectedDocs", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + total_positions_ = profile->add_nonzero_counter("SniiPrxTotalPositions", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + selected_positions_ = profile->add_nonzero_counter("SniiPrxSelectedPositions", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + fetch_ns_ = profile->add_nonzero_counter("SniiPrxFetchTime", TUnit::TIME_NS, + RuntimeProfile::ROOT_COUNTER, 1); + decode_ns_ = profile->add_nonzero_counter("SniiPrxInclusiveDecodeTime", TUnit::TIME_NS, + RuntimeProfile::ROOT_COUNTER, 1); + phrase_verify_ns_ = + profile->add_nonzero_counter("SniiPrxExclusivePhraseVerifyTime", TUnit::TIME_NS, + RuntimeProfile::ROOT_COUNTER, 1); + } + + void update(const OlapReaderStatistics& stats) const { + const SniiQueryStats& s = stats.snii_stats; + COUNTER_UPDATE(raw_frames_, s.prx_raw_frames); + COUNTER_UPDATE(zstd_frames_, s.prx_zstd_frames); + COUNTER_UPDATE(pfor_frames_, s.prx_pfor_frames); + COUNTER_UPDATE(plaintext_bytes_, s.prx_plaintext_bytes); + COUNTER_UPDATE(total_docs_, s.prx_total_docs); + COUNTER_UPDATE(selected_docs_, s.prx_selected_docs); + COUNTER_UPDATE(total_positions_, s.prx_total_positions); + COUNTER_UPDATE(selected_positions_, s.prx_selected_positions); + COUNTER_UPDATE(fetch_ns_, s.prx_fetch_ns); + COUNTER_UPDATE(decode_ns_, s.prx_decode_ns); + COUNTER_UPDATE(phrase_verify_ns_, s.prx_phrase_verify_ns); + } + +private: + RuntimeProfile::Counter* raw_frames_ = nullptr; + RuntimeProfile::Counter* zstd_frames_ = nullptr; + RuntimeProfile::Counter* pfor_frames_ = nullptr; + RuntimeProfile::Counter* plaintext_bytes_ = nullptr; + RuntimeProfile::Counter* total_docs_ = nullptr; + RuntimeProfile::Counter* selected_docs_ = nullptr; + RuntimeProfile::Counter* total_positions_ = nullptr; + RuntimeProfile::Counter* selected_positions_ = nullptr; + RuntimeProfile::Counter* fetch_ns_ = nullptr; + RuntimeProfile::Counter* decode_ns_ = nullptr; + RuntimeProfile::Counter* phrase_verify_ns_ = nullptr; +}; + +class SniiPhraseRuntimeProfileCounters { +public: + void initialize(RuntimeProfile* profile) { + candidate_docs_ = profile->add_nonzero_counter("SniiPhraseCandidateDocs", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + candidate_visits_ = profile->add_nonzero_counter("SniiPhraseCandidateVisits", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + prefix_leading_candidate_docs_ = + profile->add_nonzero_counter("SniiPhrasePrefixLeadingCandidateDocs", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + prefix_tail_candidate_visits_ = + profile->add_nonzero_counter("SniiPhrasePrefixTailCandidateVisits", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + common_grams_candidate_queries_ = profile->add_nonzero_counter( + "SniiCommonGramsCandidateQueries", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_plain_plans_ = profile->add_nonzero_counter( + "SniiCommonGramsPlainPlans", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_gram_plans_ = profile->add_nonzero_counter( + "SniiCommonGramsGramPlans", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_no_gram_ = profile->add_nonzero_counter( + "SniiCommonGramsFallbackNoGram", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_incompatible_ = + profile->add_nonzero_counter("SniiCommonGramsFallbackIncompatible", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_kill_switch_ = profile->add_nonzero_counter( + "SniiCommonGramsFallbackKillSwitch", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_cost_ = profile->add_nonzero_counter( + "SniiCommonGramsFallbackCost", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_base_analyzer_mismatch_ = + profile->add_nonzero_counter("SniiCommonGramsFallbackBaseAnalyzerMismatch", + TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_fallback_prefix_tail_empty_ = + profile->add_nonzero_counter("SniiCommonGramsFallbackPrefixTailEmpty", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + common_grams_authoritative_empty_ = profile->add_nonzero_counter( + "SniiCommonGramsAuthoritativeEmpty", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_plain_posting_bytes_ = profile->add_nonzero_counter( + "SniiCommonGramsPlainPostingBytes", TUnit::BYTES, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_gram_posting_bytes_ = profile->add_nonzero_counter( + "SniiCommonGramsGramPostingBytes", TUnit::BYTES, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_plain_estimated_candidate_df_ = + profile->add_nonzero_counter("SniiCommonGramsPlainEstimatedCandidateDf", + TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_gram_estimated_candidate_df_ = + profile->add_nonzero_counter("SniiCommonGramsGramEstimatedCandidateDf", TUnit::UNIT, + RuntimeProfile::ROOT_COUNTER, 1); + common_grams_plain_estimated_cost_ = profile->add_nonzero_counter( + "SniiCommonGramsPlainEstimatedCost", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_gram_estimated_cost_ = profile->add_nonzero_counter( + "SniiCommonGramsGramEstimatedCost", TUnit::UNIT, RuntimeProfile::ROOT_COUNTER, 1); + common_grams_planning_ns_ = profile->add_nonzero_counter( + "SniiCommonGramsPlanningTime", TUnit::TIME_NS, RuntimeProfile::ROOT_COUNTER, 1); + } + + void update(const OlapReaderStatistics& stats) const { + const SniiQueryStats& s = stats.snii_stats; + COUNTER_UPDATE(candidate_docs_, s.phrase_candidate_docs); + COUNTER_UPDATE(candidate_visits_, s.phrase_candidate_visits); + COUNTER_UPDATE(prefix_leading_candidate_docs_, s.phrase_prefix_leading_candidate_docs); + COUNTER_UPDATE(prefix_tail_candidate_visits_, s.phrase_prefix_tail_candidate_visits); + COUNTER_UPDATE(common_grams_candidate_queries_, s.common_grams_candidate_queries); + COUNTER_UPDATE(common_grams_plain_plans_, s.common_grams_plain_plans); + COUNTER_UPDATE(common_grams_gram_plans_, s.common_grams_gram_plans); + COUNTER_UPDATE(common_grams_fallback_no_gram_, s.common_grams_fallback_no_gram); + COUNTER_UPDATE(common_grams_fallback_incompatible_, s.common_grams_fallback_incompatible); + COUNTER_UPDATE(common_grams_fallback_kill_switch_, s.common_grams_fallback_kill_switch); + COUNTER_UPDATE(common_grams_fallback_cost_, s.common_grams_fallback_cost); + COUNTER_UPDATE(common_grams_fallback_base_analyzer_mismatch_, + s.common_grams_fallback_base_analyzer_mismatch); + COUNTER_UPDATE(common_grams_fallback_prefix_tail_empty_, + s.common_grams_fallback_prefix_tail_empty); + COUNTER_UPDATE(common_grams_authoritative_empty_, s.common_grams_authoritative_empty); + COUNTER_UPDATE(common_grams_plain_posting_bytes_, s.common_grams_plain_posting_bytes); + COUNTER_UPDATE(common_grams_gram_posting_bytes_, s.common_grams_gram_posting_bytes); + COUNTER_UPDATE(common_grams_plain_estimated_candidate_df_, + s.common_grams_plain_estimated_candidate_df); + COUNTER_UPDATE(common_grams_gram_estimated_candidate_df_, + s.common_grams_gram_estimated_candidate_df); + COUNTER_UPDATE(common_grams_plain_estimated_cost_, s.common_grams_plain_estimated_cost); + COUNTER_UPDATE(common_grams_gram_estimated_cost_, s.common_grams_gram_estimated_cost); + COUNTER_UPDATE(common_grams_planning_ns_, s.common_grams_planning_ns); + } + +private: + RuntimeProfile::Counter* candidate_docs_ = nullptr; + RuntimeProfile::Counter* candidate_visits_ = nullptr; + RuntimeProfile::Counter* prefix_leading_candidate_docs_ = nullptr; + RuntimeProfile::Counter* prefix_tail_candidate_visits_ = nullptr; + RuntimeProfile::Counter* common_grams_candidate_queries_ = nullptr; + RuntimeProfile::Counter* common_grams_plain_plans_ = nullptr; + RuntimeProfile::Counter* common_grams_gram_plans_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_no_gram_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_incompatible_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_kill_switch_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_cost_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_base_analyzer_mismatch_ = nullptr; + RuntimeProfile::Counter* common_grams_fallback_prefix_tail_empty_ = nullptr; + RuntimeProfile::Counter* common_grams_authoritative_empty_ = nullptr; + RuntimeProfile::Counter* common_grams_plain_posting_bytes_ = nullptr; + RuntimeProfile::Counter* common_grams_gram_posting_bytes_ = nullptr; + RuntimeProfile::Counter* common_grams_plain_estimated_candidate_df_ = nullptr; + RuntimeProfile::Counter* common_grams_gram_estimated_candidate_df_ = nullptr; + RuntimeProfile::Counter* common_grams_plain_estimated_cost_ = nullptr; + RuntimeProfile::Counter* common_grams_gram_estimated_cost_ = nullptr; + RuntimeProfile::Counter* common_grams_planning_ns_ = nullptr; +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/snii_query_stats.h b/be/src/storage/index/snii/snii_query_stats.h new file mode 100644 index 00000000000000..a7466af8d7c8de --- /dev/null +++ b/be/src/storage/index/snii/snii_query_stats.h @@ -0,0 +1,66 @@ +// 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. + +#pragma once + +#include + +namespace doris::snii { + +// Populated by snii_prx_profile.h from a per-execution query::QueryProfile, and read back out by +// SniiPrxRuntimeProfileCounters / SniiPhraseRuntimeProfileCounters into the scan node's +// RuntimeProfile. Lives in its own header, mirroring InvertedIndexStatistics +// (storage/index/inverted/inverted_index_stats.h), so OlapReaderStatistics -- shared by every +// storage format -- holds one named field here instead of one field per counter. +struct SniiQueryStats { + int64_t prx_raw_frames = 0; + int64_t prx_zstd_frames = 0; + int64_t prx_pfor_frames = 0; + int64_t prx_plaintext_bytes = 0; + int64_t prx_total_docs = 0; + int64_t prx_selected_docs = 0; + int64_t prx_total_positions = 0; + int64_t prx_selected_positions = 0; + int64_t prx_fetch_ns = 0; + int64_t prx_decode_ns = 0; + int64_t prx_phrase_verify_ns = 0; + + int64_t phrase_candidate_docs = 0; + int64_t phrase_candidate_visits = 0; + int64_t phrase_prefix_leading_candidate_docs = 0; + int64_t phrase_prefix_tail_candidate_visits = 0; + + int64_t common_grams_candidate_queries = 0; + int64_t common_grams_plain_plans = 0; + int64_t common_grams_gram_plans = 0; + int64_t common_grams_fallback_no_gram = 0; + int64_t common_grams_fallback_incompatible = 0; + int64_t common_grams_fallback_kill_switch = 0; + int64_t common_grams_fallback_cost = 0; + int64_t common_grams_fallback_base_analyzer_mismatch = 0; + int64_t common_grams_fallback_prefix_tail_empty = 0; + int64_t common_grams_authoritative_empty = 0; + int64_t common_grams_plain_posting_bytes = 0; + int64_t common_grams_gram_posting_bytes = 0; + int64_t common_grams_plain_estimated_candidate_df = 0; + int64_t common_grams_gram_estimated_candidate_df = 0; + int64_t common_grams_plain_estimated_cost = 0; + int64_t common_grams_gram_estimated_cost = 0; + int64_t common_grams_planning_ns = 0; +}; + +} // namespace doris::snii diff --git a/be/src/storage/index/snii/stats/snii_stats_provider.cpp b/be/src/storage/index/snii/stats/snii_stats_provider.cpp new file mode 100644 index 00000000000000..e12fe6d5f73d80 --- /dev/null +++ b/be/src/storage/index/snii/stats/snii_stats_provider.cpp @@ -0,0 +1,145 @@ +// 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 "storage/index/snii/stats/snii_stats_provider.h" + +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/stats_block.h" + +namespace doris::snii::stats { + +using format::DictEntry; +using format::NormsPodReader; +using format::RegionRef; + +namespace { + +// Resolves a term's DictEntry. *found=false for an absent term (OK status). +Status lookup_entry(const reader::LogicalIndexReader& idx, std::string_view term, bool* found, + DictEntry* entry) { + uint64_t frq_base = 0; + uint64_t prx_base = 0; + return idx.lookup(term, found, entry, &frq_base, &prx_base); +} + +} // namespace + +Status SniiStatsProvider::open(const reader::LogicalIndexReader* idx, SniiStatsProvider* out) { + return open_impl(idx, out, true); +} + +#ifdef BE_TEST +Status SniiStatsProvider::open_legacy_for_test(const reader::LogicalIndexReader* idx, + SniiStatsProvider* out) { + return open_impl(idx, out, false); +} +#endif + +Status SniiStatsProvider::open_impl(const reader::LogicalIndexReader* idx, SniiStatsProvider* out, + bool require_semantic_metadata) { + if (idx == nullptr || out == nullptr) { + return Status::Error("stats_provider: null argument"); + } + out->idx_ = idx; + const auto& sb = idx->stats(); + out->doc_count_ = sb.doc_count; + out->indexed_doc_count_ = sb.indexed_doc_count; + out->sum_total_term_freq_ = sb.sum_total_term_freq; + + const RegionRef& norms = idx->section_refs().norms; + const auto* metadata = idx->common_grams_metadata(); + if (metadata != nullptr) { + using namespace segment_v2::inverted_index; + RETURN_IF_ERROR(validate_snii_scoring_metadata( + metadata, sb.doc_count, sb.sum_total_term_freq, + idx->tier() == format::IndexTier::kT3, idx->has_positions(), norms.length != 0)); + out->doc_count_ = metadata->scoring_doc_count; + out->indexed_doc_count_ = metadata->scoring_doc_count; + out->sum_total_term_freq_ = metadata->scoring_token_count; + } else if (require_semantic_metadata) { + RETURN_IF_ERROR(segment_v2::inverted_index::validate_snii_scoring_metadata( + nullptr, sb.doc_count, sb.sum_total_term_freq, + idx->tier() == format::IndexTier::kT3, idx->has_positions(), norms.length != 0)); + } + if (norms.length == 0) { + out->has_norms_ = false; + return Status::OK(); + } + + RETURN_IF_ERROR(idx->open_norms(&out->norms_reader_)); + if (metadata != nullptr && out->norms_reader_.doc_count() != metadata->scoring_doc_count) { + return Status::Error( + "snii_stats: semantic norms doc count {} differs from scoring doc count {}", + out->norms_reader_.doc_count(), metadata->scoring_doc_count); + } + out->has_norms_ = true; + return Status::OK(); +} + +double SniiStatsProvider::avgdl() const { + const uint64_t denom = std::max(1, indexed_doc_count_); + return static_cast(sum_total_term_freq_) / static_cast(denom); +} + +Status SniiStatsProvider::doc_freq(std::string_view term, uint64_t* df) const { + if (df == nullptr) + return Status::Error("stats_provider: null df"); + *df = 0; + bool found = false; + DictEntry entry; + RETURN_IF_ERROR(lookup_entry(*idx_, term, &found, &entry)); + if (found) *df = entry.df; + return Status::OK(); +} + +Status SniiStatsProvider::total_term_freq(std::string_view term, uint64_t* ttf) const { + if (ttf == nullptr) + return Status::Error("stats_provider: null ttf"); + *ttf = 0; + bool found = false; + DictEntry entry; + RETURN_IF_ERROR(lookup_entry(*idx_, term, &found, &entry)); + if (!found) return Status::OK(); + // tier>=T2 entries carry the total term frequency directly in ttf_delta (the + // LogicalIndexWriter stores ttf there, not a delta from df). G16-f blocks + // (kNoTermStats: freq-dropped index) omit it -- fail with the semantic + // error instead of silently returning the default 0 (mixed old/new + // segments would otherwise disagree on the same term). + if (!entry.term_stats_present) { + return Status::Error( + "snii_stats: ttf requested but the dict block carries no term stats " + "(freq-dropped index)"); + } + *ttf = entry.ttf_delta; + return Status::OK(); +} + +Status SniiStatsProvider::encoded_norm(uint32_t docid, uint8_t* out) const { + if (out == nullptr) + return Status::Error("stats_provider: null out"); + if (!has_norms_) { + return Status::Error( + "stats_provider: index has no norms"); + } + return norms_reader_.try_encoded_norm(docid, out); +} + +} // namespace doris::snii::stats diff --git a/be/src/storage/index/snii/stats/snii_stats_provider.h b/be/src/storage/index/snii/stats/snii_stats_provider.h new file mode 100644 index 00000000000000..284a6ecc07cb71 --- /dev/null +++ b/be/src/storage/index/snii/stats/snii_stats_provider.h @@ -0,0 +1,92 @@ +// 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. + +#pragma once + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +// SniiStatsProvider -- exposes the native SNII scoring statistics required by +// BM25, sourced directly from the on-disk structures of one logical index: +// - semantic segment-level counts from CommonGrams scoring metadata. +// - per-term df / ttf from the term's DictEntry (resolved through the reader's +// lookup flow). The LogicalIndexWriter stores ttf directly in ttf_delta for +// tier>=T2 entries, so total_term_freq returns entry.ttf_delta. +// - per-doc length normalization byte (encoded_norm) from the norms POD, +// lazily loaded and validated once by LogicalIndexReader, then shared by +// every stats provider for that cached logical index. +// +// avgdl() = sum_total_term_freq / max(1, indexed_doc_count): the average document +// length used by BM25 length normalization. The provider performs no scoring; it +// only surfaces the statistics so query::Bm25Scorer can combine them. +namespace doris::snii::stats { + +class SniiStatsProvider { +public: + SniiStatsProvider() = default; + + // Binds to idx and acquires its shared validated norms view when the index + // carries scoring norms. idx must outlive this provider. Complete CommonGrams + // scoring metadata requires compatible semantic statistics and norms. + static Status open(const reader::LogicalIndexReader* idx, SniiStatsProvider* out); + +#ifdef BE_TEST + // Legacy scorer fixtures predate semantic scoring metadata. Production + // callers must use open() and fail closed when the proof is absent. + static Status open_legacy_for_test(const reader::LogicalIndexReader* idx, + SniiStatsProvider* out); +#endif + + // Segment-level semantic counts persisted by the SNII writer. + uint64_t doc_count() const { return doc_count_; } + uint64_t indexed_doc_count() const { return indexed_doc_count_; } + uint64_t sum_total_term_freq() const { return sum_total_term_freq_; } + + // Average document length: sum_total_term_freq / max(1, indexed_doc_count). + double avgdl() const; + + // Per-term document frequency. Absent term -> *df = 0 (OK status). + Status doc_freq(std::string_view term, uint64_t* df) const; + + // Per-term total term frequency (ttf = df + ttf_delta at tier>=T2). Absent + // term -> *ttf = 0 (OK status). + Status total_term_freq(std::string_view term, uint64_t* ttf) const; + + // 1-byte encoded doc-length norm for docid (raw byte from the norms POD). + // Out-of-range docid -> InvalidArgument; index without norms -> InvalidArgument. + Status encoded_norm(uint32_t docid, uint8_t* out) const; + + bool has_norms() const { return has_norms_; } + +private: + static Status open_impl(const reader::LogicalIndexReader* idx, SniiStatsProvider* out, + bool require_semantic_metadata); + + const reader::LogicalIndexReader* idx_ = nullptr; + uint64_t doc_count_ = 0; + uint64_t indexed_doc_count_ = 0; + uint64_t sum_total_term_freq_ = 0; + bool has_norms_ = false; + // Zero-copy view into the immutable bytes owned by idx_ after open_norms(). + format::NormsPodReader norms_reader_; +}; + +} // namespace doris::snii::stats diff --git a/be/src/storage/index/snii/writer/compact_posting_pool.cpp b/be/src/storage/index/snii/writer/compact_posting_pool.cpp new file mode 100644 index 00000000000000..9cef1a38b9f572 --- /dev/null +++ b/be/src/storage/index/snii/writer/compact_posting_pool.cpp @@ -0,0 +1,120 @@ +// 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 "storage/index/snii/writer/compact_posting_pool.h" + +#include +#include +#include + +namespace doris::snii::writer { + +// Gentle (~1.5x) many-level payload-capacity schedule. Starting at 5 bytes with a +// slow ramp keeps the over-allocated FINAL slice small for the millions of low-df +// terms (the dominant arena-overhead source) while still reaching multi-KiB slices +// for high-df chains in a bounded number of hops (so the per-slice 4-byte forward +// pointer stays a small fraction of a large chain's bytes). +const uint32_t CompactPostingPool::kSliceSizes[kLevelCount] = { + 5, 8, 12, 18, 27, 40, 60, 90, 135, 202, 303, 455, 683, 1024, 1536, 2304}; +const uint8_t CompactPostingPool::kNextLevel[kLevelCount] = {1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 15}; + +CompactPostingPool::CompactPostingPool() = default; + +uint32_t CompactPostingPool::kSliceSizes_level0() { + return kSliceSizes[0]; +} + +uint32_t CompactPostingPool::kSliceSize_at(int level) { + return kSliceSizes[level]; +} + +uint8_t CompactPostingPool::kNextLevel_at(int level) { + return kNextLevel[level]; +} + +void CompactPostingPool::reset() { + std::vector>().swap(blocks_); + next_offset_ = 0; + payload_bytes_ = 0; +} + +uint32_t CompactPostingPool::alloc_run(uint32_t bytes) { + const uint32_t in_block = next_offset_ & kBlockMask; + // A fresh block is needed when (a) there is no tail block yet, (b) the run does + // not fit in the current tail block's remaining space, or (c) next_offset_ sits + // exactly on a block boundary whose block has not been allocated (a previous run + // that exactly filled the tail leaves next_offset_ == blocks_.size()*kBlockSize, + // so in_block == 0 must NOT be mistaken for an empty fresh block). + const bool tail_exists = (next_offset_ >> kBlockShift) < blocks_.size(); + const bool need_block = !tail_exists || in_block + bytes > kBlockSize; + // Hard invariant (see arena_bytes()): the uint32 offset must never wrap. The spimi + // accumulator force-spills below 4 GiB, but enforce it here too -- in release as + // well as debug -- so any direct user of the pool fails loudly instead of silently + // aliasing block 0. We are a library: throw and let the caller decide how to + // handle it, rather than aborting the process. The run starts either in the + // current tail or at a new block's base; compute that start in 64 bits before the + // uint32 arithmetic can wrap. + const uint64_t run_start = + need_block ? static_cast(blocks_.size()) * kBlockSize : next_offset_; + if (run_start + bytes > UINT32_MAX) { + throw std::overflow_error( + "snii: CompactPostingPool arena exceeded the 4 GiB uint32 offset limit; " + "the caller must spill before this point"); + } + if (need_block) { + blocks_.emplace_back(kBlockSize, 0); + next_offset_ = static_cast((blocks_.size() - 1) * kBlockSize); + } + const uint32_t off = next_offset_; + next_offset_ += bytes; + return off; +} + +uint32_t CompactPostingPool::alloc_slice(int level, uint32_t* slice_end) { + const uint32_t cap = kSliceSizes[level]; + const uint32_t first = alloc_run(cap + kPtrBytes); + *slice_end = first + cap; + // Zero the forward pointer so a not-yet-extended tail slice reads next_head == 0. + std::memset(at(*slice_end), 0, kPtrBytes); + return first; +} + +uint32_t CompactPostingPool::read_ptr(uint32_t slice_end) const { + uint32_t v; + std::memcpy(&v, at(slice_end), sizeof(v)); + return v; +} + +void CompactPostingPool::write_ptr(uint32_t slice_end, uint32_t next_head) { + std::memcpy(at(slice_end), &next_head, sizeof(next_head)); +} + +uint32_t CompactPostingPool::start_chain(SliceWriter* w, uint8_t* level) { + *level = 0; + const uint32_t head = alloc_slice(0, &w->slice_end); + w->cur = head; + return head; +} + +CompactPostingPool::Cursor::Cursor(const CompactPostingPool* pool, uint32_t head, uint64_t budget) + : pool_(pool), cur_(head), level_(0), budget_(budget) { + // The first slice is level 0; its payload region ends kSliceSizes[0] bytes in. + slice_end_ = head + CompactPostingPool::kSliceSizes[0]; +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/compact_posting_pool.h b/be/src/storage/index/snii/writer/compact_posting_pool.h new file mode 100644 index 00000000000000..38adf4363b6528 --- /dev/null +++ b/be/src/storage/index/snii/writer/compact_posting_pool.h @@ -0,0 +1,333 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace doris::snii::writer { + +// SEGMENTED BYTE ARENA with per-term SLICED runs (a ByteBlockPool, after Lucene). +// +// WHY: the SPIMI accumulator's bulk memory is the per-term posting bytes. Backing +// each term with its own std::vector pays two taxes that dominate peak +// RSS at scale: (1) geometric-growth doubling slack (~1.17x of the live payload), +// and (2) a 24-32 B vector/struct header per term (hundreds of thousands of +// terms). This pool removes both: all term bytes live in a few large fixed-size +// blocks (so slack is ~one block, amortized to ~1.05x), and a term needs only two +// 32-bit cursors of live state (chain head for reads + write head for appends). +// +// HOW (slices): a term's bytes are not stored contiguously. They live in a chain +// of SLICES of geometrically growing payload capacity (the kSliceSizes schedule: +// 4, 8, 16, ... bytes of payload). Each slice is laid out as +// [ payload bytes ... ][ 4-byte forward pointer ] +// The forward pointer holds the absolute offset of the next slice's first payload +// byte (0 while the slice is still the tail of the chain). When a slice's payload +// region fills, the writer allocates a larger slice, stores its head into the old +// slice's 4 pointer bytes, and keeps appending. A reader walks the chain by +// reading payload bytes until a slice boundary, then following the pointer. +// +// Both writer and reader recompute each slice's capacity from the chain's slice +// INDEX (0, 1, 2, ...) via the deterministic schedule, so neither needs to store +// per-slice sizes. The writer carries the current slice's end offset in its +// SliceWriter handle; the reader recomputes capacities as it advances. +// +// Offsets are GLOBAL absolute byte indices into the logical concatenation of all +// blocks: offset = block_index * kBlockSize + byte_in_block. kBlockSize is a power +// of two, so offset -> (block, byte) is a shift/mask. +class CompactPostingPool { +public: + // Block size (power of two). 32 KiB blocks keep per-block tail waste tiny (it + // matters at the smaller 1M scale where the whole arena is only tens of MiB) and + // bound the outer vector header cost; at the 5M scale a few thousand + // blocks is still cheap. Empirically the lowest peak across both scales. + static constexpr uint32_t kBlockShift = 15; + static constexpr uint32_t kBlockSize = 1u << kBlockShift; // 32 KiB + static constexpr uint32_t kBlockMask = kBlockSize - 1; + + // Per-slice forward-pointer width (absolute uint32 next-slice offset). + static constexpr uint32_t kPtrBytes = 4; + + // Geometric slice payload-capacity schedule and the level transition. Level i + // slices hold kSliceSizes[i] payload bytes; on overflow the chain advances to + // kNextLevel[i] (capping at the largest level). A GENTLE (~1.5x) many-level + // schedule starting small minimizes the over-allocated final slice (the + // dominant arena overhead) while keeping the per-slice forward-pointer count + // bounded for high-df chains. + static constexpr int kLevelCount = 16; + + CompactPostingPool(); + + CompactPostingPool(const CompactPostingPool&) = delete; + CompactPostingPool& operator=(const CompactPostingPool&) = delete; + + // Payload capacity (bytes) of a fresh level-0 slice. Exposed for tests that need + // to fill exactly one slice without hardcoding the schedule. + static uint32_t kSliceSizes_level0(); + + // Payload capacity of the slice at `level`, and the level a chain advances to when + // that slice overflows. Exposed (like kSliceSizes_level0) so tests can simulate the + // arena's bump allocator exactly -- e.g. to construct an EXACT block-boundary fill -- + // without hardcoding the private schedule. `level` must be in [0, kLevelCount). + static uint32_t kSliceSize_at(int level); + static uint8_t kNextLevel_at(int level); + + // Live append handle for one term's chain. POD, 8 bytes: the absolute write + // cursor and the absolute end of the current slice's payload region. The chain's + // current slice LEVEL is kept by the caller (a uint8, packed alongside its other + // flags) so this handle stays 8 bytes -- shaving the per-term accumulator. `head` + // (the chain's first payload offset) is also stored by the CALLER (the read entry + // point); start_chain returns it. + struct SliceWriter { + uint32_t cur = 0; // next byte to write (absolute) + uint32_t slice_end = 0; // one-past-last payload byte of the current slice + }; + + // Begins a fresh chain, initializing `w` to its first (level-0) slice and + // *level to 0, and returns the chain head (absolute first payload offset). + uint32_t start_chain(SliceWriter* w, uint8_t* level); + + // Appends one payload byte to the chain described by `w` / `*level`, growing the + // chain with a new linked slice (and advancing *level) when the current slice's + // payload region is exhausted. + void append_byte(SliceWriter* w, uint8_t* level, uint8_t value); + // Encodes and appends one unsigned LEB128 value. The common case copies all + // encoded bytes into the current slice in one operation; only a slice-boundary + // value enters the cold allocation loop. + void append_varint(SliceWriter* w, uint8_t* level, uint64_t value); + + // Total live payload bytes ever written across all chains (excludes slice + // forward-pointer overhead). Drives the spill-threshold estimate only. + uint64_t payload_bytes() const { return payload_bytes_; } + + // Bytes the arena currently occupies (block_count * kBlockSize). The pool + // addresses bytes with a uint32 offset (next_offset_), so the arena MUST stay + // below 4 GiB or alloc_run wraps and silently aliases block 0. The accumulator + // watches this to force a safety spill before the wrap; alloc_run also enforces it + // directly (throws std::overflow_error on a would-be wrap) so a direct user of the + // pool fails loudly rather than silently corrupting. + // Hard invariant: a single CompactPostingPool never exceeds UINT32_MAX bytes. + uint64_t arena_bytes() const { return static_cast(blocks_.size()) << kBlockShift; } + + // Releases ALL blocks back to the OS. Called after the accumulator is fully + // drained (or before a spill's next fill) so no input-side bytes stay resident. + void reset(); + + // ---- Reader ---------------------------------------------------------------- + // Forward cursor over one term's chain, yielding its payload bytes in write + // order by walking the slice forward pointers. + // + // CONTRACT of the `budget` ctor argument (single, unambiguous meaning): + // `budget` is an UPPER BOUND on the number of bytes this cursor may yield. It + // is NOT required to equal the exact payload length: passing the exact length + // is fine, and so is passing any value >= it (the production caller passes the + // chain's write-head offset, which always bounds the payload from above). The + // cursor is SELF-TERMINATING: once it walks off the last written byte it sees + // the tail slice's zero forward pointer and stops, regardless of how much + // budget remains. So an over-large budget can never make next() read past the + // chain (no aliasing of block 0, no off-chain access) -- the budget is purely a + // secondary cap. has_next() is therefore a reliable "more bytes remain" + // predicate for ANY budget >= the true length: it becomes false at the smaller + // of (budget exhausted, chain tail reached). + class Cursor { + public: + Cursor(const CompactPostingPool* pool, uint32_t head, uint64_t budget); + + // True while the cursor can still yield a REAL payload byte: the budget is not + // spent AND the cursor has not reached the chain tail. It peeks the tail forward + // pointer at a slice boundary so it never reports a phantom trailing byte, making + // has_next()/next() a safe loop for any budget >= the true payload length. + bool has_next() const; + // Yields the next payload byte. Returns 0 (and yields no more) once the chain + // tail is reached or the budget is spent -- never reads past the chain. + uint8_t next(); + // Decodes one unsigned LEB128 value directly from the current slice, following + // a slice pointer only when the encoded value straddles a boundary. + uint64_t read_varint(); + + private: + const CompactPostingPool* pool_; + uint32_t cur_; // absolute read cursor + uint32_t slice_end_; // one-past-last payload byte of the current slice + uint32_t level_; // current slice level + uint64_t budget_; // remaining byte budget (upper bound on bytes to yield) + }; + + // Builds a cursor over the chain at `head`. `budget` is an UPPER BOUND on bytes to + // read (see Cursor's contract): the exact payload length or anything larger. The + // production caller passes the write-head offset, which always bounds the payload + // from above; the cursor self-terminates at the chain tail regardless. + Cursor cursor(uint32_t head, uint64_t budget) const { return Cursor(this, head, budget); } + +private: + static const uint32_t kSliceSizes[kLevelCount]; + static const uint8_t kNextLevel[kLevelCount]; + + uint8_t* at(uint32_t off) { return &blocks_[off >> kBlockShift][off & kBlockMask]; } + const uint8_t* at(uint32_t off) const { return &blocks_[off >> kBlockShift][off & kBlockMask]; } + + // Reads/writes the 4-byte forward pointer at the END of a slice whose payload + // region ends at `slice_end` (pointer occupies [slice_end, slice_end+4)). + uint32_t read_ptr(uint32_t slice_end) const; + void write_ptr(uint32_t slice_end, uint32_t next_head); + + // Reserves `bytes` contiguous bytes from the arena tail (a fresh block if the + // current tail cannot hold them) and returns the first reserved absolute offset. + // `bytes` must be <= kBlockSize. + uint32_t alloc_run(uint32_t bytes); + + // Allocates a slice at `level` (payload region + 4 pointer bytes), zeroes its + // forward pointer, and returns the first payload offset; sets *slice_end. + uint32_t alloc_slice(int level, uint32_t* slice_end); + + std::vector> blocks_; // fixed kBlockSize blocks + uint32_t next_offset_ = 0; // global bump pointer (absolute) into the tail block + uint64_t payload_bytes_ = 0; +}; + +// ---- Inlined per-byte hot paths -------------------------------------------- +// append_byte (one call per encoded payload byte during SPIMI ingest) and +// Cursor::has_next/next (one call per arena byte during finalize drain/merge) +// are the writer's two hottest per-byte loops. Their callers live in a DIFFERENT +// translation unit (spimi_term_buffer.cpp), and the build has no LTO/IPO/unity +// build, so an out-of-line .cpp definition forces a non-inlinable cross-TU call +// (plus stack spill of cur_/slice_end_/budget_) on every single byte. Defining +// the bodies inline HERE lets the caller's TU inline them, keeping the cursor +// state in registers and eliminating the per-byte call/ret. They are placed +// AFTER the class so every private member they touch (at, read_ptr, alloc_slice, +// write_ptr, kSliceSizes, kNextLevel, payload_bytes_) is already declared and +// visible; Cursor is a nested class, so it may use the enclosing pool's privates +// (read_ptr/at) through pool_. The cold slice-overflow helpers (alloc_slice, +// write_ptr) stay out-of-line in the .cpp -- only the per-byte work is inlined. +// Pure code move from the .cpp: behavior and produced bytes are identical. + +inline void CompactPostingPool::append_byte(SliceWriter* w, uint8_t* level, uint8_t value) { + if (w->cur == w->slice_end) { + // Current slice payload region is full: grow the chain with a larger slice and + // record the link in the old slice's trailing pointer bytes. + const uint8_t next_level = kNextLevel[*level]; + uint32_t new_end = 0; + const uint32_t new_head = alloc_slice(next_level, &new_end); + write_ptr(w->slice_end, new_head); + *level = next_level; + w->cur = new_head; + w->slice_end = new_end; + } + *at(w->cur) = value; + ++w->cur; + ++payload_bytes_; +} + +inline void CompactPostingPool::append_varint(SliceWriter* w, uint8_t* level, uint64_t value) { + std::array encoded {}; + size_t size = 0; + do { + uint8_t byte = static_cast(value & 0x7fU); + value >>= 7; + if (value != 0) { + byte |= 0x80U; + } + encoded[size++] = byte; + } while (value != 0); + + size_t copied = 0; + while (copied < size) { + if (w->cur == w->slice_end) { + const uint8_t next_level = kNextLevel[*level]; + uint32_t new_end = 0; + const uint32_t new_head = alloc_slice(next_level, &new_end); + write_ptr(w->slice_end, new_head); + *level = next_level; + w->cur = new_head; + w->slice_end = new_end; + } + const size_t available = w->slice_end - w->cur; + const size_t count = std::min(size - copied, available); + std::memcpy(at(w->cur), encoded.data() + copied, count); + w->cur += static_cast(count); + copied += count; + payload_bytes_ += count; + } +} + +inline bool CompactPostingPool::Cursor::has_next() const { + if (budget_ == 0) { + return false; + } + // At a slice boundary, the chain continues only if the forward pointer is non-zero; + // a zero pointer is the tail marker (offset 0 is never a valid next-slice head). Peek + // it so has_next() never reports a phantom byte that next() would have to fabricate. + if (cur_ == slice_end_) { + return pool_->read_ptr(slice_end_) != 0; + } + return true; +} + +inline uint8_t CompactPostingPool::Cursor::next() { + // Budget guard: the caller's stated upper bound is spent -- yield nothing more. + if (budget_ == 0) { + return 0; + } + if (cur_ == slice_end_) { + // Reached this slice's payload boundary. Follow the forward pointer to the next + // slice -- UNLESS it is zero, which marks the CHAIN TAIL (offset 0 is always the + // pool's very first slice, never a valid *next*-slice head, so a zero pointer is + // unambiguously "no more slices"). Without this tail check, an over-reading caller + // would follow the zero pointer to offset 0 and alias block 0's bytes (or read an + // unallocated block) -- UB. Stopping here makes the cursor self-terminating and + // safe regardless of how large a budget the caller passed. + const uint32_t next_head = pool_->read_ptr(slice_end_); + if (next_head == 0) { + budget_ = 0; // chain exhausted: no further bytes exist + return 0; + } + level_ = CompactPostingPool::kNextLevel[level_]; + cur_ = next_head; + slice_end_ = next_head + CompactPostingPool::kSliceSizes[level_]; + } + const uint8_t v = *pool_->at(cur_); + ++cur_; + --budget_; + return v; +} + +inline uint64_t CompactPostingPool::Cursor::read_varint() { + uint64_t result = 0; + uint32_t shift = 0; + for (;;) { + uint8_t byte; + if (budget_ != 0 && cur_ != slice_end_) { + byte = *pool_->at(cur_); + ++cur_; + --budget_; + } else { + byte = next(); + } + result |= static_cast(byte & 0x7fU) << shift; + if ((byte & 0x80U) == 0) { + return result; + } + shift += 7; + } +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/global_memory_limiter.cpp b/be/src/storage/index/snii/writer/global_memory_limiter.cpp new file mode 100644 index 00000000000000..ed8038b03d3b49 --- /dev/null +++ b/be/src/storage/index/snii/writer/global_memory_limiter.cpp @@ -0,0 +1,166 @@ +// 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 "storage/index/snii/writer/global_memory_limiter.h" + +#include +#include +#include + +#include "common/logging.h" + +namespace doris::snii::writer { + +GlobalMemoryLimiter* GlobalMemoryLimiter::instance() { + // Intentionally leaked (never destroyed): buffers un-register from their + // destructors, which may run during static teardown at process exit -- a + // destroyed registry there would be use-after-free. A leaked singleton + // makes the "limiter outlives every attached buffer" contract + // unconditional. + static auto* g_instance = new GlobalMemoryLimiter(); + return g_instance; +} + +void GlobalMemoryLimiter::register_buffer(std::atomic* spill_flag, int64_t resident_bytes, + int64_t arena_bytes) { + std::lock_guard guard(mutex_); + Entry& slot = entries_[spill_flag]; // zero-initialized on first insert + total_bytes_locked_ += resident_bytes - slot.resident; + slot.resident = resident_bytes; + slot.arena = arena_bytes; +} + +void GlobalMemoryLimiter::report(std::atomic* spill_flag, int64_t resident_bytes, + int64_t arena_bytes) { + std::lock_guard guard(mutex_); + auto it = entries_.find(spill_flag); + if (it == entries_.end()) { + // Not registered (or already unregistered): ignore rather than + // resurrect an entry nobody will remove. + return; + } + total_bytes_locked_ += resident_bytes - it->second.resident; + it->second.resident = resident_bytes; + it->second.arena = arena_bytes; + request_spills_locked(); +} + +void GlobalMemoryLimiter::unregister_buffer(std::atomic* spill_flag) { + std::lock_guard guard(mutex_); + auto it = entries_.find(spill_flag); + if (it == entries_.end()) { + return; + } + total_bytes_locked_ -= it->second.resident; + entries_.erase(it); +} + +int64_t GlobalMemoryLimiter::total_bytes() const { + std::lock_guard guard(mutex_); + return total_bytes_locked_; +} + +size_t GlobalMemoryLimiter::registered_count() const { + std::lock_guard guard(mutex_); + return entries_.size(); +} + +bool GlobalMemoryLimiter::budget_degraded() const { + std::lock_guard guard(mutex_); + return degraded_; +} + +void GlobalMemoryLimiter::request_spills_locked() { + const int64_t budget = budget_bytes_.load(std::memory_order_relaxed); + if (budget <= 0 || total_bytes_locked_ <= budget) { + return; + } + // BUDGET SANITY (graceful degradation): the budget bounds SPILLABLE + // memory only -- a forced spill reclaims a buffer's posting arena, never + // its persistent vocab / pair-map structures. When the per-writer budget + // share falls below the minimum useful working set, those persistent + // bytes alone exceed the budget and no amount of forced spilling can get + // under it; flagging would only manufacture a storm of floor-sized runs + // for zero net relief (the conc=16 wikipedia field failure). Log ONCE per + // degradation episode and stop flagging until the ratio recovers (writers + // draining/unregistering, or a raised budget). + const auto writer_count = static_cast(entries_.size()); + const int64_t min_useful = min_useful_budget_per_writer_bytes_.load(std::memory_order_relaxed); + if (writer_count > 0 && min_useful > 0 && budget / writer_count < min_useful) { + if (!degraded_) { + degraded_ = true; + LOG(WARNING) << "SNII global index-writer memory budget cannot be met by spilling: " + << "budget=" << budget << " B across " << writer_count + << " registered writers is below the minimum useful " << min_useful + << " B/writer (persistent per-writer structures dominate; the budget " + << "bounds SPILLABLE memory, not persistent memory). Forced spilling is " + << "suspended until writers drain or the budget " + << "(snii_index_writer_global_memory_bytes) is raised."; + } + return; + } + if (degraded_) { + degraded_ = false; // episode over; a relapse will log once again + LOG(INFO) << "SNII global index-writer memory budget recovered (" << budget << " B / " + << writer_count << " writers); forced spilling resumes."; + } + const int64_t overage = total_bytes_locked_ - budget; + // FORCED-SPILL FLOOR + PER-BUFFER COOLDOWN: only buffers holding at least + // the floor of RECLAIMABLE arena are eligible victims -- flagging a + // smaller arena would cut a tiny run and reclaim next to nothing, and a + // buffer that just honored a forced spill (arena ~0) stays exempt until + // its arena regrows past the floor. Never below one byte: an empty arena + // has nothing to write to a run. + const int64_t victim_floor = + std::max(min_victim_arena_bytes_.load(std::memory_order_relaxed), 1); + // Largest RECLAIMABLE consumers first: victims are ranked by their + // spillable ARENA (what the forced spill frees), not the + // persistent-dominated resident total. n is the live writer count of the + // process (at most a few hundred), and this only runs while over budget, + // so the sort under the mutex is bounded, allocation-light work. + std::vector*>> by_arena; + by_arena.reserve(entries_.size()); + for (const auto& [flag, entry] : entries_) { + if (entry.arena >= victim_floor) { + by_arena.emplace_back(entry.arena, flag); + } + } + std::sort(by_arena.begin(), by_arena.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + int64_t covered = 0; + for (const auto& [arena, flag] : by_arena) { + if (covered >= overage) { + break; + } + // An ALREADY-pending flag (set by an earlier report, owner not yet at + // its next token) counts toward the covered sum without a fresh store: + // re-flagging it would be a no-op, and skipping the store avoids + // dirtying the owner's cache line every over-budget report. + if (!flag->load(std::memory_order_relaxed)) { + flag->store(true, std::memory_order_relaxed); + } + // Count the ARENA toward coverage: it is all a forced spill of this + // victim can actually reclaim. When the persistent remainder keeps the + // sum over budget even after every eligible arena is flagged, the loop + // simply flags them all -- each still cuts a >= floor-sized run, and + // the cooldown above keeps any one buffer from being re-victimized + // before it has a floor's worth of arena again. + covered += arena; + } +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/global_memory_limiter.h b/be/src/storage/index/snii/writer/global_memory_limiter.h new file mode 100644 index 00000000000000..2f19168fc3a4a8 --- /dev/null +++ b/be/src/storage/index/snii/writer/global_memory_limiter.h @@ -0,0 +1,201 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include +#include + +namespace doris::snii::writer { + +// Process-wide SNII build-RAM limiter (G09) -- the index-build analogue of +// Doris's MemTableMemoryLimiter. Every live SPIMI accumulator registers here +// and forwards its ACCURATE resident-byte total (G08) plus its SPILLABLE +// arena bytes through its existing debounced report path; when the resident +// sum across ALL writers of the process (tablets x segments x concurrent +// loads) exceeds the configured budget, the limiter requests spills from the +// largest-ARENA eligible buffers until the flagged (reclaimable) arena sum +// covers the overage. +// +// WHY: the per-writer gate-2 cap (e.g. 512 MiB) bounds ONE writer, but a load +// keeps (tablets x concurrency) writers alive at once -- wikipedia at +// concurrency 16 held 100+ writers at 300-500 MB each (~41 GiB), none of which +// ever reached its own cap, so per-writer spilling never fired. This registry +// bounds the SUM. +// +// ASYNC-SAFE REQUESTS: the SPIMI structures are single-threaded, so the +// limiter must never spill on the reporting thread. A request is a relaxed +// atomic FLAG on the target buffer (SpimiTermBuffer::global_spill_requested_) +// that the OWNER's next add_token / maybe_spill_after_token observes and +// honors on its own thread (bypassing the G08 per-writer anti-churn floor but +// still requiring the FORCED-SPILL FLOOR of reclaimable arena -- see below -- +// so every forced run is worth its fixed costs). Flags are ADVISORY: the +// owner may have just spilled or drained -- the flag is then a (harmless) +// no-op or one extra floor-sized run. The limiter itself only ever takes its +// registry mutex and flips atomics; it never blocks a reporting thread beyond +// that mutex and never calls back into a buffer. +// +// LIFETIME: buffers un-register in their destructor. register / report / +// unregister all serialize on the registry mutex, and flags are only ever set +// UNDER that mutex, so once unregister_buffer returns no thread can touch the +// (about-to-die) flag again. The limiter must outlive every attached buffer +// (trivial for the process singleton; test-local instances are declared before +// the buffers they serve). +// +// THE BUDGET BOUNDS SPILLABLE MEMORY, NOT PERSISTENT MEMORY: a forced spill +// releases only the buffer's posting ARENA; the persistent vocab / pair-map / +// slot structures (~100-500 MB per wikipedia writer) survive it. If the +// persistent bytes of all writers alone exceed the budget, NO amount of +// spilling can reach the budget -- the budget is a back-pressure valve over +// the reclaimable arenas, not a hard cap on resident RSS. Three defenses keep +// an unreachable budget from degenerating into a forced-spill storm (the +// conc=16 wikipedia field failure: every report re-flagged every buffer, each +// honoring with one 32 KiB arena block -> thousands of tiny runs per buffer -> +// EMFILE re-opening them for the k-way merge -> failed loads): +// * VICTIMS BY ARENA: victims are selected by their reported SPILLABLE arena +// bytes -- the only bytes a forced spill can actually reclaim -- never by +// the persistent-dominated resident total, and only buffers whose arena is +// at least min_victim_arena_bytes (config snii_forced_spill_min_arena_bytes) +// are eligible. Every forced run is therefore at least floor-sized. +// * PER-BUFFER COOLDOWN: right after a buffer honors a forced spill its +// arena is ~0, below the victim floor, so it is EXEMPT from new flags +// until the arena regrows past the floor. No timer state: the eligibility +// rule IS the cooldown. +// * BUDGET SANITY: when budget / registered_count falls below a per-writer +// minimum useful share (kMinUsefulBudgetPerWriterBytes), persistent +// structures alone dominate and the budget is provably unreachable by +// spilling; the limiter LOGS ONCE (per degradation episode) and stops +// flagging until the ratio recovers (writers finishing / budget raised). +class GlobalMemoryLimiter { +public: + // Victim floor default (mirrors config snii_forced_spill_min_arena_bytes): + // a buffer is only ever asked to force-spill once its reclaimable arena + // holds at least this much, so no forced run is smaller than this. + static constexpr int64_t kDefaultMinVictimArenaBytes = 64LL << 20; // 64 MiB + // Minimum useful per-writer budget share: below budget/registered_count == + // this, spilling cannot meet the budget (persistent per-writer structures + // alone exceed the share) -- the limiter degrades to log-once-and-stop. + static constexpr int64_t kMinUsefulBudgetPerWriterBytes = 96LL << 20; // 96 MiB + + // Local instances are constructible for unit tests; production code uses + // the process singleton below. + GlobalMemoryLimiter() = default; + GlobalMemoryLimiter(const GlobalMemoryLimiter&) = delete; + GlobalMemoryLimiter& operator=(const GlobalMemoryLimiter&) = delete; + + // Process singleton (never destroyed before the writers that use it). + static GlobalMemoryLimiter* instance(); + + // Budget in bytes across every registered buffer; <= 0 disables the + // limiter (registration still tracked, but no flag is ever set). Refreshed + // from the mutable BE config at each writer init, so a config change takes + // effect for the whole registry at the next writer creation. + void set_budget_bytes(int64_t budget_bytes) { + budget_bytes_.store(budget_bytes, std::memory_order_relaxed); + } + int64_t budget_bytes() const { return budget_bytes_.load(std::memory_order_relaxed); } + + // Victim-eligibility floor over a buffer's reported SPILLABLE arena bytes + // (see the class comment). Refreshed from the mutable BE config alongside + // the budget at each writer init. Values < 1 behave as 1 (an empty arena + // is never a victim -- there would be nothing to write to the run). + void set_min_victim_arena_bytes(int64_t bytes) { + min_victim_arena_bytes_.store(bytes, std::memory_order_relaxed); + } + int64_t min_victim_arena_bytes() const { + return min_victim_arena_bytes_.load(std::memory_order_relaxed); + } + + // Per-writer minimum useful budget share for the degradation check. + // Production keeps the default (kMinUsefulBudgetPerWriterBytes); tests + // lower it to exercise selection with synthetic byte scales. + void set_min_useful_budget_per_writer_bytes(int64_t bytes) { + min_useful_budget_per_writer_bytes_.store(bytes, std::memory_order_relaxed); + } + + // True while the limiter is in the degraded log-once-and-stop state (the + // per-writer budget share fell below the useful minimum at the last + // over-budget report). Observability / tests. + bool budget_degraded() const; + + // Adds `spill_flag` (the owning buffer's advisory request flag; also the + // entry's identity) with its current resident bytes and its SPILLABLE + // arena bytes (<= resident; what a forced spill can reclaim, the victim + // selection key). Re-registering an already-registered flag just updates + // its bytes. Never sets flags itself: a single registration cannot create + // NEW overage worth reacting to before the buffer's first report. + void register_buffer(std::atomic* spill_flag, int64_t resident_bytes, + int64_t arena_bytes); + + // Updates the entry's resident and spillable-arena bytes (ABSOLUTE totals, + // not deltas -- self-healing across any missed report). When the + // registered RESIDENT sum exceeds the budget, sets the request flags of + // the largest-ARENA eligible entries (arena >= the victim floor; see the + // class comment) -- counting entries whose flag is ALREADY pending toward + // the covered sum, so an in-flight request is not amplified -- until the + // flagged ARENA bytes cover the overage or eligible entries run out. A + // report for a flag that is not registered is ignored. + void report(std::atomic* spill_flag, int64_t resident_bytes, int64_t arena_bytes); + + // Removes the entry (subtracting its bytes from the total). After this + // returns, the limiter never touches `spill_flag` again -- safe to destroy + // the owning buffer. + void unregister_buffer(std::atomic* spill_flag); + + // Registered sum / entry count, for tests and observability. + int64_t total_bytes() const; + size_t registered_count() const; + +private: + // One registered buffer's last reported byte totals. `resident` feeds the + // over-budget decision (the sum the budget bounds); `arena` -- the + // SPILLABLE subset a forced spill can actually reclaim -- feeds victim + // selection and eligibility. + struct Entry { + int64_t resident = 0; + int64_t arena = 0; + }; + + // Called with mutex_ held whenever the total may exceed the budget: after + // the budget-sanity check (degrade to log-once-and-stop when the + // per-writer share is below the useful minimum), sorts the ELIGIBLE + // entries (arena >= victim floor) by ARENA descending and flags from the + // top until the flagged arena sum covers (total - budget) or eligible + // entries run out. O(n log n) over the live writer count (at most a few + // hundred) -- bounded work under the mutex, no I/O, no callbacks. + void request_spills_locked(); + + mutable std::mutex mutex_; + std::atomic budget_bytes_ {0}; + std::atomic min_victim_arena_bytes_ {kDefaultMinVictimArenaBytes}; + std::atomic min_useful_budget_per_writer_bytes_ {kMinUsefulBudgetPerWriterBytes}; + // All below guarded by mutex_. total_bytes_locked_ is the maintained sum + // of entries_ resident values (avoids an O(n) walk per report). + int64_t total_bytes_locked_ = 0; + // Degradation episode latch: set (with ONE warning log) when the budget + // sanity check fails, cleared when a later over-budget report finds the + // ratio recovered -- so a relapse logs again, but a sustained episode + // logs exactly once. + bool degraded_ = false; + phmap::flat_hash_map*, Entry> entries_; +}; + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/logical_index_writer.cpp b/be/src/storage/index/snii/writer/logical_index_writer.cpp new file mode 100644 index 00000000000000..63cd04bd3914c9 --- /dev/null +++ b/be/src/storage/index/snii/writer/logical_index_writer.cpp @@ -0,0 +1,952 @@ +// 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 "storage/index/snii/writer/logical_index_writer.h" + +#include +#include +#include +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/bsbf.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/writer/posting_window_emitter.h" + +namespace doris::snii::writer { + +using format::BlockRef; +using format::DictBlockBuilder; +using format::DictBlockDirectoryBuilder; +using format::DictEntry; +using format::DictEntryEnc; +using format::DictEntryKind; +using format::SampledTermIndexBuilder; +using format::SectionRefs; +using segment_v2::inverted_index::CG_V1_MARKER; +using segment_v2::inverted_index::CommonGramsCoverage; +using segment_v2::inverted_index::ScoringCoverage; +using segment_v2::inverted_index::validate_common_grams_segment_metadata; + +namespace { + +// Target false-positive probability for the block-split bloom XFilter. Sizes +// the filter via Parquet OptimalNumOfBytes; L0 keeps the probe in memory and L1 +// keeps the per-query cost at one 32-byte block. +constexpr double kBsbfFpp = 0.01; +// Force-raw level for .frq dd/freq regions. Their plaintext is PFOR-bit-packed +// doc-deltas/freqs -- already high-entropy, so zstd shrinks ~30 MB of input by +// <0.1 MiB while burning ~0.4s CPU (and an extra crc pass over the compressed +// bytes) at 5M. We force raw here and keep zstd only on .prx (which compresses +// ~77%). Output stays self-describing: the region meta records zstd=false. +constexpr int kRawFrqRegion = 0; +// zstd level for whole-DICT-block compression comes from +// SniiIndexInput::dict_block_zstd_level (default 3: ~40% on the 64KiB +// front-coded blocks at ~120 MiB/s encode / ~600 MiB/s decode; higher levels +// trade import CPU for size, decode speed unchanged). G16-h made it (and the +// .prx auto level) caller-tunable. + +using format::FrqRegionMeta; + +// Fused single-pass term-level freq statistics: total_freq (running sum) and +// max_freq (running max) in ONE scan, reused by validate_term (has_prx +// position-count budget), stats_.sum_total_term_freq, and the DictEntry +// ttf_delta/max_freq. Byte-identical to the former separate SumOf/MaxOf scans: +// same left-to-right accumulation order and the same max init of 0, so a freq of +// 0 never lowers the max. Complete CommonGrams entries bypass this helper: +// their ttf is the already-known PRX position count and max_freq is not stored. +FreqStats fuse_freq_stats(const std::vector& freqs) { +#ifdef BE_TEST + testing::note_term_freq_scan(); +#endif + FreqStats fs; + for (uint32_t f : freqs) { + fs.total_freq += f; + fs.max_freq = std::max(f, fs.max_freq); + } + return fs; +} + +// Default window doc count by df: high-df windowed terms combine kFrqBaseUnit +// units into larger (kAdaptiveWindowDocs) windows. PRX limits may subsequently +// recut one of these default windows at document boundaries. +uint32_t adaptive_window_docs(uint32_t df) { + return df >= format::kAdaptiveWindowDfThreshold ? format::kAdaptiveWindowDocs + : format::kFrqBaseUnit; +} + +bool fits_prx_window_shape(uint64_t doc_count, uint64_t position_count, + const format::PrxWindowLimits& limits) { + return doc_count <= limits.max_docs && position_count <= limits.max_positions; +} + +} // namespace + +// The only encoder for TermPostingSource input. It borrows the writer's reusable +// posting buffer, streams PRX windows directly to the final sink, and stages grouped DD +// and frequency regions without retaining the complete term. +class StreamingTermEncoder { +public: + StreamingTermEncoder(LogicalIndexWriter* writer, StreamedTermPostings* postings, + bool declared_common_gram, bool term_has_freq, bool term_has_prx, + TermPostingBuffer* buffer, uint64_t frq_base, uint64_t prx_base) + : writer_(writer), + postings_(postings), + declared_common_gram_(declared_common_gram), + term_has_freq_(term_has_freq), + term_has_prx_(term_has_prx), + buffer_(buffer), + frq_base_(frq_base), + prx_base_(prx_base), + emitter_(WindowEmitterOptions { + .posting_out = writer->posting_out_, + .posting_region_offset = writer->posting_off0_, + .frq_base = frq_base, + .prx_base = prx_base, + .encoded_norms = writer->encoded_norms_, + .has_freq = term_has_freq, + .has_prx = term_has_prx, + .prx_zstd_level = writer->prx_zstd_level_, + .prx_window_limits = writer->prx_window_limits_, + .term_frequency_source = + declared_common_gram ? (postings->retain_positions + ? TermFrequencySource::kPositions + : TermFrequencySource::kDocuments) + : TermFrequencySource::kFrequenciesOrDocuments, + .memory_reporter = writer->memory_reporter_, + }) { + DCHECK(buffer_ != nullptr); + DCHECK(buffer_->empty()); + } + + ~StreamingTermEncoder() { + buffer_->clear_reuse_and_release_excess(format::kAdaptiveWindowDfThreshold); + } + + Status encode(DictEntry* entry, FreqStats* stats) { + if (postings_->source == nullptr) { + return Status::Error( + "logical_index: streamed term has a null posting source"); + } + if (writer_->has_prx_ && !postings_->retain_positions && !declared_common_gram_) { + return Status::Error( + "logical_index: only a declared CommonGrams term may omit positions"); + } +#ifdef BE_TEST + if (!declared_common_gram_) { + testing::note_term_freq_scan(); + } +#endif + bool exhausted = false; + RETURN_IF_ERROR(fill(format::kAdaptiveWindowDfThreshold, &exhausted)); + if (exhausted && buffer_->document_count() < format::kAdaptiveWindowDfThreshold) { + entry->term = std::move(postings_->term); + entry->df = total_docs_; + entry->ttf_delta = stats_.total_freq; + entry->max_freq = stats_.max_freq; + RETURN_IF_ERROR(encode_small(entry)); + *stats = stats_; + return Status::OK(); + } + + if (!buffer_->empty()) { + RETURN_IF_ERROR(encode_windowed_buffer(format::kAdaptiveWindowDocs)); + } + while (!exhausted) { + RETURN_IF_ERROR(fill(format::kAdaptiveWindowDocs, &exhausted)); + if (!buffer_->empty()) { + RETURN_IF_ERROR(encode_windowed_buffer(format::kAdaptiveWindowDocs)); + } + } + + entry->term = std::move(postings_->term); + entry->df = total_docs_; + entry->ttf_delta = stats_.total_freq; + entry->max_freq = stats_.max_freq; + RETURN_IF_ERROR(finish_windowed(entry)); + *stats = stats_; + return Status::OK(); + } + +private: + Status fill(uint32_t target_docs, bool* exhausted) { + buffer_->clear_reuse(); + position_offsets_.clear(); + *exhausted = false; + RETURN_IF_ERROR(postings_->source->fill(target_docs, buffer_, exhausted)); + const size_t count = buffer_->document_count(); + if (count > target_docs) { + return Status::Error( + "logical_index: posting source exceeded target_docs"); + } + if (!*exhausted && count != target_docs) { + return Status::Error( + "logical_index: posting source returned a short non-terminal fill"); + } + if (count == 0 && !*exhausted) { + return Status::Error( + "logical_index: posting source returned empty before EOF"); + } + return validate_and_accumulate(); + } + + Status validate_and_accumulate() { + const auto docids = buffer_->docids(); + const auto freqs = buffer_->freqs(); + const auto positions = buffer_->positions_flat(); + if (postings_->retain_positions && freqs.size() != docids.size()) { + return Status::Error( + "logical_index: positioned source must provide one freq per docid"); + } + if (!postings_->retain_positions && !freqs.empty() && freqs.size() != docids.size()) { + return Status::Error( + "logical_index: docs-only source freqs must be empty or parallel"); + } + if (postings_->retain_positions && + positions.size() > std::numeric_limits::max()) { + return Status::Error( + "logical_index: one source fill exceeds uint32 position offsets"); + } + if (postings_->retain_positions) { + position_offsets_.resize(freqs.size() + 1); + } + // One fused pass serves both the positions-count validation and the + // frequency statistics below. fill() has already capped the buffer at + // target_docs (at most the adaptive window sizes), so a uint64 sum of + // uint32 frequencies cannot overflow within one fill; the per-term + // accumulation below keeps its overflow check. + uint64_t fill_freq_sum = 0; + uint32_t fill_max_freq = 0; + for (size_t doc = 0; doc < freqs.size(); ++doc) { + const uint32_t freq = freqs[doc]; + fill_freq_sum += freq; + fill_max_freq = std::max(fill_max_freq, freq); + if (postings_->retain_positions) { + position_offsets_[doc + 1] = static_cast(fill_freq_sum); + } + } + if (postings_->retain_positions) { + if (fill_freq_sum != positions.size()) { + return Status::Error( + "logical_index: source positions count must equal sum(freqs)"); + } + } else { + if (!positions.empty()) { + return Status::Error( + "logical_index: docs-only source must not provide positions"); + } + } + + for (uint32_t docid : docids) { + if (last_docid_.has_value() && docid <= *last_docid_) { + return Status::Error( + "logical_index: source docids must be strictly ascending"); + } + if (docid >= writer_->doc_count_) { + return Status::Error( + "logical_index: source docid must be less than doc_count"); + } + last_docid_ = docid; + } + if (docids.size() > std::numeric_limits::max() - total_docs_) { + return Status::Error( + "logical_index: source document count exceeds uint32"); + } + total_docs_ += static_cast(docids.size()); + + if (declared_common_gram_) { + const uint64_t increment = + postings_->retain_positions ? positions.size() : docids.size(); + if (increment > std::numeric_limits::max() - stats_.total_freq) { + return Status::Error( + "logical_index: source total frequency overflow"); + } + stats_.total_freq += increment; + } else if (freqs.empty()) { + if (docids.size() > std::numeric_limits::max() - stats_.total_freq) { + return Status::Error( + "logical_index: source total frequency overflow"); + } + stats_.total_freq += docids.size(); + } else { + if (fill_freq_sum > std::numeric_limits::max() - stats_.total_freq) { + return Status::Error( + "logical_index: source total frequency overflow"); + } + stats_.total_freq += fill_freq_sum; + stats_.max_freq = std::max(stats_.max_freq, fill_max_freq); + } + return Status::OK(); + } + + Status encode_small(DictEntry* entry) { + if (total_docs_ >= format::kSlimDfThreshold || + (term_has_prx_ && + !fits_prx_window_shape(total_docs_, stats_.total_freq, writer_->prx_window_limits_))) { + RETURN_IF_ERROR(encode_windowed_buffer(adaptive_window_docs(total_docs_))); + return finish_windowed(entry); + } + + std::vector prx_window; + if (term_has_prx_) { + ByteSink sink; + format::PrxWindowBuildOutcome outcome = format::PrxWindowBuildOutcome::kBuilt; + RETURN_IF_ERROR(format::try_build_prx_window_flat( + buffer_->positions_flat(), buffer_->freqs(), -writer_->prx_zstd_level_, + writer_->prx_window_limits_, &sink, &outcome)); + if (outcome == format::PrxWindowBuildOutcome::kNeedsSplit) { + RETURN_IF_ERROR(encode_windowed_buffer(format::kFrqBaseUnit)); + return finish_windowed(entry); + } + prx_window = sink.take(); + } + + ByteSink frq_sink; + FrqRegionMeta dd_meta; + FrqRegionMeta freq_meta {}; + RETURN_IF_ERROR(format::build_dd_region(buffer_->docids(), /*win_base=*/0, kRawFrqRegion, + &frq_sink, &dd_meta)); + if (term_has_freq_) { + RETURN_IF_ERROR(format::build_freq_region(buffer_->freqs(), kRawFrqRegion, &frq_sink, + &freq_meta)); + } + std::vector frq_window = frq_sink.take(); + entry->enc = DictEntryEnc::kSlim; + entry->dd_meta = dd_meta; + entry->freq_meta = freq_meta; + if (frq_window.size() <= format::kDefaultInlineThreshold) { + entry->kind = DictEntryKind::kInline; + entry->inline_dd_disk_len = dd_meta.disk_len; + entry->frq_bytes = std::move(frq_window); + if (term_has_prx_) entry->prx_bytes = std::move(prx_window); + return Status::OK(); + } + + entry->kind = DictEntryKind::kPodRef; + entry->frq_docs_len = dd_meta.disk_len; + if (term_has_prx_) { + const uint64_t prx_off = writer_->posting_size(); + RETURN_IF_ERROR(writer_->posting_out_->append(Slice(prx_window))); + entry->prx_off_delta = prx_off - prx_base_; + entry->prx_len = writer_->posting_size() - prx_off; + } + const uint64_t frq_off = writer_->posting_size(); + RETURN_IF_ERROR(writer_->posting_out_->append(Slice(frq_window))); + entry->frq_off_delta = frq_off - frq_base_; + entry->frq_len = writer_->posting_size() - frq_off; + return Status::OK(); + } + + Status encode_windowed_buffer(uint32_t unit) { + DCHECK_GT(unit, 0); + for (size_t begin = 0; begin < buffer_->document_count(); begin += unit) { + const size_t count = + std::min(buffer_->document_count() - begin, static_cast(unit)); + const auto offsets = + term_has_prx_ + ? std::span(position_offsets_).subspan(begin, count + 1) + : std::span {}; + const auto positions = + term_has_prx_ ? buffer_->positions_flat().subspan( + offsets.front(), + static_cast(offsets.back() - offsets.front())) + : std::span {}; + RETURN_IF_ERROR(emitter_.emit_window(PostingRunView { + .docids = buffer_->docids().subspan(begin, count), + .freqs = buffer_->freqs().empty() ? std::span {} + : buffer_->freqs().subspan(begin, count), + .position_offsets = offsets, + .positions_flat = positions, + })); + } + return Status::OK(); + } + + Status finish_windowed(DictEntry* entry) { + TermAggregateStats emitted_stats; + RETURN_IF_ERROR(emitter_.finish_term(entry, &emitted_stats)); + DCHECK_EQ(emitted_stats.df, total_docs_); + DCHECK_EQ(emitted_stats.total_freq, stats_.total_freq); + DCHECK_EQ(emitted_stats.max_freq, stats_.max_freq); + return Status::OK(); + } + + LogicalIndexWriter* writer_; + StreamedTermPostings* postings_; + bool declared_common_gram_ = false; + bool term_has_freq_ = false; + bool term_has_prx_ = false; + TermPostingBuffer* buffer_ = nullptr; + uint64_t frq_base_ = 0; + uint64_t prx_base_ = 0; + WindowEmitter emitter_; + std::vector position_offsets_; + FreqStats stats_; + std::optional last_docid_; + uint32_t total_docs_ = 0; +}; + +namespace testing { +#ifdef BE_TEST +namespace { +// Function-local-static op-count seam backing term_freq_scans(). One atomic, +// relaxed: the writer build path is single-threaded, so only the COUNT matters, +// not ordering (the atomic keeps it race-clean if a test ever parallelizes). +std::atomic& term_freq_scan_counter() { + static std::atomic counter {0}; + return counter; +} +} // namespace +#endif + +void note_term_freq_scan() { +#ifdef BE_TEST + term_freq_scan_counter().fetch_add(1, std::memory_order_relaxed); +#endif +} +uint64_t term_freq_scans() { +#ifdef BE_TEST + return term_freq_scan_counter().load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_term_freq_scans() { +#ifdef BE_TEST + term_freq_scan_counter().store(0, std::memory_order_relaxed); +#endif +} + +// Forwards to the real fused helper so pure boundary tests exercise production +// code (not a test-local re-implementation). +FreqStats fuse_freq_stats_for_test(const std::vector& freqs) { + return fuse_freq_stats(freqs); +} +} // namespace testing + +LogicalIndexWriter::LogicalIndexWriter(const SniiIndexInput& in) + : LogicalIndexWriter(in, TrackedNullDocids(std::vector(in.null_docids))) {} + +LogicalIndexWriter::LogicalIndexWriter(const SniiIndexInput& in, TrackedNullDocids null_docids) + : index_id_(in.index_id), + index_suffix_(in.index_suffix), + index_config_(in.config), + tier_(format::tier_of(in.config)), + has_prx_(format::has_positions(in.config)), + // G16-c: the caller can drop freq layout entirely (in.write_freq == + // false) on a freq-capable tier -- see SniiIndexInput::write_freq. + has_freq_(format::tier_of(in.config) >= format::IndexTier::kT2 && in.write_freq), + has_norms_(format::has_scoring(in.config)), + doc_count_(in.doc_count), + null_docids_(std::move(null_docids)), + terms_(in.terms), + term_source_(in.term_source), + encoded_norms_(in.encoded_norms), + common_grams_metadata_(in.common_grams_metadata), + common_grams_posting_policy_(in.common_grams_posting_policy), + target_dict_block_bytes_(in.target_dict_block_bytes != 0 + ? in.target_dict_block_bytes + : format::kDefaultTargetDictBlockBytes), + dict_block_zstd_level_(in.dict_block_zstd_level), + prx_zstd_level_(in.prx_zstd_level), + prx_window_limits_(in.prx_window_limits), + memory_reporter_(in.mem_reporter), + dict_buf_(in.dict_resident_cap_bytes, "dict", in.mem_reporter), + norms_section_reservation_(in.mem_reporter == nullptr + ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation()), + null_bitmap_section_reservation_(in.mem_reporter == nullptr + ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation()), + term_hashes_reservation_(in.mem_reporter == nullptr + ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation()), + bsbf_bytes_reservation_(in.mem_reporter == nullptr + ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation()) {} + +Status LogicalIndexWriter::reserve_term_hash_for_append() { + if (memory_reporter_ == nullptr || term_hashes_.size() < term_hashes_.capacity()) { + return Status::OK(); + } + size_t target_capacity = term_hashes_.capacity() == 0 ? 1 : term_hashes_.capacity() * 2; + if (target_capacity < term_hashes_.capacity() || + target_capacity > std::numeric_limits::max() / sizeof(uint64_t)) { + return Status::Error( + "logical_index: term hash capacity overflow"); + } + MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(term_hashes_reservation_.prepare_replacement(target_capacity * sizeof(uint64_t), + &replacement)); + term_hashes_.reserve(target_capacity); + DCHECK_EQ(term_hashes_.capacity(), target_capacity); + term_hashes_reservation_ = std::move(replacement); + return Status::OK(); +} + +// Serializes the current open block, zstd-compresses it (the dict region is the +// single largest section -- term keys + entry meta + inline postings -- and the +// 64KiB blocks compress ~40%), streams the compressed bytes into the dict +// scratch file, and records a directory entry. The block-level crc32c +// (rec.checksum) covers the UNCOMPRESSED bytes, so DictBlockReader::open +// verifies integrity after the reader decompresses. A compressed block also +// shrinks the bytes a term lookup fetches from S3 -- aligning with the +// read-byte thesis. If zstd does not shrink a (tiny) block, it is stored raw so +// a lookup never pays a pointless decompress. +Status LogicalIndexWriter::flush_block(DictBlockBuilder* block, std::string first_term) { + std::vector plain_bytes = block->finish_owned(); + const Slice plain(plain_bytes); + BlockRecord rec; + rec.rel_offset = dict_buf_.size(); + rec.n_entries = block->n_entries(); + rec.checksum = crc32c(plain); // crc over UNCOMPRESSED block bytes + rec.first_term = std::move(first_term); + + std::vector comp; + Status zs = zstd_compress(plain, dict_block_zstd_level_, &comp); + if (zs.ok() && comp.size() < plain.size()) { + rec.flags = format::block_ref_flags::kZstd; + rec.uncomp_len = static_cast(plain.size()); + rec.length = static_cast(comp.size()); + RETURN_IF_ERROR(dict_buf_.append_move(std::move(comp))); + } else { + rec.flags = 0; + rec.uncomp_len = 0; + rec.length = static_cast(plain.size()); + RETURN_IF_ERROR(dict_buf_.append_move(std::move(plain_bytes))); + } + blocks_.push_back(std::move(rec)); + return Status::OK(); +} + +// Running state for the in-flight DICT block while terms stream past. +struct LogicalIndexWriter::BlockState { + explicit BlockState(MemoryReporter* memory_reporter) : transfer_buffer(memory_reporter) {} + + std::unique_ptr block; + std::string block_first_term; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + bool term_stats = true; + TermPostingBuffer transfer_buffer; +}; + +// Out-of-line so unique_ptr sees the complete type (see header). +LogicalIndexWriter::~LogicalIndexWriter() = default; + +Status LogicalIndexWriter::process_term(StreamedTermPostings& tp, BlockState* st) { + const bool is_declared_common_gram = + common_grams_metadata_.has_value() && tp.term.starts_with(CG_V1_MARKER) && + (common_grams_metadata_->common_grams_coverage == CommonGramsCoverage::kComplete || + common_grams_posting_policy_ == format::CommonGramsPostingPolicy::kHybridV1); + const bool term_has_prx = has_prx_ && tp.retain_positions; + const bool term_has_freq = has_freq_ && !is_declared_common_gram; + + if (st->block && st->term_stats != term_has_freq) { + RETURN_IF_ERROR(flush_block(st->block.get(), st->block_first_term)); + st->block.reset(); + } + if (!st->block) { + const uint64_t base = posting_size(); + st->frq_base = base; + st->prx_base = base; + st->term_stats = term_has_freq; + st->block = std::make_unique(tier_, has_prx_, st->frq_base, st->prx_base, + /*anchor_interval=*/16, + /*term_stats=*/term_has_freq); + st->block_first_term = tp.term; + } + + RETURN_IF_ERROR(reserve_term_hash_for_append()); + const uint64_t term_hash = format::bsbf_hash(tp.term); + DictEntry entry; + FreqStats stats; + StreamingTermEncoder encoder(this, &tp, is_declared_common_gram, term_has_freq, term_has_prx, + &st->transfer_buffer, st->frq_base, st->prx_base); + RETURN_IF_ERROR(encoder.encode(&entry, &stats)); + + term_hashes_.push_back(term_hash); + ++term_count_; + stats_.sum_total_term_freq += stats.total_freq; + st->block->add_entry(std::move(entry)); + if (st->block->estimated_bytes() >= target_dict_block_bytes_) { + RETURN_IF_ERROR(flush_block(st->block.get(), st->block_first_term)); + st->block.reset(); + } + return Status::OK(); +} + +Status LogicalIndexWriter::build_blocks() { + BlockState st(memory_reporter_); + if (term_source_ != nullptr) { + RETURN_IF_ERROR(term_source_->for_each_term_sorted( + [&](StreamedTermPostings&& tp) { return process_term(tp, &st); })); + } else { + for (const auto& tp : terms_) { + SpanTermPostingSource source(tp.docids, tp.freqs, tp.positions_flat); + StreamedTermPostings streamed { + .term = tp.term, .retain_positions = tp.retain_positions, .source = &source}; + RETURN_IF_ERROR(process_term(streamed, &st)); + } + } + if (st.block) RETURN_IF_ERROR(flush_block(st.block.get(), st.block_first_term)); + return Status::OK(); +} + +Status LogicalIndexWriter::prepare_build(io::FileWriter* posting_out) { + if (posting_out == nullptr) { + return Status::Error( + "logical_index: null posting sink"); + } + RETURN_IF_ERROR(format::validate_prx_window_limits(prx_window_limits_)); + if (has_norms_ && encoded_norms_.size() != doc_count_) { + return Status::Error( + "logical_index: norms length must equal doc_count"); + } + for (size_t i = 0; i < null_docids_.size(); ++i) { + if (null_docids_[i] >= doc_count_) { + return Status::Error( + "logical_index: null docid must be less than doc_count"); + } + if (i != 0 && null_docids_[i] <= null_docids_[i - 1]) { + return Status::Error( + "logical_index: null docids must be strictly ascending"); + } + } + if (common_grams_metadata_) { + RETURN_IF_ERROR(validate_common_grams_segment_metadata(*common_grams_metadata_)); + if (common_grams_metadata_->common_grams_coverage == CommonGramsCoverage::kComplete && + !has_prx_) { + return Status::Error( + "logical_index: complete CommonGrams metadata requires positions"); + } + if (common_grams_metadata_->scoring_coverage == ScoringCoverage::kComplete) { + if (!has_norms_ || !has_freq_) { + return Status::Error( + "logical_index: complete scoring metadata requires frequencies and " + "semantic norms"); + } + if (common_grams_metadata_->scoring_doc_count != doc_count_) { + return Status::Error( + "logical_index: scoring doc count must equal doc_count"); + } + } + } + if (common_grams_posting_policy_ == format::CommonGramsPostingPolicy::kHybridV1 && + (!common_grams_metadata_ || + common_grams_metadata_->common_grams_coverage != CommonGramsCoverage::kMixed || + !has_prx_)) { + return Status::Error( + "logical_index: hybrid CommonGrams postings require mixed metadata and positions"); + } + // The interleaved posting region streams STRAIGHT into the container output + // (no temp round-trip): posting_size() is the region-relative byte count, + // derived from the output offset advanced since this index's region began. + // The DICT region is staged in dict_buf_ (tiered: RAM under the cap = + // spill-only; spills above it) since it must land contiguously after the + // concurrently-streamed posting region. + posting_out_ = posting_out; + posting_off0_ = posting_out->bytes_written(); + return Status::OK(); +} + +Status LogicalIndexWriter::finalize_build() { + if (common_grams_metadata_ && + common_grams_metadata_->scoring_coverage == ScoringCoverage::kComplete) { + if (common_grams_metadata_->scoring_token_count > stats_.sum_total_term_freq) { + return Status::Error( + "logical_index: semantic scoring token count exceeds physical term frequency"); + } + if (stats_.sum_total_term_freq != 0 && common_grams_metadata_->scoring_token_count == 0) { + return Status::Error( + "logical_index: non-empty physical postings have zero semantic scoring tokens"); + } + if (common_grams_metadata_->plain_term_key_version == + ::doris::segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal && + common_grams_metadata_->common_grams_coverage == CommonGramsCoverage::kNone && + common_grams_metadata_->scoring_token_count != stats_.sum_total_term_freq) { + return Status::Error( + "logical_index: semantic plain token count must equal physical term frequency"); + } + } + // Seal the dict buffer so a spilled temp is flushed before + // stream_dict_region_into reads it back. A no-op for a RAM-resident dict. + RETURN_IF_ERROR(dict_buf_.seal()); + + stats_.doc_count = doc_count_; + stats_.indexed_doc_count = doc_count_ - static_cast(null_docids_.size()); + stats_.term_count = term_count_; + stats_.null_count = static_cast(null_docids_.size()); + + if (has_norms_) { + const size_t payload_size = varint_len(encoded_norms_.size()) + encoded_norms_.size(); + const size_t section_size = 1 + varint_len(payload_size) + payload_size + sizeof(uint32_t); + MemoryReporter::Reservation build_reservation = + memory_reporter_ == nullptr ? MemoryReporter::Reservation() + : memory_reporter_->make_reservation(); + if (memory_reporter_ != nullptr) { + RETURN_IF_ERROR(build_reservation.set_bytes(payload_size)); + RETURN_IF_ERROR(norms_section_reservation_.set_bytes(section_size)); + } + ByteSink nsink; + format::NormsPodWriter::finish(encoded_norms_, &nsink); + norms_section_ = nsink.take(); + DORIS_CHECK_EQ(norms_section_.capacity(), section_size); + if (memory_reporter_ != nullptr) { + DORIS_CHECK_EQ(norms_section_reservation_.bytes(), norms_section_.capacity()); + } + } + + if (!null_docids_.empty()) { + MemoryReporter::Reservation bitmap_build_reservation = + memory_reporter_ == nullptr ? MemoryReporter::Reservation() + : memory_reporter_->make_reservation(); + if (memory_reporter_ != nullptr) { + RETURN_IF_ERROR(bitmap_build_reservation.set_bytes( + format::NullBitmapWriter::build_memory_upper_bound( + std::span(null_docids_.data(), null_docids_.size())))); + } + format::NullBitmapWriter null_writer; + null_writer.add_many(std::span(null_docids_.data(), null_docids_.size())); + null_docids_.release(); + + format::NullBitmapSerializationSizes sizes; + RETURN_IF_ERROR(null_writer.serialization_sizes(doc_count_, &sizes)); + if (sizes.roaring_bytes > std::numeric_limits::max() - sizes.payload_bytes) { + return Status::Error( + "logical_index: null bitmap scratch size overflows"); + } + MemoryReporter::Reservation scratch_reservation = + memory_reporter_ == nullptr ? MemoryReporter::Reservation() + : memory_reporter_->make_reservation(); + if (memory_reporter_ != nullptr) { + RETURN_IF_ERROR( + scratch_reservation.set_bytes(sizes.roaring_bytes + sizes.payload_bytes)); + RETURN_IF_ERROR(null_bitmap_section_reservation_.set_bytes(sizes.framed_bytes)); + } + ByteSink null_sink; + RETURN_IF_ERROR(null_writer.finish(doc_count_, &null_sink)); + null_bitmap_section_ = null_sink.take(); + DORIS_CHECK_EQ(null_bitmap_section_.size(), sizes.framed_bytes); + if (memory_reporter_ != nullptr) { + DORIS_CHECK_LE(null_bitmap_section_.capacity(), + null_bitmap_section_reservation_.bytes()); + } + } + null_docids_.release(); + + // Build the absent-term filter (block-split bloom, Parquet-canonical) from + // the per-term keys (no retained strings) as a [28B header][bitset] blob; the + // compound writer places it as a PHYSICAL section probed one 32-byte block on + // demand. + bsbf_bytes_.clear(); + bsbf_built_ = false; + if (!term_hashes_.empty()) { + const uint32_t bitset_bytes = format::bsbf_optimal_num_bytes( + static_cast(term_hashes_.size()), kBsbfFpp); + const size_t serialized_bytes = format::kBsbfHeaderSize + bitset_bytes; + MemoryReporter::Reservation builder_reservation = + memory_reporter_ == nullptr ? MemoryReporter::Reservation() + : memory_reporter_->make_reservation(); + if (memory_reporter_ != nullptr) { + RETURN_IF_ERROR(builder_reservation.set_bytes(bitset_bytes)); + RETURN_IF_ERROR(bsbf_bytes_reservation_.set_bytes(serialized_bytes)); + } + format::BsbfBuilder bf; + RETURN_IF_ERROR(format::BsbfBuilder::create(static_cast(term_hashes_.size()), + kBsbfFpp, &bf)); + DCHECK_EQ(bf.resident_capacity_bytes(), bitset_bytes); + for (uint64_t k : term_hashes_) bf.insert(k); + ByteSink bsink; + bsink.reserve(serialized_bytes); + RETURN_IF_ERROR(bf.serialize(&bsink)); + bsbf_bytes_ = bsink.take(); + DCHECK_EQ(bsbf_bytes_.capacity(), serialized_bytes); + bsbf_built_ = true; + } + std::vector().swap(term_hashes_); // release + term_hashes_reservation_.reset(); + + return Status::OK(); +} + +void LogicalIndexWriter::release_bsbf_bytes() { + std::vector().swap(bsbf_bytes_); + bsbf_bytes_reservation_.reset(); +} + +void LogicalIndexWriter::release_null_bitmap_bytes() { + std::vector().swap(null_bitmap_section_); + null_bitmap_section_reservation_.reset(); +} + +void LogicalIndexWriter::release_norms_bytes() { + std::vector().swap(norms_section_); + norms_section_reservation_.reset(); +} + +Status LogicalIndexWriter::build(io::FileWriter* posting_out) { + // Single-session invariant: a writer that ran (or is running) a streamed + // session must not also build() -- the posting sink anchor and the dict + // buffer are one-shot. + if (stream_phase_ != StreamPhase::kIdle) { + return Status::Error( + "logical_index: build() on a writer with a streamed session"); + } + // prepare_build is pure entry validation up to its final sink-anchor + // assignments, so a failure there leaves the writer clean (still kIdle). + RETURN_IF_ERROR(prepare_build(posting_out)); + // Poison-by-default past this point: build_blocks/finalize_build may fail + // AFTER posting bytes hit the sink or term state advanced, and a dirty + // writer must never accept a later begin_streamed/build (single-session + + // crash-safety invariant 6). Only full success seals to kFinished. + stream_phase_ = StreamPhase::kFailed; + RETURN_IF_ERROR(build_blocks()); + RETURN_IF_ERROR(finalize_build()); + // Claim the (only) session so a later begin_streamed/push_term errors out. + stream_phase_ = StreamPhase::kFinished; + return Status::OK(); +} + +Status LogicalIndexWriter::begin_streamed(io::FileWriter* posting_out) { + if (stream_phase_ == StreamPhase::kFailed) { + return Status::Error( + "logical_index: begin_streamed on a failed writer (a prior session error left " + "partial state; allocate a fresh writer)"); + } + if (stream_phase_ != StreamPhase::kIdle) { + return Status::Error( + "logical_index: begin_streamed on an already-claimed writer session"); + } + RETURN_IF_ERROR(prepare_build(posting_out)); + stream_state_ = std::make_unique(memory_reporter_); + stream_phase_ = StreamPhase::kActive; + return Status::OK(); +} + +Status LogicalIndexWriter::push_term(StreamedTermPostings&& tp) { + if (stream_phase_ != StreamPhase::kActive) { + return Status::Error( + "logical_index: push_term without an active streamed session"); + } + if (has_pushed_term_ && tp.term <= last_pushed_term_) { + return Status::Error( + "logical_index: pushed terms must be strictly increasing ('{}' after '{}')", + tp.term, last_pushed_term_); + } + last_pushed_term_ = tp.term; + has_pushed_term_ = true; + Status status = process_term(tp, stream_state_.get()); + if (!status.ok()) { + stream_phase_ = StreamPhase::kFailed; + stream_state_.reset(); + } + return status; +} + +Status LogicalIndexWriter::finish_streamed() { + if (stream_phase_ == StreamPhase::kFinished) { + return Status::Error( + "logical_index: finish_streamed on an already-finished session"); + } + if (stream_phase_ == StreamPhase::kFailed) { + return Status::Error( + "logical_index: finish_streamed on a failed session (a prior push/finish error " + "poisoned the writer; the partial output must be discarded)"); + } + if (stream_phase_ != StreamPhase::kActive) { + return Status::Error( + "logical_index: finish_streamed without begin_streamed"); + } + // Poison-by-default: a failed trailing flush or finalize leaves partial + // output, so only full success below seals the session to kFinished. + stream_phase_ = StreamPhase::kFailed; + // Trailing-block flush mirrors the tail of build_blocks(); the shared + // finalize_build() then seals the dict buffer and materializes the + // stats/norms/null-bitmap/BSBF sections. + if (stream_state_->block) { + RETURN_IF_ERROR(flush_block(stream_state_->block.get(), stream_state_->block_first_term)); + } + stream_state_.reset(); + RETURN_IF_ERROR(finalize_build()); + stream_phase_ = StreamPhase::kFinished; + return Status::OK(); +} + +Status LogicalIndexWriter::finish_metadata(const SectionRefs& abs_refs, uint64_t dict_region_offset, + SerializedMetadataGroup* out) const { + if (out == nullptr) { + return Status::Error( + "logical_index: null metadata output"); + } + *out = {}; + + SampledTermIndexBuilder sti; + for (const auto& b : blocks_) sti.add_block_first_term(b.first_term); + ByteSink sti_sink; + sti.finish(&sti_sink); + + DictBlockDirectoryBuilder dir; + for (const auto& b : blocks_) { + BlockRef ref; + ref.offset = dict_region_offset + b.rel_offset; + ref.length = b.length; + ref.n_entries = b.n_entries; + ref.flags = b.flags; + ref.checksum = b.checksum; + ref.uncomp_len = b.uncomp_len; + dir.add(ref); + } + ByteSink dir_sink; + dir.finish(&dir_sink); + + ByteSink sti_blob; + RETURN_IF_ERROR( + format::encode_metadata_blob(sti_sink.view(), format::SectionType::kSampledTermIndex, + format::SectionType::kSampledTermIndexZstd, &sti_blob)); + out->sampled_term_index = sti_blob.take(); + + ByteSink dbd_blob; + RETURN_IF_ERROR( + format::encode_metadata_blob(dir_sink.view(), format::SectionType::kDictBlockDirectory, + format::SectionType::kDictBlockDirectoryZstd, &dbd_blob)); + out->dict_block_directory = dbd_blob.take(); + + format::CoreMetadata core; + core.index_config = index_config_; + core.stats = stats_; + core.section_refs = abs_refs; + core.common_grams_metadata = common_grams_metadata_; + core.common_grams_posting_policy = common_grams_posting_policy_; + ByteSink core_sink; + RETURN_IF_ERROR(format::encode_core_metadata(core, &core_sink)); + out->core = core_sink.take(); + return Status::OK(); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/logical_index_writer.h b/be/src/storage/index/snii/writer/logical_index_writer.h new file mode 100644 index 00000000000000..c93c58d8b722ac --- /dev/null +++ b/be/src/storage/index/snii/writer/logical_index_writer.h @@ -0,0 +1,459 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/metadata_blob.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/format/stats_block.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/spillable_byte_buffer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_source.h" + +// LogicalIndexWriter -- builds the per-logical-index section bytes (interleaved +// posting region + DICT block region) plus the SampledTermIndex and DICT block +// directory metadata for ONE logical index. It owns the in-memory section bytes, +// runtime statistics, and references needed by the container orchestrator +// (SniiCompoundWriter) to resolve absolute offsets and emit the Core/STI/DBD +// metadata group. +// +// This module deliberately produces ONLY relative bytes/structures: it has no +// knowledge of the absolute file position where the sections will land. The +// orchestrator stitches the absolute offsets in afterward (append-only, no +// seek-back). See snii_compound_writer.h for the precise offset contract. +// +// POSTING REGION (single interleaved sink): the former separate .frq POD and .prx +// POD are merged into ONE posting region. For each pod_ref term, in term order, the +// writer appends its prx span FIRST then its frq span, contiguously: +// posting region = concat over pod_ref terms of [prx span][frq span]. +// The prx span is empty when !has_prx (docs-only / keyword tier). INLINE terms +// append NOTHING to the posting region. +// +// Per-term encoding policy (v1): +// df >= kSlimDfThreshold (512), or a lower-df term whose positions cannot fit +// one configured reader-safe PRX window: WINDOWED pod_ref. The term's [prx +// windows] are appended to the posting region first, then its +// [prelude][dd-block][freq-block] frq span. The DictEntry records frq/prx +// off_delta+len relative to frq_base/prx_base (see below). +// Other df < kSlimDfThreshold terms: SLIM. The postings are encoded as a +// single .frq window (and .prx window). If the encoded .frq bytes are small +// (<= kDefaultInlineThreshold), they are stored INLINE inside the DictEntry +// (kind=inline); otherwise the term's [prx][frq] spans are appended to the +// posting region as a slim pod_ref (kind=pod_ref, enc=slim, no prelude). +// +// frq_base / prx_base convention (DOCUMENTED CONTRACT): +// For each DICT block, frq_base == prx_base == the running byte offset into THIS +// index's posting region at the moment the block opens (the posting-region size +// when the block's first POD-backed entry is appended). A windowed/slim pod_ref +// entry then sets frq_off_delta = (offset of its frq span within the posting +// region) - frq_base, so the reader computes the absolute file offset as +// section_refs.posting_region.offset + frq_base + frq_off_delta. +// prx_base / prx_off_delta follow the identical rule against the SAME region. +// Because [prx][frq] are written contiguously per term, a writer-side property +// holds when has_prx: frq_off_delta == prx_off_delta + prx_len. The reader does +// NOT rely on it -- each delta is resolved independently. +// Inline entries carry no off_delta (bytes live in the entry). +namespace doris::snii::writer { + +class SniiStreamedIndexSession; +class StreamingTermEncoder; + +struct SerializedMetadataGroup { + std::vector core; + std::vector sampled_term_index; + std::vector dict_block_directory; +}; + +// Inputs describing one logical index to be written. +struct SniiIndexInput { + uint64_t index_id = 0; + std::string index_suffix; + format::IndexConfig config = format::IndexConfig::kDocsPositions; + uint32_t doc_count = 0; + std::vector null_docids; + // Per-doc 1-byte encoded norm (length doc_count); only consumed when the + // config has scoring. May be empty otherwise. + std::vector encoded_norms; + // G16-h: zstd levels for the dict-block whole-block compression and the + // .prx window auto mode (both default 3 == the historical constants). + // Higher levels trade import CPU for size; decode speed is unaffected. + int dict_block_zstd_level = 3; + int prx_zstd_level = 3; + // Internal writer policy. Production callers keep the reader limits; unit + // tests may only tighten them to exercise extreme-window behavior without + // allocating hundreds of MiB. + format::PrxWindowLimits prx_window_limits = format::kReaderPrxWindowLimits; + // G16-c: whether freq-capable (tier>=T2) postings lay out freq regions at + // all. Freq bytes serve ONLY BM25 scoring (want_freq=true lives solely in + // scoring_query), so the CALLER resolves the policy -- the Doris adapter + // passes has_scoring(config) || config::snii_positions_index_write_freq, + // i.e. plain kDocsPositions indexes drop freq unless the escape hatch is + // set. Defaults to true so the core library and existing callers keep the + // full T2 layout unless they opt out. The drop is value-driven on disk + // (windowed prelude flags bit0; slim/inline zero-length freq regions), so + // readers need no index-level flag. Ignored for docs-only configs. + bool write_freq = true; + // Lexicographically sorted terms with ascending-docid postings. Used when + // `term_source` is null (callers that already hold a materialized vector, + // e.g. unit tests). The writer reads but does not retain these. + std::vector terms; + // Optional streaming term source. When non-null, the writer DRAINS it via + // SpimiTermBuffer::for_each_term_sorted so that only one term's postings is + // materialized at a time (avoiding the full TermPostings vector and its + // second-copy peak). `terms` is ignored when this is set. The buffer is + // consumed (emptied) by build(); the caller must keep it alive until build() + // returns and must not reuse it afterwards. + SpimiTermBuffer* term_source = nullptr; + // Target DICT block size in bytes; a block is cut once its estimate reaches + // this. 0 uses kDefaultTargetDictBlockBytes. Smaller values yield more blocks + // (and a finer-grained sampled-term index). + uint32_t target_dict_block_bytes = 0; + // Maximum resident capacity of the staged DICT region before it spills. + // Ordinary builds keep the default unlimited local cap and use their shared + // spill-threshold reporter. Streamed compaction sets a bounded watermark so + // reclaimable DICT blocks cannot consume the hard-capped term workspace. + uint64_t dict_resident_cap_bytes = std::numeric_limits::max(); + // Optional writer-level build-RAM reporter (one per SniiCompoundWriter = one + // segment inverted index). When non-null, the dict buffer reports its REAL + // resident-byte deltas (positive on grow, negative on spill). The SPIMI side + // (arena + slot index) reports through the SAME reporter, injected directly at + // the term_source's construction by the caller. null in bench / unit tests -> no + // reporting. NEVER report live_bytes_ (a gated estimate); report + // arena_bytes()+slot_of_+dict ram_bytes_. + MemoryReporter* mem_reporter = nullptr; + // Optional persisted CommonGrams capability and semantic scoring stats. + // Missing metadata preserves the legacy SNII image and cannot be treated as + // compatibility proof by readers. + std::optional common_grams_metadata; + // Optional per-term CommonGrams postings shape. HybridV1 requires Mixed + // coverage metadata and a positions-capable logical index. + format::CommonGramsPostingPolicy common_grams_posting_policy = + format::CommonGramsPostingPolicy::kNone; +}; + +// Move-only ownership of a NULL-docid allocation and its precharged bytes. +// Reservation is declared first so destruction always frees the vector before +// returning its charge. Move assignment is intentionally forbidden because its +// default member order would release the destination charge before its vector. +class TrackedNullDocids { +public: + explicit TrackedNullDocids(std::vector&& docids) : docids_(std::move(docids)) {} + TrackedNullDocids(MemoryReporter::Reservation&& reservation, std::vector&& docids) + : reservation_(std::move(reservation)), docids_(std::move(docids)) {} + + TrackedNullDocids(const TrackedNullDocids&) = delete; + TrackedNullDocids& operator=(const TrackedNullDocids&) = delete; + TrackedNullDocids(TrackedNullDocids&&) noexcept = default; + TrackedNullDocids& operator=(TrackedNullDocids&&) = delete; + + bool empty() const { return docids_.empty(); } + size_t size() const { return docids_.size(); } + const uint32_t* data() const { return docids_.data(); } + uint32_t operator[](size_t index) const { return docids_[index]; } + auto begin() const { return docids_.begin(); } + auto end() const { return docids_.end(); } + + void release() { + std::vector().swap(docids_); + reservation_.reset(); + } + +private: + MemoryReporter::Reservation reservation_; + std::vector docids_; +}; + +// Move-only ownership of a destination norms allocation and its precharged +// bytes. Streamed sessions adopt both together so the vector remains accounted +// for until the writer has materialized the norms section. +class TrackedEncodedNorms { +public: + explicit TrackedEncodedNorms(std::vector&& norms) : norms_(std::move(norms)) {} + TrackedEncodedNorms(MemoryReporter::Reservation&& reservation, std::vector&& norms) + : reservation_(std::move(reservation)), norms_(std::move(norms)) {} + + TrackedEncodedNorms(const TrackedEncodedNorms&) = delete; + TrackedEncodedNorms& operator=(const TrackedEncodedNorms&) = delete; + TrackedEncodedNorms(TrackedEncodedNorms&&) noexcept = default; + TrackedEncodedNorms& operator=(TrackedEncodedNorms&&) = delete; + + bool empty() const { return norms_.empty(); } + size_t size() const { return norms_.size(); } + uint8_t operator[](size_t index) const { return norms_[index]; } + auto begin() const { return norms_.begin(); } + auto end() const { return norms_.end(); } + + void release() { + std::vector().swap(norms_); + reservation_.reset(); + } + +private: + friend class SniiStreamedIndexSession; + MemoryReporter::Reservation reservation_; + std::vector norms_; +}; + +// Term-level frequency statistics. Ordinary terms compute sum(freqs) and +// max(freqs) in one fused scan. Complete CommonGrams entries derive total_freq +// from their required PRX position count and leave max_freq at 0 because their +// statless DICT block does not serialize it. +struct FreqStats { + uint64_t total_freq = 0; + uint32_t max_freq = 0; +}; + +// Builds and holds the section bytes + meta sub-sections for one logical index. +class LogicalIndexWriter { +public: + explicit LogicalIndexWriter(const SniiIndexInput& in); + // Out-of-line: stream_state_ points at the private nested BlockState, which + // is incomplete here (unique_ptr needs the complete type at destruction). + ~LogicalIndexWriter(); + + // Builds DICT blocks, the interleaved posting region, sampled-term index, dict + // directory, stats and bsbf. The posting region is written STRAIGHT into + // `posting_out` as terms are produced (no temp round-trip for the bulk); the + // orchestrator captures its absolute offset/length from posting_out->bytes_written() + // around this call. Must be called once before the accessors below. Returns + // InvalidArgument on a null sink or inconsistent input (e.g. norms/doc_count + // mismatch when scoring is enabled, or non-ascending docids). + Status build(io::FileWriter* posting_out); + + // Streamed three-phase alternative to build() (T2.1, the compaction index + // merge fast path): the CALLER produces terms one at a time (k-way merge + // over source segments) instead of handing the writer a term source. + // begin_streamed(sink) -> push_term(tp) x N -> finish_streamed() + // push_term funnels through the SAME process_term choke point build() + // drains through (bigram prune gates, shape validation, encode), and the + // setup/finalize steps are shared with build() -- so the produced bytes are + // IDENTICAL to a build() fed the same terms in the same order (the T2 + // byte-golden invariant). A writer instance runs EXACTLY ONE session: + // build() and begin_streamed are mutually exclusive, push_term after + // finish_streamed and a second finish_streamed are errors -- a half-fed + // session can never masquerade as a sealed index (crash-safety invariant 6). + // Failure poisons the session: any push_term/finish_streamed/build failure + // past entry validation (e.g. a posting-sink append error mid-encode) moves + // the writer to a terminal failed state where every subsequent + // push_term/finish_streamed/build/begin_streamed is rejected, so a caller + // that swallows an error can never seal (or re-claim) a corrupt index. + // Entry rejections themselves (term-order / postings-shape / phase checks) + // do not poison an active session EXCEPT postings-shape violations, which + // fail inside the shared process_term and conservatively poison too; only + // the term-order guard is explicitly recoverable (see the UT contract). + Status begin_streamed(io::FileWriter* posting_out); + // Consumes tp synchronously. Entry validation: terms must arrive in STRICTLY + // increasing lexicographic order (the one invariant process_term cannot see + // -- DICT blocks, the sampled term index and the reader's binary search all + // assume it; equal terms are rejected too, the upstream merge must have + // combined duplicates). The streaming encoder rejects invalid per-term + // posting shapes. Returns InvalidArgument on any violation. + Status push_term(StreamedTermPostings&& tp); + Status finish_streamed(); + + // DICT region byte length (relative; orchestrator decides its absolute offset). The + // DICT region (zstd-compressed blocks) is built into a tiered buffer during build() + // -- it must land contiguously AFTER the posting region (streamed concurrently), so + // it cannot stream directly. The buffer stays in RAM while small (spill-only build) + // and spills to a temp once it crosses the RAM cap (bounded peak RSS for a huge + // dict). Its bytes are emitted via stream_dict_region_into below. The posting region + // went straight to the output during build(), so it has no length accessor here -- + // the orchestrator measures it directly. norms stays in RAM (1 byte/doc). + uint64_t dict_region_size() const { return dict_buf_.size(); } + const std::vector& norms_bytes() const { return norms_section_; } + const std::vector& null_bitmap_bytes() const { return null_bitmap_section_; } + // Block-split bloom XFilter blob ([28B header][bitset]); empty when no terms. + const std::vector& bsbf_bytes() const { return bsbf_bytes_; } + bool has_bsbf() const { return bsbf_built_; } + void release_bsbf_bytes(); + void release_null_bitmap_bytes(); + void release_norms_bytes(); + bool has_null_bitmap() const { return !null_bitmap_section_.empty(); } + + // Streams the DICT region (RAM or spilled temp) into the append-only container + // after its posting region. + Status stream_dict_region_into(io::FileWriter* out) { + return dict_buf_.stream_into_and_release(out); + } + + bool has_prx() const { return has_prx_; } + bool has_norms() const { return has_norms_; } + format::IndexTier tier() const { return tier_; } + uint64_t index_id() const { return index_id_; } + const std::string& index_suffix() const { return index_suffix_; } + + // Builds the three mandatory v1 metadata blobs. The orchestrator writes them + // contiguously as Core -> STI -> DBD and publishes their absolute references + // only after all three appends succeed. + Status finish_metadata(const format::SectionRefs& abs_refs, uint64_t dict_region_offset, + SerializedMetadataGroup* out) const; + +private: + friend class SniiStreamedIndexSession; + LogicalIndexWriter(const SniiIndexInput& in, TrackedNullDocids null_docids); + + // One DICT block's directory record. The block's serialized bytes are appended to + // the in-RAM dict buffer as soon as the block is cut; only this compact summary + // (offset within the dict region + length + entry count + checksum) is kept to + // build the DICT block directory at finish_metadata time. The absolute file offset is + // computed as dict_region_offset + rel_offset. + struct BlockRecord { + uint64_t rel_offset = 0; // byte offset of this block within the dict region + uint64_t length = 0; // ON-DISK block length (compressed when flags&kZstd) + uint32_t n_entries = 0; + uint32_t checksum = 0; // crc32c of the UNCOMPRESSED block bytes + uint8_t flags = 0; // block_ref_flags::* (kZstd when block is compressed) + uint64_t uncomp_len = 0; // uncompressed block length (when flags&kZstd) + std::string first_term; + }; + + // Shared entry/exit of build() and the streamed session, extracted so the + // two paths are byte-identical BY CONSTRUCTION (not by parallel-maintained + // copies). prepare_build validates the sink/norms/CommonGrams identity and + // anchors the posting region; finalize_build re-checks the stats-dependent + // CommonGrams invariants, seals the dict buffer and materializes + // stats/norms/null-bitmap/BSBF. + Status prepare_build(io::FileWriter* posting_out); + Status finalize_build(); + // Iterates terms (from the streaming source or the materialized vector), + // splitting DICT blocks by target size and filling PODs + blocks_. + Status build_blocks(); + // Per-term driver shared by every producer. It validates the term, opens a + // block if needed, encodes it, and cuts the block at the target size. + struct BlockState; + Status process_term(StreamedTermPostings& tp, BlockState* st); + Status reserve_term_hash_for_append(); + // Region-relative byte count of the posting bytes written so far (the offset basis + // for frq_base/prx_base + frq_off_delta/prx_off_delta). During build() the only + // writes to posting_out_ are this index's posting region, so the count is the + // output offset advanced since the region began. + uint64_t posting_size() const { return posting_out_->bytes_written() - posting_off0_; } + // Serializes the current open block, streams its bytes into the dict scratch + // file, and records a compact directory entry (no block bytes retained). + Status flush_block(format::DictBlockBuilder* block, std::string first_term); + + uint64_t index_id_; + std::string index_suffix_; + format::IndexConfig index_config_; + format::IndexTier tier_; + bool has_prx_; + bool has_freq_; // tier >= T2: a freq region is encoded per window + bool has_norms_; + uint32_t doc_count_; + TrackedNullDocids null_docids_; + const std::vector& terms_; // materialized fallback (may be empty) + SpimiTermBuffer* term_source_; // streaming source (null => use terms_) + uint64_t term_count_ = 0; // distinct terms actually consumed + const std::vector& encoded_norms_; + std::optional common_grams_metadata_; + format::CommonGramsPostingPolicy common_grams_posting_policy_ = + format::CommonGramsPostingPolicy::kNone; + + uint32_t target_dict_block_bytes_; + // G16-h: zstd levels (dict whole-block / prx auto mode), from SniiIndexInput. + int dict_block_zstd_level_ = 3; + int prx_zstd_level_ = 3; + format::PrxWindowLimits prx_window_limits_ = format::kReaderPrxWindowLimits; + // The DICT region (zstd-compressed blocks) is staged here as blocks flush. It must + // land contiguously AFTER the posting region (which streams concurrently to the + // output), so it cannot stream directly; the orchestrator streams it into the + // container right after the posting region. It has NO independent local cap -- it + // spills to a temp via the writer's shared tracked-memory budget (the + // MemoryReporter from SniiIndexInput, null off-Doris). This reporter covers + // explicitly reserved structures; codec-local scratch remains governed by + // its format window limits rather than being represented as a whole-writer + // RSS hard cap. + MemoryReporter* memory_reporter_ = nullptr; + SpillableByteBuffer dict_buf_; + // The interleaved [prx][frq] posting region streams STRAIGHT into the container + // output during build() -- no temp. posting_out_ is the container writer (borrowed + // for the duration of build); posting_off0_ is its absolute offset when this index's + // region began, so posting_size() = bytes_written() - posting_off0_. + io::FileWriter* posting_out_ = nullptr; + uint64_t posting_off0_ = 0; + MemoryReporter::Reservation norms_section_reservation_; + std::vector norms_section_; + MemoryReporter::Reservation null_bitmap_section_reservation_; + std::vector null_bitmap_section_; + + std::vector blocks_; + MemoryReporter::Reservation term_hashes_reservation_; + MemoryReporter::Reservation bsbf_bytes_reservation_; + // One 8-byte XXH64 (seed 0) filter key per term, collected during the build pass + // so the whole-vocabulary string copy is never retained. + std::vector term_hashes_; + format::StatsBlock stats_; + std::vector bsbf_bytes_; // serialized block-split bloom XFilter section + bool bsbf_built_ = false; + + // Streamed-session state (T2.1). kIdle until build()/begin_streamed claims + // the writer; build() jumps straight to kFinished on completion so the two + // entry points can never interleave on one instance (single-session + // invariant). kFailed is the poison state: any failure PAST entry + // validation (process_term inside push_term, build_blocks inside build(), + // flush/finalize inside either finish path) may have left partial posting + // bytes in the sink or partial term state (term_hashes_/stats_), so the + // writer transitions to kFailed and every subsequent + // push/finish/build/begin is rejected -- a half-fed session can never + // masquerade as a sealed index (crash-safety invariant 6), and a dirty + // writer can never be re-claimed for a second session. Pure entry + // rejections (phase check, term-order guard, prepare_build validation) + // mutate nothing and therefore do NOT poison. + enum class StreamPhase : uint8_t { kIdle, kActive, kFinished, kFailed }; + StreamPhase stream_phase_ = StreamPhase::kIdle; + std::unique_ptr stream_state_; // live only while kActive + std::string last_pushed_term_; // strict term-order entry guard + bool has_pushed_term_ = false; + + friend class StreamingTermEncoder; +}; + +// TEST-ONLY observability seam (mirrors the reader-side decode-counter and the +// SPIMI vocab-materialization patterns). term_freq_scans() returns a +// process-global count of term-level fused freqs scans. Ordinary terms call +// fuse_freq_stats exactly once; complete CommonGrams entries call it zero times. +// note_term_freq_scan() bumps the counter; reset_term_freq_scans() zeroes it +// between tests; fuse_freq_stats_for_test() exposes the real fused helper so +// pure boundary tests exercise production code. Process-global; reset between +// tests. Not part of the production API. +namespace testing { +void note_term_freq_scan(); +uint64_t term_freq_scans(); +void reset_term_freq_scans(); +FreqStats fuse_freq_stats_for_test(const std::vector& freqs); +} // namespace testing + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/memory_reporter.h b/be/src/storage/index/snii/writer/memory_reporter.h new file mode 100644 index 00000000000000..04875118f96872 --- /dev/null +++ b/be/src/storage/index/snii/writer/memory_reporter.h @@ -0,0 +1,227 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "common/status.h" + +namespace doris::snii::writer { + +// Per-WRITER accurate byte counter for build-time RAM (one per SniiCompoundWriter = +// one per segment's inverted index). Legacy modules report resident-byte deltas +// after their allocation; hard-gated modules own a Reservation that atomically +// pre-charges before allocation. current_bytes() is their shared live total. +// consume_release mirrors successful changes into Doris's LOAD MemTracker; it is +// null off-Doris (bench / unit tests), where only the local atomic is updated. +class MemoryReporter { +public: + // The callback may be invoked concurrently and from Reservation destructors; + // it must be thread-safe and must not throw. Null off-Doris. + using ConsumeReleaseFn = std::function; + + enum class CapPolicy : uint8_t { + // Reservations fail before an allocation would cross the cap. Native + // compaction uses this policy so an over-budget merge can fall back to + // the raw-column rebuild path without exceeding its bounded workspace. + kHardLimit, + // The cap is a spill trigger, not an allocation limit. Ordinary index + // ingestion uses this policy because persistent vocabulary structures + // can exceed the reclaimable posting-arena threshold by design. + kSpillThreshold, + }; + + // Move-only ownership of bytes pre-charged against this reporter. Growing a + // reservation atomically charges before allocation. Hard-limit reporters + // reject an over-cap charge without changing state; spill-threshold reporters + // retain exact accounting above the threshold. Callers must release/shrink the + // physical buffer before lowering the reservation. A Reservation borrows its + // reporter, which must outlive it. + class Reservation { + public: + Reservation() = default; + Reservation(const Reservation&) = delete; + Reservation& operator=(const Reservation&) = delete; + Reservation(Reservation&& other) noexcept; + Reservation& operator=(Reservation&& other) noexcept; + ~Reservation(); + + Status set_bytes(uint64_t target_bytes); + // Pre-charges an independent allocation while this Reservation keeps + // covering the old one. After the physical replacement succeeds, move + // `replacement` back into this Reservation to release the old charge. + Status prepare_replacement(uint64_t target_bytes, Reservation* replacement) const; + void reset(); + uint64_t bytes() const { return bytes_; } + + private: + friend class MemoryReporter; + explicit Reservation(MemoryReporter* owner) : owner_(owner) {} + + MemoryReporter* owner_ = nullptr; + uint64_t bytes_ = 0; + }; + + // cap_bytes is the shared gate-2 threshold (0 = unlimited). Hard-limit + // reporters reject reservations before their covered allocations cross it. + // Spill-threshold reporters keep exact accounting above it so over_cap() can + // drive reclaim without turning irreducible vocabulary growth into an import + // failure. + explicit MemoryReporter(ConsumeReleaseFn consume_release = nullptr, uint64_t cap_bytes = 0, + CapPolicy cap_policy = CapPolicy::kHardLimit) + : consume_release_(std::move(consume_release)), + cap_bytes_(cap_bytes), + cap_policy_(cap_policy) {} + + MemoryReporter(const MemoryReporter&) = delete; + MemoryReporter& operator=(const MemoryReporter&) = delete; + + Reservation make_reservation() { return Reservation(this); } + + // Observe-only legacy path: delta > 0 grows, delta < 0 shrinks/frees. New + // hard-gated allocations must use Reservation instead. + void report(int64_t delta) { + if (delta == 0) return; + DCHECK_NE(delta, std::numeric_limits::min()); + int64_t current = current_.load(std::memory_order_relaxed); + while (true) { + DCHECK_GE(current, 0); + if (delta > 0) { + DCHECK_LE(delta, std::numeric_limits::max() - current); + } else { + DCHECK_GE(current, -delta); + } + const int64_t desired = current + delta; + if (current_.compare_exchange_weak(current, desired, std::memory_order_relaxed, + std::memory_order_relaxed)) { + if (consume_release_) consume_release_(delta); + return; + } + } + } + + int64_t current_bytes() const { return current_.load(std::memory_order_relaxed); } + + // True once all reported/reserved build RAM reaches the shared spill threshold. + bool over_cap() const { + const int64_t current = current_bytes(); + DCHECK_GE(current, 0); + return cap_bytes_ != 0 && static_cast(current) >= cap_bytes_; + } + uint64_t cap_bytes() const { return cap_bytes_; } + +private: + Status try_acquire(uint64_t bytes); + void release(uint64_t bytes); + + std::atomic current_ {0}; + ConsumeReleaseFn consume_release_; + uint64_t cap_bytes_ = 0; + CapPolicy cap_policy_ = CapPolicy::kHardLimit; +}; + +inline MemoryReporter::Reservation::Reservation(Reservation&& other) noexcept + : owner_(std::exchange(other.owner_, nullptr)), bytes_(std::exchange(other.bytes_, 0)) {} + +inline MemoryReporter::Reservation& MemoryReporter::Reservation::operator=( + Reservation&& other) noexcept { + if (this != &other) { + reset(); + owner_ = std::exchange(other.owner_, nullptr); + bytes_ = std::exchange(other.bytes_, 0); + } + return *this; +} + +inline MemoryReporter::Reservation::~Reservation() { + reset(); +} + +inline Status MemoryReporter::Reservation::set_bytes(uint64_t target_bytes) { + DORIS_CHECK(owner_ != nullptr); + if (target_bytes > bytes_) { + RETURN_IF_ERROR(owner_->try_acquire(target_bytes - bytes_)); + } else if (target_bytes < bytes_) { + owner_->release(bytes_ - target_bytes); + } + bytes_ = target_bytes; + return Status::OK(); +} + +inline Status MemoryReporter::Reservation::prepare_replacement(uint64_t target_bytes, + Reservation* replacement) const { + DORIS_CHECK(owner_ != nullptr); + DORIS_CHECK(replacement != nullptr); + DORIS_CHECK(replacement->owner_ == nullptr); + Reservation pending(owner_); + RETURN_IF_ERROR(pending.set_bytes(target_bytes)); + *replacement = std::move(pending); + return Status::OK(); +} + +inline void MemoryReporter::Reservation::reset() { + if (owner_ != nullptr && bytes_ != 0) { + owner_->release(bytes_); + bytes_ = 0; + } +} + +inline Status MemoryReporter::try_acquire(uint64_t bytes) { + if (bytes == 0) { + return Status::OK(); + } + int64_t current = current_.load(std::memory_order_relaxed); + while (true) { + DCHECK_GE(current, 0); + const uint64_t current_bytes = static_cast(current); + const bool exceeds_cap = cap_policy_ == CapPolicy::kHardLimit && cap_bytes_ != 0 && + (current_bytes > cap_bytes_ || bytes > cap_bytes_ - current_bytes); + const bool exceeds_counter = + bytes > static_cast(std::numeric_limits::max()) - current_bytes; + if (exceeds_cap || exceeds_counter) { + return Status::Error( + "SNII memory reservation exceeds limit: request={} current={} cap={}", bytes, + current_bytes, cap_bytes_); + } + const int64_t desired = current + static_cast(bytes); + if (current_.compare_exchange_weak(current, desired, std::memory_order_relaxed, + std::memory_order_relaxed)) { + if (consume_release_) { + consume_release_(static_cast(bytes)); + } + return Status::OK(); + } + } +} + +inline void MemoryReporter::release(uint64_t bytes) { + DCHECK_LE(bytes, static_cast(std::numeric_limits::max())); + const int64_t delta = static_cast(bytes); + const int64_t previous = current_.fetch_sub(delta, std::memory_order_relaxed); + DCHECK_GE(previous, delta); + if (consume_release_) { + consume_release_(-delta); + } +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/posting_window_emitter.cpp b/be/src/storage/index/snii/writer/posting_window_emitter.cpp new file mode 100644 index 00000000000000..709d59b72c9ee8 --- /dev/null +++ b/be/src/storage/index/snii/writer/posting_window_emitter.cpp @@ -0,0 +1,579 @@ +// 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 "storage/index/snii/writer/posting_window_emitter.h" + +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/spillable_byte_buffer.h" + +namespace doris::snii::writer { + +namespace { + +constexpr int kRawFrqRegion = 0; +constexpr uint32_t kPreludeGroupSize = 64; + +struct PostingWindowPlan { + size_t doc_begin = 0; + size_t doc_count = 0; + uint64_t position_begin = 0; + uint64_t position_count = 0; + uint32_t max_freq = 0; +}; + +bool fits_prx_window_shape(uint64_t doc_count, uint64_t position_count, + const format::PrxWindowLimits& limits) { + return doc_count <= limits.max_docs && position_count <= limits.max_positions; +} + +// Five bytes is the maximum encoded width of every uint32 field in the raw +// payload. This gate is used only after an exact build requests a split. +bool conservatively_fits_prx_window(uint64_t doc_count, uint64_t position_count, + const format::PrxWindowLimits& limits) { + return fits_prx_window_shape(doc_count, position_count, limits) && + 1 + doc_count + position_count <= limits.max_uncomp_bytes / 5; +} + +uint8_t window_max_norm(std::span norms, std::span docs) { + if (norms.empty() || docs.empty()) { + return 0; + } +#ifdef BE_TEST + testing::note_window_norm_doc_visits(docs.size()); +#endif + uint8_t best = 0xFF; + for (uint32_t docid : docs) { + DCHECK_LT(docid, norms.size()); + best = std::min(best, norms[docid]); + } + return best == 0xFF ? 0 : best; +} + +Status build_prelude(const std::vector& windows, bool has_freq, bool has_prx, + std::vector* output) { + format::FrqPreludeColumns columns; + columns.has_freq = has_freq; + columns.has_prx = has_prx; + columns.group_size = kPreludeGroupSize; + columns.windows = windows; + ByteSink sink; + RETURN_IF_ERROR(format::build_frq_prelude(columns, &sink)); + *output = sink.take(); + return Status::OK(); +} + +Status checked_add(uint64_t increment, uint64_t* value) { + if (increment > std::numeric_limits::max() - *value) { + return Status::Error( + "window emitter: term frequency overflow"); + } + *value += increment; + return Status::OK(); +} + +} // namespace + +class WindowEmitter::Impl { +public: + explicit Impl(WindowEmitterOptions options) + : options_(options), + dd_stager_(std::numeric_limits::max(), "term_dd", options.memory_reporter), + freq_stager_(std::numeric_limits::max(), "term_freq", + options.memory_reporter) { + if (options_.posting_out != nullptr && + options_.posting_out->bytes_written() >= options_.posting_region_offset) { + prx_off_ = options_.posting_out->bytes_written() - options_.posting_region_offset; + posting_offset_valid_ = true; + } + } + + Status emit_window(const PostingRunView& run) { + if (phase_ != Phase::kActive) { + return phase_error("emit_window"); + } + Status status = emit_window_impl(run); + if (!status.ok()) { + phase_ = Phase::kFailed; + } + return status; + } + + Status finish_term(format::DictEntry* entry, TermAggregateStats* stats) { + if (phase_ != Phase::kActive) { + return phase_error("finish_term"); + } + if (entry == nullptr || stats == nullptr) { + phase_ = Phase::kFailed; + return Status::Error( + "window emitter: null finish output"); + } + if (windows_.empty()) { + phase_ = Phase::kFailed; + return Status::Error( + "window emitter: cannot finish an empty term"); + } + Status status = finish_term_impl(entry); + if (!status.ok()) { + phase_ = Phase::kFailed; + return status; + } + *stats = stats_; + phase_ = Phase::kFinished; +#ifdef BE_TEST + finished_term_counter().fetch_add(1, std::memory_order_relaxed); +#endif + return Status::OK(); + } + +private: + enum class Phase : uint8_t { kActive, kFinished, kFailed }; + + Status phase_error(std::string_view operation) const { + return Status::Error( + "window emitter: {} after {}", operation, + phase_ == Phase::kFailed ? "failure" : "finish"); + } + + Status posting_size(uint64_t* size) const { + if (options_.posting_out == nullptr) { + return Status::Error( + "window emitter: null posting sink"); + } + if (!posting_offset_valid_ || + options_.posting_out->bytes_written() < options_.posting_region_offset) { + return Status::Error( + "window emitter: invalid posting region offset"); + } + *size = options_.posting_out->bytes_written() - options_.posting_region_offset; + return Status::OK(); + } + + Status validate_run(const PostingRunView& run) const { + if (options_.posting_out == nullptr) { + return Status::Error( + "window emitter: null posting sink"); + } + if (run.docids.empty()) { + return Status::Error( + "window emitter: empty posting window"); + } + if ((!run.freqs.empty() || options_.has_freq || options_.has_prx) && + run.freqs.size() != run.docids.size()) { + return Status::Error( + "window emitter: frequency shape must match documents"); + } + if (options_.term_frequency_source == TermFrequencySource::kPositions && + !options_.has_prx) { + return Status::Error( + "window emitter: position-derived statistics require PRX offsets"); + } + if (options_.has_prx) { + if (run.position_offsets.size() != run.docids.size() + 1) { + return Status::Error( + "window emitter: position offsets must have docs plus one entries"); + } + if (run.position_offsets.front() > run.position_offsets.back() || + run.position_offsets.back() - run.position_offsets.front() != + run.positions_flat.size()) { + return Status::Error( + "window emitter: position offsets differ from the position run"); + } + } else if (!run.position_offsets.empty() || !run.positions_flat.empty()) { + return Status::Error( + "window emitter: positions require a PRX term"); + } + if (last_input_docid_.has_value() && run.docids.front() <= *last_input_docid_) { + return Status::Error( + "window emitter: posting windows must be strictly ordered"); + } + return Status::OK(); + } + + Status accumulate_constant_stats(const PostingRunView& run) { + if (run.docids.size() > std::numeric_limits::max() - stats_.df) { + return Status::Error( + "window emitter: document frequency overflow"); + } + stats_.df += static_cast(run.docids.size()); + switch (options_.term_frequency_source) { + case TermFrequencySource::kDocuments: + return checked_add(run.docids.size(), &stats_.total_freq); + case TermFrequencySource::kPositions: + return checked_add(run.position_offsets.back() - run.position_offsets.front(), + &stats_.total_freq); + case TermFrequencySource::kFrequenciesOrDocuments: + if (run.freqs.empty()) { + return checked_add(run.docids.size(), &stats_.total_freq); + } + return Status::OK(); + } + __builtin_unreachable(); + } + + uint64_t position_count(const PostingRunView& run, size_t begin, size_t count) const { + if (!options_.has_prx) { + return 0; + } + return run.position_offsets[begin + count] - run.position_offsets[begin]; + } + + Status emit_window_impl(const PostingRunView& run) { + RETURN_IF_ERROR(validate_run(run)); + RETURN_IF_ERROR(accumulate_constant_stats(run)); + + const bool accumulate_frequencies = + options_.term_frequency_source == TermFrequencySource::kFrequenciesOrDocuments && + !run.freqs.empty(); + if (!options_.has_prx && !options_.has_freq && !accumulate_frequencies) { + RETURN_IF_ERROR( + emit_planned(run, make_plan(run, 0, run.docids.size(), /*max_freq=*/0))); + last_input_docid_ = run.docids.back(); + return Status::OK(); + } + + size_t window_begin = 0; + uint32_t window_max_freq = 0; + for (size_t doc = 0; doc < run.docids.size(); ++doc) { + const uint64_t document_positions = options_.has_prx ? position_count(run, doc, 1) : 0; + if (options_.has_prx && (run.position_offsets[doc + 1] < run.position_offsets[doc] || + document_positions != run.freqs[doc])) { + return Status::Error( + "window emitter: position offsets must match frequencies"); + } + if (options_.has_prx && document_positions > options_.prx_window_limits.max_positions) { + return Status::Error( + "window emitter: one document exceeds the PRX position limit"); + } + if (accumulate_frequencies) { + RETURN_IF_ERROR(checked_add(run.freqs[doc], &stats_.total_freq)); + stats_.max_freq = std::max(stats_.max_freq, run.freqs[doc]); + } + const uint64_t candidate_docs = doc - window_begin + 1; + const uint64_t candidate_positions = position_count(run, window_begin, candidate_docs); + if (doc != window_begin && options_.has_prx && + !fits_prx_window_shape(candidate_docs, candidate_positions, + options_.prx_window_limits)) { + RETURN_IF_ERROR(emit_planned( + run, make_plan(run, window_begin, doc - window_begin, window_max_freq))); + window_begin = doc; + window_max_freq = 0; + } + if (options_.has_freq) { +#ifdef BE_TEST + testing::note_window_freq_doc_visits(); +#endif + window_max_freq = std::max(window_max_freq, run.freqs[doc]); + } + } + RETURN_IF_ERROR(emit_planned( + run, + make_plan(run, window_begin, run.docids.size() - window_begin, window_max_freq))); + last_input_docid_ = run.docids.back(); + return Status::OK(); + } + + PostingWindowPlan make_plan(const PostingRunView& run, size_t begin, size_t count, + uint32_t max_freq) const { + return { + .doc_begin = begin, + .doc_count = count, + .position_begin = options_.has_prx ? run.position_offsets[begin] - + run.position_offsets.front() + : uint64_t {0}, + .position_count = position_count(run, begin, count), + .max_freq = max_freq, + }; + } + + Status emit_planned(const PostingRunView& run, const PostingWindowPlan& plan) { + format::PrxWindowBuildOutcome outcome = format::PrxWindowBuildOutcome::kBuilt; + RETURN_IF_ERROR(emit_physical_window(run, plan, &outcome)); + if (outcome == format::PrxWindowBuildOutcome::kBuilt) { + return Status::OK(); + } + + std::vector recut; + recut_window(run, plan, &recut); + for (const PostingWindowPlan& subplan : recut) { + outcome = format::PrxWindowBuildOutcome::kBuilt; + RETURN_IF_ERROR(emit_physical_window(run, subplan, &outcome)); + if (outcome == format::PrxWindowBuildOutcome::kNeedsSplit) { + return Status::Error( + "window emitter: one document exceeds the PRX byte limit"); + } + } + return Status::OK(); + } + + void recut_window(const PostingRunView& run, const PostingWindowPlan& input, + std::vector* output) const { + size_t window_begin = input.doc_begin; + uint32_t window_max_freq = 0; + const size_t input_end = input.doc_begin + input.doc_count; + for (size_t doc = input.doc_begin; doc < input_end; ++doc) { + const uint64_t candidate_docs = doc - window_begin + 1; + const uint64_t candidate_positions = position_count(run, window_begin, candidate_docs); + if (doc != window_begin && + !conservatively_fits_prx_window(candidate_docs, candidate_positions, + options_.prx_window_limits)) { + output->push_back( + make_plan(run, window_begin, doc - window_begin, window_max_freq)); + window_begin = doc; + window_max_freq = 0; + } + if (options_.has_freq) { +#ifdef BE_TEST + testing::note_window_freq_doc_visits(); +#endif + window_max_freq = std::max(window_max_freq, run.freqs[doc]); + } + } + output->push_back(make_plan(run, window_begin, input_end - window_begin, window_max_freq)); + } + + Status emit_physical_window(const PostingRunView& run, const PostingWindowPlan& plan, + format::PrxWindowBuildOutcome* outcome) { + const auto docs = run.docids.subspan(plan.doc_begin, plan.doc_count); + const auto freqs = run.freqs.empty() ? std::span {} + : run.freqs.subspan(plan.doc_begin, plan.doc_count); + format::WindowMeta window; + window.last_docid = docs.back(); + window.win_base = window_base_; + window.doc_count = static_cast(docs.size()); + window.max_freq = options_.has_freq ? plan.max_freq : 0; + window.max_norm = options_.has_freq ? window_max_norm(options_.encoded_norms, docs) : 0; + + if (options_.has_prx) { + const auto positions = + run.positions_flat.subspan(static_cast(plan.position_begin), + static_cast(plan.position_count)); + prx_scratch_.clear(); + RETURN_IF_ERROR(format::try_build_prx_window_flat( + positions, freqs, -options_.prx_zstd_level, options_.prx_window_limits, + &prx_scratch_, outcome)); + if (*outcome == format::PrxWindowBuildOutcome::kNeedsSplit) { + return Status::OK(); + } + window.prx_off = prx_total_len_; + window.prx_len = prx_scratch_.size(); + RETURN_IF_ERROR(options_.posting_out->append(prx_scratch_.view())); + prx_total_len_ += window.prx_len; + } else { + *outcome = format::PrxWindowBuildOutcome::kBuilt; + } + + ByteSink dd_sink; + format::FrqRegionMeta dd_meta; + window.dd_off = dd_stager_.size(); + RETURN_IF_ERROR( + format::build_dd_region(docs, window_base_, kRawFrqRegion, &dd_sink, &dd_meta)); + window.dd_zstd = dd_meta.zstd; + window.dd_disk_len = dd_meta.disk_len; + window.dd_uncomp_len = dd_meta.uncomp_len; + window.crc_dd = dd_meta.crc; + RETURN_IF_ERROR(dd_stager_.append_move(dd_sink.take())); + + if (options_.has_freq) { + ByteSink freq_sink; + format::FrqRegionMeta freq_meta; + window.freq_off = freq_stager_.size(); + RETURN_IF_ERROR( + format::build_freq_region(freqs, kRawFrqRegion, &freq_sink, &freq_meta)); + window.freq_zstd = freq_meta.zstd; + window.freq_disk_len = freq_meta.disk_len; + window.freq_uncomp_len = freq_meta.uncomp_len; + window.crc_freq = freq_meta.crc; + RETURN_IF_ERROR(freq_stager_.append_move(freq_sink.take())); + } + + windows_.push_back(window); + window_base_ = window.last_docid; +#ifdef BE_TEST + physical_window_counter().fetch_add(1, std::memory_order_relaxed); +#endif + return Status::OK(); + } + + Status finish_term_impl(format::DictEntry* entry) { + std::vector prelude; + RETURN_IF_ERROR(build_prelude(windows_, options_.has_freq, options_.has_prx, &prelude)); + entry->kind = format::DictEntryKind::kPodRef; + entry->enc = format::DictEntryEnc::kWindowed; + entry->has_sb = true; + entry->prelude_len = prelude.size(); + entry->frq_docs_len = entry->prelude_len + dd_stager_.size(); + + uint64_t frq_off = 0; + RETURN_IF_ERROR(posting_size(&frq_off)); + RETURN_IF_ERROR(options_.posting_out->append(Slice(prelude))); + RETURN_IF_ERROR(dd_stager_.seal()); + RETURN_IF_ERROR(dd_stager_.stream_into_and_release(options_.posting_out)); + RETURN_IF_ERROR(freq_stager_.seal()); + RETURN_IF_ERROR(freq_stager_.stream_into_and_release(options_.posting_out)); + entry->frq_off_delta = frq_off - options_.frq_base; + uint64_t end = 0; + RETURN_IF_ERROR(posting_size(&end)); + entry->frq_len = end - frq_off; + if (options_.has_prx) { + entry->prx_off_delta = prx_off_ - options_.prx_base; + entry->prx_len = prx_total_len_; + } + return Status::OK(); + } + +#ifdef BE_TEST + static std::atomic& finished_term_counter(); + static std::atomic& physical_window_counter(); +#endif + + WindowEmitterOptions options_; + SpillableByteBuffer dd_stager_; + SpillableByteBuffer freq_stager_; + std::vector windows_; + ByteSink prx_scratch_; + TermAggregateStats stats_; + std::optional last_input_docid_; + uint64_t prx_off_ = 0; + uint64_t prx_total_len_ = 0; + uint64_t window_base_ = 0; + bool posting_offset_valid_ = false; + Phase phase_ = Phase::kActive; +}; + +#ifdef BE_TEST +namespace { +std::atomic& window_norm_doc_visit_counter() { + static std::atomic counter {0}; + return counter; +} +std::atomic& window_freq_doc_visit_counter() { + static std::atomic counter {0}; + return counter; +} +std::atomic& emitter_finished_term_counter() { + static std::atomic counter {0}; + return counter; +} +std::atomic& emitter_physical_window_counter() { + static std::atomic counter {0}; + return counter; +} +} // namespace + +std::atomic& WindowEmitter::Impl::finished_term_counter() { + return emitter_finished_term_counter(); +} + +std::atomic& WindowEmitter::Impl::physical_window_counter() { + return emitter_physical_window_counter(); +} +#endif + +WindowEmitter::WindowEmitter(WindowEmitterOptions options) + : impl_(std::make_unique(options)) {} + +WindowEmitter::~WindowEmitter() = default; + +Status WindowEmitter::emit_window(const PostingRunView& window) { + return impl_->emit_window(window); +} + +Status WindowEmitter::finish_term(format::DictEntry* entry, TermAggregateStats* stats) { + return impl_->finish_term(entry, stats); +} + +namespace testing { + +void note_window_norm_doc_visits(uint64_t count) { +#ifdef BE_TEST + window_norm_doc_visit_counter().fetch_add(count, std::memory_order_relaxed); +#endif +} + +uint64_t window_norm_doc_visits() { +#ifdef BE_TEST + return window_norm_doc_visit_counter().load(std::memory_order_relaxed); +#else + return 0; +#endif +} + +void reset_window_norm_doc_visits() { +#ifdef BE_TEST + window_norm_doc_visit_counter().store(0, std::memory_order_relaxed); +#endif +} + +void note_window_freq_doc_visits() { +#ifdef BE_TEST + window_freq_doc_visit_counter().fetch_add(1, std::memory_order_relaxed); +#endif +} + +uint64_t window_freq_doc_visits() { +#ifdef BE_TEST + return window_freq_doc_visit_counter().load(std::memory_order_relaxed); +#else + return 0; +#endif +} + +void reset_window_freq_doc_visits() { +#ifdef BE_TEST + window_freq_doc_visit_counter().store(0, std::memory_order_relaxed); +#endif +} + +uint64_t window_emitter_finished_terms() { +#ifdef BE_TEST + return emitter_finished_term_counter().load(std::memory_order_relaxed); +#else + return 0; +#endif +} + +uint64_t window_emitter_physical_windows() { +#ifdef BE_TEST + return emitter_physical_window_counter().load(std::memory_order_relaxed); +#else + return 0; +#endif +} + +void reset_window_emitter_counters() { +#ifdef BE_TEST + emitter_finished_term_counter().store(0, std::memory_order_relaxed); + emitter_physical_window_counter().store(0, std::memory_order_relaxed); +#endif +} + +} // namespace testing + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/posting_window_emitter.h b/be/src/storage/index/snii/writer/posting_window_emitter.h new file mode 100644 index 00000000000000..e6b0e3444b31fd --- /dev/null +++ b/be/src/storage/index/snii/writer/posting_window_emitter.h @@ -0,0 +1,112 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/prx_pod.h" + +namespace doris::snii { +namespace format { +struct DictEntry; +} +namespace io { +class FileWriter; +} +namespace writer { + +class MemoryReporter; + +// Synchronously borrowed posting arrays for one canonical posting window. +// position_offsets has docids.size()+1 entries when positions are present. The +// offsets may start above zero when the view borrows a sub-window; positions_flat +// covers exactly position_offsets.back()-position_offsets.front() values. +struct PostingRunView { + std::span docids; + std::span freqs; + std::span position_offsets; + std::span positions_flat; +}; + +struct TermAggregateStats { + uint32_t df = 0; + uint64_t total_freq = 0; + uint32_t max_freq = 0; +}; + +// CommonGrams entries define semantic term frequency from documents or +// positions rather than the transient physical frequency array. +enum class TermFrequencySource : uint8_t { + kFrequenciesOrDocuments, + kDocuments, + kPositions, +}; + +struct WindowEmitterOptions { + io::FileWriter* posting_out = nullptr; + uint64_t posting_region_offset = 0; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + std::span encoded_norms; + bool has_freq = false; + bool has_prx = false; + int prx_zstd_level = 3; + format::PrxWindowLimits prx_window_limits = format::kReaderPrxWindowLimits; + TermFrequencySource term_frequency_source = TermFrequencySource::kFrequenciesOrDocuments; + MemoryReporter* memory_reporter = nullptr; +}; + +// The single owner of windowed DD/frequency/PRX encoding and prelude metadata. +// A failed emit poisons the instance; finish_term cannot publish a partial term. +class WindowEmitter { +public: + explicit WindowEmitter(WindowEmitterOptions options); + ~WindowEmitter(); + + WindowEmitter(const WindowEmitter&) = delete; + WindowEmitter& operator=(const WindowEmitter&) = delete; + WindowEmitter(WindowEmitter&&) = delete; + WindowEmitter& operator=(WindowEmitter&&) = delete; + + Status emit_window(const PostingRunView& window); + Status finish_term(format::DictEntry* entry, TermAggregateStats* stats); + +private: + class Impl; + std::unique_ptr impl_; +}; + +// Process-global test observability for the one emitter choke point and the +// existing bounded-work counters. Reset between tests. +namespace testing { +void note_window_norm_doc_visits(uint64_t count); +uint64_t window_norm_doc_visits(); +void reset_window_norm_doc_visits(); +void note_window_freq_doc_visits(); +uint64_t window_freq_doc_visits(); +void reset_window_freq_doc_visits(); +uint64_t window_emitter_finished_terms(); +uint64_t window_emitter_physical_windows(); +void reset_window_emitter_counters(); +} // namespace testing + +} // namespace writer +} // namespace doris::snii diff --git a/be/src/storage/index/snii/writer/snii_compound_writer.cpp b/be/src/storage/index/snii/writer/snii_compound_writer.cpp new file mode 100644 index 00000000000000..472b9e1460cad4 --- /dev/null +++ b/be/src/storage/index/snii/writer/snii_compound_writer.cpp @@ -0,0 +1,398 @@ +// 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 "storage/index/snii/writer/snii_compound_writer.h" + +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/format/bootstrap_header.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/metadata_directory.h" +#include "storage/index/snii/format/tail_pointer.h" + +namespace doris::snii::writer { + +using format::BootstrapHeader; +using format::LogicalIndexMetadataRef; +using format::SectionRefs; +using format::TailPointer; + +SniiCompoundWriter::SniiCompoundWriter(io::FileWriter* out) : out_(out) {} + +Status SniiCompoundWriter::poison(Status status) { + DCHECK(!status.ok()); + if (failed_.ok()) { + failed_ = std::move(status); + } + return failed_; +} + +Status SniiCompoundWriter::append(const std::vector& bytes) { + if (bytes.empty()) return Status::OK(); + return out_->append(Slice(bytes)); +} + +// The bootstrap header occupies offset 0 and must precede the first posting region, +// which streams straight into the output during build(). Written lazily exactly once +// (on the first add, or in finish() for an empty container). +Status SniiCompoundWriter::ensure_bootstrap() { + if (!failed_.ok()) return failed_; + if (bootstrap_written_) return Status::OK(); + const Status status = write_bootstrap(); + if (!status.ok()) return poison(status); + bootstrap_written_ = true; + return Status::OK(); +} + +Status SniiCompoundWriter::add_logical_index(const SniiIndexInput& in) { + if (out_ == nullptr) + return Status::Error("compound: null file writer"); + if (finished_) + return Status::Error("compound: add after finish"); + if (!failed_.ok()) return failed_; + if (has_active_session()) + return Status::Error( + "compound: add_logical_index while a streamed index session is active (its " + "posting region streams straight into the output; interleaving would corrupt " + "both indexes)"); + RETURN_IF_ERROR(ensure_bootstrap()); + auto liw = std::make_unique(in); + Placement p; + // The posting region streams DIRECTLY into the container during build() -- no temp + // round-trip for the bulk -- followed immediately by this index's compact DICT + // trailer (produced interleaved into a temp, but laid out right after its posting + // region, preserving the per-index [posting][dict] layout). Offsets are read off + // the output writer (the single source of truth -- no separate cursor). + p.post_off = out_->bytes_written(); + Status status = liw->build(out_); + if (!status.ok()) return poison(status); + p.post_len = out_->bytes_written() - p.post_off; + p.dict_off = out_->bytes_written(); + status = liw->stream_dict_region_into(out_); + if (!status.ok()) return poison(status); + p.dict_len = out_->bytes_written() - p.dict_off; + indexes_.push_back(std::move(liw)); + placements_.push_back(p); + return Status::OK(); +} + +SniiIndexInput SniiStreamedIndexSession::attach_encoded_norms(SniiIndexInput in, + TrackedEncodedNorms* encoded_norms, + uint64_t reserved_bytes) { + DORIS_CHECK(encoded_norms != nullptr); + DORIS_CHECK(in.encoded_norms.empty()); + if (in.mem_reporter != nullptr) { + DORIS_CHECK_EQ(reserved_bytes, encoded_norms->norms_.capacity()); + } else { + DORIS_CHECK_EQ(reserved_bytes, 0); + } + in.encoded_norms = std::move(encoded_norms->norms_); + return in; +} + +SniiStreamedIndexSession::SniiStreamedIndexSession(SniiCompoundWriter* owner, SniiIndexInput in, + TrackedNullDocids null_docids, + TrackedEncodedNorms encoded_norms) + : owner_(owner), + encoded_norms_reservation_(std::move(encoded_norms.reservation_)), + input_(attach_encoded_norms(std::move(in), &encoded_norms, + encoded_norms_reservation_.bytes())), + // input_ (a member, initialized above) owns the vectors the writer + // keeps references into -- NOT the caller's already-moved-from `in`. + writer_(new LogicalIndexWriter(input_, std::move(null_docids))), + semantic_token_count_required_( + input_.common_grams_metadata.has_value() && + input_.common_grams_metadata->scoring_coverage == + segment_v2::inverted_index::ScoringCoverage::kComplete) {} + +Status SniiStreamedIndexSession::push_term(StreamedTermPostings&& tp) { + if (!owner_->failed_.ok()) return owner_->failed_; + if (finished_) { + return Status::Error( + "compound: push_term on a finished streamed index session"); + } + const Status status = writer_->push_term(std::move(tp)); + if (!status.ok()) return owner_->poison(status); + return Status::OK(); +} + +Status SniiStreamedIndexSession::set_semantic_token_count(uint64_t token_count) { + if (!owner_->failed_.ok()) return owner_->failed_; + if (finished_) { + return Status::Error( + "compound: semantic token count on a finished streamed index session"); + } + if (!semantic_token_count_required_) { + return Status::Error( + "compound: streamed index session has no complete semantic scoring metadata"); + } + if (semantic_token_count_set_) { + return Status::Error( + "compound: semantic token count was already set"); + } + DORIS_CHECK(input_.common_grams_metadata.has_value()); + DORIS_CHECK(writer_->common_grams_metadata_.has_value()); + input_.common_grams_metadata->scoring_token_count = token_count; + writer_->common_grams_metadata_->scoring_token_count = token_count; + semantic_token_count_set_ = true; + return Status::OK(); +} + +Status SniiStreamedIndexSession::finish() { + if (!owner_->failed_.ok()) return owner_->failed_; + if (finished_) { + return Status::Error( + "compound: finish on an already-finished streamed index session"); + } + if (semantic_token_count_required_ && !semantic_token_count_set_) { + return owner_->poison(Status::Error( + "compound: semantic token count must be set before streamed index finish")); + } + return owner_->finish_streamed_index(this); +} + +void SniiStreamedIndexSession::abort(const Status& cause) { + DCHECK(!cause.ok()); + static_cast(owner_->poison(cause)); +} + +Status SniiCompoundWriter::begin_streamed_index(SniiIndexInput in, + SniiStreamedIndexSession** session) { + std::vector null_docids; + null_docids.swap(in.null_docids); + MemoryReporter::Reservation null_docids_reservation = + in.mem_reporter == nullptr ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation(); + if (in.mem_reporter != nullptr) { + RETURN_IF_ERROR(null_docids_reservation.set_bytes( + static_cast(null_docids.capacity()) * sizeof(uint32_t))); + } + return begin_streamed_index( + std::move(in), + TrackedNullDocids(std::move(null_docids_reservation), std::move(null_docids)), session); +} + +Status SniiCompoundWriter::begin_streamed_index(SniiIndexInput in, TrackedNullDocids null_docids, + SniiStreamedIndexSession** session) { + std::vector encoded_norms; + encoded_norms.swap(in.encoded_norms); + MemoryReporter::Reservation encoded_norms_reservation = + in.mem_reporter == nullptr ? MemoryReporter::Reservation() + : in.mem_reporter->make_reservation(); + if (in.mem_reporter != nullptr) { + RETURN_IF_ERROR(encoded_norms_reservation.set_bytes(encoded_norms.capacity())); + } + return begin_streamed_index( + std::move(in), std::move(null_docids), + TrackedEncodedNorms(std::move(encoded_norms_reservation), std::move(encoded_norms)), + session); +} + +Status SniiCompoundWriter::begin_streamed_index(SniiIndexInput in, TrackedNullDocids null_docids, + TrackedEncodedNorms encoded_norms, + SniiStreamedIndexSession** session) { + if (session == nullptr) + return Status::Error( + "compound: null session out parameter"); + *session = nullptr; + if (out_ == nullptr) + return Status::Error("compound: null file writer"); + if (finished_) + return Status::Error("compound: begin after finish"); + if (!failed_.ok()) return failed_; + if (has_active_session()) + return Status::Error( + "compound: a streamed index session is already active (one at a time: its " + "posting region streams straight into the container output)"); + // A streamed session takes terms ONLY via push_term; a term source or a + // materialized vector would silently be ignored by the streamed writer. + if (in.term_source != nullptr || !in.terms.empty()) + return Status::Error( + "compound: a streamed index session must not carry a term source or " + "materialized terms"); + if (!in.null_docids.empty()) + return Status::Error( + "compound: tracked streamed NULL docids must not also be present in input"); + if (!in.encoded_norms.empty()) + return Status::Error( + "compound: tracked streamed norms must not also be present in input"); + RETURN_IF_ERROR(ensure_bootstrap()); + auto s = std::unique_ptr(new SniiStreamedIndexSession( + this, std::move(in), std::move(null_docids), std::move(encoded_norms))); + s->post_off_ = out_->bytes_written(); + RETURN_IF_ERROR(s->writer_->begin_streamed(out_)); + sessions_.push_back(std::move(s)); + *session = sessions_.back().get(); + return Status::OK(); +} + +Status SniiCompoundWriter::finish_streamed_index(SniiStreamedIndexSession* session) { + // Flushes the trailing DICT block and finalizes the stats / null-bitmap / + // BSBF sections (poisoning the logical writer on failure). + Status status = session->writer_->finish_streamed(); + if (!status.ok()) return poison(status); + // finalize materialized the framed norms section. Drop the source vector and + // its transferred charge before retaining that section for compound finish. + std::vector().swap(session->input_.encoded_norms); + session->encoded_norms_reservation_.reset(); + Placement p; + p.post_off = session->post_off_; + p.post_len = out_->bytes_written() - p.post_off; + p.dict_off = out_->bytes_written(); + status = session->writer_->stream_dict_region_into(out_); + if (!status.ok()) return poison(status); + p.dict_len = out_->bytes_written() - p.dict_off; + // Only now does the index join the container. Any failure above leaves + // finished_ false, so finish() keeps failing loudly instead of sealing a + // tail that silently omits an index whose posting bytes are in the file. + indexes_.push_back(std::move(session->writer_)); + placements_.push_back(p); + session->finished_ = true; + return Status::OK(); +} + +Status SniiCompoundWriter::write_bootstrap() { + BootstrapHeader bh; + bh.tail_pointer_size = static_cast(format::tail_pointer_size()); + ByteSink sink; + RETURN_IF_ERROR(format::encode_bootstrap_header(bh, &sink)); + return append(sink.buffer()); +} + +// Writes each index's norms POD then bsbf section (in add order), after all the +// per-index [posting][dict] regions. +Status SniiCompoundWriter::write_norms() { + for (size_t i = 0; i < indexes_.size(); ++i) { + const LogicalIndexWriter& w = *indexes_[i]; + if (!w.has_norms() || w.norms_bytes().empty()) continue; + Placement& p = placements_[i]; + p.norms_off = out_->bytes_written(); + RETURN_IF_ERROR(append(w.norms_bytes())); + p.norms_len = out_->bytes_written() - p.norms_off; + indexes_[i]->release_norms_bytes(); + } + for (size_t i = 0; i < indexes_.size(); ++i) { + LogicalIndexWriter& w = *indexes_[i]; + if (!w.has_null_bitmap()) continue; + Placement& p = placements_[i]; + p.null_off = out_->bytes_written(); + RETURN_IF_ERROR(append(w.null_bitmap_bytes())); + p.null_len = out_->bytes_written() - p.null_off; + w.release_null_bitmap_bytes(); + } + for (size_t i = 0; i < indexes_.size(); ++i) { + LogicalIndexWriter& w = *indexes_[i]; + if (!w.has_bsbf()) continue; + Placement& p = placements_[i]; + p.bsbf_off = out_->bytes_written(); + RETURN_IF_ERROR(append(w.bsbf_bytes())); + p.bsbf_len = out_->bytes_written() - p.bsbf_off; + w.release_bsbf_bytes(); + } + return Status::OK(); +} + +Status SniiCompoundWriter::write_tail() { + std::vector directory_entries; + directory_entries.reserve(indexes_.size()); + for (size_t i = 0; i < indexes_.size(); ++i) { + const LogicalIndexWriter& w = *indexes_[i]; + const Placement& p = placements_[i]; + + SectionRefs refs; + refs.dict_region = {.offset = p.dict_off, .length = p.dict_len}; + refs.posting_region = {.offset = p.post_off, .length = p.post_len}; + refs.norms = {.offset = p.norms_off, .length = p.norms_len}; + refs.null_bitmap = {.offset = p.null_off, .length = p.null_len}; + refs.bsbf = {.offset = p.bsbf_off, .length = p.bsbf_len}; + + SerializedMetadataGroup group; + RETURN_IF_ERROR(w.finish_metadata(refs, p.dict_off, &group)); + + LogicalIndexMetadataRef entry; + entry.index_id = w.index_id(); + entry.index_suffix = w.index_suffix(); + entry.core_metadata = {.offset = out_->bytes_written(), .length = group.core.size()}; + RETURN_IF_ERROR(append(group.core)); + DORIS_CHECK_EQ(out_->bytes_written(), + entry.core_metadata.offset + entry.core_metadata.length); + + entry.sampled_term_index = {.offset = out_->bytes_written(), + .length = group.sampled_term_index.size()}; + DORIS_CHECK_EQ(entry.sampled_term_index.offset, + entry.core_metadata.offset + entry.core_metadata.length); + RETURN_IF_ERROR(append(group.sampled_term_index)); + DORIS_CHECK_EQ(out_->bytes_written(), + entry.sampled_term_index.offset + entry.sampled_term_index.length); + + entry.dict_block_directory = {.offset = out_->bytes_written(), + .length = group.dict_block_directory.size()}; + DORIS_CHECK_EQ(entry.dict_block_directory.offset, + entry.sampled_term_index.offset + entry.sampled_term_index.length); + RETURN_IF_ERROR(append(group.dict_block_directory)); + DORIS_CHECK_EQ(out_->bytes_written(), + entry.dict_block_directory.offset + entry.dict_block_directory.length); + directory_entries.push_back(std::move(entry)); + } + + ByteSink directory_sink; + RETURN_IF_ERROR(format::encode_metadata_directory(directory_entries, &directory_sink)); + const uint64_t directory_offset = out_->bytes_written(); + RETURN_IF_ERROR(append(directory_sink.buffer())); + const uint64_t directory_length = out_->bytes_written() - directory_offset; + + TailPointer tp; + tp.directory_offset = directory_offset; + tp.directory_length = directory_length; + tp.directory_crc32c = crc32c(directory_sink.view()); + ByteSink tail_sink; + RETURN_IF_ERROR(format::encode_tail_pointer(tp, &tail_sink)); + return append(tail_sink.buffer()); +} + +Status SniiCompoundWriter::finish() { + if (out_ == nullptr) + return Status::Error("compound: null file writer"); + if (!failed_.ok()) { + return failed_; + } + if (finished_) + return Status::Error("compound: finish called twice"); + // Crash-safety invariant 6: a begun-but-unfinished streamed session already + // streamed posting bytes into the file but recorded no placement; sealing + // the container now would silently drop that index. Fail loudly -- the + // whole compaction round must fail and be retried. + if (has_active_session()) + return Status::Error( + "compound: finish with an unfinished streamed index session; the half-fed " + "index must never be sealed away silently"); + finished_ = true; + + RETURN_IF_ERROR(ensure_bootstrap()); // empty container still gets a header + Status status = write_norms(); + if (!status.ok()) return poison(status); + status = write_tail(); + if (!status.ok()) return poison(status); + status = out_->finalize(); + if (!status.ok()) return poison(status); + return Status::OK(); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/snii_compound_writer.h b/be/src/storage/index/snii/writer/snii_compound_writer.h new file mode 100644 index 00000000000000..48243581343075 --- /dev/null +++ b/be/src/storage/index/snii/writer/snii_compound_writer.h @@ -0,0 +1,206 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/logical_index_writer.h" + +// SniiCompoundWriter -- orchestrates a single-segment SNII container for one or +// more logical indexes, written front-to-back through an append-only +// io::FileWriter (no seek-back). It resolves all back-references by writing the +// metadata groups, raw directory, and fixed tail pointer LAST. +// +// CONTAINER LAYOUT PRODUCED (this is the on-disk contract the reader matches): +// [bootstrap_header] (kBootstrapHeaderSize bytes) +// for each logical index, in add order: +// [posting region] interleaved [prx][frq] per pod_ref term, term order +// (prx span empty when !has_prx) +// [DICT blocks region] concatenated DICT blocks, split by +// target_dict_block_bytes +// for each logical index, in add order: +// [norms POD] NormsPodWriter::finish (scoring only; else absent) +// [null bitmap POD] NullBitmapWriter::finish (when nulls exist) +// for each logical index, in add order: +// [Core metadata][SampledTermIndex blob][DICT block directory blob] +// [metadata directory] raw SniiMetadataDirectoryPB bytes +// [tail_pointer] encode_tail_pointer at EOF +// +// (The posting region is streamed BEFORE the DICT region per index: postings are +// the large append-only term-ordered stream; the DICT region is the compact +// compressed trailer.) +// +// OFFSET CONVENTIONS (ABSOLUTE file offsets unless stated otherwise): +// - SectionRefs in each Core metadata record ABSOLUTE file offset+length of +// that index's posting, DICT, norms, null-bitmap, and BSBF regions. Absent +// regions are (0,0); a present-but-empty posting region (all-INLINE index) +// is (off, 0). +// - DictBlockDirectory entries record each DICT block's ABSOLUTE file offset + +// length. +// - A windowed/slim pod_ref entry's absolute .frq offset = +// section_refs.posting_region.offset + frq_base + frq_off_delta +// where frq_base is the posting-region-relative running offset captured at the +// block's open (see logical_index_writer.h). prx follows the identical rule +// against the SAME region (prx_base == frq_base). +// - tail_pointer.directory_offset/length point at the raw metadata directory. +namespace doris::snii::writer { + +class SniiCompoundWriter; + +// T2.2 (compaction index merge fast path): handle for ONE streamed logical-index +// session inside a SniiCompoundWriter, obtained from begin_streamed_index(). The +// caller pushes lexicographically sorted terms (the k-way merge output) and seals +// the index with finish(), which lays the [posting][dict] regions out exactly like +// add_logical_index. The handle is owned by the compound writer and stays valid +// for the writer's lifetime; it must not outlive it. +// +// CRASH SAFETY (invariant 6): a session that was begun but never successfully +// finished keeps the container permanently unsealable -- its posting bytes are +// already in the file, so SniiCompoundWriter::finish() fails loudly instead of +// writing a tail that silently omits the half-fed index. +class SniiStreamedIndexSession { +public: + SniiStreamedIndexSession(const SniiStreamedIndexSession&) = delete; + SniiStreamedIndexSession& operator=(const SniiStreamedIndexSession&) = delete; + SniiStreamedIndexSession(SniiStreamedIndexSession&&) = delete; + SniiStreamedIndexSession& operator=(SniiStreamedIndexSession&&) = delete; + + // Terms must arrive in strictly increasing lexicographic order with + // ascending-docid postings. Unlike the standalone LogicalIndexWriter API, + // every rejection is terminal here because posting bytes may already have + // entered the compound output; all later calls return the first error. + Status push_term(StreamedTermPostings&& tp); + // Binds the semantic (plain-token) count after the merge's single postings + // pass. Complete scoring metadata requires exactly one call, including for + // an empty destination whose count is zero. + Status set_semantic_token_count(uint64_t token_count); + // Seals this index: flushes the trailing DICT block, streams the DICT region + // right after the posting region and records the placements. A failed finish + // leaves the session unfinished (and the container unsealable) -- there is + // no retry, the whole compaction round must fail and be redone. + Status finish(); + // Makes the owning compound permanently unsealable. Merge plans call this + // on every destination when any source or sibling destination fails, because + // a successfully written prefix is not a complete logical index. + void abort(const Status& cause); + bool finished() const { return finished_; } + +private: + friend class SniiCompoundWriter; + SniiStreamedIndexSession(SniiCompoundWriter* owner, SniiIndexInput in, + TrackedNullDocids null_docids, TrackedEncodedNorms encoded_norms); + static SniiIndexInput attach_encoded_norms(SniiIndexInput in, + TrackedEncodedNorms* encoded_norms, + uint64_t reserved_bytes); + + SniiCompoundWriter* owner_; + // The reservation precedes input_ so input_.encoded_norms is destroyed + // before its charge is released. + MemoryReporter::Reservation encoded_norms_reservation_; + // Owns the input: LogicalIndexWriter keeps references into it (terms / + // encoded_norms), so it must live exactly as long as the writer. + SniiIndexInput input_; + std::unique_ptr writer_; + uint64_t post_off_ = 0; + bool semantic_token_count_required_ = false; + bool semantic_token_count_set_ = false; + bool finished_ = false; +}; + +class SniiCompoundWriter { +public: + explicit SniiCompoundWriter(io::FileWriter* out); + SniiCompoundWriter(const SniiCompoundWriter&) = delete; + SniiCompoundWriter& operator=(const SniiCompoundWriter&) = delete; + SniiCompoundWriter(SniiCompoundWriter&&) = delete; + SniiCompoundWriter& operator=(SniiCompoundWriter&&) = delete; + + // Buffers one logical index: builds its section bytes and meta sub-sections. + // The actual file writing happens in finish() (single front-to-back pass). + Status add_logical_index(const SniiIndexInput& in); + + // T2.2: begins a STREAMED logical-index session (the compaction merge fast + // path) -- the caller pushes pre-merged terms through *session instead of + // handing the writer a term source, so `in` must carry NO term_source and NO + // materialized terms. Only ONE session may be active at a time (its posting + // region streams straight into the container output, so a concurrent + // add_logical_index or second session would interleave bytes); both are + // rejected while a session is unfinished, as is finish(). The returned + // handle is owned by this writer and valid for its lifetime. + Status begin_streamed_index(SniiIndexInput in, SniiStreamedIndexSession** session); + Status begin_streamed_index(SniiIndexInput in, TrackedNullDocids null_docids, + SniiStreamedIndexSession** session); + Status begin_streamed_index(SniiIndexInput in, TrackedNullDocids null_docids, + TrackedEncodedNorms encoded_norms, + SniiStreamedIndexSession** session); + + // Writes bootstrap header + all index sections + adjacent metadata groups + + // raw directory + tail pointer, then finalizes the underlying writer. + Status finish(); + +private: + // Absolute placement of one index's sections, resolved during finish(). + struct Placement { + uint64_t dict_off = 0; + uint64_t dict_len = 0; + uint64_t post_off = 0; // interleaved [prx][frq] posting region (was frq + prx) + uint64_t post_len = 0; + uint64_t norms_off = 0; + uint64_t norms_len = 0; + uint64_t null_off = 0; + uint64_t null_len = 0; + uint64_t bsbf_off = 0; + uint64_t bsbf_len = 0; + }; + + friend class SniiStreamedIndexSession; + + Status ensure_bootstrap(); + Status write_bootstrap(); + Status write_norms(); + Status write_tail(); + Status append(const std::vector& bytes); + Status poison(Status status); + // Seals one streamed session: records its placements and adopts its + // LogicalIndexWriter into indexes_ (only on FULL success -- see the + // crash-safety note on SniiStreamedIndexSession). + Status finish_streamed_index(SniiStreamedIndexSession* session); + // An unfinished streamed session (begun but not successfully sealed): + // blocks add_logical_index, another begin_streamed_index and finish(). + bool has_active_session() const { return !sessions_.empty() && !sessions_.back()->finished(); } + + io::FileWriter* out_; + std::vector> indexes_; + // Streamed sessions in begin order (at most the last one is unfinished). + // Owned here so the raw handles returned to callers stay valid for the + // writer's lifetime; a finished session is inert (its writer moved out). + std::vector> sessions_; + // Per-index placement; post_off/post_len are filled as each index's posting region + // streams in during add_logical_index, the rest during finish(). The absolute write + // offset is out_->bytes_written() (the single source of truth -- no separate cursor). + std::vector placements_; + bool bootstrap_written_ = false; + bool finished_ = false; + Status failed_ = Status::OK(); +}; + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/spill_run_codec.cpp b/be/src/storage/index/snii/writer/spill_run_codec.cpp new file mode 100644 index 00000000000000..b85f903c26e0f4 --- /dev/null +++ b/be/src/storage/index/snii/writer/spill_run_codec.cpp @@ -0,0 +1,937 @@ +// 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 "storage/index/snii/writer/spill_run_codec.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::writer { + +namespace { + +// Flush staging at this exact bound. A large write buffer (4 MiB) collapses the +// per-flush write() syscall count by ~64x: at 64 KiB the 5M build issued +// ~8800 write()s to ext4 (~9s of syscall overhead) for ~553 MiB of runs, versus +// a raw dd of the same bytes taking ~1.2s. Wide terms are appended and flushed +// in chunks, so the staging allocation never grows with term width. +constexpr size_t kWriteFlushBytes = 1u << 22; // 4 MiB +// RunReader reads this much per disk fill; the window slides so a single record +// never needs the whole run in RAM (only the current term's encoded span). KEEP +// this small (64 KiB): a large read chunk x many open runs would inflate the +// merge-phase peak RSS at low spill thresholds (each reader holds a window). +constexpr size_t kReadChunkBytes = 1u << 16; // 64 KiB + +enum class RunPostingShape : uint8_t { + kDocsOnlyStatless = 0, + kDocsAndFreqs = 1, + kPositioned = 2, +}; + +RunPostingShape posting_shape(const TermPostings& tp) { + if (tp.retain_positions) { + return RunPostingShape::kPositioned; + } + return tp.freqs.empty() ? RunPostingShape::kDocsOnlyStatless : RunPostingShape::kDocsAndFreqs; +} + +// Writes the full byte range [data, data+len) to fd, looping over short writes. +Status write_all(int fd, const uint8_t* data, size_t len) { + size_t off = 0; + while (off < len) { + const ssize_t n = ::write(fd, data + off, len - off); + if (n < 0) { + if (errno == EINTR) continue; + return Status::Error(std::string("run write failed: ") + + std::strerror(errno)); + } + off += static_cast(n); + } + return Status::OK(); +} + +template +Status reserve_vector_for_size(std::vector* values, size_t target, + MemoryReporter* memory_reporter, + MemoryReporter::Reservation* reservation) { + if (target <= values->capacity()) { + return Status::OK(); + } + if (target > std::numeric_limits::max() / sizeof(T)) { + return Status::Error( + "run reader: vector byte capacity overflow"); + } + if (memory_reporter == nullptr) { + values->reserve(target); + return Status::OK(); + } + MemoryReporter::Reservation replacement; + const uint64_t target_bytes = static_cast(target) * sizeof(T); + RETURN_IF_ERROR(reservation->prepare_replacement(target_bytes, &replacement)); + values->reserve(target); + DCHECK_EQ(values->capacity(), target); + *reservation = std::move(replacement); + return Status::OK(); +} + +Status reserve_write_buffer_for_append(std::vector* buffer, size_t target, + MemoryReporter* memory_reporter, + MemoryReporter::Reservation* reservation) { + if (target <= buffer->capacity()) { + return Status::OK(); + } + if (target > kWriteFlushBytes) { + return Status::Error( + "run writer: staging buffer exceeds flush bound"); + } + + size_t capacity = std::max(buffer->capacity(), 1); + while (capacity < target) { + capacity = capacity > kWriteFlushBytes / 2 ? kWriteFlushBytes : capacity * 2; + } + return reserve_vector_for_size(buffer, capacity, memory_reporter, reservation); +} + +} // namespace + +// --------------------------------------------------------------------------- +// RunWriter +// --------------------------------------------------------------------------- + +RunWriter::RunWriter(MemoryReporter* memory_reporter) + : memory_reporter_(memory_reporter), + buffer_reservation_(memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()) {} + +RunWriter::~RunWriter() { + if (fd_ >= 0) ::close(fd_); +} + +Status RunWriter::open(const std::string& path) { + fd_ = ::open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0600); + if (fd_ < 0) { + return Status::Error("run open(" + path + + "): " + std::strerror(errno)); + } + buf_.clear(); + return Status::OK(); +} + +Status RunWriter::flush() { + if (buf_.empty()) return Status::OK(); + RETURN_IF_ERROR(write_all(fd_, buf_.data(), buf_.size())); + buf_.clear(); + return Status::OK(); +} + +Status RunWriter::append_bytes(const uint8_t* data, size_t size) { + while (size != 0) { + if (buf_.size() == kWriteFlushBytes) { + RETURN_IF_ERROR(flush()); + } + const size_t count = std::min(size, kWriteFlushBytes - buf_.size()); + const size_t target = buf_.size() + count; + RETURN_IF_ERROR(reserve_write_buffer_for_append(&buf_, target, memory_reporter_, + &buffer_reservation_)); + buf_.insert(buf_.end(), data, data + count); + data += count; + size -= count; + } + return Status::OK(); +} + +Status RunWriter::append_varint(uint64_t value) { + uint8_t bytes[10]; + const size_t size = encode_varint64(value, bytes); + return append_bytes(bytes, size); +} + +Status RunWriter::append_raw_u32(const uint32_t* values, size_t count) { + if (count > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "run writer: raw u32 byte count overflows size_t"); + } + return append_bytes(reinterpret_cast(values), count * sizeof(uint32_t)); +} + +void RunWriter::release_buffer() { + std::vector().swap(buf_); + buffer_reservation_.reset(); +} + +Status RunWriter::write_term(uint32_t term_id, const TermPostings& tp) { + DCHECK(tp.retain_positions || tp.positions_flat.empty()); + const RunPostingShape shape = posting_shape(tp); + const size_t doc_count = tp.document_count(); + if (shape != RunPostingShape::kDocsOnlyStatless) { + DCHECK_EQ(tp.docids.size(), tp.freqs.size()); + } + RETURN_IF_ERROR(append_varint(term_id)); + RETURN_IF_ERROR(append_varint(static_cast(shape))); + RETURN_IF_ERROR(append_varint(doc_count)); + // Docids are a RAW fixed-width u32 block (bulk memcpy), NOT per-value VInt. + // Per-value varint over ~60M docids cost ~1.5s of encode CPU on the spill feed + // side; raw is a single memcpy and the decode side becomes a memcpy too. Runs + // are PRIVATE temp files written then read back from page cache, so the modestly + // larger run (no delta packing) costs ~0 extra real I/O. Absolute docids are + // stored (the merge concatenates per-term across runs and re-deltas at encode). + RETURN_IF_ERROR(append_raw_u32(tp.docids.data(), tp.docids.size())); + if (shape != RunPostingShape::kDocsOnlyStatless) { + RETURN_IF_ERROR(append_raw_u32(tp.freqs.data(), tp.freqs.size())); + } + if (shape == RunPostingShape::kPositioned) { + const uint64_t n_pos = tp.positions_flat.size(); + RETURN_IF_ERROR(append_varint(n_pos)); + RETURN_IF_ERROR(append_raw_u32(tp.positions_flat.data(), tp.positions_flat.size())); + } + return Status::OK(); +} + +Status RunWriter::close() { + if (fd_ < 0) return Status::OK(); + RETURN_IF_ERROR(flush()); + const int fd = fd_; + fd_ = -1; + if (::close(fd) != 0) { + return Status::Error(std::string("run close: ") + + std::strerror(errno)); + } + release_buffer(); + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// RunReader +// --------------------------------------------------------------------------- + +RunReader::RunReader(MemoryReporter* memory_reporter) + : memory_reporter_(memory_reporter), + window_reservation_(memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()), + docids_reservation_(memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()), + freqs_reservation_(memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()), + positions_reservation_(memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()) { +} + +RunReader::~RunReader() { + if (fd_ >= 0) ::close(fd_); +} + +Status RunReader::open(const std::string& path, bool has_positions) { + fd_ = ::open(path.c_str(), O_RDONLY); + if (fd_ < 0) { + return Status::Error("run reopen(" + path + + "): " + std::strerror(errno)); + } + // Record the run's byte size so every length decoded from the stream can be + // bounded against it before allocating (no record holds more u32s than the whole + // file). Honors the header's "lengths validated against the file size" contract, + // turning a corrupt/truncated length into Status::Corruption rather than an + // uncaught std::bad_alloc from a giant resize(). + struct stat st {}; + if (::fstat(fd_, &st) != 0) { + return Status::Error(std::string("run fstat: ") + + std::strerror(errno)); + } + file_size_ = static_cast(st.st_size); + bytes_read_ = 0; + has_positions_ = has_positions; + exhausted_ = false; + eof_ = false; + pos_ = 0; + pos_count_ = 0; + pos_remaining_ = 0; + window_.clear(); + return advance(); +} + +// Slides consumed bytes out of the window, then appends one disk chunk. +Status RunReader::fill() { + if (pos_ > 0) { + window_.erase(window_.begin(), window_.begin() + pos_); + pos_ = 0; + } + if (eof_) return Status::OK(); + if (bytes_read_ > file_size_) { + return Status::Error( + "run reader: bytes read exceed file size"); + } + const uint64_t remaining = file_size_ - bytes_read_; + if (remaining == 0) { + eof_ = true; + return Status::OK(); + } + const size_t read_size = static_cast( + std::min(remaining, static_cast(kReadChunkBytes))); + const size_t base = window_.size(); + if (base > std::numeric_limits::max() - read_size) { + return Status::Error( + "run reader: decode window capacity overflow"); + } + RETURN_IF_ERROR(reserve_vector_for_size(&window_, base + read_size, memory_reporter_, + &window_reservation_)); + window_.resize(base + read_size); + ssize_t n; + do { + n = ::read(fd_, window_.data() + base, read_size); + } while (n < 0 && errno == EINTR); + if (n < 0) + return Status::Error(std::string("run read: ") + + std::strerror(errno)); + window_.resize(base + static_cast(n)); + bytes_read_ += static_cast(n); + if (n == 0 || bytes_read_ == file_size_) eof_ = true; + return Status::OK(); +} + +// Buffered bytes available to the decoder right now (from pos_ to window end). +// fill() may slide the window (erasing consumed bytes), so callers must compare +// THIS quantity -- not window_.size() -- to decide whether more data arrived. +size_t RunReader::available() const { + return window_.size() - pos_; +} + +Status RunReader::ensure(size_t n) { + while (available() < n) { + const size_t had = available(); + RETURN_IF_ERROR(fill()); + if (available() == had && eof_) { + return Status::Error( + "run truncated: needed more bytes than available"); + } + } + return Status::OK(); +} + +// Streamed varint: decode from the current window; if it straddles the buffered +// boundary, top up from disk and retry. A varint is at most 10 bytes, so this +// loops at most a couple of times. Bounds-safe: decode_varint64 never reads past +// `end`, and a partial varint at true eof is reported as corruption. +Status RunReader::read_varint(uint64_t* v) { + while (true) { + const uint8_t* p = window_.data() + pos_; + const uint8_t* end = window_.data() + window_.size(); + const uint8_t* next = nullptr; + Status s = decode_varint64(p, end, v, &next); + if (s.ok()) { + pos_ += static_cast(next - p); + return Status::OK(); + } + if (eof_) + return Status::Error( + "run truncated: incomplete varint"); + const size_t had = available(); + RETURN_IF_ERROR(fill()); + if (available() == had && eof_) { + return Status::Error( + "run truncated: incomplete varint at eof"); + } + } +} + +// Streams `count` raw little-endian u32s from the window into `dst` (caller-owned +// storage of at least count*4 bytes), topping up the window from disk as needed. +// Copies whatever is buffered each pass (the window may hold only part of a large +// block), so a high-df term's freqs/positions stream through in 64 KiB chunks +// without ever needing the whole block resident at once. +Status RunReader::pull_raw_u32(uint8_t* dst, size_t count) { + if (count == 0) return Status::OK(); + if (count > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "run: raw u32 byte count overflows size_t"); + } + size_t need = count * sizeof(uint32_t); + size_t written = 0; + while (need > 0) { + if (available() == 0) { + const size_t had = available(); + RETURN_IF_ERROR(fill()); + if (available() == had && eof_) { + return Status::Error( + "run truncated: needed more raw bytes than available"); + } + } + const size_t take = std::min(need, available()); + std::memcpy(dst + written, window_.data() + pos_, take); + pos_ += take; + written += take; + need -= take; + } + return Status::OK(); +} + +// Bulk-decodes `count` raw u32s into `out` (resized to count). +Status RunReader::read_raw_u32(size_t count, std::vector* out, + MemoryReporter::Reservation* reservation) { + // Bound `count` against the run's byte size BEFORE resize(): a record can never + // hold more u32s than the whole file. Rejects a corrupt/truncated length varint + // (which is otherwise an unbounded resize -> uncaught std::bad_alloc). + if (count > file_size_ / sizeof(uint32_t)) { + return Status::Error( + "run: raw u32 count exceeds file size"); + } + RETURN_IF_ERROR(reserve_vector_for_size(out, count, memory_reporter_, reservation)); + out->resize(count); + if (count == 0) return Status::OK(); + return pull_raw_u32(reinterpret_cast(out->data()), count); +} + +// Materializes the current term's deferred position block into positions_flat. +// A no-op once the positions are already drained (idempotent within a term). +Status RunReader::materialize_positions() { + if (pos_remaining_ == 0) { + current_.positions_flat.clear(); + return Status::OK(); + } + if (pos_remaining_ > std::numeric_limits::max()) { + return Status::Error( + "run: position count exceeds addressable memory"); + } + const size_t n = static_cast(pos_remaining_); + RETURN_IF_ERROR(read_raw_u32(n, ¤t_.positions_flat, &positions_reservation_)); + pos_remaining_ = 0; + return Status::OK(); +} + +// Streams the next `n` positions of the current term straight from the window. +Status RunReader::stream_positions(uint32_t* dst, size_t n) { + if (n == 0) return Status::OK(); + if (n > pos_remaining_) { + return Status::Error( + "run: stream_positions past block end"); + } + RETURN_IF_ERROR(pull_raw_u32(reinterpret_cast(dst), n)); + pos_remaining_ -= n; + return Status::OK(); +} + +// Discards any positions of the current term left unread, so the window cursor +// lands at the next record boundary before advance() reads the next term. +Status RunReader::skip_remaining_positions() { + if (pos_remaining_ == 0) return Status::OK(); + std::array scratch; + while (pos_remaining_ != 0) { + const size_t count = static_cast( + std::min(pos_remaining_, static_cast(scratch.size()))); + RETURN_IF_ERROR(pull_raw_u32(reinterpret_cast(scratch.data()), count)); + pos_remaining_ -= count; + } + return Status::OK(); +} + +Status RunReader::advance() { + // Drain any positions the owner left unread for the previous term so the window + // cursor lands at the next record boundary. + RETURN_IF_ERROR(skip_remaining_positions()); + // End-of-run detection: at a record boundary, if no bytes remain we are done. + if (available() == 0) { + RETURN_IF_ERROR(fill()); + if (available() == 0 && eof_) { + exhausted_ = true; + return Status::OK(); + } + } + uint64_t term_id = 0; + RETURN_IF_ERROR(read_varint(&term_id)); + if (term_id > UINT32_MAX) + return Status::Error( + "run term_id exceeds uint32"); + current_id_ = static_cast(term_id); + current_.term.clear(); // runs store only the id; owner resolves the string + + uint64_t encoded_shape = 0; + RETURN_IF_ERROR(read_varint(&encoded_shape)); + if (encoded_shape > static_cast(RunPostingShape::kPositioned)) { + return Status::Error( + "run: unknown posting shape"); + } + const auto shape = static_cast(encoded_shape); + if (shape == RunPostingShape::kPositioned && !has_positions_) { + return Status::Error( + "run: positioned record in docs-only run"); + } + if (shape == RunPostingShape::kDocsOnlyStatless && !has_positions_) { + return Status::Error( + "run: statless record requires a positioned mixed-shape run"); + } + + uint64_t n_docs = 0; + RETURN_IF_ERROR(read_varint(&n_docs)); + if (n_docs > file_size_ / sizeof(uint32_t) || n_docs > std::numeric_limits::max()) { + return Status::Error( + "run: document count exceeds file size or addressable memory"); + } + // Docids: RAW absolute u32 block (bulk read), matching the writer's AppendRawU32. + RETURN_IF_ERROR( + read_raw_u32(static_cast(n_docs), ¤t_.docids, &docids_reservation_)); + for (size_t i = 1; i < current_.docids.size(); ++i) { + if (current_.docids[i] <= current_.docids[i - 1]) { + return Status::Error( + "run: docids must be strictly ascending within one record"); + } + } + current_.freqs.clear(); + current_.positions_flat.clear(); + pos_count_ = 0; + pos_remaining_ = 0; + if (shape == RunPostingShape::kDocsOnlyStatless) { + current_.retain_positions = false; + return Status::OK(); + } + + // Freqs: RAW u32 block (bulk read), matching the writer's AppendRawU32. + RETURN_IF_ERROR( + read_raw_u32(static_cast(n_docs), ¤t_.freqs, &freqs_reservation_)); + uint64_t total_freq = 0; + for (uint32_t freq : current_.freqs) { + if (freq == 0) { + return Status::Error( + "run: frequency must be positive"); + } + if (freq > UINT64_MAX - total_freq) { + return Status::Error( + "run: frequency sum overflows uint64"); + } + total_freq += freq; + } + if (shape == RunPostingShape::kDocsAndFreqs) { + current_.retain_positions = false; + return Status::OK(); + } + + uint64_t n_pos = 0; + RETURN_IF_ERROR(read_varint(&n_pos)); + if (n_pos != total_freq) { + return Status::Error( + "run: position count does not match frequency sum"); + } + if (n_pos > file_size_ / sizeof(uint32_t)) { + return Status::Error( + "run: position count exceeds file size"); + } + // Positions are LAZY: record the block count and leave the window cursor parked + // at the block start. The owner picks materialize_positions() for explicit + // materialization or stream_positions() for bounded writer-owned windows. + current_.retain_positions = true; + pos_count_ = n_pos; + pos_remaining_ = n_pos; + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// K-way merge +// --------------------------------------------------------------------------- + +namespace { + +// Min-heap entry: orders by the run's current term-id's PRECOMPUTED integer +// string-rank (rank[term_id] == its lexicographic rank over the dense vocabulary), +// tie-broken by run index so equal terms are gathered run-order (keeping +// concatenated docids ascending). The rank is a lexicographic bijection on a dense +// vocab, so ordering by the dense 4 B rank array reproduces the exact dictionary +// order a vocab-string compare would -- with an integer compare and zero random +// vocab string access in the inner loop. +struct HeapItem { + uint32_t term_id; + size_t run; +}; +struct HeapGreater { + const std::vector* rank; + bool operator()(const HeapItem& a, const HeapItem& b) const { + const uint32_t ra = (*rank)[a.term_id]; + const uint32_t rb = (*rank)[b.term_id]; + if (ra != rb) { + return ra > rb; + } // smaller rank first (lexicographic min-heap) + return a.run > b.run; // same term across runs: run-order tie-break + } +}; + +// Appends src's postings onto dst (run order). Later runs only cover docids +// >= dst's last, so docids stay ascending. COALESCE the boundary doc: if a spill +// fell BETWEEN two tokens of the same doc, that doc ends one run and begins the +// next with the SAME docid -- merge them (sum freqs, splice positions) so the +// merged term has exactly one entry per docid (matching the in-memory build). +// +// Positions are FLAT: doc order, partitioned by freqs. Because both dst and src +// already store doc-ordered flat positions, the common (no-boundary-overlap) case +// is a single bulk append. The boundary-overlap case must INSERT src's first +// doc's positions right after dst's last doc's positions so flat order stays +// consistent with the merged (coalesced) freqs. +void concat(TermPostings* dst, const TermPostings& src, RunPostingShape shape) { + if (src.docids.empty()) return; + const bool has_positions = shape == RunPostingShape::kPositioned; + const bool statless = shape == RunPostingShape::kDocsOnlyStatless; + DCHECK(posting_shape(src) == shape); + size_t start = 0; + size_t src_pos_start = 0; // flat offset of src positions to append after splice + if (!dst->docids.empty() && dst->docids.back() == src.docids.front()) { + const uint32_t head_fc = statless ? 0 : src.freqs.front(); + if (has_positions && head_fc != 0) { + // Splice src's first-doc positions in right after dst's last-doc positions. + // dst's last doc owns dst->freqs.back() entries at the tail of positions_flat + // BEFORE we bump that freq, so insert at end() (last doc is the tail run). + auto& flat = dst->positions_flat; + flat.insert(flat.end(), src.positions_flat.begin(), + src.positions_flat.begin() + head_fc); + } + if (!statless) { + dst->freqs.back() += head_fc; + } + src_pos_start = head_fc; + start = 1; // boundary doc folded in; append the rest + } + dst->docids.insert(dst->docids.end(), src.docids.begin() + start, src.docids.end()); + if (!statless) { + dst->freqs.insert(dst->freqs.end(), src.freqs.begin() + start, src.freqs.end()); + } + if (has_positions) { + dst->positions_flat.insert(dst->positions_flat.end(), + src.positions_flat.begin() + src_pos_start, + src.positions_flat.end()); + } +} + +class RunTermPostingSource final : public TermPostingSource { +public: + RunTermPostingSource(std::vector>* readers, + const std::vector* matching, RunPostingShape shape) + : readers_(readers), matching_(matching), shape_(shape) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + if (out == nullptr || exhausted == nullptr || target_docs == 0 || !out->empty()) { + return Status::Error( + "run posting source: invalid fill arguments"); + } + + Cursor planned = cursor_; + size_t document_count = 0; + size_t position_count = 0; + while (document_count < target_docs) { + normalize(&planned); + if (planned.run == matching_->size()) { + break; + } + const uint32_t docid = current_docid(planned); + uint64_t frequency = 0; + do { + const TermPostings& postings = current_postings(planned); + if (shape_ != RunPostingShape::kDocsOnlyStatless) { + frequency += postings.freqs[planned.doc]; + if (frequency > std::numeric_limits::max()) { + return Status::Error( + "run: coalesced frequency exceeds uint32"); + } + } + advance(&planned); + normalize(&planned); + if (planned.run == matching_->size()) { + break; + } + const uint32_t next_docid = current_docid(planned); + if (next_docid < docid) { + return Status::Error( + "run: docids overlap across spill runs"); + } + if (next_docid != docid) { + break; + } + } while (true); + if (shape_ == RunPostingShape::kPositioned) { + if (frequency > std::numeric_limits::max() - position_count) { + return Status::Error( + "run posting source: position window exceeds size_t"); + } + position_count += static_cast(frequency); + } + ++document_count; + } + + MutableTermPostingSpan destination; + const bool has_freqs = shape_ != RunPostingShape::kDocsOnlyStatless; + RETURN_IF_ERROR( + out->grow_uninitialized(document_count, has_freqs, position_count, &destination)); + size_t position_offset = 0; + for (size_t output = 0; output < document_count; ++output) { + normalize(&cursor_); + DCHECK_LT(cursor_.run, matching_->size()); + const uint32_t docid = current_docid(cursor_); + uint64_t frequency = 0; + do { + const TermPostings& postings = current_postings(cursor_); + const uint32_t run_frequency = has_freqs ? postings.freqs[cursor_.doc] : 0; + if (shape_ == RunPostingShape::kPositioned) { + RunReader* reader = (*readers_)[(*matching_)[cursor_.run]].get(); + RETURN_IF_ERROR(reader->stream_positions( + destination.positions_flat.data() + position_offset, run_frequency)); + position_offset += run_frequency; + } + frequency += run_frequency; + advance(&cursor_); + normalize(&cursor_); + if (cursor_.run == matching_->size() || current_docid(cursor_) != docid) { + break; + } + } while (true); + destination.docids[output] = docid; + if (has_freqs) { + destination.freqs[output] = static_cast(frequency); + } + } + DCHECK_EQ(position_offset, destination.positions_flat.size()); + normalize(&cursor_); + *exhausted = cursor_.run == matching_->size(); + return Status::OK(); + } + + bool exhausted() { + normalize(&cursor_); + return cursor_.run == matching_->size(); + } + +private: + struct Cursor { + size_t run = 0; + size_t doc = 0; + }; + + const TermPostings& current_postings(const Cursor& cursor) const { + return (*readers_)[(*matching_)[cursor.run]]->current(); + } + + uint32_t current_docid(const Cursor& cursor) const { + return current_postings(cursor).docids[cursor.doc]; + } + + void normalize(Cursor* cursor) const { + while (cursor->run < matching_->size() && + cursor->doc == current_postings(*cursor).docids.size()) { + ++cursor->run; + cursor->doc = 0; + } + } + + static void advance(Cursor* cursor) { ++cursor->doc; } + + std::vector>* readers_; + const std::vector* matching_; + RunPostingShape shape_; + Cursor cursor_; +}; + +} // namespace + +Status merge_run_sources(const std::vector& run_paths, + const std::vector& vocab, + const std::vector& string_rank, bool has_positions, + const StreamedTermConsumer& fn, TermKeyMaterializer materialize_term_key, + MemoryReporter* memory_reporter) { + if (string_rank.size() != vocab.size()) { + return Status::Error( + "merge_run_sources: string_rank/vocab size mismatch"); + } + std::vector> readers; + readers.reserve(run_paths.size()); + std::priority_queue, HeapGreater> heap( + HeapGreater {&string_rank}); + for (size_t i = 0; i < run_paths.size(); ++i) { + auto reader = std::make_unique(memory_reporter); + RETURN_IF_ERROR(reader->open(run_paths[i], has_positions)); + if (!reader->exhausted()) { + if (reader->current_id() >= vocab.size()) { + return Status::Error( + "run term_id out of vocab range"); + } + heap.push({reader->current_id(), i}); + } + readers.push_back(std::move(reader)); + } + + std::vector matching; + while (!heap.empty()) { + const uint32_t id = heap.top().term_id; + matching.clear(); + while (!heap.empty() && heap.top().term_id == id) { + matching.push_back(heap.top().run); + heap.pop(); + } + DCHECK(!matching.empty()); + const RunPostingShape shape = posting_shape(readers[matching.front()]->current()); + for (size_t run : matching) { + if (posting_shape(readers[run]->current()) != shape) { + return Status::Error( + "run: posting shape differs across matching terms"); + } + } + + RunTermPostingSource source(&readers, &matching, shape); + StreamedTermPostings postings {.term = materialize_term_key + ? materialize_term_key(vocab[id]) + : std::string(vocab[id]), + .retain_positions = shape == RunPostingShape::kPositioned, + .source = &source}; + RETURN_IF_ERROR(fn(std::move(postings))); + if (!source.exhausted()) { + return Status::Error( + "run posting source: consumer returned before term exhaustion"); + } + + for (size_t run : matching) { + RunReader* reader = readers[run].get(); + RETURN_IF_ERROR(reader->advance()); + if (!reader->exhausted()) { + if (reader->current_id() >= vocab.size()) { + return Status::Error( + "run term_id out of vocab range"); + } + heap.push({reader->current_id(), run}); + } + } + } + return Status::OK(); +} + +Status compact_runs(const std::vector& run_paths, + const std::vector& string_rank, bool has_positions, + const std::string& out_path, MemoryReporter* memory_reporter) { + // Same heap machinery as merge_run_sources, but the output is a RUN (records keyed + // by term-id, ordered by string rank -- the exact invariant every run file + // carries), not a resolved term stream: no vocab strings are needed, and + // positions are always materialized because the run codec serializes + // positions_flat directly. + std::vector> readers; + readers.reserve(run_paths.size()); + std::priority_queue, HeapGreater> heap( + HeapGreater {&string_rank}); + for (size_t i = 0; i < run_paths.size(); ++i) { + auto r = std::make_unique(memory_reporter); + RETURN_IF_ERROR(r->open(run_paths[i], has_positions)); + if (!r->exhausted()) { + if (r->current_id() >= string_rank.size()) { + return Status::Error( + "run term_id out of rank range"); + } + heap.push({r->current_id(), i}); + } + readers.push_back(std::move(r)); + } + + RunWriter w(memory_reporter); + RETURN_IF_ERROR(w.open(out_path)); + std::vector matching; // run indices contributing the current term + while (!heap.empty()) { + const uint32_t id = heap.top().term_id; + MemoryReporter::Reservation merged_docids_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation merged_freqs_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation merged_positions_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + TermPostings merged; + matching.clear(); + uint64_t total_docs = 0; + uint64_t total_pos = 0; + while (!heap.empty() && heap.top().term_id == id) { + const size_t ri = heap.top().run; + heap.pop(); + const RunReader* r = readers[ri].get(); + const uint64_t run_docs = r->current().docids.size(); + const uint64_t run_positions = r->current_pos_count(); + if (run_docs > std::numeric_limits::max() - total_docs || + run_positions > std::numeric_limits::max() - total_pos) { + return Status::Error( + "run compaction: merged posting size overflows uint64"); + } + total_docs += run_docs; + total_pos += run_positions; + matching.push_back(ri); + } + DCHECK(!matching.empty()); + const RunPostingShape shape = posting_shape(readers[matching.front()]->current()); + for (size_t ri : matching) { + if (posting_shape(readers[ri]->current()) != shape) { + return Status::Error( + "run: posting shape differs across matching terms"); + } + } + const bool term_has_positions = shape == RunPostingShape::kPositioned; + const bool statless = shape == RunPostingShape::kDocsOnlyStatless; + merged.retain_positions = term_has_positions; + if (total_docs > std::numeric_limits::max() || + total_pos > std::numeric_limits::max()) { + return Status::Error( + "run compaction: merged posting exceeds addressable memory"); + } + RETURN_IF_ERROR(reserve_vector_for_size(&merged.docids, static_cast(total_docs), + memory_reporter, &merged_docids_reservation)); + if (!statless) { + RETURN_IF_ERROR(reserve_vector_for_size(&merged.freqs, static_cast(total_docs), + memory_reporter, &merged_freqs_reservation)); + } + if (term_has_positions) { + RETURN_IF_ERROR(reserve_vector_for_size(&merged.positions_flat, + static_cast(total_pos), memory_reporter, + &merged_positions_reservation)); + } + // concat (WITH boundary-doc coalescing) is deliberately the SAME + // append the final merge applies: coalescing the seam between two + // adjacent input runs here yields exactly what the final merge would + // have produced from the uncompacted pair, so compaction is invisible + // in the emitted term stream. + for (size_t ri : matching) { + RunReader* r = readers[ri].get(); + if (term_has_positions) { + RETURN_IF_ERROR(r->materialize_positions()); + } + concat(&merged, r->current(), shape); + } + RETURN_IF_ERROR(w.write_term(id, merged)); + for (size_t ri : matching) { + RunReader* r = readers[ri].get(); + RETURN_IF_ERROR(r->advance()); + if (!r->exhausted()) { + if (r->current_id() >= string_rank.size()) { + return Status::Error( + "run term_id out of rank range"); + } + heap.push({r->current_id(), ri}); + } + } + } + return w.close(); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/spill_run_codec.h b/be/src/storage/index/snii/writer/spill_run_codec.h new file mode 100644 index 00000000000000..f63bd06d6a9156 --- /dev/null +++ b/be/src/storage/index/snii/writer/spill_run_codec.h @@ -0,0 +1,236 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::writer { + +using TermKeyMaterializer = std::function; + +// On-disk SPIMI "run" codec for the spill / k-way-merge out-of-core build path. +// +// A RUN is a self-describing file holding a sequence of terms keyed by TERM-ID, +// each followed by its postings, in this exact wire layout. The file is produced +// and consumed by THIS module only (a private temp file -- the on-disk INDEX is +// unaffected), so the format is chosen for cheap I/O: docids, freqs and positions +// are ALL RAW fixed-width little-endian u32 BLOCKS (bulk memcpy on both ends, +// ~10x cheaper than per-value varint -- which cost ~1.5s of encode CPU over the +// 5M build's ~60M docids and compressed those streams poorly anyway). Decode +// still validates every length against the file size. +// +// run := record* (term-ids ordered by vocab string, +// strictly ascending within a run) +// record := +// VInt term_id (index into the shared vocabulary) +// VInt shape (0=docs-only-statless, 1=docs+freq, +// 2=positioned) +// VInt n_docs +// u32 docid * n_docs (RAW LE absolute ascending docids) +// shape=1: u32 freq * n_docs (RAW LE, each >= 1) +// shape=2: u32 freq * n_docs, VInt n_pos, u32 position * n_pos +// (n_pos == sum(freqs)) +// +// Shape 0 is the CommonGrams docs-only set representation. It writes neither +// synthetic all-one frequencies nor n_pos; run files are private temporaries, +// so this does not change the persisted SNII index format. +// +// Decode is fully STREAMED: a RunReader reads a small fixed buffer at a time and +// materializes only the CURRENT term's postings, never the whole run. The k-way +// merge keeps one heap slot per run (each holding only its current term-id + +// that term's postings), so peak memory is bounded by the widest single term +// summed across the runs that contain it -- not by total postings. The merge +// orders runs by a PRECOMPUTED integer string-rank (term-id -> its lexicographic +// rank over the shared dense vocabulary): an integer compare that reproduces the +// exact lexicographic order without touching a vocab string in the inner loops. + +// Writes a sorted sequence of terms (by id) to one run file. Term-ids must be +// handed to write_term in vocab-string ascending order (the spill caller sorts +// before spilling). RAII: the file is flushed and closed on close(); the partial +// file is left for the owning SpimiTermBuffer to delete on its temp-path list. +class RunWriter { +public: + explicit RunWriter(MemoryReporter* memory_reporter = nullptr); + ~RunWriter(); + + RunWriter(const RunWriter&) = delete; + RunWriter& operator=(const RunWriter&) = delete; + + // Opens `path` for writing (truncating). Returns IoError on failure. + Status open(const std::string& path); + + // Appends one term's postings under `term_id`. Empty freqs denotes the + // docs-only-statless shape; otherwise freqs parallels docids. Positioned + // postings additionally hold sum(freqs) positions in document order. + Status write_term(uint32_t term_id, const TermPostings& tp); + + // Flushes the buffer and closes the file. Safe to call once; idempotent. + Status close(); + +private: + Status flush(); + Status append_bytes(const uint8_t* data, size_t size); + Status append_varint(uint64_t value); + Status append_raw_u32(const uint32_t* values, size_t count); + void release_buffer(); + + MemoryReporter* memory_reporter_ = nullptr; + MemoryReporter::Reservation buffer_reservation_; + int fd_ = -1; + std::vector buf_; // bounded staging buffer; flushed in fixed-size chunks +}; + +// Streamed reader over one run file. After open() the first term is loaded; +// current()/current_id() expose it; advance() loads the next (or marks +// exhausted). Only the current term's postings live in memory at a time. The +// current record's `term` string is left EMPTY -- runs store only the id; the +// owner resolves the string via the shared vocabulary. +// +// LAZY POSITIONS (peak-RSS optimization for the widest merged term): advance() +// loads term_id / docids / freqs and the position-block COUNT, but does NOT read +// the position bytes -- it leaves the decode window cursor parked at the start of +// the position block. The owner then chooses, per term: +// * materialize_positions(): bulk-reads the block into current().positions_flat +// (the default; behaves exactly as the old eager reader). +// * stream_positions(dst, n): pulls the next n positions straight from the +// window in 64 KiB chunks, never materializing the whole block -- used by the +// k-way merge source to decode directly into each writer-owned window. +// advance() drains any positions left unread from the previous term before the +// next record, so a partly-streamed (or skipped) term still lands at the right +// record boundary. The yielded byte sequence is identical either way. +class RunReader { +public: + explicit RunReader(MemoryReporter* memory_reporter = nullptr); + ~RunReader(); + + RunReader(const RunReader&) = delete; + RunReader& operator=(const RunReader&) = delete; + + // Opens `path`, loading the first record (if any). has_positions declares + // whether this run may contain positioned or mixed statless records. + Status open(const std::string& path, bool has_positions); + + bool exhausted() const { return exhausted_; } + const TermPostings& current() const { return current_; } + uint32_t current_id() const { return current_id_; } + + // Number of positions in the current term's (lazily-loaded) position block. + uint64_t current_pos_count() const { return pos_count_; } + // True once the current term's positions have been materialized OR fully + // streamed (i.e. nothing remains to read before advance()). + bool positions_drained() const { return pos_remaining_ == 0; } + + // Materializes the current term's position block into current().positions_flat + // (bulk read). Idempotent within a term: a no-op once positions are drained. + Status materialize_positions(); + // Streams the next `n` positions of the current term into dst[0..n) directly + // from the decode window (64 KiB chunks topped up on demand). Caller must not + // request more than positions_remaining(); each call advances the cursor. + Status stream_positions(uint32_t* dst, size_t n); + uint64_t positions_remaining() const { return pos_remaining_; } + + // Loads the next record into current(); sets exhausted() at end of file. Any + // positions of the current term left unread are skipped first. + Status advance(); + +private: + size_t available() const; // buffered bytes from pos_ to window end + Status fill(); // tops up the decode window from disk + Status ensure(size_t n); // guarantees >= n buffered bytes (or eof) + Status read_varint(uint64_t* v); // bounds-checked streamed varint + // Bulk-reads `count` RAW little-endian u32s from the window into `out` (resized + // to count). Bounds-checked against the run's true length (Corruption on EOF). + Status read_raw_u32(size_t count, std::vector* out, + MemoryReporter::Reservation* reservation); + // Streams `count` raw u32s from the window into dst (caller-owned, sized by the + // caller); shared by read_raw_u32 (into a vector) and stream_positions. + Status pull_raw_u32(uint8_t* dst, size_t count); + // Drains (and discards) any remaining positions of the current term so the + // window cursor lands at the next record boundary. + Status skip_remaining_positions(); + + MemoryReporter* memory_reporter_ = nullptr; + // Reservations precede their vectors so allocations are destroyed first. + MemoryReporter::Reservation window_reservation_; + MemoryReporter::Reservation docids_reservation_; + MemoryReporter::Reservation freqs_reservation_; + MemoryReporter::Reservation positions_reservation_; + int fd_ = -1; + bool has_positions_ = false; + bool exhausted_ = false; + uint64_t file_size_ = 0; // total run byte size (fstat at open); bounds lengths + uint64_t bytes_read_ = 0; // bytes pulled from fd; never exceeds file_size_ + std::vector window_; // sliding decode window + size_t pos_ = 0; // consumed offset within window_ + bool eof_ = false; // no more bytes on disk + uint32_t current_id_ = 0; // current record's term-id + uint64_t pos_count_ = 0; // current term's total position count (from n_pos) + uint64_t pos_remaining_ = 0; // positions still unread in the current block + TermPostings current_; +}; + +// K-way merges the given run files into a single term stream ordered by a +// PRECOMPUTED integer string-rank (string_rank[term_id] == the term-id's +// lexicographic rank over the dense vocabulary), invoking `fn` once per distinct +// term-id with its postings concatenated across all runs that contain it (in run +// order -> docids stay ascending) and its `term` resolved from `vocab` once. +// Because a dense vocab maps each id to a distinct string, the rank is a +// lexicographic bijection: ordering by the dense 4 B rank array (an integer +// compare) reproduces the EXACT order a vocab-string compare would -- but never +// reads a vocab string in the inner heap/gather loops. Only one bounded posting +// window is materialized at a time. Returns IoError/Corruption on bad run data, or +// InternalError when string_rank.size() != vocab.size(). has_positions must match +// how the runs were written. `vocab` (term-id -> string) and `string_rank` +// (term-id -> rank) are both borrowed and MUST be sized to the vocabulary. +// +// The source callback is synchronous: matching run readers remain parked on the +// current term until the callback returns. The source fills writer-owned windows +// directly and coalesces equal docids at run boundaries. A successful callback +// must exhaust the source. +Status merge_run_sources(const std::vector& run_paths, + const std::vector& vocab, + const std::vector& string_rank, bool has_positions, + const StreamedTermConsumer& fn, + TermKeyMaterializer materialize_term_key = {}, + MemoryReporter* memory_reporter = nullptr); + +// G09 run-file cap support: k-way merges `run_paths` into ONE new run file at +// `out_path`, keyed and ordered exactly like merge_run_sources (heap on +// string_rank[term_id]; per-term postings concatenated across runs in run +// order, boundary docs coalesced -- the same concat the final merge applies, +// so compact-then-merge emits the identical term stream as merging the +// originals). Positions are fully materialized for each term because the run +// codec serializes positions_flat. +// Every record's term-id must index string_rank (else Corruption). On error +// `out_path` may hold a partial file the caller must delete; the input runs +// are never modified. Opens run_paths.size() read fds + 1 write fd for the +// call's duration -- the caller (SpimiTermBuffer::compact_runs) bounds that +// fan-in with its run-count cap. +Status compact_runs(const std::vector& run_paths, + const std::vector& string_rank, bool has_positions, + const std::string& out_path, MemoryReporter* memory_reporter = nullptr); + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/spillable_byte_buffer.h b/be/src/storage/index/snii/writer/spillable_byte_buffer.h new file mode 100644 index 00000000000000..2760600b00f019 --- /dev/null +++ b/be/src/storage/index/snii/writer/spillable_byte_buffer.h @@ -0,0 +1,305 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "io/fs/file_reader.h" +#include "io/fs/file_reader_writer_fwd.h" +#include "io/fs/file_writer.h" +#include "io/fs/local_file_system.h" +#include "io/fs/path.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/temp_dir.h" +#include "util/slice.h" + +namespace doris::snii::writer { + +// A tiered append buffer for one build-time section. While resident it holds the +// bytes as a CHAIN OF CHUNKS (one per append) rather than a single growing vector. +// Copying appends own a right-sized allocation. Move appends preserve the caller's +// allocation, including capacity slack, and account that retained capacity as resident +// memory. Once the resident capacity crosses `cap_bytes` the buffer SPILLS to a temp +// file (resolve_temp_dir()) and routes later appends there, so a huge section stays +// RSS-bounded at ~cap_bytes while a small one is RAM-only (zero disk, spill-only build). +// Append order/bytes are identical wherever they land; stream_into() reproduces the +// logical section bytes in order. RAII-removes the temp. (cap_bytes == UINT64_MAX +// disables spilling -> always RAM.) +class SpillableByteBuffer { +public: + // `reporter` is an OPTIONAL writer-level build-RAM reporter (null off-Doris / + // unit tests). Resident chunk capacities are hard-reserved before adoption or + // allocation. If the shared cap cannot admit the next chunk, the buffer spills + // its resident prefix first and writes that chunk directly to disk. Spilled + // bytes are not resident and hold no reservation. + SpillableByteBuffer(uint64_t cap_bytes, std::string tag, MemoryReporter* reporter = nullptr) + : cap_bytes_(cap_bytes), + tag_(std::move(tag)), + reporter_(reporter), + reservation_(reporter == nullptr ? MemoryReporter::Reservation() + : reporter->make_reservation()) {} + ~SpillableByteBuffer() { release_storage(); } + SpillableByteBuffer(const SpillableByteBuffer&) = delete; + SpillableByteBuffer& operator=(const SpillableByteBuffer&) = delete; + + // Total bytes appended so far (the offset basis for callers recording sub-offsets). + uint64_t size() const { + return consumed_ ? consumed_size_ : (spilled_ ? spilled_bytes_ : ram_bytes_); + } + + // Copying append (the Slice bytes are copied into a fresh chunk). + Status append(Slice bytes) { + if (consumed_) { + return Status::Error( + "spillable buffer: append after stream"); + } + if (spilled_) { + const ::doris::Slice s(bytes.data(), bytes.size()); + RETURN_IF_ERROR(to_snii(temp_writer_->appendv(&s, 1))); + spilled_bytes_ += bytes.size(); + return Status::OK(); + } + if (!bytes.empty()) { + bool keep_resident = true; + RETURN_IF_ERROR(reserve_resident_capacity(bytes.size(), &keep_resident)); + if (!keep_resident) { + const ::doris::Slice s(bytes.data(), bytes.size()); + RETURN_IF_ERROR(to_snii(temp_writer_->appendv(&s, 1))); + spilled_bytes_ += bytes.size(); + return Status::OK(); + } + std::vector chunk; + chunk.reserve(bytes.size()); + DCHECK_EQ(chunk.capacity(), bytes.size()); + chunk.insert(chunk.end(), bytes.data(), bytes.data() + bytes.size()); + chunks_.push_back(std::move(chunk)); + ram_bytes_ += bytes.size(); + ram_capacity_bytes_ += chunks_.back().capacity(); + } + if (over_cap()) return spill_to_disk(); + return Status::OK(); + } + + // Move append: the section ADOPTS the caller's vector without copying. The common + // dict path hands off each flushed block this way. Logical bytes remain v.size(), + // while resident accounting uses the capacity retained by the adopted vector. + Status append_move(std::vector&& v) { + if (consumed_) { + return Status::Error( + "spillable buffer: append after stream"); + } + if (spilled_) { + const ::doris::Slice s(v.data(), v.size()); + RETURN_IF_ERROR(to_snii(temp_writer_->appendv(&s, 1))); + spilled_bytes_ += v.size(); + return Status::OK(); + } + if (!v.empty()) { + const size_t logical_bytes = v.size(); + const size_t retained_capacity = v.capacity(); + bool keep_resident = true; + RETURN_IF_ERROR(reserve_resident_capacity(retained_capacity, &keep_resident)); + if (!keep_resident) { + const ::doris::Slice s(v.data(), v.size()); + RETURN_IF_ERROR(to_snii(temp_writer_->appendv(&s, 1))); + spilled_bytes_ += v.size(); + return Status::OK(); + } + chunks_.push_back(std::move(v)); + ram_bytes_ += logical_bytes; + ram_capacity_bytes_ += chunks_.back().capacity(); + DCHECK_EQ(chunks_.back().capacity(), retained_capacity); + } + if (over_cap()) return spill_to_disk(); + return Status::OK(); + } + + // Must be called once after the last append, before stream_into(): flushes the temp + // (if spilled) so it can be read back. A no-op for a RAM-resident buffer. + Status seal() { + if (consumed_) { + return Status::Error( + "spillable buffer: seal after stream"); + } + if (spilled_ && !sealed_) { + RETURN_IF_ERROR(to_snii(temp_writer_->close())); + sealed_ = true; + } + return Status::OK(); + } + + // Streams the whole section (RAM chunks or sealed temp) into `out`, in append order. + Status stream_into(io::FileWriter* out) const { + if (consumed_) { + return Status::Error( + "spillable buffer: stream called twice"); + } + if (!spilled_) { + for (const auto& c : chunks_) { + if (!c.empty()) RETURN_IF_ERROR(out->append(Slice(c))); + } + return Status::OK(); + } + ::doris::io::FileReaderSPtr reader; + RETURN_IF_ERROR( + to_snii(::doris::io::global_local_filesystem()->open_file(temp_path_, &reader))); + constexpr uint64_t kMaxChunk = 1u << 20; // bounded copy window (no whole-section reload) + size_t read_capacity = static_cast(std::min(kMaxChunk, spilled_bytes_)); + MemoryReporter::Reservation read_reservation = reporter_ == nullptr + ? MemoryReporter::Reservation() + : reporter_->make_reservation(); + if (reporter_ != nullptr) { + Status reserve_status = read_reservation.set_bytes(read_capacity); + while (!reserve_status.ok() && reserve_status.is() && + read_capacity > 1) { + read_capacity = std::max(1, read_capacity / 2); + reserve_status = read_reservation.set_bytes(read_capacity); + } + RETURN_IF_ERROR(reserve_status); + } + std::vector buf; + buf.reserve(read_capacity); + DCHECK_EQ(buf.capacity(), read_capacity); + for (uint64_t off = 0; off < spilled_bytes_; off += read_capacity) { + const uint64_t n = std::min(read_capacity, spilled_bytes_ - off); + buf.resize(static_cast(n)); + size_t bytes_read = 0; + RETURN_IF_ERROR(to_snii(reader->read_at( + off, ::doris::Slice(buf.data(), static_cast(n)), &bytes_read))); + if (bytes_read != n) { + return Status::Error( + "short read from spill scratch file"); + } + RETURN_IF_ERROR(out->append(Slice(buf.data(), static_cast(n)))); + } + return Status::OK(); + } + + // The staged bytes have no readers after they are copied into the compound + // output. Releasing them here, rather than in the logical-writer destructor, + // keeps multi-destination compaction from retaining one DICT image per + // completed index until rowset close. + Status stream_into_and_release(io::FileWriter* out) { + RETURN_IF_ERROR(stream_into(out)); + release_storage(); + return Status::OK(); + } + + bool spilled() const { return spilled_; } + +private: + Status reserve_resident_capacity(size_t additional_capacity, bool* keep_resident) { + DCHECK(keep_resident != nullptr); + *keep_resident = true; + if (reporter_ == nullptr) return Status::OK(); + if (additional_capacity > std::numeric_limits::max() - ram_capacity_bytes_) { + return Status::Error( + "spillable buffer: resident capacity overflows uint64"); + } + const Status reserved = reservation_.set_bytes(ram_capacity_bytes_ + additional_capacity); + if (reserved.ok()) return Status::OK(); + if (!reserved.is()) return reserved; + RETURN_IF_ERROR(spill_to_disk()); + *keep_resident = false; + return Status::OK(); + } + + // Gate-2 spill condition (UNIFIED): spill when the writer's TOTAL build RAM crosses + // the one shared cap (reporter_->over_cap()), with the local cap_bytes_ kept only as + // a defensive per-buffer hard ceiling (e.g. when no reporter is attached). + bool over_cap() const { + return (reporter_ != nullptr && reporter_->over_cap()) || ram_capacity_bytes_ >= cap_bytes_; + } + // Bridge a Doris IO Status into SNII's Status. R01 (status migration) is not done yet, + // so this buffer still returns Status; this mirrors snii_doris_adapter's + // to_snii_status (ok -> OK, otherwise IoError carrying the Doris message). + static Status to_snii(const Status& s) { + if (s.ok()) return Status::OK(); + return Status::Error(s.to_string_no_stack()); + } + Status spill_to_disk() { + temp_path_ = resolve_temp_dir() + "/snii_" + tag_ + "_" + std::to_string(::getpid()) + "_" + + std::to_string(reinterpret_cast(this)) + ".tmp"; + RETURN_IF_ERROR(to_snii( + ::doris::io::global_local_filesystem()->create_file(temp_path_, &temp_writer_))); + for (const auto& c : chunks_) { + if (!c.empty()) { + const ::doris::Slice s(c.data(), c.size()); + RETURN_IF_ERROR(to_snii(temp_writer_->appendv(&s, 1))); + } + } + spilled_bytes_ = ram_bytes_; + std::vector>().swap(chunks_); // reclaim the RAM immediately + ram_bytes_ = 0; + ram_capacity_bytes_ = 0; + reservation_.reset(); + spilled_ = true; + return Status::OK(); + } + + void release_storage() { + if (consumed_) { + return; + } + consumed_size_ = spilled_ ? spilled_bytes_ : ram_bytes_; + std::vector>().swap(chunks_); + ram_bytes_ = 0; + ram_capacity_bytes_ = 0; + reservation_.reset(); + + // A sealed temp writer is already closed; an unsealed error-path writer + // aborts on reset. In both cases remove the scratch path best-effort. + temp_writer_.reset(); + if (!temp_path_.empty()) { + std::remove(temp_path_.c_str()); + temp_path_.clear(); + } + spilled_bytes_ = 0; + consumed_ = true; + } + + uint64_t cap_bytes_; + std::string tag_; + MemoryReporter* reporter_ = nullptr; // optional build-RAM reporter (null off-Doris) + MemoryReporter::Reservation reservation_; // resident inner-vector capacities + std::vector> chunks_; // resident tier: one chunk per append + uint64_t ram_bytes_ = 0; // logical section bytes retained in RAM + uint64_t ram_capacity_bytes_ = 0; // resident capacities of the retained chunks + bool spilled_ = false; + bool sealed_ = false; + ::doris::io::FileWriterPtr temp_writer_; // Doris local writer for the spill scratch file + std::string temp_path_; + uint64_t spilled_bytes_ = 0; + uint64_t consumed_size_ = 0; + bool consumed_ = false; +}; + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/spimi_term_buffer.cpp b/be/src/storage/index/snii/writer/spimi_term_buffer.cpp new file mode 100644 index 00000000000000..2ee6d90bc20b81 --- /dev/null +++ b/be/src/storage/index/snii/writer/spimi_term_buffer.cpp @@ -0,0 +1,2097 @@ +// 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 "storage/index/snii/writer/spimi_term_buffer.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/exception.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/writer/global_memory_limiter.h" +#include "storage/index/snii/writer/spill_run_codec.h" +#include "storage/index/snii/writer/temp_dir.h" + +#if defined(__GLIBC__) +#include +#endif + +namespace doris::snii::writer { + +namespace { + +constexpr size_t kCommonGramPairKeySize = 10; +constexpr char kCommonGramPairKeyTag = 'P'; + +std::array EncodeCommonGramPairKey(PlainTermId left, + PlainTermId right) { + std::array key {}; + key[0] = segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN.front(); + key[1] = kCommonGramPairKeyTag; + for (size_t byte = 0; byte < sizeof(uint32_t); ++byte) { + const size_t shift = (sizeof(uint32_t) - byte - 1) * 8; + key[2 + byte] = static_cast((left.value >> shift) & 0xffU); + key[6 + byte] = static_cast((right.value >> shift) & 0xffU); + } + return key; +} + +bool is_common_gram_pair_key(std::string_view key) { + return key.size() == kCommonGramPairKeySize && + key.front() == segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN.front() && + key[1] == kCommonGramPairKeyTag; +} + +struct CommonGramPairIds { + PlainTermId left; + PlainTermId right; +}; + +uint32_t decode_big_endian_uint32_unchecked(const char* bytes) { + uint32_t value = 0; + for (size_t byte = 0; byte < sizeof(uint32_t); ++byte) { + value = (value << 8) | static_cast(bytes[byte]); + } + return value; +} + +#ifdef BE_TEST +std::atomic g_common_gram_pair_unchecked_decodes {0}; +std::atomic g_common_gram_trusted_plain_decodes {0}; +std::atomic g_common_gram_pair_cache_probes {0}; +std::atomic g_common_gram_pair_cache_pair_hits {0}; +std::atomic g_common_gram_pair_cache_same_doc_hits {0}; +std::atomic g_common_gram_native_pair_probes {0}; +std::atomic g_common_gram_native_pair_hits {0}; +std::atomic g_common_gram_native_pair_inserts {0}; +std::atomic g_common_gram_logical_validations {0}; +std::atomic g_common_gram_plain_cache_probes {0}; +std::atomic g_common_gram_plain_cache_hits {0}; +std::atomic g_common_gram_plain_intern_table_probes {0}; +std::atomic g_owned_term_full_byte_comparisons {0}; +std::atomic g_fail_next_owned_term_reserve {false}; +std::atomic g_fail_next_owned_term_emplace {false}; +std::atomic g_spill_gate_checks {0}; +std::atomic g_compact_chain_varint_decodes {0}; +#endif + +CommonGramPairIds decode_common_gram_pair_key_unchecked(std::string_view key) { + DCHECK(is_common_gram_pair_key(key)); +#ifdef BE_TEST + g_common_gram_pair_unchecked_decodes.fetch_add(1, std::memory_order_relaxed); +#endif + return { + .left = PlainTermId {.value = decode_big_endian_uint32_unchecked(key.data() + 2)}, + .right = PlainTermId {.value = decode_big_endian_uint32_unchecked(key.data() + 6)}, + }; +} + +class LogicalPlainKeyView { +public: + explicit LogicalPlainKeyView(std::string_view physical) : physical_(physical) { + escaped_ = !physical.empty() && + physical.front() == segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX; + if (escaped_) { + DCHECK_GE(physical.size(), 2U); + DCHECK(physical[1] == 'E' || physical[1] == 'G'); + } + } + + size_t size() const { return physical_.size() - (escaped_ ? 1 : 0); } + + uint8_t operator[](size_t index) const { + DCHECK_LT(index, size()); + if (!escaped_) { + return static_cast(physical_[index]); + } + if (index == 0) { + return static_cast( + physical_[1] == 'E' ? segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX : '\x1f'); + } + return static_cast(physical_[index + 1]); + } + +private: + std::string_view physical_; + bool escaped_ = false; +}; + +std::string_view decode_logical_plain_term_trusted(std::string_view physical, + std::string* scratch) { + DCHECK(scratch != nullptr); +#ifdef BE_TEST + g_common_gram_trusted_plain_decodes.fetch_add(1, std::memory_order_relaxed); +#endif + scratch->clear(); + if (physical.empty() || physical.front() != segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX) { + DCHECK(!segment_v2::inverted_index::is_internal_term_key(physical)); + return physical; + } + + DCHECK_GE(physical.size(), 2U); + DCHECK(physical[1] == 'E' || physical[1] == 'G'); + scratch->reserve(physical.size() - 1); + scratch->push_back(physical[1] == 'E' + ? segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX + : segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN.front()); + scratch->append(physical.substr(2)); + return std::string_view(*scratch); +} + +int compare_logical_plain_keys(const LogicalPlainKeyView& left, const LogicalPlainKeyView& right) { + const size_t common = std::min(left.size(), right.size()); + for (size_t i = 0; i < common; ++i) { + if (left[i] != right[i]) { + return left[i] < right[i] ? -1 : 1; + } + } + if (left.size() == right.size()) { + return 0; + } + return left.size() < right.size() ? -1 : 1; +} + +} // namespace + +struct SpimiTermBuffer::CommonGramPairCache { + struct Entry { + uint64_t pair = 0; + uint32_t term_id = std::numeric_limits::max(); + uint32_t last_docid = 0; + }; + static_assert(sizeof(Entry) == 16); + + static constexpr size_t kEntryCount = 1024; + static constexpr uint64_t kHashMultiplier = 11400714819323198485ULL; + static constexpr uint32_t kInvalidTermId = std::numeric_limits::max(); + + static size_t index(uint64_t pair) { + return static_cast((pair * kHashMultiplier) >> 54); + } + + std::array entries; + static_assert(sizeof(std::array) == 16 * 1024); +}; + +struct SpimiTermBuffer::CommonGramPlainTermCache { + struct Entry { + uint32_t fingerprint = 0; + uint32_t term_id = std::numeric_limits::max(); + }; + static_assert(sizeof(Entry) == 8); + + static constexpr size_t kSetCount = 1024; + static constexpr uint32_t kInvalidTermId = std::numeric_limits::max(); + using Set = std::array; + static_assert((kSetCount & (kSetCount - 1)) == 0); + + static size_t index(size_t term_hash) { return term_hash & (kSetCount - 1); } + + static uint32_t fingerprint(size_t term_hash) { + const auto hash = static_cast(term_hash); + return static_cast(hash ^ (hash >> 32)); + } + + static bool matches(const Entry& entry, uint32_t expected_fingerprint, std::string_view term, + const std::vector& vocab) { + if (entry.term_id == kInvalidTermId || entry.fingerprint != expected_fingerprint) { + return false; + } + DCHECK_LT(entry.term_id, vocab.size()); + return std::string_view(vocab[entry.term_id]) == term; + } + + uint32_t find(size_t term_hash, std::string_view term, const std::vector& vocab) { + Set& set = sets[index(term_hash)]; + const uint32_t expected_fingerprint = fingerprint(term_hash); + if (matches(set[0], expected_fingerprint, term, vocab)) { + return set[0].term_id; + } + if (matches(set[1], expected_fingerprint, term, vocab)) { + std::swap(set[0], set[1]); + return set[0].term_id; + } + return kInvalidTermId; + } + + void remember(size_t term_hash, uint32_t term_id) { + Set& set = sets[index(term_hash)]; + set[1] = set[0]; + set[0] = Entry {.fingerprint = fingerprint(term_hash), .term_id = term_id}; + } + + std::array sets {}; + static_assert(sizeof(std::array) == 16 * 1024); +}; + +bool SpimiTermBuffer::OwnedVocabEq::operator()(uint32_t stored, + std::string_view probe) const noexcept { +#ifdef BE_TEST + g_owned_term_full_byte_comparisons.fetch_add(1, std::memory_order_relaxed); +#endif + DCHECK_LT(stored, vocab->size()); + return std::string_view((*vocab)[stored]) == probe; +} + +bool SpimiTermBuffer::OwnedVocabEq::operator()(std::string_view probe, + uint32_t stored) const noexcept { + return (*this)(stored, probe); +} + +#ifdef BE_TEST +size_t SpimiTermBuffer::owned_term_key_size_for_test() { + return sizeof(decltype(intern_)::key_type); +} + +void SpimiTermBuffer::set_owned_term_hash_mask_for_test(size_t mask) { + DORIS_CHECK(intern_.empty()); + intern_ = decltype(intern_)(0, OwnedVocabHash {.vocab = &owned_vocab_, .hash_mask = mask}, + OwnedVocabEq {&owned_vocab_}); +} +#endif + +namespace { + +// Returns freed heap arenas to the OS (glibc only). The spill encode churns many +// small allocations whose freed chunks glibc retains in its arenas; trimming +// before the peak-RSS-defining merge phase recovers that retention. No-op (and +// harmless) on non-glibc libcs. +void trim_malloc() { +#if defined(__GLIBC__) + ::malloc_trim(0); +#endif +} + +// Process-unique temp path for a spill run under `dir` (pid + monotonic counter so +// parallel builds / multiple buffers never collide). +std::string make_run_path(const std::string& dir) { + static std::atomic counter {0}; + const uint64_t n = counter.fetch_add(1); + return dir + "/snii_spill_" + std::to_string(::getpid()) + "_" + std::to_string(n) + ".run"; +} + +// TEST-ONLY seam backing testing::vocab_string_materialization_count(). Bumped once +// per DISTINCT interned term (owned_vocab_.emplace_back), never per token. Relaxed: +// the build path is single-threaded, so only the COUNT matters, not ordering. +#ifdef BE_TEST +std::atomic g_vocab_materializations {0}; +#endif + +// G09 seam: spills that consumed a pending process-wide forced-spill request +// (the limiter flagged this buffer as one of the largest registered consumers +// while the global total exceeded the budget). Incremented under BE_TEST only +// (per-token path shared by concurrent writers). +std::atomic g_global_forced_spills {0}; + +// G09 run-file cap seam: merge-compactions of a buffer's run list (always-on: +// at most one per cap-many spills, contention-free). +std::atomic g_run_compactions {0}; + +// Test seam for complete-vocabulary rank rebuilds. The increment is compiled +// out of production because ensure_string_rank() may run on the import path. +#ifdef BE_TEST +std::atomic g_string_rank_rebuilds {0}; +std::atomic g_dense_rank_inversions {0}; +std::atomic g_rank_comparison_sorts {0}; +#endif + +// G11 bench seam: when set (BE_TEST paths only), the add-path prefetch hints +// are skipped so the locality bench can A/B them in one process. Production +// builds never read it (the hint compiles in unconditionally there). +std::atomic g_bench_disable_g11_prefetch {false}; + +// G11 add-path prefetch gate: always-on in production; toggleable under +// BE_TEST for the in-process A/B bench. The branch is perfectly predicted, so +// the bench's OFF arm measures the pre-G11 code path faithfully. +inline bool g11_prefetch_enabled() { +#ifdef BE_TEST + return !g_bench_disable_g11_prefetch.load(std::memory_order_relaxed); +#else + return true; +#endif +} + +// G08: heap payload of one owned-vocab string -- 0 while it fits the SSO buffer +// (those bytes live inside the 32 B header owned_vocab_.capacity() charges), else +// the allocated buffer (capacity + NUL). The SSO capacity is probed from the +// running stdlib so the classification is exact, not hardcoded. +uint64_t string_heap_bytes(const std::string& s) { + static const size_t kSsoCapacity = std::string().capacity(); + return s.capacity() > kSsoCapacity ? static_cast(s.capacity()) + 1 : 0; +} + +void order_ids_by_dense_rank(std::vector* ids, const std::vector& rank) { + if (ids->size() == rank.size()) { + // Touched ids are unique. Equal cardinality therefore means the run covers + // the complete vocabulary, so invert the dense rank in linear time. + for (uint32_t term_id = 0; term_id < rank.size(); ++term_id) { + (*ids)[rank[term_id]] = term_id; + } +#ifdef BE_TEST + g_dense_rank_inversions.fetch_add(1, std::memory_order_relaxed); +#endif + return; + } + + std::ranges::sort(*ids, [&](uint32_t a, uint32_t b) { return rank[a] < rank[b]; }); +#ifdef BE_TEST + g_rank_comparison_sorts.fetch_add(1, std::memory_order_relaxed); +#endif +} + +} // namespace + +namespace testing { +void set_bench_disable_g11_prefetch(bool disabled) { + g_bench_disable_g11_prefetch.store(disabled, std::memory_order_relaxed); +} +uint64_t vocab_string_materialization_count() { +#ifdef BE_TEST + return g_vocab_materializations.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_vocab_string_materialization_count() { +#ifdef BE_TEST + g_vocab_materializations.store(0, std::memory_order_relaxed); +#endif +} +uint64_t global_forced_spills() { + return g_global_forced_spills.load(std::memory_order_relaxed); +} +void reset_global_forced_spills() { + g_global_forced_spills.store(0, std::memory_order_relaxed); +} +uint64_t run_compactions() { + return g_run_compactions.load(std::memory_order_relaxed); +} +void reset_run_compactions() { + g_run_compactions.store(0, std::memory_order_relaxed); +} +uint64_t string_rank_rebuilds() { +#ifdef BE_TEST + return g_string_rank_rebuilds.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_string_rank_rebuilds() { +#ifdef BE_TEST + g_string_rank_rebuilds.store(0, std::memory_order_relaxed); +#endif +} +uint64_t dense_rank_inversions() { +#ifdef BE_TEST + return g_dense_rank_inversions.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t rank_comparison_sorts() { +#ifdef BE_TEST + return g_rank_comparison_sorts.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_rank_ordering_counts() { +#ifdef BE_TEST + g_dense_rank_inversions.store(0, std::memory_order_relaxed); + g_rank_comparison_sorts.store(0, std::memory_order_relaxed); +#endif +} +uint64_t common_gram_pair_unchecked_decode_count() { +#ifdef BE_TEST + return g_common_gram_pair_unchecked_decodes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_trusted_plain_decode_count() { +#ifdef BE_TEST + return g_common_gram_trusted_plain_decodes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_common_gram_pair_fast_path_counts() { +#ifdef BE_TEST + g_common_gram_pair_unchecked_decodes.store(0, std::memory_order_relaxed); + g_common_gram_trusted_plain_decodes.store(0, std::memory_order_relaxed); +#endif +} +uint64_t common_gram_pair_cache_probes() { +#ifdef BE_TEST + return g_common_gram_pair_cache_probes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_pair_cache_pair_hits() { +#ifdef BE_TEST + return g_common_gram_pair_cache_pair_hits.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_pair_cache_same_doc_hits() { +#ifdef BE_TEST + return g_common_gram_pair_cache_same_doc_hits.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_common_gram_pair_cache_counts() { +#ifdef BE_TEST + g_common_gram_pair_cache_probes.store(0, std::memory_order_relaxed); + g_common_gram_pair_cache_pair_hits.store(0, std::memory_order_relaxed); + g_common_gram_pair_cache_same_doc_hits.store(0, std::memory_order_relaxed); +#endif +} +uint64_t common_gram_native_pair_probes() { +#ifdef BE_TEST + return g_common_gram_native_pair_probes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_native_pair_hits() { +#ifdef BE_TEST + return g_common_gram_native_pair_hits.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_native_pair_inserts() { +#ifdef BE_TEST + return g_common_gram_native_pair_inserts.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_common_gram_native_pair_intern_counts() { +#ifdef BE_TEST + g_common_gram_native_pair_probes.store(0, std::memory_order_relaxed); + g_common_gram_native_pair_hits.store(0, std::memory_order_relaxed); + g_common_gram_native_pair_inserts.store(0, std::memory_order_relaxed); +#endif +} +uint64_t common_gram_logical_validation_count() { +#ifdef BE_TEST + return g_common_gram_logical_validations.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_common_gram_logical_validation_count() { +#ifdef BE_TEST + g_common_gram_logical_validations.store(0, std::memory_order_relaxed); +#endif +} +uint64_t common_gram_plain_cache_probes() { +#ifdef BE_TEST + return g_common_gram_plain_cache_probes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_plain_cache_hits() { +#ifdef BE_TEST + return g_common_gram_plain_cache_hits.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +uint64_t common_gram_plain_intern_table_probes() { +#ifdef BE_TEST + return g_common_gram_plain_intern_table_probes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_common_gram_plain_cache_counts() { +#ifdef BE_TEST + g_common_gram_plain_cache_probes.store(0, std::memory_order_relaxed); + g_common_gram_plain_cache_hits.store(0, std::memory_order_relaxed); + g_common_gram_plain_intern_table_probes.store(0, std::memory_order_relaxed); +#endif +} +uint64_t owned_term_full_byte_comparison_count() { +#ifdef BE_TEST + return g_owned_term_full_byte_comparisons.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_owned_term_full_byte_comparison_count() { +#ifdef BE_TEST + g_owned_term_full_byte_comparisons.store(0, std::memory_order_relaxed); +#endif +} +void fail_next_owned_term_reserve() { +#ifdef BE_TEST + g_fail_next_owned_term_reserve.store(true, std::memory_order_relaxed); +#endif +} +void fail_next_owned_term_emplace() { +#ifdef BE_TEST + g_fail_next_owned_term_emplace.store(true, std::memory_order_relaxed); +#endif +} +uint64_t spill_gate_check_count() { +#ifdef BE_TEST + return g_spill_gate_checks.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_spill_gate_check_count() { +#ifdef BE_TEST + g_spill_gate_checks.store(0, std::memory_order_relaxed); +#endif +} +uint64_t compact_chain_varint_decode_count() { +#ifdef BE_TEST + return g_compact_chain_varint_decodes.load(std::memory_order_relaxed); +#else + return 0; +#endif +} +void reset_compact_chain_varint_decode_count() { +#ifdef BE_TEST + g_compact_chain_varint_decodes.store(0, std::memory_order_relaxed); +#endif +} +} // namespace testing + +SpimiTermBuffer::SpimiTermBuffer(const std::vector* vocab, bool has_positions, + size_t spill_threshold_bytes, MemoryReporter* reporter) + : vocab_(vocab), + // Bind the equality functor to &owned_vocab_ even in borrowed mode: + // add_token(string_view) rejects before the functor can dereference it, + // and binding unconditionally keeps both constructors symmetric. + // Initialized in the member-init list (NOT the body): the functors are + // NESTED types, whose default-constructibility is not yet established at + // the point the flat set's default ctor would be needed. The + // (bucket_count, hash, equal) constructor sidesteps that entirely. + // owned_vocab_ is constructed before intern_ (declaration order) and the + // buffer is non-movable, so &owned_vocab_ is stable for the buffer's life. + intern_(0, OwnedVocabHash {.vocab = &owned_vocab_}, OwnedVocabEq {&owned_vocab_}), + has_positions_(has_positions), + spill_threshold_bytes_(spill_threshold_bytes), + mem_reporter_(reporter) { + // Borrowed-vocab mode: only the 4 B/id slot-index array is sized to the + // vocabulary; the Term pool (slots_) grows with the LIVE touched count, so an + // all-but-empty vocabulary costs ~4 B/id instead of ~80 B/id. + slot_of_.assign(vocab_->size(), 0); + // The vocab-sized slot index is resident immediately and survives spills; report + // its initial positive delta now. + report_arena_delta(); +} + +SpimiTermBuffer::SpimiTermBuffer(bool has_positions, size_t spill_threshold_bytes, + MemoryReporter* reporter) + : vocab_(&owned_vocab_), + // Owned-vocab mode: bind both functors to the sole vocabulary so stored + // ids can rehash and string probes resolve full term equality. + // Initialized in the member-init list (NOT the body): + // the functors are NESTED types whose default-constructibility is not yet + // established where the flat set's default ctor (whose noexcept spec inspects + // the functors) would be needed for a body assignment, so the + // (bucket_count, hash, equal) constructor is used instead. owned_vocab_ is + // constructed before intern_ (declaration order) and the buffer is + // non-movable, so &owned_vocab_ is stable for the buffer's life. + intern_(0, OwnedVocabHash {.vocab = &owned_vocab_}, OwnedVocabEq {&owned_vocab_}), + has_positions_(has_positions), + spill_threshold_bytes_(spill_threshold_bytes), + mem_reporter_(reporter) { + report_arena_delta(); +} + +SpimiTermBuffer::~SpimiTermBuffer() { + // G09: leave the process-wide registry FIRST. unregister_buffer removes the + // entry (and its bytes) under the registry mutex -- the same mutex every + // flag store is made under -- so once it returns, no other thread can touch + // global_spill_requested_ while this buffer dies. + if (global_limiter_ != nullptr) { + global_limiter_->unregister_buffer(&global_spill_requested_); + global_limiter_ = nullptr; + } + // Balance the writer-level / Doris tracker on the error path: if the buffer is + // destroyed while resident bytes were reported but not yet freed-and-reported + // (e.g. a build aborts before draining), return them here so nothing leaks. + if (mem_reporter_ != nullptr && reported_resident_ != 0) { + mem_reporter_->report(-reported_resident_); + reported_resident_ = 0; + } + cleanup_runs(); +} + +void SpimiTermBuffer::attach_global_limiter(GlobalMemoryLimiter* limiter) { + // At-most-once: a re-attach would leave a stale registry entry behind (the + // dtor un-registers only the current limiter). + if (limiter == nullptr || global_limiter_ != nullptr) { + return; + } + global_limiter_ = limiter; + // Race-safe vs report: registration and every report run on the OWNER's + // thread, strictly ordered; the registry serializes them against other + // buffers' calls internally. Register with the CURRENT resident total AND + // the current spillable arena bytes (the victim-selection key) so the + // registry is exact from the first moment (a borrowed-vocab buffer + // already holds its vocab-sized slot index here). + global_limiter_->register_buffer(&global_spill_requested_, + static_cast(resident_bytes()), + static_cast(pool_.arena_bytes())); +} + +void SpimiTermBuffer::report_arena_delta() { + if (mem_reporter_ == nullptr && global_limiter_ == nullptr) { + return; + } + // Diff the REAL resident bytes (resident_bytes()) against the last reported + // total; emit the signed delta exactly once. + const auto now = static_cast(resident_bytes()); + // Per-token zero-delta debounce: skip the locked fetch_add when resident is + // unchanged (the common case -- arena_bytes() grows only ~every 32 KiB block and + // the other charged structures grow by geometric capacity steps / per new term + // only, so most tokens see delta==0). A + // delta==0 report() is a no-op (current_.fetch_add(0) plus a mirrored + // consume_release(0)) and leaves reported_resident_ == now, so current_bytes(), + // every over_cap() result, and the gate-2 spill timing stay bit-for-bit identical. + // The spill gate still evaluates the writer-level UNIFIED total whenever the + // arena is large enough to reclaim, even if this buffer's local delta is 0: + // the shared dict buffer may have crossed the cap independently. + if (now == reported_resident_) { + return; + } + if (mem_reporter_ != nullptr) { + mem_reporter_->report(now - reported_resident_); + } + // G09: forward the same debounced total -- as an ABSOLUTE, self-healing + // value -- to the process-wide registry, together with the current + // SPILLABLE arena bytes (the victim-selection key: only the arena is + // reclaimable by a forced spill; the persistent vocab/pair structures are + // not). This is the limiter's decision point: report() flags the + // largest-arena eligible buffers (possibly this one) while the global sum + // exceeds the budget. It only ever takes the registry mutex and flips + // advisory atomics; no lock is held here while spilling (any spill this + // buffer performs happens AFTER this returns, back in + // maybe_spill_after_token, on this thread). + if (global_limiter_ != nullptr) { + global_limiter_->report(&global_spill_requested_, now, + static_cast(pool_.arena_bytes())); + } + reported_resident_ = now; +} + +size_t SpimiTermBuffer::unique_terms() const { + return live_term_count_; +} + +uint64_t SpimiTermBuffer::resident_bytes() const { + // Everything live is charged by CAPACITY (the reserved tail is resident RSS + // and survives spills). All reads are O(1), since this runs once per token. + uint64_t b = pool_.arena_bytes(); // posting chains: docs + prx payload + b += static_cast(slot_of_.capacity()) * sizeof(uint32_t); // vocab-sized slot index + b += static_cast(slots_.capacity()) * sizeof(Term); // live Term pool + b += static_cast(free_slots_.capacity()) * sizeof(uint32_t); + b += static_cast(touched_ids_.capacity()) * sizeof(uint32_t); + // Owned-vocab machinery (all zero in borrowed mode): string headers by vector + // capacity, heap payloads via the incrementally-maintained counter, and the + // intern set's entries at a fixed per-entry estimate (kept at the + // pre-G10 node-set value so the gate-2 spill points are unchanged; see the + // constant's comment). + b += static_cast(owned_vocab_.capacity()) * sizeof(std::string); + b += owned_vocab_heap_bytes_; + b += static_cast(common_word_classification_.capacity()) * + sizeof(CommonWordClassification); + b += static_cast(intern_.size() + common_gram_pair_intern_.size()) * + kInternEntryEstimateBytes; + b += common_gram_pair_cache_bytes_; + b += common_gram_plain_term_cache_bytes_; + // Cached lexicographic ranks survive spills and are included by capacity. + b += static_cast(string_rank_.capacity()) * sizeof(uint32_t); + return b; +} + +// Returns the live Term for `term_id`, claiming a pool slot on first touch (1 == +// new). Reuses a freed slot from free_slots_ when available; otherwise appends a +// fresh Term to slots_. slot_of_[term_id] holds (slot index + 1); 0 means empty. +SpimiTermBuffer::Term& SpimiTermBuffer::term_slot(uint32_t term_id, bool* new_term) { + uint32_t enc = slot_of_[term_id]; + if (enc != 0) { + *new_term = false; + return slots_[enc - 1]; + } + *new_term = true; + uint32_t slot; + if (!free_slots_.empty()) { + slot = free_slots_.back(); + free_slots_.pop_back(); + } else { + slot = static_cast(slots_.size()); + slots_.emplace_back(); + } + slot_of_[term_id] = slot + 1; + return slots_[slot]; +} + +void SpimiTermBuffer::put_varint(Term* t, uint64_t v) { + if (t->head == kNoChain) { + t->head = pool_.start_chain(&t->w, &t->level); + } + if (v < 0x80U) { + pool_.append_byte(&t->w, &t->level, static_cast(v)); + return; + } + pool_.append_varint(&t->w, &t->level, v); +} + +void SpimiTermBuffer::accumulate_without_spill_gate(uint32_t term_id, uint32_t docid, uint32_t pos, + PostingChainShape shape) { + const bool retain_positions = shape == PostingChainShape::kTaggedPositioned; + const bool statless_common_gram = shape == PostingChainShape::kStatlessDocsOnly; + DCHECK(!retain_positions || has_positions_); + bool new_term = false; + Term& t = term_slot(term_id, &new_term); + if (new_term) { + t.shape = shape; + touched_ids_.push_back(term_id); + ++live_term_count_; + } else { + DCHECK(t.shape == shape); + } + // Docs-only accelerator postings are sets. Tokens for one input document are + // contiguous on the writer path, so discard repeated occurrences before they + // allocate arena bytes or enter spill/sort/posting encoding. + if (!retain_positions && t.started && t.cur_docid == docid) { + ++total_tokens_; + return; + } + // A token starts a new doc unless it continues the most-recent doc for this term. + const bool first_token = !t.started; + const bool new_doc = first_token || t.cur_docid != docid; + // A statless CommonGram singleton owns no chain. On its second distinct doc, + // backfill the first absolute docid before appending the current delta. + if (statless_common_gram && !first_token && t.head == kNoChain) { + DCHECK_EQ(t.ntok, 1U); + DCHECK_EQ(t.ndocs, 1U); + put_varint(&t, zigzag_encode(static_cast(t.cur_docid))); + } + + // Positioned and ordinary docs-only terms retain the tagged token stream used + // to reconstruct frequency. A statless CommonGram is already deduplicated per + // document and has no frequency, so its new_doc tag would be the constant 1; + // omit it and store only the document delta. + if (!statless_common_gram) { + // Widen to 64-bit so a full 32-bit position survives the shift. + const uint64_t tagged = retain_positions + ? ((static_cast(pos) << 1) | (new_doc ? 1U : 0U)) + : (new_doc ? 1U : 0U); + put_varint(&t, tagged); + } else { + DCHECK(new_doc); + } + if (new_doc) { + // Out-of-order docids are tolerated (zigzag delta is signed) and reordered at + // finalize; flag them so to_postings sorts. The delta base is the previous + // distinct doc (cur_docid), which is 0 for the very first doc (started==false). + const int64_t base = t.started ? static_cast(t.cur_docid) : 0; + if (t.started && docid < t.cur_docid) { + t.sorted = false; + } + const int64_t delta = static_cast(docid) - base; + if (!first_token || !statless_common_gram) { + put_varint(&t, zigzag_encode(delta)); + } + t.cur_docid = docid; + t.started = true; + // Exact new-doc group count; out-of-order coalescing can only shrink it. + ++t.ndocs; + } + ++t.ntok; + ++total_tokens_; +} + +void SpimiTermBuffer::accumulate(uint32_t term_id, uint32_t docid, uint32_t pos, + bool retain_positions) { + accumulate_without_spill_gate(term_id, docid, pos, + retain_positions ? PostingChainShape::kTaggedPositioned + : PostingChainShape::kTaggedDocsOnly); + maybe_spill_after_token(); +} + +// Per-input-token gate-2 tail. Ordinary adds invoke it after one posting; the +// fused CommonGrams path invokes it after its gram and right plain posting. It +// reports the token's REAL resident growth FIRST so the writer's unified total +// (reporter_->current_bytes()) reflects it before the gate check (single-source +// diff; cheap: a subtraction + relaxed atomic add), then evaluates the spill triggers: +// * Gate-2 (UNIFIED): with a reporter attached, trigger on the writer's TOTAL +// build RAM (arena + vocab structures + dict) crossing the one +// configured cap -- the same total and cap every buffer of this writer +// shares, not a per-buffer threshold. Off Doris (no reporter) fall back to +// the local spill_threshold_bytes_ against resident_bytes(). +// * G08 anti-churn floor: a gate-2 spill reclaims ONLY the posting arena +// (pool_.reset()); the vocab / slot structures resident_bytes() +// now also charges SURVIVE it. Once those persistent bytes alone exceed the +// cap, an unconditioned +// trigger would spill EVERY subsequent token -- one-block runs, k-way-merge +// and spill-fixed-cost blowup. Honor the cap only when at least a quarter of +// it is reclaimable arena: peak stays bounded at persistent + cap/4 and no +// run is smaller than cap/4, while the one-block minimum keeps small caps +// (tests, tiny configs) spilling on the first block exactly as before. +// * Hard arena safety stop, active even in unlimited mode and BYPASSING the +// floor: when the arena nears the 4 GiB uint32-offset limit, spill now -- +// without it a single >4 GiB in-memory segment wraps alloc_run and silently +// corrupts data. A forced spill + final k-way merge stays byte-identical +// regardless of when it fires. +// spill_to_run() resets the arena and reports its negative internally, so the +// unified total drops (and the trigger self-rearms) after each spill. +void SpimiTermBuffer::maybe_spill_after_token() { +#ifdef BE_TEST + g_spill_gate_checks.fetch_add(1, std::memory_order_relaxed); +#endif + constexpr uint64_t kArenaSpillCap = 0xE0000000ULL; // 3.5 GiB, < UINT32_MAX margin + const bool global_requested = global_spill_requested_.load(std::memory_order_relaxed); + const bool arena_near_limit = pool_.arena_bytes() >= kArenaSpillCap; + report_arena_delta(); + const uint64_t gate_cap = + mem_reporter_ != nullptr ? mem_reporter_->cap_bytes() : spill_threshold_bytes_; + const bool arena_worth_spilling = + pool_.arena_bytes() >= std::max(CompactPostingPool::kBlockSize, gate_cap / 4); + // G09: the process-wide limiter flagged this buffer (one of the + // largest-ARENA eligible consumers while the global total exceeded the + // budget). Honored HERE, on the owner's own thread -- never on the + // reporting thread that set the flag. The G08 anti-churn floor (cap/4) is + // deliberately BYPASSED (each victim's arena is below cap/4 by + // construction: it never reached its per-writer gate -- that is exactly + // why the global sum grew), but the FORCED-SPILL FLOOR + // (snii_forced_spill_min_arena_bytes, >= one arena block so a run is + // writable) still applies: a forced spill reclaims ONLY the arena, so + // honoring below the floor would cut a tiny run for near-zero relief. + // Below the floor the request is a NO-OP that stays PENDING -- it is NOT + // retried as a spill each token -- and is honored once the arena regrows + // past the floor (the limiter's victim selection applies the same floor, + // so a below-floor flag only arises from a floor/config race or a test + // seam). A request that finds the owner already drained is never observed + // again -- an advisory no-op (the dtor un-registers) -- and a stale + // re-request after a spill costs at most one extra floor-sized run + // (double-spill is harmless, byte-identical output). + const bool global_spill_now = + global_requested && + pool_.arena_bytes() >= std::max(CompactPostingPool::kBlockSize, + forced_spill_min_arena_bytes_); + const bool over_cap = !global_spill_now && !arena_near_limit && arena_worth_spilling && + (mem_reporter_ != nullptr ? mem_reporter_->over_cap() + : (spill_threshold_bytes_ != 0 && + resident_bytes() >= spill_threshold_bytes_)); + if ((over_cap || global_spill_now || arena_near_limit) && spill_status_.ok()) { + if (global_requested) { + // Consume the request BEFORE spilling: this spill releases exactly + // the arena a forced spill would, so it satisfies the request no + // matter which trigger won the OR above. + global_spill_requested_.store(false, std::memory_order_relaxed); +#ifdef BE_TEST + // Seam under BE_TEST only: per-token path shared by every + // concurrent writer. + g_global_forced_spills.fetch_add(1, std::memory_order_relaxed); +#endif + } + spill_status_ = spill_to_run(); + } +} + +void SpimiTermBuffer::add_token(uint32_t term_id, uint32_t docid, uint32_t pos) { + add_token(term_id, docid, pos, has_positions_); +} + +void SpimiTermBuffer::add_token(uint32_t term_id, uint32_t docid, uint32_t pos, + bool retain_positions) { + // Hot path: a pooled slot lookup + a couple of pushes. No hashing, no string + // construction per token. Reject (and latch) an out-of-range id. + if (term_id >= slot_of_.size()) { + if (spill_status_.ok()) { + spill_status_ = Status::Error( + "spimi: term_id out of vocab range"); + } + return; + } + accumulate(term_id, docid, pos, retain_positions); +} + +void SpimiTermBuffer::add_token(std::string_view term, uint32_t docid, uint32_t pos) { + add_token(term, docid, pos, has_positions_); +} + +void SpimiTermBuffer::add_token(std::string_view term, uint32_t docid, uint32_t pos, + bool retain_positions) { + // Compatibility path: intern the term into the owned vocabulary on first + // occurrence, then accumulate by its id. ONLY valid in OWNED-vocab mode. In + // BORROWED-vocab mode vocab_ points at the caller's vector, NOT &owned_vocab_: + // interning here would grow owned_vocab_ / intern_ / slot_of_ out of step with + // the active (borrowed) vocab, so the new id indexes the WRONG string and writes + // a slot_of_ entry the borrowed-vocab build never reconciles -- silent + // corruption. Reject (and latch) instead of forwarding by a bogus id. + if (vocab_ != &owned_vocab_) { + if (spill_status_.ok()) { + spill_status_ = Status::Error( + "spimi: add_token(string_view) requires owned-vocab mode"); + } + return; + } + DCHECK(!common_gram_pair_keys_); + const uint32_t term_id = find_or_intern_owned_term(term); + accumulate(term_id, docid, pos, retain_positions); +} + +PlainTermId SpimiTermBuffer::intern_plain_term(std::string_view physical_plain_term) { + DCHECK(common_gram_pair_keys_); + DCHECK(vocab_ == &owned_vocab_); + DCHECK(!physical_plain_term.empty()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(physical_plain_term)); + const size_t term_hash = intern_.hash(physical_plain_term); + uint32_t term_id = find_interned_plain_term(physical_plain_term, term_hash); + if (term_id == CommonGramPlainTermCache::kInvalidTermId) { + term_id = intern_owned_term(std::string(physical_plain_term), term_hash); + remember_plain_term(term_hash, term_id); + } + return PlainTermId {.value = term_id}; +} + +PlainTermId SpimiTermBuffer::intern_plain_term(std::string_view physical_plain_term, + std::string_view logical_plain_term) { + DCHECK(common_gram_pair_keys_); + DCHECK(vocab_ == &owned_vocab_); + DCHECK(!physical_plain_term.empty()); + DCHECK(!logical_plain_term.empty()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(physical_plain_term)); + + const size_t term_hash = intern_.hash(physical_plain_term); + uint32_t term_id = find_interned_plain_term(physical_plain_term, term_hash); + if (term_id != CommonGramPlainTermCache::kInvalidTermId) { + return PlainTermId {.value = term_id}; + } + +#ifdef BE_TEST + g_common_gram_logical_validations.fetch_add(1, std::memory_order_relaxed); +#endif + auto validation = segment_v2::inverted_index::validate_common_grams_logical_term( + logical_plain_term, "input token"); + if (!validation.ok()) { + throw Exception(validation); + } + if (physical_plain_term.size() > segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES) { + throw Exception(Status::Error( + "CommonGrams escaped plain term would exceed the 16383-byte key limit; " + "set enable_common_grams_index_build=false and retry the import in a new " + "transaction")); + } + term_id = intern_owned_term(std::string(physical_plain_term), term_hash); + remember_plain_term(term_hash, term_id); + return PlainTermId {.value = term_id}; +} + +ClassifiedPlainTerm SpimiTermBuffer::intern_classified_plain_term( + std::string_view physical_plain_term, std::string_view logical_plain_term, + const segment_v2::inverted_index::CommonWordSet& common_words) { + DCHECK(common_gram_pair_keys_); + const PlainTermId id = intern_plain_term(physical_plain_term, logical_plain_term); + DCHECK_EQ(common_word_classification_.size(), owned_vocab_.size()); + DCHECK_LT(id.value, common_word_classification_.size()); + CommonWordClassification& classification = common_word_classification_[id.value]; + if (classification == CommonWordClassification::kUnknown) { + classification = common_words.contains(logical_plain_term) + ? CommonWordClassification::kCommon + : CommonWordClassification::kNotCommon; + } + return ClassifiedPlainTerm { + .id = id, + .is_common = classification == CommonWordClassification::kCommon, + }; +} + +void SpimiTermBuffer::add_plain_token(PlainTermId term_id, uint32_t docid, uint32_t pos) { + DCHECK(common_gram_pair_keys_); + DCHECK_LT(term_id.value, owned_vocab_.size()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[term_id.value])); + accumulate(term_id.value, docid, pos, has_positions_); +} + +void SpimiTermBuffer::add_common_gram_without_spill_gate(PlainTermId left, PlainTermId right, + uint32_t docid, uint32_t pos, + bool retain_positions) { + DCHECK(common_gram_pair_keys_); + DCHECK_LT(left.value, owned_vocab_.size()); + DCHECK_LT(right.value, owned_vocab_.size()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[left.value])); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[right.value])); + DCHECK(common_gram_pair_cache_ != nullptr); + const PostingChainShape shape = retain_positions ? PostingChainShape::kTaggedPositioned + : PostingChainShape::kStatlessDocsOnly; + const uint64_t pair = (static_cast(left.value) << 32) | right.value; + CommonGramPairCache::Entry& entry = + common_gram_pair_cache_->entries[CommonGramPairCache::index(pair)]; +#ifdef BE_TEST + g_common_gram_pair_cache_probes.fetch_add(1, std::memory_order_relaxed); +#endif + if (entry.term_id != CommonGramPairCache::kInvalidTermId && entry.pair == pair) { +#ifdef BE_TEST + g_common_gram_pair_cache_pair_hits.fetch_add(1, std::memory_order_relaxed); +#endif + DCHECK_LT(entry.term_id, slot_of_.size()); + if (!retain_positions && entry.last_docid == docid) { +#ifdef BE_TEST + g_common_gram_pair_cache_same_doc_hits.fetch_add(1, std::memory_order_relaxed); +#endif + ++total_tokens_; + return; + } + entry.last_docid = docid; + accumulate_without_spill_gate(entry.term_id, docid, pos, shape); + return; + } + + const uint32_t term_id = find_or_intern_common_gram_pair(left, right, pair); + DCHECK_NE(term_id, CommonGramPairCache::kInvalidTermId); + entry = CommonGramPairCache::Entry {.pair = pair, .term_id = term_id, .last_docid = docid}; + accumulate_without_spill_gate(term_id, docid, pos, shape); +} + +void SpimiTermBuffer::add_common_gram(PlainTermId left, PlainTermId right, uint32_t docid, + uint32_t pos, bool retain_positions) { + add_common_gram_without_spill_gate(left, right, docid, pos, retain_positions); + maybe_spill_after_token(); +} + +void SpimiTermBuffer::add_common_gram_and_plain(PlainTermId left, PlainTermId right, uint32_t docid, + uint32_t gram_pos, uint32_t plain_pos, + bool retain_gram_positions) { + add_common_gram_without_spill_gate(left, right, docid, gram_pos, retain_gram_positions); + DCHECK_LT(right.value, owned_vocab_.size()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[right.value])); + accumulate_without_spill_gate(right.value, docid, plain_pos, + has_positions_ ? PostingChainShape::kTaggedPositioned + : PostingChainShape::kTaggedDocsOnly); + maybe_spill_after_token(); +} + +void SpimiTermBuffer::enable_common_gram_pair_keys() { + DORIS_CHECK(vocab_ == &owned_vocab_); + DORIS_CHECK_EQ(total_tokens_, 0); + DORIS_CHECK(owned_vocab_.empty()); + DORIS_CHECK(!common_gram_pair_keys_); + DORIS_CHECK(common_word_classification_.empty()); + auto pair_cache = std::make_unique(); + auto plain_term_cache = std::make_unique(); + common_gram_pair_keys_ = true; + common_gram_pair_cache_ = std::move(pair_cache); + common_gram_pair_cache_bytes_ = sizeof(CommonGramPairCache); + common_gram_plain_term_cache_ = std::move(plain_term_cache); + common_gram_plain_term_cache_bytes_ = sizeof(CommonGramPlainTermCache); + report_arena_delta(); +} + +uint32_t SpimiTermBuffer::find_or_intern_owned_term(std::string_view term) { + static_assert(std::is_same_v); + DCHECK_LE(term.size(), std::numeric_limits::max()); + const size_t term_hash = intern_.hash(term); + const auto found = intern_.find(term, term_hash); + if (found != intern_.end()) { + const uint32_t term_id = *found; + if (g11_prefetch_enabled()) { + __builtin_prefetch(slot_of_.data() + term_id); + } + return term_id; + } + return intern_owned_term(std::string(term), term_hash); +} + +uint32_t SpimiTermBuffer::find_or_intern_common_gram_pair(PlainTermId left, PlainTermId right, + uint64_t pair) { + DCHECK(common_gram_pair_keys_); + DCHECK_LT(left.value, owned_vocab_.size()); + DCHECK_LT(right.value, owned_vocab_.size()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[left.value])); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(owned_vocab_[right.value])); +#ifdef BE_TEST + g_common_gram_native_pair_probes.fetch_add(1, std::memory_order_relaxed); +#endif + const auto found = common_gram_pair_intern_.find(pair); + if (found != common_gram_pair_intern_.end()) { +#ifdef BE_TEST + g_common_gram_native_pair_hits.fetch_add(1, std::memory_order_relaxed); +#endif + return found->second; + } + + const auto key = EncodeCommonGramPairKey(left, right); + const uint32_t term_id = append_owned_vocab_term(std::string(key.data(), key.size())); + const auto [inserted_it, inserted] = common_gram_pair_intern_.try_emplace(pair, term_id); + DCHECK(inserted); + DCHECK_EQ(inserted_it->second, term_id); +#ifdef BE_TEST + g_common_gram_native_pair_inserts.fetch_add(1, std::memory_order_relaxed); +#endif + return term_id; +} + +uint32_t SpimiTermBuffer::find_interned_plain_term(std::string_view term, size_t term_hash) { + DCHECK(common_gram_pair_keys_); + DCHECK(common_gram_plain_term_cache_ != nullptr); +#ifdef BE_TEST + g_common_gram_plain_cache_probes.fetch_add(1, std::memory_order_relaxed); +#endif + const uint32_t cached_term_id = + common_gram_plain_term_cache_->find(term_hash, term, owned_vocab_); + if (cached_term_id != CommonGramPlainTermCache::kInvalidTermId) { +#ifdef BE_TEST + g_common_gram_plain_cache_hits.fetch_add(1, std::memory_order_relaxed); +#endif + if (g11_prefetch_enabled()) { + __builtin_prefetch(slot_of_.data() + cached_term_id); + } + return cached_term_id; + } + +#ifdef BE_TEST + g_common_gram_plain_intern_table_probes.fetch_add(1, std::memory_order_relaxed); +#endif + const auto found = intern_.find(term, term_hash); + if (found == intern_.end()) { + return CommonGramPlainTermCache::kInvalidTermId; + } + const uint32_t term_id = *found; + remember_plain_term(term_hash, term_id); + if (g11_prefetch_enabled()) { + __builtin_prefetch(slot_of_.data() + term_id); + } + return term_id; +} + +void SpimiTermBuffer::remember_plain_term(size_t term_hash, uint32_t term_id) { + DCHECK(common_gram_pair_keys_); + DCHECK(common_gram_plain_term_cache_ != nullptr); + DCHECK_LT(term_id, owned_vocab_.size()); + common_gram_plain_term_cache_->remember(term_hash, term_id); +} + +bool SpimiTermBuffer::transient_term_less(uint32_t left_id, uint32_t right_id) const { + const std::vector& v = vocab(); + const std::string_view left = v[left_id]; + const std::string_view right = v[right_id]; + if (!common_gram_pair_keys_) { + return left < right; + } + + const bool left_is_pair = is_common_gram_pair_key(left); + const bool right_is_pair = is_common_gram_pair_key(right); + if (left_is_pair != right_is_pair) { + const std::string_view plain = left_is_pair ? right : left; + DCHECK(!plain.empty()); + DCHECK(!segment_v2::inverted_index::is_internal_term_key(plain)); + const bool pair_sorts_first = + static_cast( + segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN.front()) < + static_cast(plain.front()); + return left_is_pair ? pair_sorts_first : !pair_sorts_first; + } + if (!left_is_pair) { + return left < right; + } + + const CommonGramPairIds left_ids = decode_common_gram_pair_key_unchecked(left); + const CommonGramPairIds right_ids = decode_common_gram_pair_key_unchecked(right); + DCHECK_LT(left_ids.left.value, v.size()); + DCHECK_LT(left_ids.right.value, v.size()); + DCHECK_LT(right_ids.left.value, v.size()); + DCHECK_LT(right_ids.right.value, v.size()); + + const LogicalPlainKeyView left_left_key(v[left_ids.left.value]); + const LogicalPlainKeyView left_right_key(v[left_ids.right.value]); + const LogicalPlainKeyView right_left_key(v[right_ids.left.value]); + const LogicalPlainKeyView right_right_key(v[right_ids.right.value]); + if (left_left_key.size() != right_left_key.size()) { + return left_left_key.size() < right_left_key.size(); + } + const int left_component_order = compare_logical_plain_keys(left_left_key, right_left_key); + if (left_component_order != 0) { + return left_component_order < 0; + } + return compare_logical_plain_keys(left_right_key, right_right_key) < 0; +} + +std::string SpimiTermBuffer::materialize_transient_term(std::string_view term) const { + if (!is_common_gram_pair_key(term)) { + return std::string(term); + } + + DCHECK(common_gram_pair_keys_); + const CommonGramPairIds ids = decode_common_gram_pair_key_unchecked(term); + DCHECK_LT(ids.left.value, owned_vocab_.size()); + DCHECK_LT(ids.right.value, owned_vocab_.size()); + DCHECK(!is_common_gram_pair_key(owned_vocab_[ids.left.value])); + DCHECK(!is_common_gram_pair_key(owned_vocab_[ids.right.value])); + + std::string left_scratch; + std::string right_scratch; + const std::string_view left = + decode_logical_plain_term_trusted(owned_vocab_[ids.left.value], &left_scratch); + const std::string_view right = + decode_logical_plain_term_trusted(owned_vocab_[ids.right.value], &right_scratch); + std::string output; + [[maybe_unused]] const bool encoded = + segment_v2::inverted_index::try_encode_common_gram_prevalidated(left, right, output); + DCHECK(encoded); + return output; +} + +// Prepared first-time insertion stores the string before emplace so every +// stored id remains resolvable during later growth rehashes. +uint32_t SpimiTermBuffer::intern_owned_term(std::string&& term_str, size_t term_hash) { + const size_t next_vocab_size = owned_vocab_.size() + 1; + DCHECK_LE(next_vocab_size, std::numeric_limits::max()); + + size_t target_capacity = owned_vocab_.capacity(); + if (target_capacity < next_vocab_size) { + target_capacity = target_capacity <= std::numeric_limits::max() / 2 + ? std::max(next_vocab_size, target_capacity * 2) + : next_vocab_size; + } + + // Prepare append-only vectors geometrically before publishing a vocabulary + // id. A later reserve may throw after an earlier vector already changed + // capacity, so the catch path must settle that resident delta before + // propagating the failure. + try { + owned_vocab_.reserve(target_capacity); +#ifdef BE_TEST + if (g_fail_next_owned_term_reserve.exchange(false, std::memory_order_relaxed)) { + throw std::bad_alloc(); + } +#endif + slot_of_.reserve(target_capacity); + if (common_gram_pair_keys_) { + common_word_classification_.reserve(target_capacity); + } + } catch (...) { + report_arena_delta(); + throw; + } + report_arena_delta(); + + const uint32_t term_id = append_owned_vocab_term(std::move(term_str)); + static_assert(std::is_nothrow_copy_constructible_v); + + const auto rollback_append = [&]() { +#ifdef BE_TEST + g_vocab_materializations.fetch_sub(1, std::memory_order_relaxed); +#endif + owned_vocab_heap_bytes_ -= string_heap_bytes(owned_vocab_.back()); + slot_of_.pop_back(); + if (common_gram_pair_keys_) { + common_word_classification_.pop_back(); + } + owned_vocab_.pop_back(); + report_arena_delta(); + }; + + // phmap allocates a growth table before publishing the prepared slot, and + // constructing this trivial key cannot throw. If allocation fails, the old + // table is intact and only the preceding vocabulary append needs rollback. + const auto [it, inserted] = [&]() { + try { +#ifdef BE_TEST + if (g_fail_next_owned_term_emplace.exchange(false, std::memory_order_relaxed)) { + throw std::bad_alloc(); + } +#endif + return intern_.emplace_with_hash(term_hash, term_id); + } catch (...) { + rollback_append(); + throw; + } + }(); + if (!inserted) { + rollback_append(); + } + DCHECK(inserted); + DCHECK_EQ(*it, term_id); + return term_id; +} + +uint32_t SpimiTermBuffer::append_owned_vocab_term(std::string&& term_str) { + const uint32_t term_id = static_cast(owned_vocab_.size()); + owned_vocab_.emplace_back(std::move(term_str)); + if (common_gram_pair_keys_) { + common_word_classification_.push_back(CommonWordClassification::kUnknown); + DCHECK_EQ(common_word_classification_.size(), owned_vocab_.size()); + } + slot_of_.push_back(0); // vocab grows: new id starts with no live slot + // G08: credit the stored string's heap payload (0 for SSO); the header is + // charged via owned_vocab_.capacity(). + owned_vocab_heap_bytes_ += string_heap_bytes(owned_vocab_[term_id]); +#ifdef BE_TEST + g_vocab_materializations.fetch_add(1, std::memory_order_relaxed); +#endif + return term_id; +} + +namespace { + +// Reorders a term's flat arrays into ascending-docid order, COALESCING any +// same-docid groups so the result has exactly one entry per docid -- matching the +// k-way-merge path's boundary-doc coalescing and the writer's strictly-ascending +// precondition. Only invoked for the rare term that received out-of-order docids +// (the common ascending path leaves t.sorted true and skips it). +// +// A docid may REVISIT (e.g. feed 5,1,5): the chain holds two separate doc-groups +// for doc 5. A STABLE sort keeps equal-docid groups in arrival order, then the +// coalesce pass sums their freqs and concatenates their positions in that same +// (document/arrival) order -- so the merged positions stay consistent with the +// merged freqs, exactly as the run-order merge would have produced. +template +Status reserve_tracked_vector(std::vector* values, size_t target, + MemoryReporter* memory_reporter, + MemoryReporter::Reservation* reservation) { + if (target <= values->capacity()) { + return Status::OK(); + } + if (target > std::numeric_limits::max() / sizeof(T)) { + return Status::Error( + "spimi materialization: vector byte capacity overflow"); + } + if (memory_reporter == nullptr) { + values->reserve(target); + return Status::OK(); + } + MemoryReporter::Reservation replacement; + RETURN_IF_ERROR(reservation->prepare_replacement(static_cast(target) * sizeof(T), + &replacement)); + values->reserve(target); + DCHECK_EQ(values->capacity(), target); + *reservation = std::move(replacement); + return Status::OK(); +} + +Status sort_by_docid(std::vector* docids, std::vector* freqs, + std::vector* positions_flat, bool has_positions, + MemoryReporter* memory_reporter, + MemoryReporter::Reservation* docids_reservation, + MemoryReporter::Reservation* freqs_reservation, + MemoryReporter::Reservation* positions_reservation) { + const size_t n = docids->size(); + MemoryReporter::Reservation order_reservation = memory_reporter == nullptr + ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation pos_off_reservation = memory_reporter == nullptr + ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation sorted_docids_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation sorted_freqs_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + MemoryReporter::Reservation sorted_positions_reservation = + memory_reporter == nullptr ? MemoryReporter::Reservation() + : memory_reporter->make_reservation(); + std::vector order; + RETURN_IF_ERROR(reserve_tracked_vector(&order, n, memory_reporter, &order_reservation)); + order.resize(n); + std::iota(order.begin(), order.end(), 0); + // The original index breaks equal-doc ties, preserving arrival order without + // stable_sort's implementation-owned allocation. + std::ranges::sort(order, [&](size_t a, size_t b) { + if ((*docids)[a] != (*docids)[b]) { + return (*docids)[a] < (*docids)[b]; + } + return a < b; + }); + + std::vector pos_off; + if (has_positions) { + RETURN_IF_ERROR(reserve_tracked_vector(&pos_off, n, memory_reporter, &pos_off_reservation)); + pos_off.resize(n); + uint32_t running = 0; + for (size_t i = 0; i < n; ++i) { + pos_off[i] = running; + running += (*freqs)[i]; + } + } + std::vector nd, nf, np; + RETURN_IF_ERROR(reserve_tracked_vector(&nd, n, memory_reporter, &sorted_docids_reservation)); + RETURN_IF_ERROR(reserve_tracked_vector(&nf, n, memory_reporter, &sorted_freqs_reservation)); + if (has_positions) { + RETURN_IF_ERROR(reserve_tracked_vector(&np, positions_flat->size(), memory_reporter, + &sorted_positions_reservation)); + } + for (size_t k : order) { + // Coalesce a revisited docid into the previous entry (it sorts adjacent now): + // sum freqs and append this group's positions right after the prior group's, + // so flat doc order stays partitioned by the merged freqs. + if (!nd.empty() && nd.back() == (*docids)[k]) { + if (has_positions) { + nf.back() += (*freqs)[k]; + } + } else { + nd.push_back((*docids)[k]); + nf.push_back((*freqs)[k]); + } + if (has_positions) { + np.insert(np.end(), positions_flat->begin() + pos_off[k], + positions_flat->begin() + pos_off[k] + (*freqs)[k]); + } + } + docids->swap(nd); + freqs->swap(nf); + std::swap(*docids_reservation, sorted_docids_reservation); + std::swap(*freqs_reservation, sorted_freqs_reservation); + if (has_positions) { + positions_flat->swap(np); + std::swap(*positions_reservation, sorted_positions_reservation); + } + return Status::OK(); +} + +} // namespace + +namespace { + +// Decodes one varint from a pool chain cursor. The chain was written by +// encode_varint*, so the same LEB128 continuation-bit loop reconstructs it. +uint64_t decode_chain_varint(CompactPostingPool::Cursor* c) { +#ifdef BE_TEST + g_compact_chain_varint_decodes.fetch_add(1, std::memory_order_relaxed); +#endif + return c->read_varint(); +} + +} // namespace + +// Decodes the compact tagged chain directly into caller-owned posting windows. +class SpimiTermBuffer::ArenaTermPostingSource final : public TermPostingSource { +public: + ArenaTermPostingSource(const CompactPostingPool* pool, const Term& term) + : shape_(term.shape), + remaining_docs_(term.ndocs), + remaining_tokens_(term.ntok), + inline_docid_(term.cur_docid) { + if (term.head != kNoChain) { + doc_cursor_.emplace(pool->cursor(term.head, term.w.cur)); + } + } + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + if (out == nullptr || exhausted == nullptr || target_docs == 0 || !out->empty()) { + return Status::Error( + "spimi arena source: invalid fill arguments"); + } + const uint32_t count = std::min(target_docs, remaining_docs_); + if (count == 0) { + *exhausted = true; + return Status::OK(); + } + + if (shape_ == PostingChainShape::kStatlessDocsOnly) { + MutableTermPostingSpan destination; + RETURN_IF_ERROR(out->grow_uninitialized(count, /*has_freqs=*/false, + /*position_count=*/0, &destination)); + for (uint32_t i = 0; i < count; ++i) { + if (!doc_cursor_) { + DCHECK_EQ(remaining_docs_, 1U); + destination.docids[i] = inline_docid_; + } else { + absolute_docid_ += zigzag_decode(decode_chain_varint(&*doc_cursor_)); + destination.docids[i] = static_cast(absolute_docid_); + } + } + remaining_tokens_ -= count; + } else { + RETURN_IF_ERROR(fill_tagged(count, out)); + } + + remaining_docs_ -= count; + *exhausted = remaining_docs_ == 0; + if (*exhausted) { + DCHECK_EQ(remaining_tokens_, 0U); + DCHECK(!pending_new_doc_); + } + return Status::OK(); + } + + bool exhausted() const { return remaining_docs_ == 0; } + +private: + Status fill_tagged(uint32_t count, TermPostingBuffer* out) { + const bool has_positions = shape_ == PostingChainShape::kTaggedPositioned; + const bool terminal_fill = count == remaining_docs_; + const size_t position_count = has_positions && terminal_fill ? remaining_tokens_ : 0; + MutableTermPostingSpan documents; + RETURN_IF_ERROR( + out->grow_uninitialized(count, /*has_freqs=*/true, position_count, &documents)); + size_t position_index = 0; + for (uint32_t i = 0; i < count; ++i) { + uint64_t tagged = 0; + if (pending_new_doc_) { + tagged = pending_tagged_; + pending_new_doc_ = false; + } else { + DCHECK_GT(remaining_tokens_, 0U); + tagged = decode_chain_varint(&*doc_cursor_); + } + DCHECK_NE(tagged & 1U, 0U); + absolute_docid_ += zigzag_decode(decode_chain_varint(&*doc_cursor_)); + documents.docids[i] = static_cast(absolute_docid_); + uint32_t frequency = 0; + while (true) { + if (has_positions) { + if (terminal_fill) { + documents.positions_flat[position_index++] = + static_cast(tagged >> 1); + } else { + RETURN_IF_ERROR(out->append_position(static_cast(tagged >> 1))); + } + } + ++frequency; + --remaining_tokens_; + if (remaining_tokens_ == 0) { + break; + } + tagged = decode_chain_varint(&*doc_cursor_); + if ((tagged & 1U) != 0) { + pending_tagged_ = tagged; + pending_new_doc_ = true; + break; + } + } + documents.freqs[i] = frequency; + } + DCHECK_EQ(position_index, documents.positions_flat.size()); + return Status::OK(); + } + + PostingChainShape shape_; + std::optional doc_cursor_; + uint32_t remaining_docs_ = 0; + uint32_t remaining_tokens_ = 0; + uint32_t inline_docid_ = 0; + int64_t absolute_docid_ = 0; + uint64_t pending_tagged_ = 0; + bool pending_new_doc_ = false; +}; + +Status SpimiTermBuffer::to_postings(std::string term, Term&& t, + TrackedTermPostings* tracked) const { + DCHECK(tracked != nullptr); + TermPostings& postings = tracked->postings; + DCHECK(postings.docids.empty()); + DCHECK(postings.freqs.empty()); + DCHECK(postings.positions_flat.empty()); + postings.term = std::move(term); + postings.retain_positions = t.shape == PostingChainShape::kTaggedPositioned; + if (t.ntok == 0) { + return Status::OK(); + } + + RETURN_IF_ERROR(reserve_tracked_vector(&postings.docids, t.ndocs, mem_reporter_, + &tracked->docids_reservation)); + if (t.shape != PostingChainShape::kStatlessDocsOnly) { + RETURN_IF_ERROR(reserve_tracked_vector(&postings.freqs, t.ndocs, mem_reporter_, + &tracked->freqs_reservation)); + } + if (t.shape == PostingChainShape::kTaggedPositioned) { + RETURN_IF_ERROR(reserve_tracked_vector(&postings.positions_flat, t.ntok, mem_reporter_, + &tracked->positions_reservation)); + } + + ArenaTermPostingSource source(&pool_, t); + TermPostingBuffer buffer(mem_reporter_); + bool exhausted = false; + while (!exhausted) { + buffer.clear_reuse(); + RETURN_IF_ERROR(source.fill(format::kAdaptiveWindowDocs, &buffer, &exhausted)); + postings.docids.insert(postings.docids.end(), buffer.docids().begin(), + buffer.docids().end()); + postings.freqs.insert(postings.freqs.end(), buffer.freqs().begin(), buffer.freqs().end()); + postings.positions_flat.insert(postings.positions_flat.end(), + buffer.positions_flat().begin(), + buffer.positions_flat().end()); + } + if (!t.sorted && t.shape == PostingChainShape::kStatlessDocsOnly) { + std::ranges::sort(postings.docids); + postings.docids.erase(std::unique(postings.docids.begin(), postings.docids.end()), + postings.docids.end()); + } else if (!t.sorted) { + RETURN_IF_ERROR(sort_by_docid(&postings.docids, &postings.freqs, &postings.positions_flat, + postings.retain_positions, mem_reporter_, + &tracked->docids_reservation, &tracked->freqs_reservation, + &tracked->positions_reservation)); + } + return Status::OK(); +} + +void SpimiTermBuffer::ensure_string_rank() const { + const std::vector& v = vocab(); + if (string_rank_.size() == v.size()) { + return; // already built for the current append-only vocabulary + } + // Build the complete rank required by the first spill and by k-way merge + // paths. Ordinary spills with a stale rank deliberately do not call here. + if (!common_gram_pair_keys_) { + std::vector order(v.size()); + std::iota(order.begin(), order.end(), 0U); + std::ranges::sort(order, [&](uint32_t a, uint32_t b) { return transient_term_less(a, b); }); + string_rank_.assign(v.size(), 0U); + for (uint32_t rank = 0; rank < order.size(); ++rank) { + string_rank_[order[rank]] = rank; + } + } else { + size_t pair_count = 0; + for (const std::string& term : v) { + pair_count += is_common_gram_pair_key(term); + } + + std::vector plain_order; + std::vector pair_order; + plain_order.reserve(v.size() - pair_count); + pair_order.reserve(pair_count); + for (uint32_t term_id = 0; term_id < v.size(); ++term_id) { + if (is_common_gram_pair_key(v[term_id])) { + pair_order.push_back(term_id); + } else { + plain_order.push_back(term_id); + } + } + + // EscapedV1 preserves logical byte order: 0x1e maps to 0x1eE and 0x1f + // maps to 0x1eG. One physical sort therefore supplies both the final plain + // order and the logical component rank used by a gram's right term. + std::ranges::sort(plain_order, + [&](uint32_t left, uint32_t right) { return v[left] < v[right]; }); + string_rank_.assign(v.size(), 0U); + for (uint32_t rank = 0; rank < plain_order.size(); ++rank) { + string_rank_[plain_order[rank]] = rank; + } + + // Decode each transient pair exactly once. The low word remains its term id; + // the high word temporarily carries the left plain id, while the pair's + // unused rank slot carries its right component's logical rank. + for (uint64_t& decorated_pair : pair_order) { + const uint32_t pair_term_id = static_cast(decorated_pair); + const CommonGramPairIds ids = decode_common_gram_pair_key_unchecked(v[pair_term_id]); + DCHECK_LT(ids.left.value, v.size()); + DCHECK_LT(ids.right.value, v.size()); + DCHECK(!is_common_gram_pair_key(v[ids.left.value])); + DCHECK(!is_common_gram_pair_key(v[ids.right.value])); + string_rank_[pair_term_id] = string_rank_[ids.right.value]; + decorated_pair = (static_cast(ids.left.value) << 32) | pair_term_id; + } + + // The physical gram key orders its left component by fixed-width encoded + // length, then by logical bytes. Stable per-length offsets convert the + // already-logically-sorted plain ids into that compound dense rank without + // another comparison sort. + std::vector next_length_rank( + segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES + 1, 0U); + for (uint32_t plain_id : plain_order) { + const size_t logical_size = LogicalPlainKeyView(v[plain_id]).size(); + DCHECK_LT(logical_size, next_length_rank.size()); + ++next_length_rank[logical_size]; + } + uint32_t next_rank = 0; + for (uint32_t& length_count : next_length_rank) { + const uint32_t count = length_count; + length_count = next_rank; + next_rank += count; + } + DCHECK_EQ(next_rank, plain_order.size()); + for (uint32_t plain_id : plain_order) { + const size_t logical_size = LogicalPlainKeyView(v[plain_id]).size(); + string_rank_[plain_id] = next_length_rank[logical_size]++; + } + for (uint64_t& decorated_pair : pair_order) { + const uint32_t left_plain_id = static_cast(decorated_pair >> 32); + const uint32_t pair_term_id = static_cast(decorated_pair); + decorated_pair = + (static_cast(string_rank_[left_plain_id]) << 32) | pair_term_id; + } + + std::ranges::sort(pair_order, [&](uint64_t left, uint64_t right) { + const uint32_t left_component_rank = static_cast(left >> 32); + const uint32_t right_component_rank = static_cast(right >> 32); + if (left_component_rank != right_component_rank) { + return left_component_rank < right_component_rank; + } + return string_rank_[static_cast(left)] < + string_rank_[static_cast(right)]; + }); + + // No EscapedV1 plain key enters 0x1f, so every materialized gram forms one + // contiguous namespace group between the two physical-plain ranges. + const auto pair_position = std::lower_bound( + plain_order.begin(), plain_order.end(), segment_v2::inverted_index::CG_V1_MARKER, + [&](uint32_t plain_id, std::string_view marker) { return v[plain_id] < marker; }); + uint32_t final_rank = 0; + for (auto it = plain_order.begin(); it != pair_position; ++it) { + string_rank_[*it] = final_rank++; + } + for (uint64_t decorated_pair : pair_order) { + string_rank_[static_cast(decorated_pair)] = final_rank++; + } + for (auto it = pair_position; it != plain_order.end(); ++it) { + string_rank_[*it] = final_rank++; + } + DCHECK_EQ(final_rank, v.size()); + } +#ifdef BE_TEST + g_string_rank_rebuilds.fetch_add(1, std::memory_order_relaxed); +#endif +} + +std::vector SpimiTermBuffer::sorted_ids() const { + std::vector ids = touched_ids_; + const std::vector& v = vocab(); + if (string_rank_.empty()) { + // Preserve the fixed-vocabulary fast path: the first spill pays once for + // a complete rank, then every later spill is integer-only until vocab grows. + ensure_string_rank(); + } + if (string_rank_.size() == v.size()) { + order_ids_by_dense_rank(&ids, string_rank_); + } else { + // Vocabulary grew after the last complete rank. A run needs only its touched + // terms in lexical order; defer the O(vocab log vocab) rebuild until a k-way + // merge needs rank lookups for arbitrary ids. Reserve the same persistent + // rank capacity the old rebuild allocated so resident accounting and later + // spill-trigger timing remain unchanged. + string_rank_.reserve(v.size()); + std::ranges::sort(ids, [&](uint32_t a, uint32_t b) { return transient_term_less(a, b); }); + } + return ids; +} + +void SpimiTermBuffer::release_term(uint32_t term_id) { + const uint32_t enc = slot_of_[term_id]; + DCHECK_NE(enc, 0U); + const uint32_t slot = enc - 1; + slots_[slot] = Term(); // free this term's arrays; the empty Term slot is reusable + free_slots_.push_back(slot); + slot_of_[term_id] = 0; + --live_term_count_; +} + +Status SpimiTermBuffer::drain_sorted_streamed(const StreamedTermConsumer& fn) { + const std::vector& v = vocab(); + ensure_string_rank(); + report_arena_delta(); + order_ids_by_dense_rank(&touched_ids_, string_rank_); + intern_ = decltype(intern_)(0, OwnedVocabHash {.vocab = &owned_vocab_}, + OwnedVocabEq {&owned_vocab_}); + common_gram_pair_intern_ = decltype(common_gram_pair_intern_)(); + std::vector().swap(common_word_classification_); + std::vector().swap(string_rank_); + report_arena_delta(); + + constexpr size_t kSlotIndexPrefetchDistance = 32; + constexpr size_t kTermPrefetchDistance = 16; + Status callback_status = Status::OK(); + for (size_t ordinal = 0; ordinal < touched_ids_.size(); ++ordinal) { + if (ordinal + kSlotIndexPrefetchDistance < touched_ids_.size()) { + const uint32_t future_id = touched_ids_[ordinal + kSlotIndexPrefetchDistance]; + __builtin_prefetch(slot_of_.data() + future_id); + } + if (ordinal + kTermPrefetchDistance < touched_ids_.size()) { + const uint32_t future_id = touched_ids_[ordinal + kTermPrefetchDistance]; + const uint32_t future_enc = slot_of_[future_id]; + DCHECK_NE(future_enc, 0U); + __builtin_prefetch(slots_.data() + future_enc - 1); + __builtin_prefetch(v.data() + future_id); + } + const uint32_t id = touched_ids_[ordinal]; + const uint32_t enc = slot_of_[id]; + DCHECK_NE(enc, 0U); + Term term = slots_[enc - 1]; + slots_[enc - 1] = Term(); + slot_of_[id] = 0; + --live_term_count_; + + std::string output_term = materialize_transient_term(v[id]); + if (term.sorted) { + ArenaTermPostingSource source(&pool_, term); + StreamedTermPostings postings { + .term = std::move(output_term), + .retain_positions = term.shape == PostingChainShape::kTaggedPositioned, + .source = &source}; + callback_status = fn(std::move(postings)); + if (callback_status.ok() && !source.exhausted()) { + callback_status = Status::Error( + "spimi arena source: consumer returned before term exhaustion"); + } + } else { + TrackedTermPostings materialized(mem_reporter_); + callback_status = to_postings(std::move(output_term), std::move(term), &materialized); + if (callback_status.ok()) { + SpanTermPostingSource source(materialized.postings.docids, + materialized.postings.freqs, + materialized.postings.positions_flat); + StreamedTermPostings postings { + .term = std::move(materialized.postings.term), + .retain_positions = materialized.postings.retain_positions, + .source = &source}; + callback_status = fn(std::move(postings)); + if (callback_status.ok() && !source.exhausted()) { + callback_status = Status::Error( + "spimi span source: consumer returned before term exhaustion"); + } + } + } + if (!callback_status.ok()) { + break; + } + } + + pool_.reset(); + std::vector().swap(slots_); + std::vector().swap(free_slots_); + std::vector().swap(slot_of_); + std::vector().swap(touched_ids_); + live_term_count_ = 0; + std::vector().swap(owned_vocab_); + owned_vocab_heap_bytes_ = 0; + common_gram_pair_cache_.reset(); + common_gram_pair_cache_bytes_ = 0; + common_gram_plain_term_cache_.reset(); + common_gram_plain_term_cache_bytes_ = 0; + trim_malloc(); + report_arena_delta(); + return callback_status; +} + +Status SpimiTermBuffer::drain_to_writer(RunWriter* w) { + Status st = Status::OK(); + const std::vector& v = vocab(); + // Spill writes by term-id (no string IO). Iterate touched ids in vocab-string + // order so each run is sorted; the k-way merge re-orders runs by the same key. + for (uint32_t id : sorted_ids()) { + const uint32_t enc = slot_of_[id]; + DCHECK_NE(enc, 0U); + Term term = slots_[enc - 1]; + release_term(id); + if (st.ok()) { + TrackedTermPostings materialized(mem_reporter_); + st = to_postings(v[id], std::move(term), &materialized); + if (st.ok()) { + st = w->write_term(id, materialized.postings); + } + } + } + touched_ids_.clear(); + pool_.reset(); // all chains decoded into the run; free the arena for the refill + // The spill returns the arena to 0; slot_of_ keeps its capacity (survives + // the spill). Report the arena-drop negative now so the gate-2 spill is balanced + // immediately, not deferred to the next token. + report_arena_delta(); + return st; +} + +Status SpimiTermBuffer::compact_runs() { + if (run_paths_.size() < 2) { + return Status::OK(); + } + // The compaction heap can encounter any id held by an earlier run, so it + // requires a complete rank for the current vocabulary. New append-only ids + // can shift existing lexicographic ranks, hence the explicit refresh here. + ensure_string_rank(); + const std::string out_path = make_run_path(resolve_temp_dir()); + Status s = + writer::compact_runs(run_paths_, string_rank_, has_positions_, out_path, mem_reporter_); + if (!s.ok()) { + std::remove(out_path.c_str()); // drop the partial output; inputs intact + return s; + } + // The compacted run REPLACES its inputs at the FRONT of the run order: + // it holds exactly runs [0..n) merged in run order, and any later run only + // covers strictly-later docids, so per-term run-order concatenation (the + // k-way merge invariant) is preserved. + for (const std::string& p : run_paths_) { + std::remove(p.c_str()); + } + run_paths_.clear(); + run_paths_.push_back(out_path); + g_run_compactions.fetch_add(1, std::memory_order_relaxed); + return Status::OK(); +} + +Status SpimiTermBuffer::spill_to_run() { + // G09 run-file cap: a buffer must never accumulate unbounded run files -- + // the final k-way merge (re)opens ALL of them simultaneously and holds + // the fds for its whole duration, so unbounded runs across ~100 + // concurrent writers exhausted the BE nofile rlimit ('Too many open + // files' at run reopen). At the cap, merge-compact the existing runs into + // one before cutting the new run: the merge fan-in (and its fd count) is + // bounded by cap + 1 per buffer. + if (max_run_files_ != 0 && run_paths_.size() >= max_run_files_) { + RETURN_IF_ERROR(compact_runs()); + } + const std::string dir = resolve_temp_dir(); + // Best-effort space pre-check: fail with a clear, early error rather than a + // mid-write IoError that leaves a half-written run. Best-effort only (TOCTOU; on + // tmpfs this reports RAM). The ARENA -- not full resident_bytes(), which since + // G08 also charges vocabulary structures a run never contains -- is what the + // run re-encodes, and its block slack makes it a conservative over-estimate of + // the run's on-disk size. + const uint64_t arena = pool_.arena_bytes(); + const uint64_t avail = temp_dir_available_bytes(dir); + if (avail < arena) { + return Status::Error( + "spimi: insufficient temp space in '" + dir + "' to spill ~" + + std::to_string(arena) + " B (~" + std::to_string(avail) + + " B free); set SNII_TEMP_DIR/TMPDIR to a larger disk"); + } + const std::string path = make_run_path(dir); + RunWriter w(mem_reporter_); + RETURN_IF_ERROR(w.open(path)); + run_paths_.push_back(path); // tracked for cleanup even if a later step fails + RETURN_IF_ERROR(drain_to_writer(&w)); + // The drain emptied touched_ids_ and released every live slot while retaining + // capacity for the next fill. + return w.close(); +} + +Status SpimiTermBuffer::prepare_run_merge(TermKeyMaterializer* materializer) { + if (!touched_ids_.empty()) { + Status status = spill_to_run(); + if (!status.ok() && spill_status_.ok()) { + spill_status_ = status; + } + } + if (!spill_status_.ok()) { + return spill_status_; + } + + std::vector().swap(slots_); + std::vector().swap(free_slots_); + std::vector().swap(slot_of_); + std::vector().swap(touched_ids_); + common_gram_pair_cache_.reset(); + common_gram_pair_cache_bytes_ = 0; + common_gram_plain_term_cache_.reset(); + common_gram_plain_term_cache_bytes_ = 0; + common_gram_pair_intern_ = decltype(common_gram_pair_intern_)(); + std::vector().swap(common_word_classification_); + trim_malloc(); + report_arena_delta(); + + ensure_string_rank(); + report_arena_delta(); + intern_ = decltype(intern_)(0, OwnedVocabHash {.vocab = &owned_vocab_}, + OwnedVocabEq {&owned_vocab_}); + report_arena_delta(); + if (common_gram_pair_keys_) { + *materializer = [this](std::string_view term) { return materialize_transient_term(term); }; + } + return Status::OK(); +} + +void SpimiTermBuffer::finish_run_merge() { + std::vector().swap(owned_vocab_); + owned_vocab_heap_bytes_ = 0; + std::vector().swap(string_rank_); + report_arena_delta(); + trim_malloc(); +} + +Status SpimiTermBuffer::merge_runs_streamed(const StreamedTermConsumer& fn) { + TermKeyMaterializer materializer; + RETURN_IF_ERROR(prepare_run_merge(&materializer)); + Status status = merge_run_sources(run_paths_, vocab(), string_rank_, has_positions_, fn, + std::move(materializer), mem_reporter_); + finish_run_merge(); + return status; +} + +Status SpimiTermBuffer::for_each_term_sorted(const StreamedTermConsumer& fn) { + // Single-drain contract: a second call would re-merge the (still-present) run + // files and re-emit every term, or emit nothing in the in-memory path. Return + // an error and emit NOTHING rather than produce a wrong second stream. + if (drained_) { + return Status::Error( + "spimi: already drained (single-drain contract)"); + } + drained_ = true; + if (run_paths_.empty() && spill_status_.ok()) { + return drain_sorted_streamed(fn); + } + return merge_runs_streamed(fn); +} + +std::vector SpimiTermBuffer::finalize_sorted() { + std::vector out; + out.reserve(touched_ids_.size()); + Status status = for_each_term_sorted([&out](StreamedTermPostings&& streamed) { + TermPostings materialized; + materialized.term = std::move(streamed.term); + materialized.retain_positions = streamed.retain_positions; + TermPostingBuffer buffer(nullptr); + bool exhausted = false; + while (!exhausted) { + buffer.clear_reuse(); + RETURN_IF_ERROR( + streamed.source->fill(format::kAdaptiveWindowDocs, &buffer, &exhausted)); + materialized.docids.insert(materialized.docids.end(), buffer.docids().begin(), + buffer.docids().end()); + materialized.freqs.insert(materialized.freqs.end(), buffer.freqs().begin(), + buffer.freqs().end()); + materialized.positions_flat.insert(materialized.positions_flat.end(), + buffer.positions_flat().begin(), + buffer.positions_flat().end()); + } + out.push_back(std::move(materialized)); + return Status::OK(); + }); + if (!status.ok() && spill_status_.ok()) { + spill_status_ = status; + std::vector().swap(out); + } + return out; +} + +void SpimiTermBuffer::cleanup_runs() { + for (const std::string& p : run_paths_) { + std::remove(p.c_str()); + } + run_paths_.clear(); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/spimi_term_buffer.h b/be/src/storage/index/snii/writer/spimi_term_buffer.h new file mode 100644 index 00000000000000..b5c637e161b581 --- /dev/null +++ b/be/src/storage/index/snii/writer/spimi_term_buffer.h @@ -0,0 +1,730 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/compact_posting_pool.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/term_posting_source.h" + +namespace doris::segment_v2::inverted_index { +class CommonWordSet; +} + +namespace doris::snii::writer { + +using StreamedTermConsumer = std::function; + +// G11: compiled-in marker for the per-token prefetch candidate (the locality +// bench keys its in-process A/B test off this). +#define SNII_G11_PREFETCH 1 + +class GlobalMemoryLimiter; // G09 process-wide build-RAM registry (see below) + +struct PlainTermId { + uint32_t value = 0; +}; + +struct ClassifiedPlainTerm { + PlainTermId id; + bool is_common = false; +}; + +// One term's posting list: docids ascending, with parallel freqs and (when +// positions are enabled) a single FLAT positions buffer. +// +// positions_flat holds every position for the term in document order, partitioned +// by freqs: doc i owns the next freqs[i] entries. This is the SAME layout the +// accumulator stores natively, so no per-doc vector-of-vectors is ever built on +// the build/merge hot path (that vector-of-vectors was the dominant peak-RSS +// driver for high-df terms). doc_positions(i) returns a non-owning span view of +// doc i's positions for consumers that want per-doc access (e.g. the prx window +// builder, tests). positions_flat is empty when positions are disabled. +struct TermPostings { + std::string term; + std::vector docids; // absolute docids + std::vector freqs; + std::vector positions_flat; // empty when positions disabled + // Per-term posting shape. A positioned logical index may mix ordinary terms + // with docs-only accelerator terms; the latter carry one posting per doc and + // deliberately omit frequencies/positions from the final index. + bool retain_positions = true; + + size_t document_count() const { return docids.size(); } + + // Byte offset of doc i's first position within positions_flat (prefix sum of + // freqs). O(i) -- callers iterating all docs should track a running offset. + size_t pos_offset(size_t doc_index) const { + size_t off = 0; + for (size_t i = 0; i < doc_index; ++i) { + off += freqs[i]; + } + return off; + } + // Non-owning view of doc i's positions (length freqs[i]) into positions_flat. + std::span doc_positions(size_t doc_index) const { + const size_t off = pos_offset(doc_index); + return {positions_flat.data() + off, freqs[doc_index]}; + } + + // Rebuilds the per-doc position lists (for callers/tests wanting per-doc access) + // from positions_flat partitioned by freqs. O(total positions); allocates. + std::vector> positions_per_doc() const { + std::vector> out(freqs.size()); + size_t off = 0; + for (size_t i = 0; i < freqs.size(); ++i) { + out[i].assign(positions_flat.begin() + off, positions_flat.begin() + off + freqs[i]); + off += freqs[i]; + } + return out; + } + + // Sets the flat positions from per-doc lists (convenience for tests / callers + // that produce per-doc positions). Does NOT touch freqs; the caller is expected + // to keep freqs[i] == per_doc[i].size() consistent (the writer validates this). + void set_positions_per_doc(const std::vector>& per_doc) { + positions_flat.clear(); + for (const auto& d : per_doc) { + positions_flat.insert(positions_flat.end(), d.begin(), d.end()); + } + } +}; + +// In-memory SPIMI (Single-Pass In-Memory Indexing) accumulator for one logical +// index. Records term occurrences and produces lexicographically sorted terms +// with ascending-docid posting lists. +// +// TERM-ID ACCUMULATION (no per-token string work): tokens are accumulated by an +// INTEGER term-id, not by hashing/constructing a std::string per token. The +// caller supplies a VOCABULARY mapping term-id -> term string; the buffer keeps +// a DENSE std::vector indexed by term-id, so the hot add_token path is a +// vector index + a couple of pushes -- no hashing, no allocation per token. The +// vocabulary is resolved to strings only once per distinct term at finalize. +// +// Two construction modes: +// * BORROWED vocab (the fast path): pass a non-null `vocab` that the caller +// owns and keeps alive; add_token(term_id, ...) indexes straight into it. +// * OWNED vocab (compatibility): pass a null `vocab`; the string-keyed +// add_token(string_view, ...) interns each new term into an internal owned +// vocabulary (assigning ids in first-seen order) and forwards to the id +// path. Existing callers that feed strings keep working unchanged. +// +// SPILL / K-WAY MERGE (out-of-core, bounds input RAM): when a non-zero +// spill_threshold_bytes is set, the REAL resident accumulator size (see +// resident_bytes(): the posting arena PLUS every live vocab / slot / rank +// structure, G08) is compared against the threshold as tokens arrive. Once it +// crosses the threshold and enough reclaimable posting arena has accumulated, +// the buffer SORTS its current terms, +// writes a self-describing sorted RUN to a temp file, and CLEARS memory. Each +// run record is keyed by the TERM-ID (varint); the k-way merge orders runs by +// the id's VOCAB STRING so the merged stream stays lexicographic. Because +// tokens arrive in globally ascending docid order, a term that reappears in a +// later run only covers strictly-later docids, so concatenating its postings in +// run order during the final merge keeps docids ascending. for_each_term_sorted +// flushes the residual buffer as a final run, then k-way merges all runs +// materializing only ONE merged term at a time -> peak memory stays bounded by +// the threshold (plus the widest single term), NOT by total postings. With the +// default threshold 0 (unlimited) the path is exactly the in-memory behavior. +// +// Internal representation is a COMPACT TAGGED VARINT byte stream per term, held in +// a shared SEGMENTED ARENA (CompactPostingPool), NOT per-term uint32 vectors. Each +// term owns ONE arena chain holding a stream of per-TOKEN entries in arrival +// order: positioned and ordinary docs-only tokens contribute +// varint((pos << 1) | new_doc_bit); when new_doc_bit is set, a +// zigzag-varint(docid - prev_docid) immediately follows. Statless CommonGrams are +// deduplicated per document and store only that document delta, omitting the +// constant new_doc tag. Frequencies are otherwise recovered as the count of +// consecutive same-doc tokens. This drops +// the entire freq stream and the second (positions) chain versus a freq/prox split, +// so the payload is ~3.4x smaller than raw uint32 docids/freqs/positions, and the +// shared arena removes per-vector doubling slack and per-term vector headers. Each +// positioned and ordinary docs-only tokens append straight into the chain. +// Stateless CommonGrams keep a singleton doc id inline and backfill it only when a +// second document arrives. The other live per-term state is the current doc id (to +// detect a doc change) and the delta base. +// The production writer drains each chain through a bounded TermPostingSource. +// to_postings() remains only for explicit materialization in run maintenance and +// test/finalize helpers. positions_flat stays empty (and pos is tagged as 0) when +// positions are disabled; freq still counts. +// +// Duplicate vocab strings: the vocab is assumed to map each id to a DISTINCT +// string (a dense vocabulary). If two ids share a string they sort adjacently +// but are emitted as two separate terms; callers must not rely on coalescing. +class SpimiTermBuffer { +public: + // BORROWED-vocab constructor: `vocab` maps term-id -> term string and is + // borrowed (NOT owned) -- the caller must keep it alive for the buffer's + // lifetime. add_token(term_id, ...) accumulates by id with no string work. + // spill_threshold_bytes is the gate-2 internal buffer cap (e.g. 512 MiB), + // sourced from config; == 0 means unlimited (pure in-memory, default). A + // positive value is a soft spill threshold for the REAL resident accumulator + // size (resident_bytes(): arena + every live vocab/slot/rank structure, G08), + // triggering a spill once enough reclaimable arena has accumulated -- NOT a + // hard cap on persistent vocabulary memory or the old per-token estimate. + // `reporter` is the OPTIONAL writer-level build-RAM reporter (null off-Doris / + // unit tests). When non-null, the accumulator reports its REAL resident-byte + // deltas -- resident_bytes() diffs -- positive on grow, negative on every + // reset/free, exactly once. NEVER reports live_bytes_ (a gated estimate that + // feeds only the spill threshold). + explicit SpimiTermBuffer(const std::vector* vocab, bool has_positions, + size_t spill_threshold_bytes = 0, MemoryReporter* reporter = nullptr); + + // OWNED-vocab (compatibility) constructor: no external vocab. The string-keyed + // add_token interns terms into an internal vocabulary on first occurrence. + explicit SpimiTermBuffer(bool has_positions, size_t spill_threshold_bytes = 0, + MemoryReporter* reporter = nullptr); + + ~SpimiTermBuffer(); + + SpimiTermBuffer(const SpimiTermBuffer&) = delete; + SpimiTermBuffer& operator=(const SpimiTermBuffer&) = delete; + + // Records one token by TERM-ID: term `term_id` occurs in `docid` at `pos`. + // `term_id` must be in [0, vocab_size). An out-of-range id latches an + // InvalidArgument into status() and is ignored. For a given term, docids are + // expected to arrive in non-decreasing order, and positions within a docid in + // ascending order; out-of-order docids (INCLUDING a REVISITED docid -- the same + // docid appearing again after a different one) are tolerated and reordered at + // finalize: sort_by_docid stably sorts by docid and COALESCES same-docid groups + // (summing freqs, concatenating positions in document order), so the emitted + // postings have exactly ONE strictly-ascending entry per docid -- matching the + // k-way merge path and the writer's strictly-ascending precondition. + void add_token(uint32_t term_id, uint32_t docid, uint32_t pos); + void add_token(uint32_t term_id, uint32_t docid, uint32_t pos, bool retain_positions); + + // Compatibility overload: records one token by TERM STRING. Valid ONLY on an + // OWNED-vocab buffer before enable_common_gram_pair_keys(); interns `term` into + // the internal vocabulary on first occurrence, then forwards by id. Pair-key + // mode must use the typed plain/gram APIs below so a physical gram and its + // transient pair key cannot become two ids for the same logical term. Called on + // a BORROWED-vocab buffer it is REJECTED (latches InvalidArgument, token ignored) + // -- interning would grow the owned vocab out of step with the borrowed one and + // corrupt the build. Interning probes a heterogeneous (string_view-keyed) set, + // so a repeat token for an already-seen term allocates NOTHING; a std::string is + // materialized only on a term's FIRST occurrence (stored once in owned_vocab_). + // The id overload remains the hot path (no hashing at all); prefer that and + // reserve this for tests / legacy string-fed callers. + void add_token(std::string_view term, uint32_t docid, uint32_t pos); + void add_token(std::string_view term, uint32_t docid, uint32_t pos, bool retain_positions); + + // SNII CommonGrams fast path. Plain terms are interned once and returned as + // stable ids; each gram occurrence hashes a fixed 10-byte pair of those ids + // instead of constructing and hashing the variable-length physical gram key. + PlainTermId intern_plain_term(std::string_view physical_plain_term); + // Production CommonGrams path. A physical-key hit proves the injectively mapped + // logical term was validated previously; a miss validates exactly once before + // materializing the vocabulary entry. + PlainTermId intern_plain_term(std::string_view physical_plain_term, + std::string_view logical_plain_term); + ClassifiedPlainTerm intern_classified_plain_term( + std::string_view physical_plain_term, std::string_view logical_plain_term, + const segment_v2::inverted_index::CommonWordSet& common_words); + void add_plain_token(PlainTermId term_id, uint32_t docid, uint32_t pos); + void add_common_gram(PlainTermId left, PlainTermId right, uint32_t docid, uint32_t pos, + bool retain_positions); + void add_common_gram_and_plain(PlainTermId left, PlainTermId right, uint32_t docid, + uint32_t gram_pos, uint32_t plain_pos, + bool retain_gram_positions); + void enable_common_gram_pair_keys(); + + // G09: joins the PROCESS-WIDE build-RAM registry. Registers this buffer's + // current resident bytes with `limiter` and forwards every subsequent + // (debounced, see report_arena_delta) resident total to it; the destructor + // un-registers. When the registered sum across ALL of the process's live + // buffers exceeds the limiter's budget, the limiter may set this buffer's + // ADVISORY spill-request flag from ANOTHER thread; the flag is observed -- + // and the forced spill run ON THIS BUFFER'S OWN THREAD -- by the next + // add_token's maybe_spill_after_token (see there for the honor rule). + // Call at most once, right after construction (extra calls are ignored); + // `limiter` must outlive this buffer. Null / never attached = the G08 + // per-writer behavior, byte-identical. + void attach_global_limiter(GlobalMemoryLimiter* limiter); + + // TEST-ONLY: G09 advisory-flag observability -- read the pending flag, and + // plant a request directly (what the limiter does cross-thread) so the + // owner-honors-at-next-token contract is testable without a registry. + bool global_spill_requested_for_test() const { + return global_spill_requested_.load(std::memory_order_relaxed); + } + void request_global_spill_for_test() { + global_spill_requested_.store(true, std::memory_order_relaxed); + } + + // G09 forced-spill floor (config snii_forced_spill_min_arena_bytes): a + // pending process-wide forced-spill request is honored only once the + // reclaimable posting arena holds at least this much (never below one + // arena block, so a run is always writable). A request planted while the + // arena is below the floor is a NO-OP that stays PENDING -- it is NOT + // retried as a spill every token -- and is honored when the arena regrows + // past the floor. Without the floor, an unreachable global budget (the + // persistent vocabulary/slot structures of all writers alone exceeding it) + // re-flagged every buffer on every report and each honored with a single + // 32 KiB arena block: thousands of tiny runs per buffer, EMFILE at the + // k-way merge reopen, failed loads (the conc=16 wikipedia field storm). + static constexpr uint64_t kDefaultForcedSpillMinArenaBytes = 64ULL << 20; // 64 MiB + void set_forced_spill_min_arena_bytes(uint64_t bytes) { forced_spill_min_arena_bytes_ = bytes; } + uint64_t forced_spill_min_arena_bytes() const { return forced_spill_min_arena_bytes_; } + + // G09 run-file cap (config snii_spill_max_run_files_per_buffer): when a + // new spill would grow the accumulated run-file count past this cap, the + // existing runs are first MERGE-COMPACTED into one (a k-way merge of the + // run files back into a single fresh run; term stream byte-identical, the + // old files deleted) so the buffer never holds more than the cap + 1 run + // files. Bounds both the final k-way merge's fan-in and -- decisively -- + // its OPEN FILE DESCRIPTORS: every run of a buffer is (re)opened + // simultaneously and held open for the whole merge, so unbounded run + // counts across ~100 concurrent writers exhausted the BE nofile rlimit + // ('Too many open files' at run reopen). 0 disables the cap. + static constexpr size_t kDefaultMaxRunFilesPerBuffer = 64; + void set_max_run_files(size_t cap) { max_run_files_ = cap; } + size_t max_run_files() const { return max_run_files_; } + + // Number of DISTINCT terms accumulated so far (touched ids still resident). + size_t unique_terms() const; + uint64_t total_tokens() const { return total_tokens_; } + bool has_positions() const { return has_positions_; } + + // OK unless an add_token validation error (out-of-range term-id, wrong vocab + // mode) was latched. for_each_term_sorted now returns its own I/O Status + // directly; callers that use add_token's latch-and-report pattern MUST check + // this after draining to surface input-side validation errors. + [[nodiscard]] Status status() const { return spill_status_; } + + // TEST-ONLY: number of spill run files currently HELD (== 0 in pure + // in-memory mode). Lets tests assert that a gate-2 spill actually fired + // once the REAL resident size crossed the configured cap. NOTE: a G09 + // run-cap merge-compaction (see set_max_run_files) collapses the list to + // ONE file, so the count is not monotonic. Not part of the production API. + size_t run_count_for_test() const { return run_paths_.size(); } + + // TEST-ONLY: the REAL resident accumulator bytes the gate-2 trigger and the + // MemoryReporter see (resident_bytes()). Lets the G08 accounting tests assert + // coverage and monotonicity without widening access to the private + // accounting. Not part of + // the production API. + uint64_t resident_bytes_for_test() const { return resident_bytes(); } + size_t string_rank_capacity_for_test() const { return string_rank_.capacity(); } +#ifdef BE_TEST + static size_t hash_term_bytes_for_test(std::string_view term) { return hash_term_bytes(term); } + static size_t owned_term_key_size_for_test(); + void set_owned_term_hash_mask_for_test(size_t mask); +#endif + + // Materializes all terms sorted lexicographically; each term's docids are + // ascending. Convenience wrapper around for_each_term_sorted that keeps the + // whole result alive at once. Prefer for_each_term_sorted for low peak memory. + // The returned vectors are caller-owned compatibility output and are not + // charged to this buffer's internal MemoryReporter after the callback returns. + // MUST be called at most once: it drains internal state. A SECOND drain (a + // repeat call, or a finalize_sorted after a for_each_term_sorted, or vice versa) + // returns EMPTY and latches an error into status() rather than re-emitting. + std::vector finalize_sorted(); + + // Streams terms to `fn` in lexicographic order. Each source is borrowed for + // the synchronous callback and fills the writer-owned transfer buffer. The + // callback must exhaust the source before returning success. + // MUST be called at most once: it drains internal state. A SECOND drain invokes + // `fn` zero times and returns an Internal error (a re-merge of the still-present + // run files would otherwise re-emit every term). Returns non-OK on spill/merge + // I/O or corruption errors, or if a prior add_token latched a validation error + // into status(). + Status for_each_term_sorted(const StreamedTermConsumer& fn); + +private: + struct CommonGramPairCache; + struct CommonGramPlainTermCache; + + enum class PostingChainShape : uint8_t { + kTaggedPositioned, + kTaggedDocsOnly, + kStatlessDocsOnly, + }; + + // Compact per-term accumulator: ONE tagged-varint arena chain plus a few cursors. + // A statless CommonGram keeps its first distinct doc inline in cur_docid and + // starts a chain only when a second doc arrives. For other posting shapes, a + // sentinel chain head marks an empty term. ntok / ndocs bound the decode loop + // and size reserves. + // Total 28 B per live term. + static constexpr uint32_t kNoChain = 0xFFFFFFFFU; + struct Term { + uint32_t head = kNoChain; // chain read entry point + CompactPostingPool::SliceWriter w; // chain cursor (8 B) + uint32_t ntok = 0; // total tokens (entries) in the chain + uint32_t cur_docid = 0; // most-recent doc id: detects doc change AND + // is the zigzag delta base for the next doc + // Exact count of new-doc groups in the chain (one per new_doc tag). It + // bounds the decode reserves and equals the distinct-doc count while the + // input remains sorted; a later out-of-order coalesce can only shrink it. + uint32_t ndocs = 0; + PostingChainShape shape = PostingChainShape::kTaggedPositioned; + uint8_t level = 0; // current slice level of w (packed here, not in w) + bool started = false; // false until the first token is accumulated + bool sorted = true; // false if a docid arrived out of ascending order + }; + static_assert(sizeof(CompactPostingPool::SliceWriter) == 8, + "SliceWriter must stay 8 bytes to keep Term compact"); + static_assert(sizeof(Term) == 28, "Term must stay compact for high-cardinality imports"); + + struct TrackedTermPostings { + explicit TrackedTermPostings(MemoryReporter* reporter) + : docids_reservation(reporter == nullptr ? MemoryReporter::Reservation() + : reporter->make_reservation()), + freqs_reservation(reporter == nullptr ? MemoryReporter::Reservation() + : reporter->make_reservation()), + positions_reservation(reporter == nullptr ? MemoryReporter::Reservation() + : reporter->make_reservation()) {} + + TrackedTermPostings(const TrackedTermPostings&) = delete; + TrackedTermPostings& operator=(const TrackedTermPostings&) = delete; + + // Reservations precede the posting vectors so physical allocations are + // destroyed before their charges are released. + MemoryReporter::Reservation docids_reservation; + MemoryReporter::Reservation freqs_reservation; + MemoryReporter::Reservation positions_reservation; + TermPostings postings; + }; + + // The active vocabulary (term-id -> string): either the borrowed pointer or, + // in owned mode, &owned_vocab_. Always non-null after construction. + const std::vector& vocab() const { return *vocab_; } + + // Accumulates one already-validated token into the per-id Term and checks the + // spill gate once for that input token. + void accumulate(uint32_t term_id, uint32_t docid, uint32_t pos, bool retain_positions); + void accumulate_without_spill_gate(uint32_t term_id, uint32_t docid, uint32_t pos, + PostingChainShape shape); + void add_common_gram_without_spill_gate(PlainTermId left, PlainTermId right, uint32_t docid, + uint32_t pos, bool retain_positions); + + // Per-token gate-2 tail of accumulate(): reports the token's resident growth, + // then spills when the unified cap / local threshold fires with a worthwhile + // reclaimable arena (the G08 anti-churn floor), when the G09 process-wide + // limiter's advisory request flag is pending (honored here, on the owner's + // own thread; bypasses the G08 floor but requires one allocated arena block + // so a run is writable), or when the arena nears its hard 4 GiB offset + // limit. Every public add path invokes this gate once; the fused CommonGrams + // path invokes it after appending both the gram and its right plain token. + void maybe_spill_after_token(); + + Status to_postings(std::string term, Term&& t, TrackedTermPostings* tracked) const; + class ArenaTermPostingSource; + + // Returns the touched term-ids sorted by their vocab string (lexicographic). + // The first spill builds the full integer string-rank. Later spills reuse it + // while the append-only vocabulary is unchanged. If the vocabulary grew, an + // ordinary spill sorts only this run's touched ids by string; run compaction + // and final merge rebuild the full rank when they actually require it. + std::vector sorted_ids() const; + // Builds string_rank_ (term-id -> lexicographic rank) for the current complete + // vocabulary. Idempotent until the append-only vocabulary grows. + void ensure_string_rank() const; + Status drain_sorted_streamed(const StreamedTermConsumer& fn); + // Spills the current buffer to a fresh sorted run file and clears memory. + Status spill_to_run(); + // G09 run-file cap enforcement (see set_max_run_files): merge-compacts the + // current run files into ONE fresh run (same term stream, ids ordered by + // the current string rank), deletes the old + // files and replaces run_paths_ with the compacted one. Called by + // spill_to_run before opening a new run once the cap is reached. + Status compact_runs(); + // Writes all current terms (sorted) to an already-open RunWriter, draining. + Status drain_to_writer(class RunWriter* w); + // REAL resident accumulator bytes -- the single source of truth for the gate-2 + // spill trigger and every MemoryReporter delta. G08: sums EVERY live input-side + // structure -- the posting arena (docs+prx payload) + // plus the vocab-sized slot index, the Term slot pool + free/touched lists, the + // owned vocabulary (headers by capacity + string heap payloads via + // owned_vocab_heap_bytes_) and its intern set, plus the cached string ranks. + // Capacity, not size, throughout: the reserved tail is resident RSS and + // survives spills. + uint64_t resident_bytes() const; + // Reports the signed change in REAL resident bytes (resident_bytes()) to + // mem_reporter_ since the previous call, then caches the new total. + // Single-source diff: every grow/reset/free emits EXACTLY ONE delta + // (self-balancing -> impossible to double-count or miss a negative). No-op when + // mem_reporter_ is null. + void report_arena_delta(); + Status merge_runs_streamed(const StreamedTermConsumer& fn); + Status prepare_run_merge(std::function* materializer); + void finish_run_merge(); + // Deletes every temp run file; called from the destructor (RAII cleanup). + void cleanup_runs(); + // Frees a drained term's accumulator (id leaves the touched set). + void release_term(uint32_t term_id); + + // Stores a first-seen owned-vocabulary term under a stable id. + uint32_t append_owned_vocab_term(std::string&& term_str); + uint32_t intern_owned_term(std::string&& term_str, size_t term_hash); + uint32_t find_or_intern_owned_term(std::string_view term); + uint32_t find_or_intern_common_gram_pair(PlainTermId left, PlainTermId right, uint64_t pair); + uint32_t find_interned_plain_term(std::string_view term, size_t term_hash); + void remember_plain_term(size_t term_hash, uint32_t term_id); + bool transient_term_less(uint32_t left_id, uint32_t right_id) const; + std::string materialize_transient_term(std::string_view term) const; + + const std::vector* vocab_; // active vocab (borrowed or &owned_) + std::vector owned_vocab_; // owned mode: interned term strings + + enum class CommonWordClassification : uint8_t { + kUnknown, + kNotCommon, + kCommon, + }; + // Stable semantic classification keyed by owned term id. Pair ids remain + // kUnknown; physical plain ids are classified once from their logical bytes. + std::vector common_word_classification_; + + // G08: running sum of the owned vocab strings' HEAP payloads (0 for SSO + // strings -- their bytes live inside the headers owned_vocab_.capacity() + // already charges; capacity+1 for heap strings). Maintained incrementally by + // intern_owned_term so resident_bytes() stays O(1); terminal drains zero it + // when owned_vocab_ is released. + uint64_t owned_vocab_heap_bytes_ = 0; + + // G08: fixed per-entry estimate for one intern-set entry. Sized for the + // pre-G10 NODE-based set (16 B next-ptr+id node, its malloc chunk rounding, + // and an amortized bucket-array share) and deliberately UNCHANGED by the G10 + // swap to the flat set: resident_bytes() feeds the gate-2 spill trigger, so + // keeping the constant keeps the resident-byte sequence -- and therefore + // every spill point and the drained output -- bit-identical to the prior + // build. It still OVER-approximates the 4-byte flat key plus control bytes + // and load-factor slack, which can only fire the gate earlier, never overshoot. + // Deterministic so the accounting tests can reason about it, and ZERO for an + // empty set so an untouched (borrowed-mode) buffer charges nothing for it. + static constexpr uint64_t kInternEntryEstimateBytes = 48; + + // The table slot stores only the stable vocabulary id. String probes hash + // their bytes once; stored ids rehash through the sole owned vocabulary. + // Equality always resolves term identity from the complete bytes, so hash + // collisions are harmless and only add comparisons. + static size_t hash_term_bytes(std::string_view s) noexcept { + return std::hash {}(s); + } + + struct OwnedVocabHash { + using is_transparent = void; + const std::vector* vocab = nullptr; + size_t hash_mask = std::numeric_limits::max(); + size_t operator()(std::string_view term) const noexcept { + return hash_term_bytes(term) & hash_mask; + } + size_t operator()(uint32_t term_id) const noexcept { + return (*this)(std::string_view((*vocab)[term_id])); + } + }; + struct OwnedVocabEq { + using is_transparent = void; + const std::vector* vocab = nullptr; + bool operator()(uint32_t left, uint32_t right) const noexcept { return left == right; } + bool operator()(uint32_t stored, std::string_view probe) const noexcept; + bool operator()(std::string_view probe, uint32_t stored) const noexcept; + }; + // One flat table is the sole ordinary-term admission path. Heterogeneous + // probes avoid temporary strings; prepared insertion materializes only a miss. + // A failed table insertion rolls back the preceding vocabulary append, and no + // iterator survives a mutation. + phmap::flat_hash_set intern_; + + // CommonGram pair ids are already a canonical, collision-free key. Keep them + // out of the string-content intern table: an L0 cache miss probes this native + // map and materializes the 10-byte transient vocabulary key only for a new + // pair. The map survives ordinary spills because the persistent vocabulary + // and term ids do; terminal drains release it before output reservations. + phmap::flat_hash_map common_gram_pair_intern_; + + bool has_positions_; + bool common_gram_pair_keys_ = false; + std::unique_ptr common_gram_pair_cache_; + uint64_t common_gram_pair_cache_bytes_ = 0; + std::unique_ptr common_gram_plain_term_cache_; + uint64_t common_gram_plain_term_cache_bytes_ = 0; + size_t spill_threshold_bytes_; // 0 => unlimited (no spilling) + uint64_t total_tokens_ = 0; + + // POOLED accumulators (replaces a dense vocab-sized std::vector, which + // cost ~80 B per vocab id even for the ~empty majority -- the single largest + // input-phase memory line). slot_of_ is the only vocab-sized array: a 4 B index + // per id (0 == no live Term; otherwise slot index + 1). slots_ holds ONE Term + // per CURRENTLY-LIVE id, so its size tracks the live touched count, not the + // vocabulary. On first touch an id claims a slot (reusing a freed one from + // free_slots_ when available, else appending). release_term frees the slot back + // to the pool and clears slot_of_[id]. touched_ids_ lists every live id so + // finalize/spill iterate touched ids without scanning the whole vocabulary. + // present_[id] is now (slot_of_[id] != 0). The hot add path is still a vector + // index + a couple of pushes: no hashing, no per-token allocation. + std::vector slot_of_; // vocab-sized: id -> slot index + 1 (0=empty) + std::vector slots_; // live Term pool (size ~ live touched count) + std::vector free_slots_; // recycled slot indices (drained terms) + std::vector touched_ids_; + size_t live_term_count_ = 0; // present (non-drained) terms; == unique_terms() + + // Shared arena backing every live term's DOC and POS varint byte chains. Holds + // the bulk of the accumulator's memory in a few large blocks (no per-term vector + // headers, no per-vector doubling slack) -- the compact-RSS win. + CompactPostingPool pool_; + + // Optional writer-level build-RAM reporter (null off-Doris / unit tests) and the + // last resident-byte total it was told about. report_arena_delta() diffs the live + // total (arena_bytes() + slot_of_.capacity()*4) against reported_resident_. + MemoryReporter* mem_reporter_ = nullptr; + int64_t reported_resident_ = 0; + + // ---- G09 process-wide limiter hookup (null / false = feature off) -------- + // The registry this buffer joined via attach_global_limiter (borrowed; must + // outlive the buffer), and the ADVISORY forced-spill request flag the + // limiter sets from other threads (only ever under the registry mutex; the + // owner reads it relaxed on its own thread each token). The flag pointer + // doubles as the buffer's registry identity. + GlobalMemoryLimiter* global_limiter_ = nullptr; + std::atomic global_spill_requested_ {false}; + // G09 forced-spill floor / run-file cap (see the public setters above). + uint64_t forced_spill_min_arena_bytes_ = kDefaultForcedSpillMinArenaBytes; + size_t max_run_files_ = kDefaultMaxRunFilesPerBuffer; + + // Returns the live Term for `term_id`, claiming a pool slot on first touch. + Term& term_slot(uint32_t term_id, bool* new_term); + + // Appends one varint to a term's chain, lazily starting the chain on first use + // (so an untouched term costs no arena bytes). + void put_varint(Term* t, uint64_t v); + + std::vector run_paths_; // spilled run temp files (deleted in dtor) + Status spill_status_; // first spill / range error, at finalize + bool drained_ = false; // set once finalize_sorted/for_each_term_sorted has run; + // a second drain would (spilled path) re-merge the run + // files and re-emit every term, or (in-memory path) emit + // nothing -- both wrong. Guard against the double-drain. + + // Lazily-built vocab-sized map: term-id -> its lexicographic rank among all + // vocab strings. `size() == vocab().size()` means the rank is current; a + // smaller non-zero size is a stale rank retained after vocabulary growth. + // Its capacity is still advanced on stale ordinary spills so resident-byte + // accounting and spill-trigger timing keep the prior full-rank charge. + mutable std::vector string_rank_; +}; + +// TEST-ONLY observability seam (mirrors the reader-side decode-counter pattern). +// Counts how many times a vocabulary string is MATERIALIZED into owned_vocab_ during +// owned-mode interning. With single-store interning this is bumped EXACTLY ONCE per +// DISTINCT term (the owned_vocab_.emplace_back) and NEVER per token -- so feeding the +// same term M times still materializes it once, and the per-token temporary probe +// string is gone entirely. Writer tests use it for deterministic allocation +// assertions (count == distinct terms). Process-global; reset between tests. Not part +// of the production API. +namespace testing { +// G11 bench seam (honored under BE_TEST only): disables the add-path +// prefetch hints so the locality bench can A/B them within ONE process. +// Production builds prefetch unconditionally. +void set_bench_disable_g11_prefetch(bool disabled); + +uint64_t vocab_string_materialization_count(); +void reset_vocab_string_materialization_count(); + +// G09 process-wide limiter seam: spills that observed -- and cleared -- a +// PENDING global forced-spill request at the moment they fired (whether or not +// the per-writer gate would also have spilled that token; the request was +// consumed either way). Incremented under BE_TEST only because the check sits +// on the per-token path of every +// concurrent writer). Deterministic on the single-threaded build path; reset +// between tests. Not part of the production API. +uint64_t global_forced_spills(); +void reset_global_forced_spills(); + +// G09 run-file cap seam: merge-compactions of a buffer's accumulated spill +// runs (each collapses the whole run list into one file). Always-on relaxed +// atomic (a compaction is rare -- at most once per cap-many spills -- so +// contention is a non-issue, unlike the per-token seams above). Deterministic +// on the single-threaded build path; reset between tests. Not part of the +// production API. +uint64_t run_compactions(); +void reset_run_compactions(); + +// Number of complete-vocabulary lexicographic rank rebuilds. Ordinary spills +// with a stale rank must not increment it; run compaction and final merge may. +uint64_t string_rank_rebuilds(); +void reset_string_rank_rebuilds(); + +// Complete touched vocabularies invert the dense rank in O(N); partial runs +// retain comparison sorting. Both counters compile out of production. +uint64_t dense_rank_inversions(); +uint64_t rank_comparison_sorts(); +void reset_rank_ordering_counts(); + +// CommonGrams pair-key terminal-ordering seam. Both counters compile out of +// production: tests use them to prove terminal sorting/materialization takes the +// trusted fixed-key path instead of re-running generic key validation. +uint64_t common_gram_pair_unchecked_decode_count(); +uint64_t common_gram_trusted_plain_decode_count(); +void reset_common_gram_pair_fast_path_counts(); + +// CommonGrams pair direct-cache seam. The counters compile out of production; +// tests use them to prove repeated pairs bypass key encoding and the intern table. +// Docs-only pairs additionally suppress same-document repeats, while positioned +// pairs reuse the cached term id and still accumulate every position. +uint64_t common_gram_pair_cache_probes(); +uint64_t common_gram_pair_cache_pair_hits(); +uint64_t common_gram_pair_cache_same_doc_hits(); +void reset_common_gram_pair_cache_counts(); + +// Native CommonGram pair-interner seam. Normal production pair ingestion must +// never route a transient pair key through the generic string-content table. +uint64_t common_gram_native_pair_probes(); +uint64_t common_gram_native_pair_hits(); +uint64_t common_gram_native_pair_inserts(); +void reset_common_gram_native_pair_intern_counts(); + +uint64_t common_gram_logical_validation_count(); +void reset_common_gram_logical_validation_count(); + +// CommonGrams plain-term hot-cache seam: total cache probes, cache hits, and +// fallbacks that reached the global intern table. +uint64_t common_gram_plain_cache_probes(); +uint64_t common_gram_plain_cache_hits(); +uint64_t common_gram_plain_intern_table_probes(); +void reset_common_gram_plain_cache_counts(); + +// Counts equality checks that must dereference owned vocabulary bytes after the +// inline length and prefix checks. Short terms must never increment this counter. +uint64_t owned_term_full_byte_comparison_count(); +void reset_owned_term_full_byte_comparison_count(); +void fail_next_owned_term_reserve(); +void fail_next_owned_term_emplace(); + +uint64_t spill_gate_check_count(); +void reset_spill_gate_check_count(); + +// Counts compact-chain varint decodes during arena source consumption. Tests use +// this to prevent positioned sources from replaying the token chain. +uint64_t compact_chain_varint_decode_count(); +void reset_compact_chain_varint_decode_count(); +} // namespace testing + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/temp_dir.cpp b/be/src/storage/index/snii/writer/temp_dir.cpp new file mode 100644 index 00000000000000..23da77c6a2ecff --- /dev/null +++ b/be/src/storage/index/snii/writer/temp_dir.cpp @@ -0,0 +1,41 @@ +// 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 "storage/index/snii/writer/temp_dir.h" + +#include +#include +#include + +#include "runtime/exec_env.h" +#include "storage/index/index_writer.h" // segment_v2::TmpFileDirs (full definition) + +namespace doris::snii::writer { + +std::string resolve_temp_dir() { + // Use Doris's configured spill/scratch dirs (the same source the inverted-index + // writer uses; see index_file_writer.cpp). SNII spills/section temp files live in + // a dedicated "snii" subdirectory so they do not crowd the tmp root alongside + // every other component's files. + auto dir = ExecEnv::GetInstance()->get_tmp_file_dirs()->get_tmp_file_dir(); + dir /= "snii"; + std::error_code ec; + std::filesystem::create_directories(dir, ec); + return dir.native(); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/temp_dir.h b/be/src/storage/index/snii/writer/temp_dir.h new file mode 100644 index 00000000000000..1b249087cbd213 --- /dev/null +++ b/be/src/storage/index/snii/writer/temp_dir.h @@ -0,0 +1,47 @@ +// 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. + +#pragma once + +#include + +#include +#include + +namespace doris::snii::writer { + +// Scratch directory for spill runs and section temp files. Uses Doris's configured +// tmp_file_dirs (ExecEnv::get_tmp_file_dirs) in production; when those are not +// initialized (unit tests / standalone) it falls back to SNII_TEMP_DIR -> TMPDIR -> +// /tmp. Defined in temp_dir.cpp to keep exec_env.h out of this header. +// +// The fallback (SNII_TEMP_DIR / TMPDIR) should point at a REAL disk (SSD/NVMe): +// /tmp is often tmpfs (RAM-backed), where spilling does NOT reduce RSS. +std::string resolve_temp_dir(); + +// Best-effort free bytes on the filesystem backing `dir`. Returns UINT64_MAX when +// statvfs fails, so a caller's space pre-check never false-positives on an +// unstattable path. CAVEATS: this is best-effort only -- it is subject to TOCTOU +// (free space can drop before/while the write runs), and on tmpfs it reports +// RAM-backed space (use the temp-dir config to avoid tmpfs in the first place). +inline uint64_t temp_dir_available_bytes(const std::string& dir) { + struct statvfs vfs; + if (::statvfs(dir.c_str(), &vfs) != 0) return UINT64_MAX; + return static_cast(vfs.f_bavail) * static_cast(vfs.f_frsize); +} + +} // namespace doris::snii::writer diff --git a/be/src/storage/index/snii/writer/term_posting_source.h b/be/src/storage/index/snii/writer/term_posting_source.h new file mode 100644 index 00000000000000..9de6b5405fee05 --- /dev/null +++ b/be/src/storage/index/snii/writer/term_posting_source.h @@ -0,0 +1,392 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/memory_reporter.h" + +namespace doris::snii::writer { + +struct MutableTermPostingSpan { + std::span docids; + std::span freqs; + std::span positions_flat; +}; + +// Reusable, reservation-backed transfer storage for one source fill. A source +// may append multiple runs during one fill, but every run must agree on whether +// transient frequency statistics are present. clear_reuse() preserves capacity +// and its memory charge. +class TermPostingBuffer { +public: + explicit TermPostingBuffer(MemoryReporter* memory_reporter) + : memory_reporter_(memory_reporter), + capacity_reservation_(memory_reporter == nullptr + ? MemoryReporter::Reservation() + : memory_reporter->make_reservation()) {} + + TermPostingBuffer(const TermPostingBuffer&) = delete; + TermPostingBuffer& operator=(const TermPostingBuffer&) = delete; + TermPostingBuffer(TermPostingBuffer&&) = delete; + TermPostingBuffer& operator=(TermPostingBuffer&&) = delete; + + size_t document_count() const { return docids_.size(); } + bool empty() const { return docids_.empty(); } + + void clear_reuse() { + docids_.clear(); + freqs_.clear(); + positions_flat_.clear(); + has_freqs_.reset(); + } + + void clear_reuse_and_release_excess(size_t max_retained_capacity) { + clear_reuse(); + bool released = false; + if (docids_.capacity() > max_retained_capacity) { + std::vector().swap(docids_); + released = true; + } + if (freqs_.capacity() > max_retained_capacity) { + std::vector().swap(freqs_); + released = true; + } + if (positions_flat_.capacity() > max_retained_capacity) { + std::vector().swap(positions_flat_); + released = true; + } + if (released && memory_reporter_ != nullptr) { + uint64_t retained_bytes = 0; + DORIS_CHECK(capacity_bytes(docids_.capacity(), freqs_.capacity(), + positions_flat_.capacity(), &retained_bytes) + .ok()); + DORIS_CHECK(capacity_reservation_.set_bytes(retained_bytes).ok()); + } + } + + Status append(std::span docids, std::span freqs, + std::span positions_flat) { + const bool has_freqs = !freqs.empty(); + if (has_freqs && freqs.size() != docids.size()) { + return Status::Error( + "term posting buffer: freqs length must equal docids"); + } + if (!has_freqs && !positions_flat.empty()) { + return Status::Error( + "term posting buffer: positions require parallel freqs"); + } + uint64_t expected_positions = 0; + for (uint32_t freq : freqs) { + if (freq > std::numeric_limits::max() - expected_positions) { + return Status::Error( + "term posting buffer: position count overflow"); + } + expected_positions += freq; + } + if (!positions_flat.empty() && expected_positions != positions_flat.size()) { + return Status::Error( + "term posting buffer: positions count must equal sum(freqs)"); + } + if (has_freqs_.has_value() && *has_freqs_ != has_freqs) { + return Status::Error( + "term posting buffer: frequency shape changed within one fill"); + } + + MutableTermPostingSpan destination; + RETURN_IF_ERROR( + grow_uninitialized(docids.size(), has_freqs, positions_flat.size(), &destination)); + std::ranges::copy(docids, destination.docids.begin()); + std::ranges::copy(freqs, destination.freqs.begin()); + std::ranges::copy(positions_flat, destination.positions_flat.begin()); + return Status::OK(); + } + + // Extends the current fill once and exposes the new tail for direct decode. + // The caller must initialize every returned element before returning from + // TermPostingSource::fill. Shape and capacity changes are committed only + // after all reservations succeed. + Status grow_uninitialized(size_t document_count, bool has_freqs, size_t position_count, + MutableTermPostingSpan* destination) { + if (destination == nullptr) { + return Status::Error( + "term posting buffer: null writable span destination"); + } + if (!has_freqs && position_count != 0) { + return Status::Error( + "term posting buffer: positions require parallel freqs"); + } + if (has_freqs_.has_value() && *has_freqs_ != has_freqs) { + return Status::Error( + "term posting buffer: frequency shape changed within one fill"); + } + + const size_t doc_begin = docids_.size(); + const size_t freq_begin = freqs_.size(); + const size_t position_begin = positions_flat_.size(); + size_t target_docids = 0; + size_t target_freqs = 0; + size_t target_positions = 0; + RETURN_IF_ERROR(checked_size(doc_begin, document_count, &target_docids)); + RETURN_IF_ERROR(checked_size(freq_begin, has_freqs ? document_count : 0, &target_freqs)); + RETURN_IF_ERROR(checked_size(position_begin, position_count, &target_positions)); + RETURN_IF_ERROR(reserve_for_append(target_docids, target_freqs, target_positions)); + + docids_.resize(target_docids); + freqs_.resize(target_freqs); + positions_flat_.resize(target_positions); + destination->docids = std::span(docids_).subspan(doc_begin, document_count); + destination->freqs = std::span(freqs_).subspan(freq_begin, has_freqs ? document_count : 0); + destination->positions_flat = + std::span(positions_flat_).subspan(position_begin, position_count); + if (document_count != 0) { + has_freqs_ = has_freqs; + } + return Status::OK(); + } + + // Appends one position while a source decodes a frequency-bearing fill. + // The common path writes into retained capacity; growth keeps replacement + // reservation accounting atomic. + Status append_position(uint32_t position) { + if (!has_freqs_.value_or(false)) { + return Status::Error( + "term posting buffer: incremental positions require parallel freqs"); + } + if (positions_flat_.size() == positions_flat_.capacity()) { + size_t target_positions = 0; + RETURN_IF_ERROR(checked_size(positions_flat_.size(), 1, &target_positions)); + RETURN_IF_ERROR(reserve_for_append(docids_.size(), freqs_.size(), target_positions)); + } + positions_flat_.push_back(position); + return Status::OK(); + } + + std::span docids() const { return docids_; } + std::span freqs() const { return freqs_; } + std::span positions_flat() const { return positions_flat_; } + +private: + static Status checked_size(size_t current, size_t additional, size_t* target) { + if (additional > std::numeric_limits::max() - current) { + return Status::Error( + "term posting buffer: capacity overflow"); + } + *target = current + additional; + if (*target > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "term posting buffer: byte capacity overflow"); + } + return Status::OK(); + } + + static Status growth_capacity(size_t required, size_t current, size_t* target) { + *target = required; + if (current != 0 && current <= std::numeric_limits::max() / 2) { + *target = std::max(required, current * 2); + } + if (*target > std::numeric_limits::max() / sizeof(uint32_t)) { + return Status::Error( + "term posting buffer: growth capacity overflow"); + } + return Status::OK(); + } + + static Status capacity_bytes(size_t docids_capacity, size_t freqs_capacity, + size_t positions_capacity, uint64_t* bytes) { + const uint64_t docids_bytes = docids_capacity * sizeof(uint32_t); + const uint64_t freqs_bytes = freqs_capacity * sizeof(uint32_t); + const uint64_t positions_bytes = positions_capacity * sizeof(uint32_t); + if (freqs_bytes > std::numeric_limits::max() - docids_bytes || + positions_bytes > std::numeric_limits::max() - docids_bytes - freqs_bytes) { + return Status::Error( + "term posting buffer: aggregate capacity overflow"); + } + *bytes = docids_bytes + freqs_bytes + positions_bytes; + return Status::OK(); + } + + Status reserve_for_append(size_t target_docids, size_t target_freqs, size_t target_positions) { + const bool grow_docids = target_docids > docids_.capacity(); + const bool grow_freqs = target_freqs > freqs_.capacity(); + const bool grow_positions = target_positions > positions_flat_.capacity(); + if (!grow_docids && !grow_freqs && !grow_positions) { + return Status::OK(); + } + size_t docids_capacity = docids_.capacity(); + size_t freqs_capacity = freqs_.capacity(); + size_t positions_capacity = positions_flat_.capacity(); + if (grow_docids) { + RETURN_IF_ERROR(growth_capacity(target_docids, docids_.capacity(), &docids_capacity)); + } + if (grow_freqs) { + RETURN_IF_ERROR(growth_capacity(target_freqs, freqs_.capacity(), &freqs_capacity)); + } + if (grow_positions) { + RETURN_IF_ERROR(growth_capacity(target_positions, positions_flat_.capacity(), + &positions_capacity)); + } + if (memory_reporter_ == nullptr) { + if (grow_docids) docids_.reserve(docids_capacity); + if (grow_freqs) freqs_.reserve(freqs_capacity); + if (grow_positions) positions_flat_.reserve(positions_capacity); + return Status::OK(); + } + + uint64_t previous_bytes = 0; + uint64_t final_bytes = 0; + RETURN_IF_ERROR(capacity_bytes(docids_.capacity(), freqs_.capacity(), + positions_flat_.capacity(), &previous_bytes)); + RETURN_IF_ERROR( + capacity_bytes(docids_capacity, freqs_capacity, positions_capacity, &final_bytes)); + RETURN_IF_ERROR(capacity_reservation_.set_bytes(final_bytes)); + + const uint64_t overlap_bytes = + std::max({grow_docids ? docids_.capacity() * sizeof(uint32_t) : 0, + grow_freqs ? freqs_.capacity() * sizeof(uint32_t) : 0, + grow_positions ? positions_flat_.capacity() * sizeof(uint32_t) : 0}); + MemoryReporter::Reservation overlap_reservation = memory_reporter_->make_reservation(); + Status overlap_status = overlap_reservation.set_bytes(overlap_bytes); + if (!overlap_status.ok()) { + DORIS_CHECK(capacity_reservation_.set_bytes(previous_bytes).ok()); + return overlap_status; + } + + if (grow_docids) { + docids_.reserve(docids_capacity); + DCHECK_EQ(docids_.capacity(), docids_capacity); + } + if (grow_freqs) { + freqs_.reserve(freqs_capacity); + DCHECK_EQ(freqs_.capacity(), freqs_capacity); + } + if (grow_positions) { + positions_flat_.reserve(positions_capacity); + DCHECK_EQ(positions_flat_.capacity(), positions_capacity); + } + return Status::OK(); + } + + MemoryReporter* memory_reporter_ = nullptr; + // The reservation precedes vectors so their allocations are destroyed first. + MemoryReporter::Reservation capacity_reservation_; + std::vector docids_; + std::vector freqs_; + std::vector positions_flat_; + std::optional has_freqs_; +}; + +class TermPostingSource { +public: + virtual ~TermPostingSource() = default; + + // out is empty on entry. Unless this call reaches the term end, it must + // return exactly target_docs postings. exhausted means no postings remain + // after this call. The source and output are borrowed synchronously. + virtual Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) = 0; +}; + +// Synchronous non-owning adapter for callers that already hold one materialized +// posting list. It slices the arrays into the writer's requested document +// windows without copying the whole term into an intermediate object. +class SpanTermPostingSource final : public TermPostingSource { +public: + SpanTermPostingSource(std::span docids, std::span freqs, + std::span positions_flat) + : docids_(docids), freqs_(freqs), positions_flat_(positions_flat) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + if (out == nullptr || exhausted == nullptr || target_docs == 0) { + return Status::Error( + "span posting source: invalid fill arguments"); + } + if (!out->empty()) { + return Status::Error( + "span posting source: output must be empty"); + } + if (!freqs_.empty() && freqs_.size() != docids_.size()) { + return Status::Error( + "span posting source: freqs length must equal docids"); + } + if (freqs_.empty() && !positions_flat_.empty()) { + return Status::Error( + "span posting source: positions require parallel freqs"); + } + + const size_t count = + std::min(static_cast(target_docs), docids_.size() - doc_offset_); + size_t position_count = 0; + if (!positions_flat_.empty()) { + for (size_t i = 0; i < count; ++i) { + RETURN_IF_ERROR( + checked_add(position_count, freqs_[doc_offset_ + i], &position_count)); + } + if (position_count > positions_flat_.size() - position_offset_) { + return Status::Error( + "span posting source: positions shorter than sum(freqs)"); + } + } + + RETURN_IF_ERROR(out->append( + docids_.subspan(doc_offset_, count), + freqs_.empty() ? std::span {} : freqs_.subspan(doc_offset_, count), + positions_flat_.subspan(position_offset_, position_count))); + doc_offset_ += count; + position_offset_ += position_count; + *exhausted = doc_offset_ == docids_.size(); + if (*exhausted && !positions_flat_.empty() && position_offset_ != positions_flat_.size()) { + return Status::Error( + "span posting source: positions longer than sum(freqs)"); + } + return Status::OK(); + } + + bool exhausted() const { return doc_offset_ == docids_.size(); } + +private: + static Status checked_add(size_t current, uint32_t additional, size_t* result) { + if (additional > std::numeric_limits::max() - current) { + return Status::Error( + "span posting source: position count overflow"); + } + *result = current + additional; + return Status::OK(); + } + + std::span docids_; + std::span freqs_; + std::span positions_flat_; + size_t doc_offset_ = 0; + size_t position_offset_ = 0; +}; + +struct StreamedTermPostings { + std::string term; + bool retain_positions = true; + TermPostingSource* source = nullptr; +}; + +} // namespace doris::snii::writer diff --git a/be/src/storage/olap_common.h b/be/src/storage/olap_common.h index c13963b3bff449..da6ecc0b391e8e 100644 --- a/be/src/storage/olap_common.h +++ b/be/src/storage/olap_common.h @@ -40,6 +40,7 @@ #include "common/exception.h" #include "io/io_common.h" #include "storage/index/inverted/inverted_index_stats.h" +#include "storage/index/snii/snii_query_stats.h" #include "storage/olap_define.h" #include "storage/rowset/rowset_fwd.h" #include "util/hash_util.hpp" @@ -392,6 +393,8 @@ struct OlapReaderStatistics { int64_t inverted_index_query_timer = 0; int64_t inverted_index_query_cache_hit = 0; int64_t inverted_index_query_cache_miss = 0; + int64_t inverted_index_query_cache_lookup = 0; + int64_t inverted_index_query_cache_insert = 0; int64_t inverted_index_query_null_bitmap_timer = 0; int64_t inverted_index_query_bitmap_copy_timer = 0; int64_t inverted_index_searcher_open_timer = 0; @@ -403,6 +406,8 @@ struct OlapReaderStatistics { int64_t inverted_index_downgrade_count = 0; int64_t inverted_index_analyzer_timer = 0; int64_t inverted_index_lookup_timer = 0; + // See snii_query_stats.h: one field here instead of one per SNII counter. + snii::SniiQueryStats snii_stats; InvertedIndexStatistics inverted_index_stats; int64_t ann_index_load_ns = 0; diff --git a/be/src/storage/predicate_collector.cpp b/be/src/storage/predicate_collector.cpp deleted file mode 100644 index fa8fc0117ce34f..00000000000000 --- a/be/src/storage/predicate_collector.cpp +++ /dev/null @@ -1,314 +0,0 @@ -// 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 "storage/predicate_collector.h" - -#include - -#include - -#include "exec/common/variant_util.h" -#include "exprs/vexpr.h" -#include "exprs/vexpr_context.h" -#include "exprs/vliteral.h" -#include "exprs/vsearch.h" -#include "exprs/vslot_ref.h" -#include "gen_cpp/Exprs_types.h" -#include "storage/index/index_reader_helper.h" -#include "storage/index/inverted/analyzer/analyzer.h" -#include "storage/index/inverted/util/string_helper.h" -#include "storage/tablet/tablet_schema.h" - -namespace doris { - -using namespace segment_v2; - -VSlotRef* PredicateCollector::find_slot_ref(const VExprSPtr& expr) const { - if (!expr) { - return nullptr; - } - - auto cur = VExpr::expr_without_cast(expr); - if (cur->node_type() == TExprNodeType::SLOT_REF) { - return static_cast(cur.get()); - } - - for (const auto& ch : cur->children()) { - if (auto* s = find_slot_ref(ch)) { - return s; - } - } - - return nullptr; -} - -std::string PredicateCollector::build_field_name(int32_t col_unique_id, - const std::string& suffix_path) const { - std::string field_name = std::to_string(col_unique_id); - if (!suffix_path.empty()) { - field_name += "." + suffix_path; - } - return field_name; -} - -Status MatchPredicateCollector::collect(RuntimeState* state, const TabletSchemaSPtr& tablet_schema, - const VExprSPtr& expr, CollectInfoMap* collect_infos) { - DCHECK(collect_infos != nullptr); - - auto* left_slot_ref = find_slot_ref(expr->children()[0]); - if (left_slot_ref == nullptr) { - return Status::Error( - "Index statistics collection failed: Cannot find slot reference in match predicate " - "left expression"); - } - - auto* right_literal = static_cast(expr->children()[1].get()); - DCHECK(right_literal != nullptr); - - const auto* sd = state->desc_tbl().get_slot_descriptor(left_slot_ref->slot_id()); - if (sd == nullptr) { - return Status::Error( - "Index statistics collection failed: Cannot find slot descriptor for slot_id={}", - left_slot_ref->slot_id()); - } - - int32_t col_idx = tablet_schema->field_index(left_slot_ref->column_name()); - if (col_idx == -1) { - return Status::Error( - "Index statistics collection failed: Cannot find column index for column={}", - left_slot_ref->column_name()); - } - - const auto& column = tablet_schema->column(col_idx); - auto index_metas = tablet_schema->inverted_indexs(column); - std::vector> owned_index_metas; - std::string index_suffix_path = column.suffix_path(); - - // Schema-only fallback for variant sub-columns. Collector runs at tablet - // level without segment context, so we cannot do nested-group inference - // or inherit_index runtime-type dispatch. Two paths cover what is - // resolvable from schema alone: - // 1. field_pattern templates (MATCH_NAME / MATCH_NAME_GLOB) via - // generate_sub_column_info. - // 2. Plain parent inverted index when the schema column is the dynamic - // path's VARIANT placeholder produced by _init_variant_columns. In - // that state inverted_indexs(column) misses because - // _path_set_info_map.subcolumn_indexes is only populated for typed - // paths / field_pattern outputs, not for plain parent indexes added - // by ALTER. Clone the parent's non-field-pattern indexes with the - // variant path as suffix so segment-side BM25 statistics can be - // collected. - if (index_metas.empty() && column.is_extracted_column()) { - TabletSchema::SubColumnInfo sub_column_info; - const std::string relative_path = column.path_info_ptr()->copy_pop_front().get_path(); - if (variant_util::generate_sub_column_info(*tablet_schema, column.parent_unique_id(), - relative_path, &sub_column_info) && - !sub_column_info.indexes.empty()) { - index_suffix_path = sub_column_info.column.suffix_path(); - for (auto& idx : sub_column_info.indexes) { - index_metas.push_back(idx.get()); - owned_index_metas.emplace_back(std::move(idx)); - } - } else if (column.is_variant_type()) { - const auto parent_indexes = tablet_schema->inverted_indexs(column.parent_unique_id()); - for (const auto* index : parent_indexes) { - if (!index->field_pattern().empty()) { - continue; - } - auto index_ptr = std::make_shared(*index); - index_ptr->set_escaped_escaped_index_suffix_path( - column.path_info_ptr()->get_path()); - index_metas.push_back(index_ptr.get()); - owned_index_metas.emplace_back(std::move(index_ptr)); - } - } - } - -#ifndef BE_TEST - if (index_metas.empty()) { - return Status::Error( - "Index statistics collection failed: Score query is not supported without inverted " - "index for column={}", - left_slot_ref->column_name()); - } -#endif - - for (const auto* index_meta : index_metas) { - if (!InvertedIndexAnalyzer::should_analyzer(index_meta->properties())) { - continue; - } - - if (!IndexReaderHelper::is_need_similarity_score(expr->op(), index_meta)) { - continue; - } - - auto options = DataTypeSerDe::get_default_format_options(); - options.timezone = &state->timezone_obj(); - auto term_infos = InvertedIndexAnalyzer::get_analyse_result(right_literal->value(options), - index_meta->properties()); - - std::string field_name = - build_field_name(index_meta->col_unique_ids()[0], index_suffix_path); - std::wstring ws_field_name = StringHelper::to_wstring(field_name); - - auto iter = collect_infos->find(ws_field_name); - if (iter == collect_infos->end()) { - CollectInfo collect_info; - collect_info.term_infos.insert(term_infos.begin(), term_infos.end()); - collect_info.index_meta = index_meta; - for (const auto& owned_index_meta : owned_index_metas) { - if (owned_index_meta.get() == index_meta) { - collect_info.owned_index_meta = owned_index_meta; - break; - } - } - (*collect_infos)[ws_field_name] = std::move(collect_info); - } else { - iter->second.term_infos.insert(term_infos.begin(), term_infos.end()); - } - } - - return Status::OK(); -} - -Status SearchPredicateCollector::collect(RuntimeState* state, const TabletSchemaSPtr& tablet_schema, - const VExprSPtr& expr, CollectInfoMap* collect_infos) { - DCHECK(collect_infos != nullptr); - - auto* search_expr = dynamic_cast(expr.get()); - if (search_expr == nullptr) { - return Status::InternalError("SearchPredicateCollector: expr is not VSearchExpr type"); - } - - const TSearchParam& search_param = search_expr->get_search_param(); - - RETURN_IF_ERROR(collect_from_clause(search_param.root, state, tablet_schema, collect_infos)); - - return Status::OK(); -} - -Status SearchPredicateCollector::collect_from_clause(const TSearchClause& clause, - RuntimeState* state, - const TabletSchemaSPtr& tablet_schema, - CollectInfoMap* collect_infos) { - const std::string& clause_type = clause.clause_type; - ClauseTypeCategory category = get_clause_type_category(clause_type); - - if (category == ClauseTypeCategory::COMPOUND) { - if (clause.__isset.children) { - for (const auto& child_clause : clause.children) { - RETURN_IF_ERROR( - collect_from_clause(child_clause, state, tablet_schema, collect_infos)); - } - } - return Status::OK(); - } - - return collect_from_leaf(clause, state, tablet_schema, collect_infos); -} - -Status SearchPredicateCollector::collect_from_leaf(const TSearchClause& clause, RuntimeState* state, - const TabletSchemaSPtr& tablet_schema, - CollectInfoMap* collect_infos) { - if (!clause.__isset.field_name || !clause.__isset.value) { - return Status::InvalidArgument("Search clause missing field_name or value"); - } - - const std::string& field_name = clause.field_name; - const std::string& value = clause.value; - const std::string& clause_type = clause.clause_type; - - if (!is_score_query_type(clause_type)) { - return Status::OK(); - } - - int32_t col_idx = tablet_schema->field_index(field_name); - if (col_idx == -1) { - return Status::OK(); - } - - const auto& column = tablet_schema->column(col_idx); - - auto index_metas = tablet_schema->inverted_indexs(column.unique_id(), column.suffix_path()); - if (index_metas.empty()) { - return Status::OK(); - } - - ClauseTypeCategory category = get_clause_type_category(clause_type); - for (const auto* index_meta : index_metas) { - std::set term_infos; - - if (category == ClauseTypeCategory::TOKENIZED) { - if (InvertedIndexAnalyzer::should_analyzer(index_meta->properties())) { - auto analyzed_terms = - InvertedIndexAnalyzer::get_analyse_result(value, index_meta->properties()); - term_infos.insert(analyzed_terms.begin(), analyzed_terms.end()); - } else { - term_infos.insert(TermInfo(value)); - } - } else if (category == ClauseTypeCategory::NON_TOKENIZED) { - if (clause_type == "TERM" && - InvertedIndexAnalyzer::should_analyzer(index_meta->properties())) { - auto analyzed_terms = - InvertedIndexAnalyzer::get_analyse_result(value, index_meta->properties()); - term_infos.insert(analyzed_terms.begin(), analyzed_terms.end()); - } else { - term_infos.insert(TermInfo(value)); - } - } - - std::string lucene_field_name = - build_field_name(index_meta->col_unique_ids()[0], column.suffix_path()); - std::wstring ws_field_name = StringHelper::to_wstring(lucene_field_name); - - auto iter = collect_infos->find(ws_field_name); - if (iter == collect_infos->end()) { - CollectInfo collect_info; - collect_info.term_infos = std::move(term_infos); - collect_info.index_meta = index_meta; - (*collect_infos)[ws_field_name] = std::move(collect_info); - } else { - iter->second.term_infos.insert(term_infos.begin(), term_infos.end()); - } - } - - return Status::OK(); -} - -bool SearchPredicateCollector::is_score_query_type(const std::string& clause_type) const { - return clause_type == "TERM" || clause_type == "EXACT" || clause_type == "PHRASE" || - clause_type == "MATCH" || clause_type == "ANY" || clause_type == "ALL"; -} - -SearchPredicateCollector::ClauseTypeCategory SearchPredicateCollector::get_clause_type_category( - const std::string& clause_type) const { - if (clause_type == "AND" || clause_type == "OR" || clause_type == "NOT" || - clause_type == "OCCUR_BOOLEAN") { - return ClauseTypeCategory::COMPOUND; - } else if (clause_type == "TERM" || clause_type == "EXACT") { - return ClauseTypeCategory::NON_TOKENIZED; - } else if (clause_type == "PHRASE" || clause_type == "MATCH" || clause_type == "ANY" || - clause_type == "ALL") { - return ClauseTypeCategory::TOKENIZED; - } else { - LOG(WARNING) << "Unknown clause type '" << clause_type - << "', defaulting to NON_TOKENIZED category"; - return ClauseTypeCategory::NON_TOKENIZED; - } -} - -} // namespace doris diff --git a/be/src/storage/rowset/beta_rowset.cpp b/be/src/storage/rowset/beta_rowset.cpp index 4419fc3c0d470e..08e58945c6a38d 100644 --- a/be/src/storage/rowset/beta_rowset.cpp +++ b/be/src/storage/rowset/beta_rowset.cpp @@ -831,6 +831,9 @@ Status BetaRowset::show_nested_index_file(rapidjson::Value* rowset_value, case InvertedIndexStorageFormatPB::V3: format_str = "V3"; break; + case InvertedIndexStorageFormatPB::SNII: + format_str = "SNII"; + break; default: return Status::InternalError("inverted index storage format error"); break; @@ -840,6 +843,19 @@ Status BetaRowset::show_nested_index_file(rapidjson::Value* rowset_value, rowset_value->AddMember("index_storage_format", rapidjson::Value(format_str.c_str(), allocator), allocator); rapidjson::Value segments(rapidjson::kArrayType); + auto add_file_info_to_json = [&](const std::string& path, + rapidjson::Value& json_value) -> Status { + json_value.AddMember("idx_file_path", rapidjson::Value(path.c_str(), allocator), allocator); + int64_t idx_file_size = 0; + auto st = fs->file_size(path, &idx_file_size); + if (st != Status::OK()) { + LOG(WARNING) << "show nested index file get file size error, file: " << path + << ", error: " << st.msg(); + return st; + } + json_value.AddMember("idx_file_size", rapidjson::Value(idx_file_size).Move(), allocator); + return Status::OK(); + }; for (int seg_id = 0; seg_id < num_segments(); ++seg_id) { rapidjson::Value segment(rapidjson::kObjectType); segment.AddMember("segment_id", rapidjson::Value(seg_id).Move(), allocator); @@ -850,24 +866,20 @@ Status BetaRowset::show_nested_index_file(rapidjson::Value* rowset_value, fs, std::string(index_file_path_prefix), storage_format, InvertedIndexFileInfo(), _rowset_meta->tablet_id()); RETURN_IF_ERROR(index_file_reader->init()); + if (storage_format == InvertedIndexStorageFormatPB::SNII) { + rapidjson::Value index_file(rapidjson::kObjectType); + auto index_file_path = + InvertedIndexDescriptor::get_index_file_path_v2(index_file_path_prefix); + RETURN_IF_ERROR(add_file_info_to_json(index_file_path, index_file)); + segment.AddMember("index_files", rapidjson::Value(rapidjson::kArrayType).Move(), + allocator); + auto& index_files = segment["index_files"]; + index_files.PushBack(index_file, allocator); + segments.PushBack(segment, allocator); + continue; + } auto dirs = index_file_reader->get_all_directories(); - auto add_file_info_to_json = [&](const std::string& path, - rapidjson::Value& json_value) -> Status { - json_value.AddMember("idx_file_path", rapidjson::Value(path.c_str(), allocator), - allocator); - int64_t idx_file_size = 0; - auto st = fs->file_size(path, &idx_file_size); - if (st != Status::OK()) { - LOG(WARNING) << "show nested index file get file size error, file: " << path - << ", error: " << st.msg(); - return st; - } - json_value.AddMember("idx_file_size", rapidjson::Value(idx_file_size).Move(), - allocator); - return Status::OK(); - }; - auto process_files = [&allocator, &index_file_reader](auto& index_meta, rapidjson::Value& indices, rapidjson::Value& index) -> Status { diff --git a/be/src/storage/segment/column_reader.cpp b/be/src/storage/segment/column_reader.cpp index db8f2d6819b04d..b32ae470b08e7d 100644 --- a/be/src/storage/segment/column_reader.cpp +++ b/be/src/storage/segment/column_reader.cpp @@ -56,6 +56,7 @@ #include "storage/index/index_reader.h" #include "storage/index/inverted/analyzer/analyzer.h" #include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/snii/snii_index_reader.h" #include "storage/index/zone_map/zone_map_index.h" #include "storage/iterators.h" #include "storage/olap_common.h" @@ -722,6 +723,17 @@ Status ColumnReader::_load_index(const std::shared_ptr& index_f } IndexReaderPtr index_reader; + if (index_file_reader->get_storage_format() == InvertedIndexStorageFormatPB::SNII) { + if (!is_string_type(type)) { + return Status::Error( + "SNII inverted index storage format does not support BKD index type {}", type); + } + auto reader_type = should_analyzer ? InvertedIndexReaderType::FULLTEXT + : InvertedIndexReaderType::STRING_TYPE; + index_reader = SniiIndexReader::create_shared(index_meta, index_file_reader, reader_type); + _index_readers[index_meta->index_id()] = index_reader; + return Status::OK(); + } if (is_string_type(type)) { if (should_analyzer) { diff --git a/be/src/storage/segment/column_writer.cpp b/be/src/storage/segment/column_writer.cpp index 7d604897494f7d..349749925e6b0b 100644 --- a/be/src/storage/segment/column_writer.cpp +++ b/be/src/storage/segment/column_writer.cpp @@ -565,6 +565,12 @@ Status ScalarColumnWriter::init() { RETURN_IF_ERROR(IndexColumnWriter::create( get_column(), &_inverted_index_builders[i], _opts.index_file_writer, _opts.inverted_indexes[i])); + // After create() (which runs the writer's init()) and before any + // value lands: forward the write type UNCONDITIONALLY so SNII + // can select its PRX zstd level -- and the + // documented "forwarded to every created IndexColumnWriter" + // contract holds for both values. + _inverted_index_builders[i]->set_direct_load(_opts.is_direct_load); } } while (false); } @@ -1056,6 +1062,9 @@ Status ArrayColumnWriter::init() { RETURN_IF_ERROR(IndexColumnWriter::create(get_column(), &_inverted_index_writer, _opts.index_file_writer, _opts.inverted_indexes[0])); + // Same unconditional forwarding as the scalar path: after + // create()/init(), before any array value is added. + _inverted_index_writer->set_direct_load(_opts.is_direct_load); } } if (_opts.need_ann_index) { diff --git a/be/src/storage/segment/column_writer.h b/be/src/storage/segment/column_writer.h index 50822c945c1bb7..2ddb71b6edbd69 100644 --- a/be/src/storage/segment/column_writer.h +++ b/be/src/storage/segment/column_writer.h @@ -78,6 +78,13 @@ struct ColumnWriterOptions { BloomFilterOptions bf_options; std::vector inverted_indexes; IndexFileWriter* index_file_writer = nullptr; + // The owning segment serves a direct load (stream/broker load, + // DataWriteType::TYPE_DIRECT) rather than compaction / schema change. Set + // once by the segment writer and propagated to variant subcolumn writers; + // forwarded to every created IndexColumnWriter via set_direct_load() so + // SNII can select its direct-load PRX zstd level without plumbing + // DataWriteType itself down here. + bool is_direct_load = false; SegmentFooterPB* footer = nullptr; io::FileWriter* file_writer = nullptr; diff --git a/be/src/storage/segment/count_on_index_fastpath.h b/be/src/storage/segment/count_on_index_fastpath.h new file mode 100644 index 00000000000000..30b0ccf7218157 --- /dev/null +++ b/be/src/storage/segment/count_on_index_fastpath.h @@ -0,0 +1,196 @@ +// 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. + +#pragma once + +#include +#include + +// G02 count-only fast-path caller guard (functional core, unit-testable +// without a SegmentIterator). +// +// COUNT_ON_INDEX counts the rows of THIS segment that match the pushed-down +// predicates MINUS deleted rows: SegmentIterator seeds _row_bitmap with +// [0, num_rows), intersects the index result bitmap into it, then subtracts +// the MOW delete bitmap and applies delete predicates / row ranges, and the +// scan emits |_row_bitmap| default-valued rows that the agg counts. The SNII +// fast path replaces the index result with a FABRICATED [0, df) bitmap whose +// cardinality is exact but whose row ids are not real. That is only equal in +// observable behavior when, for this segment iterator: +// 1. only the COUNT matters (COUNT_ON_INDEX agg pushdown), +// 2. the single pushed-down MATCH predicate is the ONLY filter (nothing else +// may intersect _row_bitmap before or after the index apply), +// 3. no rows are deleted (mirror of the V3 handling: _lazy_init subtracts +// the per-segment delete bitmap AFTER the index apply, and delete +// predicates filter later -- a fabricated id range cannot participate in +// either subtraction), +// 4. nothing consumes REAL row ids (rowid recording, ANN topn, BM25 +// scoring, virtual columns), and +// 5. rows are emitted as defaults without reading column data at the +// fabricated ids (the COUNT_ON_INDEX no-read-data contract must be +// active: enable_no_need_read_data_opt + DUP_KEYS or MOW). +// +// SegmentIterator fills the facts from its state right before applying the +// index and only sets IndexQueryContext::count_on_index_fastpath when this +// predicate holds. +namespace doris::segment_v2 { + +struct CountOnIndexFastpathFacts { + // (1) count-only context. + bool is_count_on_index_agg = false; + // (2) single MATCH predicate, no other conjuncts. + bool has_column_predicates = false; + size_t common_expr_count = 0; + bool single_expr_is_match_pred = false; + bool has_virtual_column_exprs = false; + // (3) deletes must be absent for this segment. + bool has_delete_predicates = false; + bool segment_delete_bitmap_empty = false; + // (2 cont.) nothing else may prune the row space around the index apply. + bool has_col_id_predicates = false; + bool has_topn_filters = false; + bool has_external_row_ranges = false; + bool row_bitmap_is_full = false; + // (4) consumers of real row ids. + bool record_rowids = false; + bool has_ann_topn = false; + bool has_score_runtime = false; + // (5) rows must be emitted as defaults (no data read at fabricated ids). + bool no_need_read_data_opt_enabled = false; + bool keys_type_supported = false; +}; + +inline bool count_on_index_fastpath_safe(const CountOnIndexFastpathFacts& f) { + return f.is_count_on_index_agg && !f.has_column_predicates && f.common_expr_count == 1 && + f.single_expr_is_match_pred && !f.has_virtual_column_exprs && !f.has_delete_predicates && + f.segment_delete_bitmap_empty && !f.has_col_id_predicates && !f.has_topn_filters && + !f.has_external_row_ranges && f.row_bitmap_is_full && !f.record_rowids && + !f.has_ann_topn && !f.has_score_runtime && f.no_need_read_data_opt_enabled && + f.keys_type_supported; +} + +// G03 count-emission shortcut guard (functional core, unit-testable without a +// SegmentIterator). +// +// After the G02 fast path answered the single MATCH predicate with a +// count-shaped bitmap, the only remaining work of the scan is to emit +// |_row_bitmap| default-valued rows batch by batch: the per-batch rowid +// iteration over the fabricated bitmap and the per-column no-read checks are +// pure overhead. The shortcut replaces them with a countdown that fills the +// block columns with defaults directly, in VStatisticsIterator-sized batches. +// +// That replacement is byte-for-byte equal to today's emission only when, at +// the end of _lazy_init: +// 1. the reader ACTUALLY answered from df (count_fastpath_hit) -- a mere +// guard pass with a row-accurate decode keeps today's path untouched, +// 2. no evaluation stage survives (vec/short-circuit/expr eval, leftover +// column predicates or common exprs, delete predicates, lazy +// materialization), +// 3. nothing consumes real row ids or per-row values (virtual columns, +// rowid recording), +// 4. batch accounting is a pure countdown (no read limit, no reverse +// key-ordered read, no condition-cache writes), and +// 5. the block is exactly the read schema and EVERY column would take a +// defaults fill in _read_columns_by_index (no real column read, no +// storage->schema cast, no version/lsn/tso rewrite) -- checked +// per-column by the iterator and summarized in one fact. +// +// Every fact is re-verified from live iterator state even though the G02 +// facts guard already implies most of them: the shortcut independently +// refuses on any drift, falling through to today's emission (which is always +// count-exact). +struct CountEmitShortcutFacts { + // (1) reader answered with a fabricated count bitmap. + bool count_fastpath_hit = false; + // (2) nothing evaluates or filters rows after the index apply. + bool needs_vec_eval = false; + bool needs_short_eval = false; + bool needs_expr_eval = false; + bool has_remaining_col_predicates = false; + bool has_remaining_common_exprs = false; + bool has_delete_predicates = false; + bool lazy_materialization_read = false; + // (3) consumers of real row ids / per-row values. + bool has_virtual_columns = false; + bool record_rowids = false; + // (4) batch accounting must be a pure countdown. + bool has_read_limit = false; + bool read_orderby_key_reverse = false; + bool has_condition_cache_digest = false; + // (5) emitted block == read schema, all columns defaults-fillable. + bool block_shape_matches_schema = false; + bool all_columns_emit_defaults = false; +}; + +inline bool count_emit_shortcut_safe(const CountEmitShortcutFacts& f) { + return f.count_fastpath_hit && !f.needs_vec_eval && !f.needs_short_eval && !f.needs_expr_eval && + !f.has_remaining_col_predicates && !f.has_remaining_common_exprs && + !f.has_delete_predicates && !f.lazy_materialization_read && !f.has_virtual_columns && + !f.record_rowids && !f.has_read_limit && !f.read_orderby_key_reverse && + !f.has_condition_cache_digest && f.block_shape_matches_schema && + f.all_columns_emit_defaults; +} + +} // namespace doris::segment_v2 + +// Deterministic seam for the G03 count-emission shortcut, mirroring the SNII +// query seam (storage/index/snii/query/internal/query_test_counters.h): +// - count_emit_shortcut_hits : engage decisions that ADMITTED the +// shortcut (incremented inside +// _should_engage_count_emit_shortcut, which +// _lazy_init calls once per iterator). Guard +// fall-throughs leave it unchanged. +// - count_emit_shortcut_batches : default-rows batches emitted by the +// shortcut (== ceil(count / 65535) per +// engaged iterator with a non-zero count). +// +// Active only under BE_TEST (library-wide define of doris_be_test); in a +// release build the macro expands to ((void)0): zero overhead, no global +// mutable state on the production path. The singleton is intentionally +// unsynchronized -- single-threaded test-only seam; reset between cases with +// `count_emit_test_counters() = {}`. +#if defined(BE_TEST) && !defined(SNII_COUNT_EMIT_TEST_COUNTERS) +#define SNII_COUNT_EMIT_TEST_COUNTERS +#endif + +#ifdef SNII_COUNT_EMIT_TEST_COUNTERS + +namespace doris::segment_v2::internal { + +struct CountEmitTestCounters { + uint64_t count_emit_shortcut_hits = 0; + uint64_t count_emit_shortcut_batches = 0; +}; + +// `inline` gives a single shared instance across all TUs that include this +// header, so counter increments made in segment_iterator.cpp are visible to +// the test that reads them. +inline CountEmitTestCounters& count_emit_test_counters() { + static CountEmitTestCounters counters; + return counters; +} + +} // namespace doris::segment_v2::internal + +#define SNII_COUNT_EMIT_COUNT(field) \ + (++::doris::segment_v2::internal::count_emit_test_counters().field) + +#else + +#define SNII_COUNT_EMIT_COUNT(field) ((void)0) + +#endif diff --git a/be/src/storage/segment/row_ranges.h b/be/src/storage/segment/row_ranges.h index 0f44f599695b8c..9d98f02e245b15 100644 --- a/be/src/storage/segment/row_ranges.h +++ b/be/src/storage/segment/row_ranges.h @@ -247,7 +247,7 @@ class RowRanges { size_t count() { return _count; } - bool is_empty() { return _count == 0; } + bool is_empty() const { return _count == 0; } bool contain(rowid_t from, rowid_t to) { // binary search diff --git a/be/src/storage/segment/segment_iterator.cpp b/be/src/storage/segment/segment_iterator.cpp index bc04ffa2aca276..1297dbc35905bb 100644 --- a/be/src/storage/segment/segment_iterator.cpp +++ b/be/src/storage/segment/segment_iterator.cpp @@ -103,6 +103,7 @@ #include "storage/segment/column_reader.h" #include "storage/segment/column_reader_cache.h" #include "storage/segment/condition_cache.h" +#include "storage/segment/count_on_index_fastpath.h" #include "storage/segment/row_ranges.h" #include "storage/segment/segment.h" #include "storage/segment/segment_prefetcher.h" @@ -614,6 +615,16 @@ Status SegmentIterator::_lazy_init(Block* block) { _init_segment_prefetchers(); + // G03: engage the count-emission shortcut. All inputs are final here (the + // index apply ran, _row_bitmap saw every subtraction/intersection above, + // _vec_init_lazy_materialization fixed the eval flags), so the post-apply + // cardinality IS the exact row count today's batch loop would emit; the + // shortcut only changes how fast those default rows are produced. + _count_emit_shortcut = _should_engage_count_emit_shortcut(block); + if (_count_emit_shortcut) { + _count_emit_rows_remaining = _row_bitmap.cardinality(); + } + return Status::OK(); } @@ -825,6 +836,18 @@ Status SegmentIterator::_get_row_ranges_by_column_conditions() { (has_index_in_iterators() || !_common_expr_ctxs_push_down.empty())) { SCOPED_RAW_TIMER(&_opts.stats->inverted_index_filter_timer); size_t input_rows = _row_bitmap.cardinality(); + // G02 count-only fast path handshake: only while the single + // pushed-down MATCH predicate of a provably filter-free + // COUNT_ON_INDEX scan is evaluated may a reader answer with a + // count-shaped bitmap (see count_on_index_fastpath.h). The reply + // direction (did the reader actually fabricate one?) is captured + // into _count_fastpath_hit and both flags are reset on every exit + // path so no later read_from_index call can observe or forge them. + if (_index_query_context != nullptr) { + _index_query_context->count_on_index_fastpath = _count_on_index_fastpath_safe(); + _index_query_context->count_on_index_fastpath_hit = false; + } + DEFER({ _capture_count_fastpath_hit(); }); // Only apply column-level inverted index if we have iterators if (has_index_in_iterators()) { RETURN_IF_ERROR(_apply_inverted_index()); @@ -1337,6 +1360,146 @@ Status SegmentIterator::_apply_index_expr() { return Status::OK(); } +bool SegmentIterator::_count_on_index_fastpath_safe() const { + CountOnIndexFastpathFacts facts; + facts.is_count_on_index_agg = _opts.push_down_agg_type_opt == TPushAggOp::COUNT_ON_INDEX; + facts.has_column_predicates = !_col_predicates.empty(); + facts.common_expr_count = _common_expr_ctxs_push_down.size(); + facts.single_expr_is_match_pred = + _common_expr_ctxs_push_down.size() == 1 && + _common_expr_ctxs_push_down.front()->root() != nullptr && + _common_expr_ctxs_push_down.front()->root()->node_type() == TExprNodeType::MATCH_PRED; + facts.has_virtual_column_exprs = !_virtual_column_exprs.empty(); + facts.has_delete_predicates = _opts.delete_condition_predicates != nullptr && + _opts.delete_condition_predicates->num_of_column_predicate() > 0; + // Mirror of the _lazy_init delete-bitmap subtraction: the fast path is only + // sound when there is nothing to subtract for THIS segment. + const auto delete_bitmap_it = _opts.delete_bitmap.find(segment_id()); + facts.segment_delete_bitmap_empty = delete_bitmap_it == _opts.delete_bitmap.end() || + delete_bitmap_it->second == nullptr || + delete_bitmap_it->second->isEmpty(); + facts.has_col_id_predicates = !_opts.col_id_to_predicates.empty(); + facts.has_topn_filters = !_opts.topn_filter_source_node_ids.empty(); + facts.has_external_row_ranges = !_opts.row_ranges.is_empty(); + // Catches every earlier pruning source in one check (condition cache, key + // ranges): the fabricated [0, df) range only counts correctly against a + // full [0, num_rows) bitmap. + facts.row_bitmap_is_full = _row_bitmap.cardinality() == uint64_t(num_rows()); + facts.record_rowids = _opts.record_rowids; + facts.has_ann_topn = _opts.ann_topn_runtime != nullptr; + facts.has_score_runtime = _score_runtime != nullptr; + // Mirror of the _need_read_data preamble: rows must be emitted as defaults, + // never materialized from the fabricated row ids. + facts.no_need_read_data_opt_enabled = + _opts.runtime_state == nullptr || + _opts.runtime_state->query_options().enable_no_need_read_data_opt; + facts.keys_type_supported = _opts.tablet_schema->keys_type() == KeysType::DUP_KEYS || + (_opts.tablet_schema->keys_type() == KeysType::UNIQUE_KEYS && + _opts.enable_unique_key_merge_on_write); + return count_on_index_fastpath_safe(facts); +} + +void SegmentIterator::_capture_count_fastpath_hit() { + if (_index_query_context == nullptr) { + return; + } + _count_fastpath_hit = _index_query_context->count_on_index_fastpath_hit; + _index_query_context->count_on_index_fastpath = false; + _index_query_context->count_on_index_fastpath_hit = false; +} + +bool SegmentIterator::_column_emits_defaults_for_count(ColumnId cid) { + // Mirror of the per-batch fill in _read_columns_by_index: a column's batch + // content is reproducible by the emission shortcut iff the column takes + // the _no_need_read_key_data defaults fill or the _prune_column defaults + // fill. Anything else -- a real column read, a storage->schema cast in + // _init_current_block/_convert_to_expected_type (whose CAST(default) need + // not equal the schema-type default), or a version/lsn/tso rewrite source + // (those columns are never no-read, so the type/read checks below already + // veto them) -- must keep today's path. + if (_is_pred_column[cid]) { + return false; + } + const auto* column_desc = _schema->column(cid); + if (column_desc == nullptr) { + return false; + } + const auto& file_column_type = _storage_name_and_type[cid].second; + DataTypePtr expected_type = Schema::get_data_type_ptr(*column_desc); + if (file_column_type == nullptr || expected_type == nullptr || + !file_column_type->equals(*expected_type)) { + return false; + } + return _no_need_read_key_data_eligible(cid) || !_need_read_data(cid); +} + +bool SegmentIterator::_should_engage_count_emit_shortcut(const Block* block) { + if (!_count_fastpath_hit) { + return false; + } + CountEmitShortcutFacts facts; + facts.count_fastpath_hit = _count_fastpath_hit; + facts.needs_vec_eval = _is_need_vec_eval; + facts.needs_short_eval = _is_need_short_eval; + facts.needs_expr_eval = _is_need_expr_eval; + facts.has_remaining_col_predicates = !_col_predicates.empty(); + facts.has_remaining_common_exprs = !_common_expr_ctxs_push_down.empty(); + facts.has_delete_predicates = _opts.delete_condition_predicates != nullptr && + _opts.delete_condition_predicates->num_of_column_predicate() > 0; + facts.lazy_materialization_read = _lazy_materialization_read; + facts.has_virtual_columns = !_virtual_column_exprs.empty(); + facts.record_rowids = _opts.record_rowids || _record_rowids; + facts.has_read_limit = _opts.read_limit > 0; + facts.read_orderby_key_reverse = _opts.read_orderby_key_reverse; + facts.has_condition_cache_digest = _opts.condition_cache_digest != 0; + facts.block_shape_matches_schema = block->columns() == _schema->num_column_ids(); + facts.all_columns_emit_defaults = true; + for (size_t i = 0; i < _schema->num_column_ids() && facts.all_columns_emit_defaults; ++i) { + facts.all_columns_emit_defaults = _column_emits_defaults_for_count(_schema->column_id(i)); + } + const bool engage = count_emit_shortcut_safe(facts); + if (engage) { + SNII_COUNT_EMIT_COUNT(count_emit_shortcut_hits); + } + return engage; +} + +Status SegmentIterator::_emit_count_shortcut_batch(Block* block) { + if (_count_emit_rows_remaining == 0) { + // EOF twin of _process_eof: shortcut batches never move the block's + // columns into _current_return_columns, so there is nothing to + // restore; deliver the empty block and release iterator memory the + // same way. + block->clear_column_data(); + _column_iterators.clear(); + _index_iterators.clear(); + return Status::EndOfFile("no more data in segment"); + } + const auto rows = static_cast( + std::min(_count_emit_rows_remaining, kCountEmitBatchRows)); + block->clear_column_data(_schema->num_column_ids()); + for (size_t i = 0; i < _schema->num_column_ids(); ++i) { + MutableColumnPtr column = std::move(*block->get_by_position(i).column).mutate(); + // Same fill as the defaults branches of _read_columns_by_index + // (_no_need_read_key_data / _prune_column): NOT-NULL defaults. + // ColumnNullable::insert_many_defaults would insert NULLs and break + // count(col) parity. + if (is_column_nullable(*column)) { + auto* nullable_col_ptr = reinterpret_cast(column.get()); + nullable_col_ptr->get_null_map_column().insert_many_defaults(rows); + nullable_col_ptr->get_nested_column_ptr()->insert_many_defaults(rows); + } else { + column->insert_many_defaults(rows); + } + block->replace_by_position(i, std::move(column)); + } + _count_emit_rows_remaining -= rows; + _opts.stats->blocks_load += 1; + _opts.stats->raw_rows_read += rows; + SNII_COUNT_EMIT_COUNT(count_emit_shortcut_batches); + return Status::OK(); +} + bool SegmentIterator::_downgrade_without_index(Status res, bool need_remaining) { bool is_fallback = _opts.runtime_state->query_options().enable_fallback_on_missing_inverted_index; @@ -2891,6 +3054,14 @@ Status SegmentIterator::_next_batch_internal(Block* block) { SCOPED_RAW_TIMER(&_opts.stats->block_load_ns); + // G03: a count-fastpath scan emits the remaining count as default rows in + // statistics-sized batches; the row-bitmap iterator and the per-rowid + // machinery below are never touched. Engaged only when read_limit == 0, + // so the limit check below stays unreachable for shortcut scans. + if (_count_emit_shortcut) { + return _emit_count_shortcut_batch(block); + } + if (_opts.read_limit > 0 && _rows_returned >= _opts.read_limit) { return _process_eof(block); } @@ -3492,8 +3663,7 @@ void SegmentIterator::_calculate_common_expr_index_exec_status() { } } -bool SegmentIterator::_no_need_read_key_data(ColumnId cid, MutableColumnPtr& column, - size_t nrows_read) { +bool SegmentIterator::_no_need_read_key_data_eligible(ColumnId cid) { if (_opts.runtime_state && !_opts.runtime_state->query_options().enable_no_need_read_data_opt) { return false; } @@ -3520,6 +3690,14 @@ bool SegmentIterator::_no_need_read_key_data(ColumnId cid, MutableColumnPtr& col return false; } + return true; +} + +bool SegmentIterator::_no_need_read_key_data(ColumnId cid, MutableColumnPtr& column, + size_t nrows_read) { + if (!_no_need_read_key_data_eligible(cid)) { + return false; + } insert_many_not_null_defaults(column, nrows_read); return true; } diff --git a/be/src/storage/segment/segment_iterator.h b/be/src/storage/segment/segment_iterator.h index 99e40f2a5d1214..5a16bc6c8f8612 100644 --- a/be/src/storage/segment/segment_iterator.h +++ b/be/src/storage/segment/segment_iterator.h @@ -198,6 +198,32 @@ class SegmentIterator : public RowwiseIterator { bool* continue_apply); [[nodiscard]] Status _apply_ann_topn_predicate(); [[nodiscard]] Status _apply_index_expr(); + // G02: true iff answering the single pushed-down MATCH predicate by its + // match COUNT alone is indistinguishable from the row-accurate bitmap for + // this COUNT_ON_INDEX scan (no deletes, no other filters, full row bitmap, + // no row-id consumers). Gates IndexQueryContext::count_on_index_fastpath; + // the decision predicate itself lives in count_on_index_fastpath.h. + bool _count_on_index_fastpath_safe() const; + // G03: teardown of the G02 handshake. Captures whether the reader answered + // with a fabricated count bitmap into _count_fastpath_hit and clears both + // context flags so no later read_from_index call can observe or forge + // them. Runs on every exit of the index-apply scope. + void _capture_count_fastpath_hit(); + // G03: true iff the per-batch defaults fill of _read_columns_by_index + // would apply to `cid` (the _no_need_read_key_data or _prune_column + // branch) AND the block column needs no storage->schema cast, i.e. the + // emission shortcut can reproduce the column's batch content exactly. + bool _column_emits_defaults_for_count(ColumnId cid); + // G03: fills CountEmitShortcutFacts from live iterator state at the end of + // _lazy_init and returns the pure-guard verdict; the decision predicate + // itself lives in count_on_index_fastpath.h. + bool _should_engage_count_emit_shortcut(const Block* block); + // G03: one emission-shortcut batch: min(remaining, kCountEmitBatchRows) + // default rows filled straight into the block (NOT-NULL defaults for + // nullable columns, mirroring _prune_column), then EOF once the countdown + // reaches zero. Replaces the whole per-rowid _next_batch_internal body for + // engaged scans. + Status _emit_count_shortcut_batch(Block* block); bool _column_has_fulltext_index(int32_t cid); bool _column_has_ann_index(int32_t cid); @@ -314,6 +340,10 @@ class SegmentIterator : public RowwiseIterator { Status _convert_to_expected_type(const std::vector& col_ids); bool _no_need_read_key_data(ColumnId cid, MutableColumnPtr& column, size_t nrows_read); + // Side-effect-free eligibility half of _no_need_read_key_data (no column + // fill); shared by the per-batch fill and the G03 engage-time per-column + // proof so the two can never drift. + bool _no_need_read_key_data_eligible(ColumnId cid); bool _has_delete_predicate(ColumnId cid); bool _can_skip_reading_extra_column(ColumnId cid); @@ -477,6 +507,23 @@ class SegmentIterator : public RowwiseIterator { IndexQueryContextPtr _index_query_context; + // G03 count-emission shortcut state (see count_on_index_fastpath.h). + // _count_fastpath_hit: the reader answered the single MATCH predicate with + // a fabricated count bitmap (captured from the G02 handshake reply). + // _count_emit_shortcut: engaged at the end of _lazy_init when + // count_emit_shortcut_safe holds; every subsequent batch is emitted by + // _emit_count_shortcut_batch from _count_emit_rows_remaining (initialized + // to the post-apply _row_bitmap cardinality) without touching the row + // bitmap iterator. + bool _count_fastpath_hit = false; + bool _count_emit_shortcut = false; + uint64_t _count_emit_rows_remaining = 0; + // Batch size for shortcut emission: VStatisticsIterator's + // MAX_ROW_SIZE_IN_COUNT, the largest default-rows block shape already + // proven through every consumer above the segment iterator by the plain + // COUNT pushdown (rowset reader, collect iterator, block reader, scanner). + static constexpr uint64_t kCountEmitBatchRows = 65535; + // key is column uid, value is the sparse column cache std::unordered_map _variant_sparse_column_cache; diff --git a/be/src/storage/segment/segment_writer.cpp b/be/src/storage/segment/segment_writer.cpp index aa480f218c58a7..5ba6910f5bd3dc 100644 --- a/be/src/storage/segment/segment_writer.cpp +++ b/be/src/storage/segment/segment_writer.cpp @@ -227,6 +227,8 @@ Status SegmentWriter::_create_column_writer(uint32_t cid, const TabletColumn& co if (_opts.write_type == DataWriteType::TYPE_DIRECT && schema->skip_write_index_on_load()) { skip_inverted_index = true; } + // Let SNII select the direct-load PRX zstd level. + opts.is_direct_load = _opts.write_type == DataWriteType::TYPE_DIRECT; // indexes for this column if (!skip_inverted_index) { auto inverted_indexs = schema->inverted_indexs(column); diff --git a/be/src/storage/segment/variant/variant_column_writer_impl.cpp b/be/src/storage/segment/variant/variant_column_writer_impl.cpp index 1e8610f42d886f..165d9ed1ca9293 100644 --- a/be/src/storage/segment/variant/variant_column_writer_impl.cpp +++ b/be/src/storage/segment/variant/variant_column_writer_impl.cpp @@ -455,6 +455,8 @@ Status prepare_materialized_subcolumn_writer( opts.rowset_ctx = base_opts.rowset_ctx; opts.file_writer = base_opts.file_writer; opts.storage_format = base_opts.storage_format; + // keep the segment writer's direct-load marking for subcolumn index writers + opts.is_direct_load = base_opts.is_direct_load; std::unique_ptr writer; variant_util::inherit_column_attributes(parent_column, tablet_column); @@ -1482,6 +1484,8 @@ Status prepare_subcolumn_writer_target( opts.rowset_ctx = base_opts.rowset_ctx; opts.file_writer = base_opts.file_writer; opts.storage_format = base_opts.storage_format; + // keep the segment writer's direct-load marking for subcolumn index writers + opts.is_direct_load = base_opts.is_direct_load; variant_util::inherit_column_attributes(parent_column, tablet_column); bool need_record_none_null_value_size = diff --git a/be/src/storage/segment/vertical_segment_writer.cpp b/be/src/storage/segment/vertical_segment_writer.cpp index f5b60ebadb756c..9f03745d3e8810 100644 --- a/be/src/storage/segment/vertical_segment_writer.cpp +++ b/be/src/storage/segment/vertical_segment_writer.cpp @@ -237,6 +237,8 @@ Status VerticalSegmentWriter::_create_column_writer(uint32_t cid, const TabletCo tablet_schema->skip_write_index_on_load()) { skip_inverted_index = true; } + // Let SNII select the direct-load PRX zstd level. + opts.is_direct_load = _opts.write_type == DataWriteType::TYPE_DIRECT; if (!skip_inverted_index) { auto inverted_indexs = tablet_schema->inverted_indexs(column); if (!inverted_indexs.empty()) { diff --git a/be/src/storage/tablet/tablet_meta.cpp b/be/src/storage/tablet/tablet_meta.cpp index 557fe02743b178..c5b191481f1d26 100644 --- a/be/src/storage/tablet/tablet_meta.cpp +++ b/be/src/storage/tablet/tablet_meta.cpp @@ -504,6 +504,9 @@ void TabletMeta::init_schema_from_thrift(const TTabletSchema& tablet_schema, case TInvertedIndexFileStorageFormat::V3: tablet_schema_pb->set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V3); break; + case TInvertedIndexFileStorageFormat::SNII: + tablet_schema_pb->set_inverted_index_storage_format(InvertedIndexStorageFormatPB::SNII); + break; default: tablet_schema_pb->set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V3); break; diff --git a/be/src/storage/task/index_builder.cpp b/be/src/storage/task/index_builder.cpp index ef49626e143ab5..c082e452f2ee8c 100644 --- a/be/src/storage/task/index_builder.cpp +++ b/be/src/storage/task/index_builder.cpp @@ -338,6 +338,13 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta if (_is_drop_op) { const auto& output_rs_tablet_schema = output_rowset_meta->tablet_schema(); + if (output_rs_tablet_schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + LOG(INFO) << "skip physical SNII inverted index rewrite for drop index. tablet_id=" + << _tablet->tablet_id() + << " rowset_id=" << output_rowset_meta->rowset_id().to_string(); + return Status::OK(); + } if (output_rs_tablet_schema->get_inverted_index_storage_format() != InvertedIndexStorageFormatPB::V1) { const auto& fs = output_rowset_meta->fs(); @@ -421,6 +428,11 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta _olap_data_convertor->reserve(_alter_inverted_indexes.size()); std::unique_ptr index_file_writer = nullptr; + if (output_rowset_schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + return Status::Error( + "BUILD INDEX is not supported for SNII inverted index storage format yet"); + } if (output_rowset_schema->get_inverted_index_storage_format() >= InvertedIndexStorageFormatPB::V2) { auto idx_file_reader_iter = _index_file_readers.find( @@ -516,8 +528,10 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta if (inverted_index_builder) { auto writer_sign = std::make_pair(seg_ptr->id(), index_id); - _index_column_writers.insert( + auto [index_column_writer_it, inserted] = _index_column_writers.insert( std::make_pair(writer_sign, std::move(inverted_index_builder))); + DORIS_CHECK(inserted); + DORIS_CHECK(index_column_writer_it->second != nullptr); inverted_index_writer_signs.emplace_back(writer_sign); } } @@ -545,8 +559,10 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta if (index_writer) { auto writer_sign = std::make_pair(seg_ptr->id(), index_id); - _index_column_writers.insert( + auto [index_column_writer_it, inserted] = _index_column_writers.insert( std::make_pair(writer_sign, std::move(index_writer))); + DORIS_CHECK(inserted); + DORIS_CHECK(index_column_writer_it->second != nullptr); inverted_index_writer_signs.emplace_back(writer_sign); } } @@ -554,7 +570,10 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta } // DO NOT forget index_file_writer for the segment, otherwise, original inverted index will be deleted. - _index_file_writers.emplace(seg_ptr->id(), std::move(index_file_writer)); + auto [index_file_writer_it, inserted] = + _index_file_writers.emplace(seg_ptr->id(), std::move(index_file_writer)); + DORIS_CHECK(inserted); + DORIS_CHECK(index_file_writer_it->second != nullptr); if (return_columns.empty()) { // no columns to read continue; @@ -605,8 +624,7 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta "handle_single_rowset_write_inverted_index_data_error"); }) if (!status.ok()) { - return Status::Error( - "failed to write block."); + return status; } block->clear_column_data(); } @@ -614,9 +632,10 @@ Status IndexBuilder::handle_single_rowset(RowsetMetaSharedPtr output_rowset_meta // finish write inverted index, flush data to compound file for (auto& writer_sign : inverted_index_writer_signs) { try { - if (_index_column_writers[writer_sign]) { - RETURN_IF_ERROR(_index_column_writers[writer_sign]->finish()); - } + auto index_column_writer_it = _index_column_writers.find(writer_sign); + DORIS_CHECK(index_column_writer_it != _index_column_writers.end()); + DORIS_CHECK(index_column_writer_it->second != nullptr); + RETURN_IF_ERROR(index_column_writer_it->second->finish()); DBUG_EXECUTE_IF("IndexBuilder::handle_single_rowset_index_build_finish_error", { _CLTHROWA(CL_ERR_IO, "debug point: handle_single_rowset_index_build_finish_error"); @@ -714,6 +733,11 @@ Status IndexBuilder::_add_nullable(const std::string& column_name, const std::pair& index_writer_sign, const TabletColumn* column, const uint8_t* null_map, const uint8_t** ptr, size_t num_rows) { + auto index_column_writer_it = _index_column_writers.find(index_writer_sign); + DORIS_CHECK(index_column_writer_it != _index_column_writers.end()); + DORIS_CHECK(index_column_writer_it->second != nullptr); + auto* index_column_writer = index_column_writer_it->second.get(); + // TODO: need to process null data for inverted index if (column->type() == FieldType::OLAP_FIELD_TYPE_ARRAY) { DCHECK(column->get_subtype_count() == 1); @@ -725,15 +749,14 @@ Status IndexBuilder::_add_nullable(const std::string& column_name, try { auto data = *(data_ptr + 2); auto nested_null_map = *(data_ptr + 3); - RETURN_IF_ERROR(_index_column_writers[index_writer_sign]->add_array_values( + RETURN_IF_ERROR(index_column_writer->add_array_values( field_type_size(column->get_sub_column(0).type()), reinterpret_cast(data), reinterpret_cast(nested_null_map), offsets_ptr, num_rows)); DBUG_EXECUTE_IF("IndexBuilder::_add_nullable_add_array_values_error", { _CLTHROWA(CL_ERR_IO, "debug point: _add_nullable_add_array_values_error"); }) - RETURN_IF_ERROR( - _index_column_writers[index_writer_sign]->add_array_nulls(null_map, num_rows)); + RETURN_IF_ERROR(index_column_writer->add_array_nulls(null_map, num_rows)); } catch (const std::exception& e) { return Status::Error( "CLuceneError occurred: {}", e.what()); @@ -757,11 +780,9 @@ Status IndexBuilder::_add_nullable(const std::string& column_name, do { auto step = next_run_step(); if (null_map[offset]) { - RETURN_IF_ERROR(_index_column_writers[index_writer_sign]->add_nulls( - static_cast(step))); + RETURN_IF_ERROR(index_column_writer->add_nulls(static_cast(step))); } else { - RETURN_IF_ERROR(_index_column_writers[index_writer_sign]->add_values(column_name, - *ptr, step)); + RETURN_IF_ERROR(index_column_writer->add_values(column_name, *ptr, step)); } *ptr += field_type_size(column->type()) * step; offset += step; @@ -779,6 +800,11 @@ Status IndexBuilder::_add_nullable(const std::string& column_name, Status IndexBuilder::_add_data(const std::string& column_name, const std::pair& index_writer_sign, const TabletColumn* column, const uint8_t** ptr, size_t num_rows) { + auto index_column_writer_it = _index_column_writers.find(index_writer_sign); + DORIS_CHECK(index_column_writer_it != _index_column_writers.end()); + DORIS_CHECK(index_column_writer_it->second != nullptr); + auto* index_column_writer = index_column_writer_it->second.get(); + try { if (column->type() == FieldType::OLAP_FIELD_TYPE_ARRAY) { DCHECK(column->get_subtype_count() == 1); @@ -791,14 +817,13 @@ Status IndexBuilder::_add_data(const std::string& column_name, if (element_cnt > 0) { auto data = *(data_ptr + 2); auto nested_null_map = *(data_ptr + 3); - RETURN_IF_ERROR(_index_column_writers[index_writer_sign]->add_array_values( + RETURN_IF_ERROR(index_column_writer->add_array_values( field_type_size(column->get_sub_column(0).type()), reinterpret_cast(data), reinterpret_cast(nested_null_map), offsets_ptr, num_rows)); } } else { - RETURN_IF_ERROR(_index_column_writers[index_writer_sign]->add_values(column_name, *ptr, - num_rows)); + RETURN_IF_ERROR(index_column_writer->add_values(column_name, *ptr, num_rows)); } DBUG_EXECUTE_IF("IndexBuilder::_add_data_throw_exception", { _CLTHROWA(CL_ERR_IO, "debug point: _add_data_throw_exception"); }) diff --git a/be/test/exec/scan/scanner_context_test.cpp b/be/test/exec/scan/scanner_context_test.cpp index 8be021454dd7b6..5af5869dd65563 100644 --- a/be/test/exec/scan/scanner_context_test.cpp +++ b/be/test/exec/scan/scanner_context_test.cpp @@ -38,6 +38,10 @@ #include "exec/scan/scanner_scheduler.h" #include "runtime/descriptors.h" #include "runtime/query_context.h" +#include "storage/options.h" +#include "storage/storage_engine.h" +#include "storage/tablet/tablet.h" +#include "storage/tablet/tablet_meta.h" #include "testutil/mock/mock_runtime_state.h" namespace doris { @@ -184,6 +188,165 @@ TEST_F(ScannerContextTest, test_init) { ASSERT_TRUE(st.ok()); } +TEST_F(ScannerContextTest, inverted_index_profile_collection_is_additive_and_idempotent) { + auto engine = std::make_unique(EngineOptions {}); + auto tablet_meta = std::make_shared(1, 2, 15673, 15674, 4, 5, TTabletSchema {}, 6, + std::unordered_map {{7, 8}}, + UniqueId(9, 10), TTabletType::TABLET_TYPE_DISK, + TCompressionType::LZ4F); + auto tablet = std::make_shared(*engine, std::move(tablet_meta), nullptr); + const int parallel_tasks = 1; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto local_state = OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + const std::vector scan_ranges; + const std::map, + std::vector>>> + shared_state_map; + LocalStateInfo local_state_info {profile.get(), scan_ranges, nullptr, shared_state_map, 0}; + const Status init_status = local_state->init(state.get(), local_state_info); + ASSERT_TRUE(init_status.ok()) << init_status.to_string(); + + auto make_scanner = [&]() { + OlapScanner::Params params; + params.state = state.get(); + params.profile = profile.get(); + params.version = 0; + params.limit = -1; + params.aggregation = false; + return OlapScanner::create_shared(local_state.get(), std::move(params)); + }; + auto scanner1 = make_scanner(); + auto scanner2 = make_scanner(); + scanner1->_tablet_reader_params.tablet = tablet; + scanner2->_tablet_reader_params.tablet = tablet; + scanner1->_tablet_reader = std::make_unique(); + scanner2->_tablet_reader = std::make_unique(); + auto* stats1 = scanner1->_tablet_reader->mutable_stats(); + stats1->snii_stats.prx_raw_frames = 1; + stats1->snii_stats.prx_plaintext_bytes = 10; + stats1->snii_stats.prx_decode_ns = 100; + stats1->snii_stats.phrase_candidate_docs = 3; + stats1->snii_stats.common_grams_gram_plans = 1; + stats1->snii_stats.common_grams_fallback_kill_switch = 5; + stats1->snii_stats.common_grams_plain_posting_bytes = 10; + stats1->snii_stats.common_grams_gram_posting_bytes = 20; + stats1->snii_stats.common_grams_plain_estimated_candidate_df = 30; + stats1->snii_stats.common_grams_gram_estimated_candidate_df = 40; + stats1->snii_stats.common_grams_plain_estimated_cost = 50; + stats1->snii_stats.common_grams_gram_estimated_cost = 60; + stats1->snii_stats.common_grams_fallback_base_analyzer_mismatch = 61; + stats1->snii_stats.common_grams_fallback_prefix_tail_empty = 62; + stats1->snii_stats.common_grams_planning_ns = 65; + auto* stats2 = scanner2->_tablet_reader->mutable_stats(); + stats2->snii_stats.prx_raw_frames = 2; + stats2->snii_stats.prx_plaintext_bytes = 20; + stats2->snii_stats.prx_decode_ns = 200; + stats2->snii_stats.phrase_candidate_docs = 4; + stats2->snii_stats.common_grams_gram_plans = 2; + stats2->snii_stats.common_grams_fallback_kill_switch = 6; + stats2->snii_stats.common_grams_plain_posting_bytes = 1; + stats2->snii_stats.common_grams_gram_posting_bytes = 2; + stats2->snii_stats.common_grams_plain_estimated_candidate_df = 3; + stats2->snii_stats.common_grams_gram_estimated_candidate_df = 4; + stats2->snii_stats.common_grams_plain_estimated_cost = 5; + stats2->snii_stats.common_grams_gram_estimated_cost = 6; + stats2->snii_stats.common_grams_fallback_base_analyzer_mismatch = 7; + stats2->snii_stats.common_grams_fallback_prefix_tail_empty = 8; + stats2->snii_stats.common_grams_planning_ns = 11; + + RuntimeProfile* index_filter = local_state->_index_filter_profile.get(); + ASSERT_NE(index_filter, nullptr); + auto* raw_frames = index_filter->get_counter("SniiPrxRawFrames"); + auto* plaintext_bytes = index_filter->get_counter("SniiPrxPlaintextBytes"); + auto* decode_time = index_filter->get_counter("SniiPrxInclusiveDecodeTime"); + auto* phrase_candidate_docs = index_filter->get_counter("SniiPhraseCandidateDocs"); + auto* common_grams_gram_plans = index_filter->get_counter("SniiCommonGramsGramPlans"); + auto* common_grams_fallback_kill_switch = + index_filter->get_counter("SniiCommonGramsFallbackKillSwitch"); + struct ExpectedSniiCounter { + const char* name; + RuntimeProfile::Counter* counter; + int64_t scanner1_value; + int64_t combined_value; + }; + const ExpectedSniiCounter snii_counters[] = { + {"SniiCommonGramsPlainPostingBytes", + index_filter->get_counter("SniiCommonGramsPlainPostingBytes"), 10, 11}, + {"SniiCommonGramsGramPostingBytes", + index_filter->get_counter("SniiCommonGramsGramPostingBytes"), 20, 22}, + {"SniiCommonGramsPlainEstimatedCandidateDf", + index_filter->get_counter("SniiCommonGramsPlainEstimatedCandidateDf"), 30, 33}, + {"SniiCommonGramsGramEstimatedCandidateDf", + index_filter->get_counter("SniiCommonGramsGramEstimatedCandidateDf"), 40, 44}, + {"SniiCommonGramsPlainEstimatedCost", + index_filter->get_counter("SniiCommonGramsPlainEstimatedCost"), 50, 55}, + {"SniiCommonGramsGramEstimatedCost", + index_filter->get_counter("SniiCommonGramsGramEstimatedCost"), 60, 66}, + {"SniiCommonGramsFallbackBaseAnalyzerMismatch", + index_filter->get_counter("SniiCommonGramsFallbackBaseAnalyzerMismatch"), 61, 68}, + {"SniiCommonGramsFallbackPrefixTailEmpty", + index_filter->get_counter("SniiCommonGramsFallbackPrefixTailEmpty"), 62, 70}, + {"SniiCommonGramsPlanningTime", + index_filter->get_counter("SniiCommonGramsPlanningTime"), 65, 76}, + }; + + std::vector zero_nodes; + index_filter->to_thrift(&zero_nodes); + ASSERT_EQ(zero_nodes.size(), 1U); + for (const auto& expected : snii_counters) { + bool serialized = false; + for (const auto& thrift_counter : zero_nodes.front().counters) { + serialized |= thrift_counter.name == expected.name; + } + EXPECT_FALSE(serialized) << expected.name; + } + ASSERT_NE(raw_frames, nullptr); + ASSERT_NE(plaintext_bytes, nullptr); + ASSERT_NE(decode_time, nullptr); + ASSERT_NE(phrase_candidate_docs, nullptr); + ASSERT_NE(common_grams_gram_plans, nullptr); + ASSERT_NE(common_grams_fallback_kill_switch, nullptr); + for (const auto& expected : snii_counters) { + ASSERT_NE(expected.counter, nullptr) << expected.name; + EXPECT_NE(dynamic_cast(expected.counter), nullptr) + << expected.name; + } + + scanner1->_collect_profile_before_close(); + EXPECT_EQ(raw_frames->value(), 1); + EXPECT_EQ(plaintext_bytes->value(), 10); + EXPECT_EQ(decode_time->value(), 100); + EXPECT_EQ(phrase_candidate_docs->value(), 3); + EXPECT_EQ(common_grams_gram_plans->value(), 1); + EXPECT_EQ(common_grams_fallback_kill_switch->value(), 5); + for (const auto& expected : snii_counters) { + EXPECT_EQ(expected.counter->value(), expected.scanner1_value) << expected.name; + } + + scanner1->_collect_profile_before_close(); + EXPECT_EQ(raw_frames->value(), 1); + EXPECT_EQ(plaintext_bytes->value(), 10); + EXPECT_EQ(decode_time->value(), 100); + EXPECT_EQ(phrase_candidate_docs->value(), 3); + EXPECT_EQ(common_grams_gram_plans->value(), 1); + EXPECT_EQ(common_grams_fallback_kill_switch->value(), 5); + for (const auto& expected : snii_counters) { + EXPECT_EQ(expected.counter->value(), expected.scanner1_value) << expected.name; + } + + scanner2->_collect_profile_before_close(); + EXPECT_EQ(raw_frames->value(), 3); + EXPECT_EQ(plaintext_bytes->value(), 30); + EXPECT_EQ(decode_time->value(), 300); + EXPECT_EQ(phrase_candidate_docs->value(), 7); + EXPECT_EQ(common_grams_gram_plans->value(), 3); + EXPECT_EQ(common_grams_fallback_kill_switch->value(), 11); + for (const auto& expected : snii_counters) { + EXPECT_EQ(expected.counter->value(), expected.combined_value) << expected.name; + } +} + TEST_F(ScannerContextTest, test_serial_run) { const int parallel_tasks = 1; auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, diff --git a/be/test/exprs/function/function_is_null_test.cpp b/be/test/exprs/function/function_is_null_test.cpp index c92c89645710de..efa3e938f40972 100644 --- a/be/test/exprs/function/function_is_null_test.cpp +++ b/be/test/exprs/function/function_is_null_test.cpp @@ -76,11 +76,18 @@ class FunctionIsNullTest : public ::testing::Test { _inverted_index_query_cache = std::unique_ptr( InvertedIndexQueryCache::create_global_cache(inverted_index_cache_limit, 1)); + // Both caches are owned by this fixture, so the previous globals must come back in + // TearDown -- otherwise ExecEnv keeps pointing at them after the fixture is destroyed and + // the next test that reaches InvertedIndexQueryCache::instance() reads freed memory. + _previous_searcher_cache = ExecEnv::GetInstance()->get_inverted_index_searcher_cache(); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); ExecEnv::GetInstance()->set_inverted_index_searcher_cache( _inverted_index_searcher_cache.get()); - ExecEnv::GetInstance()->_inverted_index_query_cache = _inverted_index_query_cache.get(); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_inverted_index_query_cache.get()); } void TearDown() override { + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_previous_searcher_cache); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); EXPECT_TRUE(io::global_local_filesystem()->delete_directory(_tablet->tablet_path()).ok()); EXPECT_TRUE(io::global_local_filesystem()->delete_directory(_absolute_dir).ok()); EXPECT_TRUE(io::global_local_filesystem()->delete_directory(tmp_dir).ok()); @@ -134,6 +141,8 @@ class FunctionIsNullTest : public ::testing::Test { std::string _absolute_dir; std::string _curreent_dir; TabletSchemaPB _schema_pb; + InvertedIndexSearcherCache* _previous_searcher_cache = nullptr; + InvertedIndexQueryCache* _previous_query_cache = nullptr; std::unique_ptr _inverted_index_searcher_cache; std::unique_ptr _inverted_index_query_cache; }; diff --git a/be/test/exprs/function/function_search_test.cpp b/be/test/exprs/function/function_search_test.cpp index ac57847f32e4de..74cd72dcb79ea6 100644 --- a/be/test/exprs/function/function_search_test.cpp +++ b/be/test/exprs/function/function_search_test.cpp @@ -21,17 +21,22 @@ #include #include +#include #include #include #include +#include #include #include +#include #include "core/block/block.h" #include "core/data_type/data_type_array.h" #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_string.h" #include "core/data_type/primitive_type.h" +#include "runtime/exec_env.h" +#include "runtime/index_policy/index_policy_mgr.h" #include "storage/index/index_file_reader.h" #include "storage/index/index_iterator.h" #include "storage/index/inverted/inverted_index_iterator.h" @@ -40,6 +45,8 @@ #include "storage/index/inverted/query_v2/phrase_query/multi_phrase_weight.h" #include "storage/index/inverted/query_v2/phrase_query/phrase_query.h" #include "storage/segment/variant/nested_group_provider.h" +#include "util/defer_op.h" +#include "util/thrift_util.h" namespace doris { @@ -115,6 +122,27 @@ class RecordingIndexIterator : public segment_v2::IndexIterator { Int32 last_int_value = 0; }; +class RecordingDirectInvertedIndexIterator final : public segment_v2::InvertedIndexIterator { +public: + Status read_from_index(const segment_v2::IndexParam& param) override { + ++read_calls; + auto* inverted_param = std::get_if(¶m); + DORIS_CHECK(inverted_param != nullptr); + DORIS_CHECK(*inverted_param != nullptr); + DORIS_CHECK((*inverted_param)->roaring != nullptr); + (*inverted_param)->roaring->add(3); + return Status::OK(); + } + + Status read_null_bitmap(segment_v2::InvertedIndexQueryCacheHandle* /*cache_handle*/) override { + return Status::OK(); + } + + Result has_null() override { return false; } + + int read_calls = 0; +}; + class DummyInvertedIndexReader final : public segment_v2::InvertedIndexReader { public: explicit DummyInvertedIndexReader(const TabletIndex* index_meta) @@ -151,6 +179,185 @@ class DummyInvertedIndexReader final : public segment_v2::InvertedIndexReader { segment_v2::InvertedIndexReaderType _reader_type = segment_v2::InvertedIndexReaderType::BKD; }; +class RejectingCluceneIndexFileReader final : public segment_v2::IndexFileReader { +public: + explicit RejectingCluceneIndexFileReader( + InvertedIndexStorageFormatPB storage_format = InvertedIndexStorageFormatPB::SNII, + const std::string& index_path = "/tmp/search_snii_native_idx") + : segment_v2::IndexFileReader(nullptr, index_path, storage_format) {} + + Status init(int32_t /*read_buffer_size*/, const io::IOContext* /*io_ctx*/) override { + ++init_calls; + return Status::OK(); + } + + Result> open( + const TabletIndex* /*index_meta*/, const io::IOContext* /*io_ctx*/) const override { + ++open_calls; + return ResultError(Status::InternalError("unexpected CLucene open for SNII search")); + } + + int init_calls = 0; + mutable int open_calls = 0; +}; + +class RecordingNativeInvertedIndexReader final : public segment_v2::InvertedIndexReader { +public: + RecordingNativeInvertedIndexReader( + const TabletIndex* index_meta, + const std::shared_ptr& index_file_reader, + segment_v2::InvertedIndexReaderType reader_type = + segment_v2::InvertedIndexReaderType::FULLTEXT) + : segment_v2::InvertedIndexReader(index_meta, index_file_reader), + _reader_type(reader_type), + _null_cache(1024 * 1024, 1), + _null_cache_key {"/tmp/search_snii_native_null", "", + segment_v2::InvertedIndexQueryType::UNKNOWN_QUERY, + std::to_string(index_meta->index_id())} { + set_has_null(false); + } + + Status new_iterator(std::unique_ptr* /*iterator*/) override { + return Status::OK(); + } + + Status query(const segment_v2::IndexQueryContextPtr& /*context*/, + const std::string& column_name, const Field& query_value, + segment_v2::InvertedIndexQueryType query_type, + std::shared_ptr& bit_map, + const InvertedIndexAnalyzerCtx* analyzer_ctx = nullptr) override { + ++query_calls; + last_column_name = column_name; + last_query_type = query_type; + last_query_value_type = query_value.get_type(); + last_analyzer_ctx = analyzer_ctx; + if (last_query_value_type == TYPE_STRING) { + last_query_value = query_value.get(); + } + + bit_map = std::make_shared(); + auto result_it = query_results.find(last_query_value); + if (result_it != query_results.end()) { + *bit_map = result_it->second; + } + return Status::OK(); + } + + Status try_query(const segment_v2::IndexQueryContextPtr& /*context*/, + const std::string& /*column_name*/, const Field& /*query_value*/, + segment_v2::InvertedIndexQueryType /*query_type*/, + size_t* /*count*/) override { + return Status::OK(); + } + + Status read_null_bitmap(const segment_v2::IndexQueryContextPtr& /*context*/, + segment_v2::InvertedIndexQueryCacheHandle* cache_handle, + lucene::store::Directory* /*dir*/ = nullptr) override { + ++null_bitmap_calls; + _null_cache.insert(_null_cache_key, std::make_shared(_null_bitmap), + cache_handle); + return Status::OK(); + } + + segment_v2::InvertedIndexReaderType type() override { return _reader_type; } + + void set_query_result(const std::string& pattern, roaring::Roaring result) { + query_results[pattern] = std::move(result); + } + + void set_null_bitmap(roaring::Roaring null_bitmap) { + _null_bitmap = std::move(null_bitmap); + set_has_null(!_null_bitmap.isEmpty()); + } + + int query_calls = 0; + int null_bitmap_calls = 0; + std::string last_column_name; + std::string last_query_value; + PrimitiveType last_query_value_type = PrimitiveType::TYPE_NULL; + segment_v2::InvertedIndexQueryType last_query_type = + segment_v2::InvertedIndexQueryType::UNKNOWN_QUERY; + const InvertedIndexAnalyzerCtx* last_analyzer_ctx = nullptr; + std::unordered_map query_results; + +private: + segment_v2::InvertedIndexReaderType _reader_type; + roaring::Roaring _null_bitmap; + segment_v2::InvertedIndexQueryCache _null_cache; + segment_v2::InvertedIndexQueryCache::CacheKey _null_cache_key; +}; + +class ScopedInvertedIndexQueryCache final { +public: + ScopedInvertedIndexQueryCache() + : _previous(ExecEnv::GetInstance()->get_inverted_index_query_cache()), + _cache(segment_v2::InvertedIndexQueryCache::create_global_cache(1024 * 1024, 1)) { + ExecEnv::GetInstance()->set_inverted_index_query_cache(_cache.get()); + } + + ~ScopedInvertedIndexQueryCache() { + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous); + } + + segment_v2::InvertedIndexQueryCache* get() const { return _cache.get(); } + +private: + segment_v2::InvertedIndexQueryCache* _previous; + std::unique_ptr _cache; +}; + +static roaring::Roaring make_bitmap(std::initializer_list docs) { + roaring::Roaring bitmap; + for (uint32_t doc : docs) { + bitmap.add(doc); + } + return bitmap; +} + +static void expect_bitmap_eq(const roaring::Roaring& actual, + std::initializer_list expected_docs) { + auto expected = make_bitmap(expected_docs); + EXPECT_EQ(expected.cardinality(), actual.cardinality()); + EXPECT_TRUE(actual == expected); +} + +static roaring::Roaring collect_docs( + const segment_v2::inverted_index::query_v2::ScorerPtr& scorer) { + roaring::Roaring docs; + for (uint32_t doc = scorer->doc(); doc != segment_v2::inverted_index::query_v2::TERMINATED; + doc = scorer->advance()) { + docs.add(doc); + } + return docs; +} + +static TSearchClause make_leaf_clause(const std::string& clause_type, const std::string& value) { + TSearchClause clause; + clause.clause_type = clause_type; + clause.field_name = "body"; + clause.value = value; + clause.__isset.field_name = true; + clause.__isset.value = true; + return clause; +} + +static Status insert_search_dsl_cache( + segment_v2::InvertedIndexQueryCache* cache, + const std::shared_ptr& index_file_reader, + const TSearchParam& search_param, roaring::Roaring bitmap) { + ThriftSerializer serializer(false, 1024); + TSearchParam copy = search_param; + std::string signature; + RETURN_IF_ERROR(serializer.serialize(©, &signature)); + + segment_v2::InvertedIndexQueryCache::CacheKey key { + index_file_reader->get_index_path_prefix(), "__search_dsl__", + segment_v2::InvertedIndexQueryType::SEARCH_DSL_QUERY, std::move(signature)}; + segment_v2::InvertedIndexQueryCacheHandle handle; + cache->insert(key, std::make_shared(std::move(bitmap)), &handle); + return Status::OK(); +} + static TabletIndex make_test_inverted_index( int64_t index_id, const std::map& properties = {}) { TabletIndex index_meta; @@ -166,6 +373,32 @@ static TabletIndex make_test_inverted_index( return index_meta; } +static Status resolve_non_variant_binding_with_mismatched_analyzer(const DataTypePtr& column_type) { + std::map index_properties; + index_properties[INVERTED_INDEX_PARSER_KEY] = INVERTED_INDEX_PARSER_STANDARD; + auto index_meta = make_test_inverted_index(13, index_properties); + auto reader = std::make_shared( + &index_meta, nullptr, segment_v2::InvertedIndexReaderType::FULLTEXT); + + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace("content", IndexFieldNameAndTypePair {"content", column_type}); + std::unordered_map iterators; + iterators["content"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "content"; + field_binding.index_properties[INVERTED_INDEX_PARSER_KEY] = INVERTED_INDEX_PARSER_ENGLISH; + field_binding.__isset.index_properties = true; + + auto context = std::make_shared(); + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + FieldReaderBinding binding; + return resolver.resolve("content", InvertedIndexQueryType::MATCH_ANY_QUERY, &binding); +} + TEST_F(FunctionSearchTest, TestGetName) { EXPECT_EQ("search", function_search->get_name()); } @@ -1716,6 +1949,7 @@ TEST_F(FunctionSearchTest, TestBuildLeafQueryPhrase) { binding.stored_field_wstr = L"content"; binding.index_properties["parser"] = "unicode"; binding.query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY; + binding.execution_mode = SearchFieldExecutionMode::CLUCENE; auto* dummy_reader = reinterpret_cast(0x1); binding.lucene_reader = std::shared_ptr( @@ -1736,6 +1970,80 @@ TEST_F(FunctionSearchTest, TestBuildLeafQueryPhrase) { EXPECT_NE(phrase_query, nullptr); } +TEST_F(FunctionSearchTest, TestBuildLeafQueryPhraseUsesPlainTerms) { + auto* exec_env = ExecEnv::GetInstance(); + auto* previous_policy_mgr = exec_env->index_policy_mgr(); + IndexPolicyMgr scoped_policy_mgr; + exec_env->_index_policy_mgr = &scoped_policy_mgr; + DEFER(exec_env->_index_policy_mgr = previous_policy_mgr); + + auto* policy_mgr = exec_env->index_policy_mgr(); + ASSERT_NE(policy_mgr, nullptr); + + TIndexPolicy tokenizer; + tokenizer.id = 910020; + tokenizer.name = "function_search_cg_tokenizer"; + tokenizer.type = TIndexPolicyType::TOKENIZER; + tokenizer.properties["type"] = "char_group"; + tokenizer.properties["tokenize_on_chars"] = "[whitespace]"; + + TIndexPolicy common_grams; + common_grams.id = 910021; + common_grams.name = "function_search_cg_filter"; + common_grams.type = TIndexPolicyType::TOKEN_FILTER; + common_grams.properties["type"] = "common_grams"; + + TIndexPolicy analyzer; + analyzer.id = 910022; + analyzer.name = "function_search_cg_analyzer"; + analyzer.type = TIndexPolicyType::ANALYZER; + analyzer.properties["tokenizer"] = tokenizer.name; + analyzer.properties["token_filter"] = "lowercase," + common_grams.name; + policy_mgr->apply_policy_changes({tokenizer, common_grams, analyzer}, {}); + + TSearchClause clause; + clause.clause_type = "PHRASE"; + clause.field_name = "content"; + clause.value = "man of the year"; + clause.__isset.field_name = true; + clause.__isset.value = true; + + auto context = std::make_shared(); + std::unordered_map data_type_with_names; + data_type_with_names.emplace("content", IndexFieldNameAndTypePair {"content", nullptr}); + std::unordered_map iterators; + FieldReaderResolver resolver(data_type_with_names, iterators, context); + + FieldReaderBinding binding; + binding.logical_field_name = "content"; + binding.stored_field_name = "content"; + binding.stored_field_wstr = L"content"; + binding.index_properties["analyzer"] = analyzer.name; + binding.query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY; + binding.execution_mode = SearchFieldExecutionMode::CLUCENE; + auto* dummy_reader = reinterpret_cast(0x1); + binding.lucene_reader = std::shared_ptr( + dummy_reader, [](lucene::index::IndexReader* /*ptr*/) {}); + binding.binding_key = + resolver.binding_key_for("content", InvertedIndexQueryType::MATCH_PHRASE_QUERY); + resolver._cache[binding.binding_key] = binding; + + inverted_index::query_v2::QueryPtr query; + std::string binding_key; + ASSERT_TRUE(function_search + ->build_leaf_query(clause, context, resolver, &query, &binding_key, "OR", 0) + .ok()); + + auto phrase = std::dynamic_pointer_cast(query); + ASSERT_NE(phrase, nullptr); + ASSERT_EQ(phrase->_term_infos.size(), 4); + EXPECT_EQ(phrase->_term_infos[0].get_single_term(), "man"); + EXPECT_EQ(phrase->_term_infos[1].get_single_term(), "of"); + EXPECT_EQ(phrase->_term_infos[2].get_single_term(), "the"); + EXPECT_EQ(phrase->_term_infos[3].get_single_term(), "year"); + policy_mgr->apply_policy_changes({}, {tokenizer.id, common_grams.id, analyzer.id}); +} + TEST_F(FunctionSearchTest, TestBuildLeafQueryVariantMissingFieldReturnsUnknown) { TSearchClause clause; clause.clause_type = "TERM"; @@ -1872,6 +2180,64 @@ TEST_F(FunctionSearchTest, EXPECT_EQ(ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND, status.code()); } +TEST_F(FunctionSearchTest, + TestFieldReaderResolverNonVariantStringBindingRejectsMismatchedAnalyzer) { + auto status = resolve_non_variant_binding_with_mismatched_analyzer( + std::make_shared()); + + ASSERT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_BYPASS, status.code()); +} + +TEST_F(FunctionSearchTest, + TestFieldReaderResolverNonVariantArrayStringBindingRejectsMismatchedAnalyzer) { + auto column_type = + std::make_shared(make_nullable(std::make_shared())); + auto status = resolve_non_variant_binding_with_mismatched_analyzer(column_type); + + ASSERT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_BYPASS, status.code()); +} + +TEST_F(FunctionSearchTest, TestFieldReaderResolverExactIgnoresAnalyzedBindingHint) { + std::map analyzed_properties; + analyzed_properties[INVERTED_INDEX_PARSER_KEY] = INVERTED_INDEX_PARSER_STANDARD; + auto analyzed_index = make_test_inverted_index(13, analyzed_properties); + auto keyword_index = make_test_inverted_index(14); + auto index_file_reader = std::make_shared( + nullptr, "/tmp/search_exact_multi_index", InvertedIndexStorageFormatPB::SNII); + auto analyzed_reader = std::make_shared( + &analyzed_index, index_file_reader, segment_v2::InvertedIndexReaderType::FULLTEXT); + auto keyword_reader = std::make_shared( + &keyword_index, index_file_reader, segment_v2::InvertedIndexReaderType::STRING_TYPE); + + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, analyzed_reader); + iterator.add_reader(segment_v2::InvertedIndexReaderType::STRING_TYPE, keyword_reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "content", IndexFieldNameAndTypePair {"content", std::make_shared()}); + std::unordered_map iterators; + iterators["content"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "content"; + field_binding.index_properties = analyzed_properties; + field_binding.__isset.index_properties = true; + + auto context = std::make_shared(); + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + FieldReaderBinding binding; + auto status = resolver.resolve("content", InvertedIndexQueryType::EQUAL_QUERY, &binding); + + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(binding.inverted_reader, nullptr); + EXPECT_EQ(binding.inverted_reader->get_index_id(), 14); + EXPECT_EQ(binding.query_type, InvertedIndexQueryType::EQUAL_QUERY); + EXPECT_TRUE(binding.index_properties.empty()); +} + TEST_F(FunctionSearchTest, TestFieldReaderResolverVariantBkdDirectReader) { auto context = std::make_shared(); @@ -1913,6 +2279,345 @@ TEST_F(FunctionSearchTest, TestFieldReaderResolverVariantBkdDirectReader) { EXPECT_TRUE(cache.begin()->second.use_direct_index_reader()); } +TEST_F(FunctionSearchTest, TestFieldReaderResolverBindsSniiWithoutOpeningClucene) { + auto context = std::make_shared(); + auto index_meta = make_test_inverted_index( + 14, {{INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}); + auto index_file_reader = std::make_shared(); + auto reader = std::make_shared( + &index_meta, index_file_reader, segment_v2::InvertedIndexReaderType::FULLTEXT); + + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = index_meta.properties(); + field_binding.__isset.index_properties = true; + + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + FieldReaderBinding binding; + auto status = resolver.resolve("body", InvertedIndexQueryType::MATCH_ANY_QUERY, &binding); + + ASSERT_TRUE(status.ok()) << status.to_string(); + EXPECT_EQ(0, index_file_reader->init_calls); + EXPECT_EQ(0, index_file_reader->open_calls); + EXPECT_EQ(reader, binding.inverted_reader); + EXPECT_EQ(nullptr, binding.lucene_reader); + EXPECT_TRUE(binding.use_snii_native_reader()); + EXPECT_FALSE(binding.use_direct_index_reader()); + EXPECT_EQ(SearchFieldExecutionMode::SNII_NATIVE, binding.execution_mode); +} + +TEST_F(FunctionSearchTest, TestBuildLeafQueryExecutesSelectedSniiWildcardReader) { + auto context = std::make_shared(); + std::map decoy_properties { + {INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_ENGLISH}, + {INVERTED_INDEX_PARSER_LOWERCASE_KEY, INVERTED_INDEX_PARSER_TRUE}}; + std::map selected_properties { + {INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}, + {INVERTED_INDEX_PARSER_LOWERCASE_KEY, INVERTED_INDEX_PARSER_TRUE}}; + auto decoy_meta = make_test_inverted_index(15, decoy_properties); + auto selected_meta = make_test_inverted_index(16, selected_properties); + auto decoy_file_reader = std::make_shared( + InvertedIndexStorageFormatPB::SNII, "/tmp/search_snii_decoy_idx"); + auto selected_file_reader = std::make_shared( + InvertedIndexStorageFormatPB::SNII, "/tmp/search_snii_selected_idx"); + auto decoy_reader = + std::make_shared(&decoy_meta, decoy_file_reader); + auto selected_reader = std::make_shared( + &selected_meta, selected_file_reader); + selected_reader->set_query_result("*lpha", make_bitmap({0, 2})); + selected_reader->set_null_bitmap(make_bitmap({3})); + + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, decoy_reader); + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, selected_reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"stored_body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = selected_properties; + field_binding.__isset.index_properties = true; + + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + auto clause = make_leaf_clause("WILDCARD", "*LPHA"); + inverted_index::query_v2::QueryPtr query; + std::string binding_key; + auto status = function_search->build_leaf_query(clause, context, resolver, &query, &binding_key, + "OR", 0, 4); + + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_NE(nullptr, query); + EXPECT_EQ(0, decoy_reader->query_calls); + EXPECT_EQ(1, selected_reader->query_calls); + EXPECT_EQ("stored_body", selected_reader->last_column_name); + EXPECT_EQ(TYPE_STRING, selected_reader->last_query_value_type); + EXPECT_EQ("*lpha", selected_reader->last_query_value); + EXPECT_EQ(InvertedIndexQueryType::WILDCARD_QUERY, selected_reader->last_query_type); + EXPECT_EQ(nullptr, selected_reader->last_analyzer_ctx); + EXPECT_EQ(0, decoy_file_reader->open_calls); + EXPECT_EQ(0, selected_file_reader->open_calls); + EXPECT_EQ(0, decoy_reader->null_bitmap_calls); + EXPECT_EQ(1, selected_reader->null_bitmap_calls); + const auto& bindings = resolver.binding_cache(); + ASSERT_EQ(1U, bindings.size()); + EXPECT_EQ(InvertedIndexQueryType::MATCH_ANY_QUERY, bindings.begin()->second.query_type); + + auto weight = query->weight(true); + ASSERT_NE(nullptr, weight); + inverted_index::query_v2::QueryExecutionContext exec_ctx; + exec_ctx.segment_num_rows = 4; + auto scorer = weight->scorer(exec_ctx, binding_key); + ASSERT_NE(nullptr, scorer); + EXPECT_EQ(0U, scorer->doc()); + EXPECT_FLOAT_EQ(1.0F, scorer->score()); + expect_bitmap_eq(collect_docs(scorer), {0, 2}); + ASSERT_TRUE(scorer->has_null_bitmap()); + const auto* null_bitmap = scorer->get_null_bitmap(); + ASSERT_NE(nullptr, null_bitmap); + expect_bitmap_eq(*null_bitmap, {3}); +} + +TEST_F(FunctionSearchTest, TestSniiWildcardPreservesThreeValuedBooleanAndFieldExists) { + auto context = std::make_shared(); + std::map properties { + {INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}; + auto index_meta = make_test_inverted_index(17, properties); + auto index_file_reader = std::make_shared(); + auto reader = + std::make_shared(&index_meta, index_file_reader); + reader->set_query_result("*lpha", make_bitmap({0})); + reader->set_query_result("beta*", make_bitmap({1})); + reader->set_null_bitmap(make_bitmap({3})); + + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = properties; + field_binding.__isset.index_properties = true; + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + + TSearchClause or_clause; + or_clause.clause_type = "OR"; + or_clause.children = {make_leaf_clause("WILDCARD", "*lpha"), + make_leaf_clause("WILDCARD", "beta*")}; + or_clause.__isset.children = true; + + TSearchClause not_clause; + not_clause.clause_type = "NOT"; + not_clause.children = {make_leaf_clause("WILDCARD", "*lpha")}; + not_clause.__isset.children = true; + auto exists_clause = make_leaf_clause("WILDCARD", "*"); + + auto verify_result = [&](const TSearchClause& root, + std::initializer_list expected_docs, + std::initializer_list expected_nulls) { + inverted_index::query_v2::QueryPtr query; + std::string binding_key; + auto status = function_search->build_query_recursive(root, context, resolver, &query, + &binding_key, "OR", 0, 4); + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_NE(nullptr, query); + auto weight = query->weight(false); + ASSERT_NE(nullptr, weight); + auto scorer = weight->scorer( + build_variant_search_query_execution_context(4, resolver, nullptr), binding_key); + ASSERT_NE(nullptr, scorer); + expect_bitmap_eq(collect_docs(scorer), expected_docs); + ASSERT_TRUE(scorer->has_null_bitmap()); + const auto* null_bitmap = scorer->get_null_bitmap(); + ASSERT_NE(nullptr, null_bitmap); + expect_bitmap_eq(*null_bitmap, expected_nulls); + }; + + verify_result(or_clause, {0, 1}, {3}); + verify_result(not_clause, {1, 2}, {3}); + verify_result(exists_clause, {0, 1, 2}, {3}); + EXPECT_EQ(3, reader->query_calls); +} + +TEST_F(FunctionSearchTest, TestSniiNativeRejectsNonWildcardClause) { + auto context = std::make_shared(); + auto index_meta = make_test_inverted_index( + 18, {{INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}); + auto index_file_reader = std::make_shared(); + auto reader = + std::make_shared(&index_meta, index_file_reader); + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = index_meta.properties(); + field_binding.__isset.index_properties = true; + FieldReaderResolver resolver(data_type_with_names, iterators, context, {field_binding}); + + auto clause = make_leaf_clause("TERM", "alpha"); + inverted_index::query_v2::QueryPtr query; + std::string binding_key; + auto status = function_search->build_leaf_query(clause, context, resolver, &query, &binding_key, + "OR", 0, 4); + + ASSERT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::NOT_IMPLEMENTED_ERROR, status.code()); + EXPECT_NE(std::string::npos, status.to_string().find("TERM")); + EXPECT_NE(std::string::npos, status.to_string().find("WILDCARD")); + EXPECT_EQ(0, reader->query_calls); + EXPECT_EQ(0, index_file_reader->open_calls); +} + +TEST_F(FunctionSearchTest, TestSearchDslCacheIsDisabledForSniiNativeExecution) { + ScopedInvertedIndexQueryCache cache_guard; + auto index_meta = make_test_inverted_index( + 19, {{INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}); + auto index_file_reader = std::make_shared(); + auto reader = + std::make_shared(&index_meta, index_file_reader); + reader->set_query_result("*lpha", make_bitmap({0})); + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = index_meta.properties(); + field_binding.__isset.index_properties = true; + TSearchParam search_param; + search_param.original_dsl = "body:*lpha"; + search_param.root = make_leaf_clause("WILDCARD", "*lpha"); + search_param.field_bindings = {field_binding}; + ASSERT_TRUE(insert_search_dsl_cache(cache_guard.get(), index_file_reader, search_param, + make_bitmap({3})) + .ok()); + + InvertedIndexResultBitmap result; + auto status = function_search->evaluate_inverted_index_with_search_param( + search_param, data_type_with_names, iterators, 4, result, true); + + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_NE(nullptr, result.get_data_bitmap()); + expect_bitmap_eq(*result.get_data_bitmap(), {0}); + EXPECT_EQ(1, reader->query_calls); +} + +TEST_F(FunctionSearchTest, TestSearchDslCacheIsDisabledWhenScoring) { + ScopedInvertedIndexQueryCache cache_guard; + auto index_meta = make_test_inverted_index( + 20, {{INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}); + auto index_file_reader = std::make_shared( + InvertedIndexStorageFormatPB::V2, "/tmp/search_scoring_v2_idx"); + auto reader = std::make_shared( + &index_meta, index_file_reader, segment_v2::InvertedIndexReaderType::FULLTEXT); + segment_v2::InvertedIndexIterator iterator; + iterator.add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &iterator; + TSearchFieldBinding field_binding; + field_binding.field_name = "body"; + field_binding.index_properties = index_meta.properties(); + field_binding.__isset.index_properties = true; + TSearchParam search_param; + search_param.original_dsl = "body:alpha"; + search_param.root = make_leaf_clause("TERM", "alpha"); + search_param.field_bindings = {field_binding}; + ASSERT_TRUE(insert_search_dsl_cache(cache_guard.get(), index_file_reader, search_param, + make_bitmap({3})) + .ok()); + + auto scoring_context = std::make_shared(); + scoring_context->collection_similarity = std::make_shared(); + InvertedIndexResultBitmap result; + std::unordered_map field_name_to_column_id; + auto status = function_search->evaluate_inverted_index_with_search_param( + search_param, data_type_with_names, iterators, 4, result, true, nullptr, + field_name_to_column_id, scoring_context); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(1, index_file_reader->init_calls); + EXPECT_EQ(1, index_file_reader->open_calls); +} + +TEST_F(FunctionSearchTest, TestSearchDslCacheRemainsEnabledForUnreferencedSniiField) { + ScopedInvertedIndexQueryCache cache_guard; + auto text_index_meta = make_test_inverted_index( + 21, {{INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_STANDARD}}); + auto number_index_meta = make_test_inverted_index(22); + auto text_file_reader = std::make_shared( + InvertedIndexStorageFormatPB::SNII, "/tmp/search_mixed_cache_idx"); + auto number_file_reader = std::make_shared( + InvertedIndexStorageFormatPB::V2, "/tmp/search_mixed_cache_idx"); + auto text_reader = std::make_shared( + &text_index_meta, text_file_reader, InvertedIndexReaderType::FULLTEXT); + auto number_reader = std::make_shared( + &number_index_meta, number_file_reader, InvertedIndexReaderType::BKD); + + InvertedIndexIterator text_iterator; + text_iterator.add_reader(InvertedIndexReaderType::FULLTEXT, text_reader); + RecordingDirectInvertedIndexIterator number_iterator; + number_iterator.add_reader(InvertedIndexReaderType::BKD, number_reader); + + std::unordered_map data_type_with_names; + data_type_with_names.emplace( + "body", IndexFieldNameAndTypePair {"body", std::make_shared()}); + data_type_with_names.emplace( + "age", IndexFieldNameAndTypePair {"age", std::make_shared()}); + std::unordered_map iterators; + iterators["body"] = &text_iterator; + iterators["age"] = &number_iterator; + + TSearchClause age_clause = make_leaf_clause("TERM", "42"); + age_clause.field_name = "age"; + TSearchParam search_param; + search_param.original_dsl = "age:42"; + search_param.root = age_clause; + ASSERT_TRUE(insert_search_dsl_cache(cache_guard.get(), number_file_reader, search_param, + make_bitmap({1})) + .ok()); + + InvertedIndexResultBitmap result; + auto status = function_search->evaluate_inverted_index_with_search_param( + search_param, data_type_with_names, iterators, 4, result, true); + + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_NE(nullptr, result.get_data_bitmap()); + expect_bitmap_eq(*result.get_data_bitmap(), {1}); + EXPECT_EQ(0, text_reader->query_calls); + EXPECT_EQ(0, number_iterator.read_calls); +} + TEST_F(FunctionSearchTest, TestBuildLeafQueryDirectUnknownClauseUsesLeafMapper) { TSearchClause clause; clause.clause_type = "PHRASE"; @@ -1942,6 +2647,7 @@ TEST_F(FunctionSearchTest, TestBuildLeafQueryDirectUnknownClauseUsesLeafMapper) binding.column_type = bool_type; binding.query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY; binding.state = SearchFieldBindingState::BOUND; + binding.execution_mode = SearchFieldExecutionMode::DIRECT_INDEX; TabletIndex index_meta; binding.inverted_reader = std::make_shared(&index_meta); @@ -2012,6 +2718,7 @@ TEST_F(FunctionSearchTest, TestBuildLeafQueryVariantBoolUsesDirectIndexReader) { binding.column_type = bool_type; binding.query_type = InvertedIndexQueryType::MATCH_ANY_QUERY; binding.state = SearchFieldBindingState::BOUND; + binding.execution_mode = SearchFieldExecutionMode::DIRECT_INDEX; TabletIndex index_meta; binding.inverted_reader = std::make_shared(&index_meta); @@ -2071,6 +2778,7 @@ TEST_F(FunctionSearchTest, TestBuildLeafQueryVariantNestedIntUsesDirectIndexRead binding.column_type = int_type; binding.query_type = InvertedIndexQueryType::MATCH_ANY_QUERY; binding.state = SearchFieldBindingState::BOUND; + binding.execution_mode = SearchFieldExecutionMode::DIRECT_INDEX; TabletIndex index_meta; binding.inverted_reader = std::make_shared(&index_meta); diff --git a/be/test/exprs/vsearch_expr_test.cpp b/be/test/exprs/vsearch_expr_test.cpp index cc3fa5820fca0c..8b3b4762652c38 100644 --- a/be/test/exprs/vsearch_expr_test.cpp +++ b/be/test/exprs/vsearch_expr_test.cpp @@ -33,7 +33,11 @@ #include "exprs/vexpr_context.h" #include "exprs/vliteral.h" #include "exprs/vsearch.h" +#include "storage/index/index_file_reader.h" #include "storage/index/index_iterator.h" +#include "storage/index/inverted/inverted_index_iterator.h" +#include "storage/index/inverted/inverted_index_parser.h" +#include "storage/index/inverted/inverted_index_reader.h" #include "storage/segment/variant/nested_group_provider.h" #if defined(__clang__) @@ -46,6 +50,7 @@ #if defined(__clang__) #pragma clang diagnostic pop #endif +#include "exprs/vcompound_pred.h" namespace doris { @@ -525,7 +530,7 @@ TEST_F(VSearchExprTest, TestEvaluateInvertedIndexEmptyDSL) { VExprContext context(dummy_expr); auto status = vsearch_expr->evaluate_inverted_index(&context, 100); EXPECT_FALSE(status.ok()); - EXPECT_TRUE(status.code() == ErrorCode::INVALID_ARGUMENT); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); EXPECT_TRUE(status.to_string().find("search DSL is empty") != std::string::npos); } @@ -1277,6 +1282,7 @@ TEST_F(VSearchExprTest, TestEvaluateInvertedIndexWithEmptyDSL) { auto status = vsearch_expr->evaluate_inverted_index(&context, 100); EXPECT_FALSE(status.ok()); // Should return error due to empty DSL EXPECT_EQ(status.code(), ErrorCode::INVALID_ARGUMENT); + EXPECT_NE(status.to_string().find("search DSL is empty"), std::string::npos); } TEST_F(VSearchExprTest, FastExecuteReturnsPrecomputedColumn) { @@ -1319,6 +1325,7 @@ TEST_F(VSearchExprTest, EvaluateInvertedIndexFailsWithoutStorageType) { auto status = expr->evaluate_inverted_index(context.get(), 128); EXPECT_FALSE(status.ok()); EXPECT_EQ(ErrorCode::INTERNAL_ERROR, status.code()); + EXPECT_NE(status.to_string().find("storage_name_type not found"), std::string::npos); } TEST_F(VSearchExprTest, EvaluateInvertedIndexWithUnsupportedChildReturnsError) { @@ -1337,6 +1344,20 @@ TEST_F(VSearchExprTest, EvaluateInvertedIndexWithUnsupportedChildReturnsError) { auto status = expr->evaluate_inverted_index(context.get(), 64); EXPECT_FALSE(status.ok()); EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(status.to_string().find("Unsupported child node type"), std::string::npos); + + TExprNode compound_node; + compound_node.__set_type(test_node.type); + compound_node.__set_node_type(TExprNodeType::COMPOUND_PRED); + compound_node.__set_opcode(TExprOpcode::COMPOUND_AND); + compound_node.__set_num_children(1); + compound_node.__set_is_nullable(false); + auto compound = VCompoundPred::create_shared(compound_node); + compound->add_child(expr); + status = compound->evaluate_inverted_index(context.get(), 64); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(status.to_string().find("Unsupported child node type"), std::string::npos); } TEST_F(VSearchExprTest, EvaluateInvertedIndexHandlesMissingIterators) { @@ -1393,7 +1414,7 @@ TEST_F(VSearchExprTest, EvaluateInvertedIndexNestedFallbackReturnsNotSupportedIn ASSERT_TRUE(provider != nullptr); if (!provider->should_enable_nested_group_read_path()) { EXPECT_FALSE(status.ok()); - EXPECT_EQ(ErrorCode::NOT_IMPLEMENTED_ERROR, status.code()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_NOT_SUPPORTED, status.code()); EXPECT_TRUE(status.to_string().find("NestedGroup support") != std::string::npos); } else { EXPECT_FALSE(status.ok()); @@ -1419,8 +1440,97 @@ TEST_F(VSearchExprTest, EvaluateInvertedIndexPropagatesFunctionFailure) { auto status = expr->evaluate_inverted_index(context.get(), 256); EXPECT_FALSE(status.ok()); - EXPECT_EQ(ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND, status.code()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_NOT_SUPPORTED, status.code()); + EXPECT_FALSE(status_map[0][expr.get()]); +} + +TEST_F(VSearchExprTest, EvaluateInvertedIndexRejectsSearchFallback) { + test_node.search_param.root.clause_type = "WILDCARD"; + test_node.search_param.root.value = "hello*"; + test_node.search_param.field_bindings[0].index_properties[INVERTED_INDEX_PARSER_KEY] = + INVERTED_INDEX_PARSER_ENGLISH; + test_node.search_param.field_bindings[0].__isset.index_properties = true; + + auto expr = VSearchExpr::create_shared(test_node); + expr->add_child(create_slot_ref(0, "title")); + + TabletIndexPB index_pb; + index_pb.set_index_type(IndexType::INVERTED); + index_pb.set_index_id(1); + index_pb.set_index_name("search_fallback_index"); + index_pb.add_col_unique_id(1); + (*index_pb.mutable_properties())[INVERTED_INDEX_PARSER_KEY] = INVERTED_INDEX_PARSER_STANDARD; + TabletIndex index_meta; + index_meta.init_from_pb(index_pb); + auto index_file_reader = std::make_shared( + nullptr, "/tmp/search_fallback_idx", InvertedIndexStorageFormatPB::V2); + auto reader = std::make_shared(&index_meta, index_file_reader); + auto iterator = std::make_unique(); + iterator->add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, reader); + + std::vector col_ids = {0}; + std::vector> index_iterators; + index_iterators.emplace_back(std::move(iterator)); + std::vector storage_types; + storage_types.emplace_back("title", std::make_shared()); + std::unordered_map> status_map; + status_map[0][expr.get()] = false; + + auto inverted_ctx = make_inverted_context(col_ids, index_iterators, storage_types, status_map); + auto context = std::make_shared(expr); + context->set_index_context(inverted_ctx); + + auto status = expr->evaluate_inverted_index(context.get(), 256); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_NOT_SUPPORTED, status.code()); + EXPECT_NE(status.to_string().find("cannot fall back"), std::string::npos); + EXPECT_NE(status.to_string().find("No inverted index found for analyzer"), std::string::npos); EXPECT_FALSE(status_map[0][expr.get()]); + + for (TExprOpcode::type opcode : {TExprOpcode::COMPOUND_AND, TExprOpcode::COMPOUND_OR}) { + TExprNode compound_node; + compound_node.__set_type(test_node.type); + compound_node.__set_node_type(TExprNodeType::COMPOUND_PRED); + compound_node.__set_opcode(opcode); + compound_node.__set_num_children(1); + compound_node.__set_is_nullable(false); + auto compound = VCompoundPred::create_shared(compound_node); + compound->add_child(expr); + + status = compound->evaluate_inverted_index(context.get(), 256); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_NOT_SUPPORTED, status.code()); + EXPECT_NE(status.to_string().find("No inverted index found for analyzer"), + std::string::npos); + } + + test_node.search_param.root.clause_type = "TERM"; + test_node.search_param.root.value = "hello"; + test_node.search_param.field_bindings[0].index_properties[INVERTED_INDEX_PARSER_KEY] = + INVERTED_INDEX_PARSER_STANDARD; + auto snii_expr = VSearchExpr::create_shared(test_node); + snii_expr->add_child(create_slot_ref(0, "title")); + + auto snii_file_reader = std::make_shared( + nullptr, "/tmp/search_fallback_snii_idx", InvertedIndexStorageFormatPB::SNII); + auto snii_reader = + std::make_shared(&index_meta, snii_file_reader); + auto snii_iterator = std::make_unique(); + snii_iterator->add_reader(segment_v2::InvertedIndexReaderType::FULLTEXT, snii_reader); + + std::vector> snii_index_iterators; + snii_index_iterators.emplace_back(std::move(snii_iterator)); + std::unordered_map> snii_status_map; + snii_status_map[0][snii_expr.get()] = false; + auto snii_inverted_ctx = + make_inverted_context(col_ids, snii_index_iterators, storage_types, snii_status_map); + auto snii_context = std::make_shared(snii_expr); + snii_context->set_index_context(snii_inverted_ctx); + + status = snii_expr->evaluate_inverted_index(snii_context.get(), 256); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(ErrorCode::INVERTED_INDEX_NOT_SUPPORTED, status.code()); + EXPECT_NE(status.to_string().find("supports only WILDCARD"), std::string::npos); } // Note: Full testing with actual IndexExecContext and real iterators diff --git a/be/test/io/cache/block_file_cache_profile_reporter_test.cpp b/be/test/io/cache/block_file_cache_profile_reporter_test.cpp index 975f1e8eddd88b..246cc539952f35 100644 --- a/be/test/io/cache/block_file_cache_profile_reporter_test.cpp +++ b/be/test/io/cache/block_file_cache_profile_reporter_test.cpp @@ -64,6 +64,10 @@ io::FileCacheStatistics make_file_cache_stats(int64_t multiplier) { stats.num_peer_race_s3_win = multiplier * 38; stats.num_peer_lazy_fetch = multiplier * 39; stats.peer_lazy_fetch_timer = multiplier * 40; + stats.inverted_index_request_bytes = multiplier * 41; + stats.inverted_index_read_bytes = multiplier * 42; + stats.inverted_index_range_read_count = multiplier * 43; + stats.inverted_index_serial_read_rounds = multiplier * 44; return stats; } @@ -114,6 +118,10 @@ void expect_file_cache_stats_eq(const io::FileCacheStatistics& actual, EXPECT_EQ(actual.num_peer_race_s3_win, expected.num_peer_race_s3_win); EXPECT_EQ(actual.num_peer_lazy_fetch, expected.num_peer_lazy_fetch); EXPECT_EQ(actual.peer_lazy_fetch_timer, expected.peer_lazy_fetch_timer); + EXPECT_EQ(actual.inverted_index_request_bytes, expected.inverted_index_request_bytes); + EXPECT_EQ(actual.inverted_index_read_bytes, expected.inverted_index_read_bytes); + EXPECT_EQ(actual.inverted_index_range_read_count, expected.inverted_index_range_read_count); + EXPECT_EQ(actual.inverted_index_serial_read_rounds, expected.inverted_index_serial_read_rounds); } } // namespace @@ -163,6 +171,14 @@ TEST(FileCacheProfileReporterTest, ReporterAggregatesDeltaReportsToExactFinalTot after_second_report.cross_cg_peer_io_timer); EXPECT_EQ(profile->get_counter("PeerLazyFetchTime")->value(), after_second_report.peer_lazy_fetch_timer); + EXPECT_EQ(profile->get_counter("InvertedIndexRequestBytes")->value(), + after_second_report.inverted_index_request_bytes); + EXPECT_EQ(profile->get_counter("InvertedIndexReadBytes")->value(), + after_second_report.inverted_index_read_bytes); + EXPECT_EQ(profile->get_counter("InvertedIndexRangeReadCount")->value(), + after_second_report.inverted_index_range_read_count); + EXPECT_EQ(profile->get_counter("InvertedIndexSerialReadRounds")->value(), + after_second_report.inverted_index_serial_read_rounds); } } // namespace doris diff --git a/be/test/runtime/fragment_mgr_cross_cluster_cancel_test.cpp b/be/test/runtime/fragment_mgr_cross_cluster_cancel_test.cpp index 70dd2e874bafc0..4b9d0f8c699de9 100644 --- a/be/test/runtime/fragment_mgr_cross_cluster_cancel_test.cpp +++ b/be/test/runtime/fragment_mgr_cross_cluster_cancel_test.cpp @@ -18,10 +18,14 @@ #include #include +#include "common/config.h" #include "runtime/descriptor_helper.h" #include "runtime/exec_env.h" #include "runtime/fragment_mgr.h" +#include "runtime/index_policy/index_policy_mgr.h" #include "runtime/workload_group/workload_group_manager.h" +#include "storage/id_manager.h" +#include "util/defer_op.h" namespace doris { diff --git a/be/test/runtime/index_policy/index_policy_mgr_test.cpp b/be/test/runtime/index_policy/index_policy_mgr_test.cpp index 923690ef3612c8..4485217652bb11 100644 --- a/be/test/runtime/index_policy/index_policy_mgr_test.cpp +++ b/be/test/runtime/index_policy/index_policy_mgr_test.cpp @@ -19,10 +19,34 @@ #include +#include +#include +#include +#include + +#include "common/config.h" #include "runtime/exec_env.h" #include "storage/index/inverted/analysis_factory_mgr.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/analyzer/segment_analyzer_context.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "util/defer_op.h" namespace doris { +namespace { + +TIndexPolicy common_grams_policy(int64_t id, std::string name) { + TIndexPolicy policy; + policy.id = id; + policy.name = std::move(name); + policy.type = TIndexPolicyType::TOKEN_FILTER; + policy.properties["type"] = "common_grams"; + return policy; +} + +} // namespace class IndexPolicyMgrTest : public testing::Test { protected: @@ -176,4 +200,220 @@ TEST_F(IndexPolicyMgrTest, TestTokenFilterProcessing) { ASSERT_NE(emptyAnalyzer, nullptr); } +TEST_F(IndexPolicyMgrTest, CommonGramsProviderRetainsImmutablePurposeConfiguration) { + TIndexPolicy tokenizer; + tokenizer.id = 20; + tokenizer.name = "cg_tokenizer"; + tokenizer.type = TIndexPolicyType::TOKENIZER; + tokenizer.properties["type"] = "char_group"; + tokenizer.properties["tokenize_on_chars"] = "[whitespace]"; + + auto common_grams = common_grams_policy(21, "cg_filter"); + + TIndexPolicy analyzer_policy; + analyzer_policy.id = 22; + analyzer_policy.name = "cg_analyzer"; + analyzer_policy.type = TIndexPolicyType::ANALYZER; + analyzer_policy.properties["tokenizer"] = "cg_tokenizer"; + analyzer_policy.properties["token_filter"] = "lowercase,cg_filter"; + mgr.apply_policy_changes({tokenizer, common_grams, analyzer_policy}, {}); + + auto provider = mgr.get_analyzer_provider_by_name("cg_analyzer"); + ASSERT_NE(provider, nullptr); + auto analyze = [](const segment_v2::inverted_index::AnalyzerProviderPtr& analyzer_provider, + segment_v2::inverted_index::AnalysisPurpose purpose, std::string_view text) { + auto analyzer = analyzer_provider->get_analyzer(purpose); + auto reader = std::make_shared>(); + reader->init(text.data(), static_cast(text.size()), true); + return segment_v2::inverted_index::InvertedIndexAnalyzer::get_analyse_result( + reader, analyzer.get()); + }; + + auto index = analyze(provider, segment_v2::inverted_index::AnalysisPurpose::kIndex, + "Man of the Year"); + auto plain = analyze(provider, segment_v2::inverted_index::AnalysisPurpose::kPlainQuery, + "Man of the Year"); + auto exact = analyze(provider, segment_v2::inverted_index::AnalysisPurpose::kExactPhraseQuery, + "Man of the Year"); + auto prefix = analyze(provider, segment_v2::inverted_index::AnalysisPurpose::kPhrasePrefixQuery, + "the wo"); + EXPECT_EQ(index.size(), 7); + EXPECT_EQ(plain.size(), 4); + EXPECT_EQ(exact.size(), 3); + EXPECT_EQ(prefix.size(), 1); + auto prefix_gram = segment_v2::inverted_index::encode_common_gram("the", "wo"); + ASSERT_TRUE(prefix_gram.has_value()) << prefix_gram.error(); + EXPECT_EQ(prefix.front().get_single_term(), prefix_gram.value()); + + auto single_index = mgr.get_analyzer_by_name( + "cg_analyzer", segment_v2::inverted_index::AnalysisPurpose::kIndex); + auto single_reader = std::make_shared>(); + const std::string single_input = "Man of the Year"; + single_reader->init(single_input.data(), static_cast(single_input.size()), true); + EXPECT_EQ(segment_v2::inverted_index::InvertedIndexAnalyzer::get_analyse_result( + single_reader, single_index.get()) + .size(), + 7); +} + +TEST_F(IndexPolicyMgrTest, AnalyzerProviderPreservesPurposeInsensitiveNormalizers) { + auto builtin = mgr.get_analyzer_provider_by_name("lowercase"); + auto builtin_analyzer = + builtin->get_analyzer(segment_v2::inverted_index::AnalysisPurpose::kPlainQuery); + EXPECT_EQ(builtin->get_analyzer(segment_v2::inverted_index::AnalysisPurpose::kIndex), + builtin_analyzer); + + TIndexPolicy normalizer; + normalizer.id = 23; + normalizer.name = "test_normalizer"; + normalizer.type = TIndexPolicyType::NORMALIZER; + normalizer.properties["token_filter"] = "lowercase"; + mgr.apply_policy_changes({normalizer}, {}); + + auto configured = mgr.get_analyzer_provider_by_name("test_normalizer"); + auto configured_analyzer = + configured->get_analyzer(segment_v2::inverted_index::AnalysisPurpose::kPlainQuery); + EXPECT_EQ(configured->get_analyzer(segment_v2::inverted_index::AnalysisPurpose::kIndex), + configured_analyzer); +} + +TEST_F(IndexPolicyMgrTest, FindsFreshAnalyzerProviderBySegmentBaseFingerprint) { + using segment_v2::inverted_index::AnalysisPurpose; + using segment_v2::inverted_index::InvertedIndexAnalyzer; + + const std::map slash_to_space = { + {INVERTED_INDEX_PARSER_CHAR_FILTER_TYPE, INVERTED_INDEX_CHAR_FILTER_CHAR_REPLACE}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_PATTERN, "/"}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_REPLACEMENT, " "}}; + auto configured = mgr.get_analyzer_provider_by_name("analyzer1", slash_to_space); + const std::string segment_fingerprint(configured->base_analyzer_fingerprint()); + + auto first = mgr.get_analyzer_provider_by_base_fingerprint(segment_fingerprint, slash_to_space); + ASSERT_NE(first, nullptr); + EXPECT_EQ(first->base_analyzer_fingerprint(), segment_fingerprint); + EXPECT_EQ(mgr.get_analyzer_provider_by_base_fingerprint(segment_fingerprint), nullptr); + EXPECT_EQ(mgr.get_analyzer_provider_by_base_fingerprint("unknown", slash_to_space), nullptr); + + auto second = + mgr.get_analyzer_provider_by_base_fingerprint(segment_fingerprint, slash_to_space); + ASSERT_NE(second, nullptr); + EXPECT_NE(first, second); + EXPECT_NE(first->get_analyzer(AnalysisPurpose::kPlainQuery), + second->get_analyzer(AnalysisPurpose::kPlainQuery)); + + mgr.apply_policy_changes({}, {5}); + auto reader = std::make_shared>(); + const std::string input = "ASCII TERM"; + reader->init(input.data(), static_cast(input.size()), true); + const auto terms = InvertedIndexAnalyzer::get_analyse_result( + reader, first->get_analyzer(AnalysisPurpose::kPlainQuery).get()); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].get_single_term(), "ascii"); + EXPECT_EQ(terms[1].get_single_term(), "term"); +} + +TEST_F(IndexPolicyMgrTest, RebuildsQueryContextForPersistedSegmentAnalyzer) { + const std::map slash_to_space = { + {INVERTED_INDEX_PARSER_CHAR_FILTER_TYPE, INVERTED_INDEX_CHAR_FILTER_CHAR_REPLACE}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_PATTERN, "/"}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_REPLACEMENT, " "}}; + auto segment_provider = mgr.get_analyzer_provider_by_name("analyzer1", slash_to_space); + const std::string segment_fingerprint(segment_provider->base_analyzer_fingerprint()); + + segment_v2::inverted_index::Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + segment_v2::inverted_index::CustomAnalyzerConfig::Builder request_builder; + request_builder.with_tokenizer_config("char_group", tokenizer_settings); + auto request_provider = std::make_shared( + request_builder.build()); + ASSERT_NE(request_provider->base_analyzer_fingerprint(), segment_fingerprint); + + InvertedIndexAnalyzerCtx request_context; + request_context.analyzer_name = "request_analyzer"; + request_context.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + request_context.char_filter_map = {{"stale", "filter"}}; + request_context.analyzer = request_provider->get_analyzer( + segment_v2::inverted_index::AnalysisPurpose::kPlainQuery); + request_context.analyzer_provider = request_provider; + request_context.common_grams_identity = segment_v2::inverted_index::CommonGramsQueryIdentity { + .common_grams_dictionary_identity = "stale-dictionary", + .base_analyzer_fingerprint = std::string(request_provider->base_analyzer_fingerprint()), + .common_grams_fingerprint = "stale-common-grams"}; + + const std::map physical_properties = slash_to_space; + auto rebuilt = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + &request_context, segment_fingerprint, physical_properties, &mgr); + ASSERT_TRUE(rebuilt.has_value()) << rebuilt.error(); + ASSERT_TRUE(rebuilt->has_value()); + const auto& effective = rebuilt->value(); + EXPECT_EQ(effective.analyzer_name, request_context.analyzer_name); + EXPECT_EQ(effective.parser_type, request_context.parser_type); + EXPECT_EQ(effective.char_filter_map, slash_to_space); + EXPECT_EQ(effective.analyzer, nullptr); + ASSERT_NE(effective.analyzer_provider, nullptr); + EXPECT_EQ(effective.analyzer_provider->base_analyzer_fingerprint(), segment_fingerprint); + EXPECT_FALSE(effective.common_grams_identity.has_value()); + EXPECT_NE(effective.analyzer_provider, segment_provider); + + InvertedIndexAnalyzerCtx matching_context = request_context; + matching_context.analyzer_provider = segment_provider; + auto unchanged = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + &matching_context, segment_fingerprint, physical_properties, &mgr); + ASSERT_TRUE(unchanged.has_value()) << unchanged.error(); + EXPECT_FALSE(unchanged->has_value()); + + auto unavailable = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + &request_context, "missing-segment-fingerprint", physical_properties, &mgr); + ASSERT_FALSE(unavailable.has_value()); + EXPECT_EQ(unavailable.error().code(), ErrorCode::INVERTED_INDEX_BYPASS); +} + +TEST_F(IndexPolicyMgrTest, SegmentAnalyzerAdmissionKeepsLegacyRequestWithoutMetadata) { + auto admitted = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + nullptr, std::optional {}, {}, + nullptr); + + ASSERT_TRUE(admitted.has_value()) << admitted.error(); + EXPECT_FALSE(admitted->has_value()); +} + +TEST_F(IndexPolicyMgrTest, SegmentAnalyzerAdmissionBypassesTypedMetadataWithoutBaseFingerprint) { + segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + auto admitted = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + nullptr, std::optional {metadata}, {}, &mgr); + + ASSERT_FALSE(admitted.has_value()); + EXPECT_EQ(admitted.error().code(), ErrorCode::INVERTED_INDEX_BYPASS); +} + +TEST_F(IndexPolicyMgrTest, SegmentAnalyzerAdmissionRebuildsTypedMetadata) { + const std::map slash_to_space = { + {INVERTED_INDEX_PARSER_CHAR_FILTER_TYPE, INVERTED_INDEX_CHAR_FILTER_CHAR_REPLACE}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_PATTERN, "/"}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_REPLACEMENT, " "}}; + auto segment_provider = mgr.get_analyzer_provider_by_name("analyzer1", slash_to_space); + const std::string segment_fingerprint(segment_provider->base_analyzer_fingerprint()); + + segment_v2::inverted_index::Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + segment_v2::inverted_index::CustomAnalyzerConfig::Builder request_builder; + request_builder.with_tokenizer_config("char_group", tokenizer_settings); + auto request_provider = std::make_shared( + request_builder.build()); + + InvertedIndexAnalyzerCtx request_context; + request_context.analyzer_provider = request_provider; + request_context.analyzer = request_provider->get_analyzer( + segment_v2::inverted_index::AnalysisPurpose::kPlainQuery); + segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.base_analyzer_fingerprint = segment_fingerprint; + + auto rebuilt = segment_v2::inverted_index::maybe_rebuild_segment_analyzer_context( + &request_context, std::optional {metadata}, slash_to_space, &mgr); + + ASSERT_TRUE(rebuilt.has_value()) << rebuilt.error(); + ASSERT_TRUE(rebuilt->has_value()); + EXPECT_EQ(rebuilt->value().analyzer_provider->base_analyzer_fingerprint(), segment_fingerprint); +} + } // namespace doris \ No newline at end of file diff --git a/be/test/storage/compaction/collection_statistics_test.cpp b/be/test/storage/compaction/collection_statistics_test.cpp index b78355b316ebdd..3f632ce7da7fc3 100644 --- a/be/test/storage/compaction/collection_statistics_test.cpp +++ b/be/test/storage/compaction/collection_statistics_test.cpp @@ -15,14 +15,19 @@ // specific language governing permissions and limitations // under the License. -#include "storage/compaction/collection_statistics.h" +#include "storage/index/inverted/similarity/collection_statistics.h" #include #include #include +#include +#include +#include #include #include +#include +#include #include "common/exception.h" #include "core/data_type/data_type_string.h" @@ -30,14 +35,27 @@ #include "exprs/vexpr.h" #include "exprs/vexpr_context.h" #include "exprs/vliteral.h" +#include "exprs/vsearch.h" #include "exprs/vslot_ref.h" #include "io/fs/local_file_system.h" -#include "storage/compaction/collection_statistics.cpp" +#include "runtime/exec_env.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/index_file_writer.h" +#include "storage/index/index_writer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/inverted_index_desc.h" +#include "storage/index/inverted/similarity/collection_statistics.cpp" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" #include "storage/rowset/rowset.h" #include "storage/rowset/rowset_meta.h" #include "storage/rowset/rowset_reader.h" #include "storage/tablet/tablet_schema.h" #include "testutil/mock/mock_runtime_state.h" +#include "util/slice.h" namespace doris { @@ -48,6 +66,13 @@ class MockVExpr : public VExpr { MockVExpr(TExprNodeType::type node_type) : _mock_node_type(node_type) { if (node_type == TExprNodeType::MATCH_PRED) { _opcode = TExprOpcode::MATCH_PHRASE; + InvertedIndexAnalyzerConfig config; + config.parser_type = InvertedIndexParserType::PARSER_STANDARD; + config.stop_words = "none"; + _analyzer_ctx = std::make_shared(); + _analyzer_ctx->analyzer_provider = + segment_v2::inverted_index::InvertedIndexAnalyzer::create_analyzer_provider( + &config); } } @@ -80,8 +105,37 @@ class MockVExpr : public VExpr { std::string debug_string() const override { return "MockVExpr"; } + const InvertedIndexAnalyzerCtx* query_analyzer_ctx() const override { + return _analyzer_ctx.get(); + } + + void set_analyzer_ctx(InvertedIndexAnalyzerCtxSPtr analyzer_ctx) { + _analyzer_ctx = std::move(analyzer_ctx); + } + + void set_opcode(TExprOpcode::type opcode) { _opcode = opcode; } + private: TExprNodeType::type _mock_node_type; + InvertedIndexAnalyzerCtxSPtr _analyzer_ctx; +}; + +class FixedFingerprintAnalyzerProvider final : public segment_v2::inverted_index::AnalyzerProvider { +public: + FixedFingerprintAnalyzerProvider(std::shared_ptr analyzer, + std::string fingerprint) + : _analyzer(std::move(analyzer)), _fingerprint(std::move(fingerprint)) {} + + std::shared_ptr get_analyzer( + segment_v2::inverted_index::AnalysisPurpose) const override { + return _analyzer; + } + + std::string_view base_analyzer_fingerprint() const override { return _fingerprint; } + +private: + std::shared_ptr _analyzer; + std::string _fingerprint; }; class MockVSlotRef : public VSlotRef { @@ -172,6 +226,7 @@ class MockRowset : public Rowset { int64_t num_segments() const override { return _num_segments; } Result segment_path(int64_t seg_id) override { + _segment_path_requests.push_back(seg_id); if (_segment_paths.find(seg_id) != _segment_paths.end()) { return _segment_paths.at(seg_id); } @@ -184,9 +239,12 @@ class MockRowset : public Rowset { void set_num_segments(int64_t num) { _num_segments = num; } + const std::vector& segment_path_requests() const { return _segment_path_requests; } + private: int64_t _num_segments; std::map _segment_paths; + std::vector _segment_path_requests; }; class MockRowsetReader : public RowsetReader { @@ -283,11 +341,23 @@ class CollectionStatisticsTest : public ::testing::Test { return tablet_schema; } - VExprContextSPtrs create_match_expr_contexts(const std::string& search_term = "search term") { + VExprContextSPtrs create_match_expr_contexts( + const std::string& search_term = "search term", + const std::string& base_analyzer_fingerprint = "") { VExprContextSPtrs contexts; auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + if (!base_analyzer_fingerprint.empty()) { + auto analyzer = match_expr->query_analyzer_ctx()->analyzer_provider->get_analyzer( + segment_v2::inverted_index::AnalysisPurpose::kPlainQuery); + auto provider = + std::make_shared( + std::move(analyzer), base_analyzer_fingerprint); + auto analyzer_ctx = std::make_shared(); + analyzer_ctx->analyzer_provider = std::move(provider); + match_expr->set_analyzer_ctx(std::move(analyzer_ctx)); + } auto slot_ref = std::make_shared("content", SlotId(1)); auto literal = std::make_shared(search_term); @@ -320,6 +390,389 @@ class CollectionStatisticsTest : public ::testing::Test { return splits; } + TabletSchemaSPtr create_legacy_v3_schema() { + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(DUP_KEYS); + schema_pb.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::V3); + auto tablet_schema = std::make_shared(); + tablet_schema->init_from_pb(schema_pb); + + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + tablet_schema->append_column(column); + + TabletIndex index; + index._index_id = 1; + index._index_type = IndexType::INVERTED; + index._col_unique_ids.push_back(1); + index._properties["parser"] = "standard"; + index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(index)); + return tablet_schema; + } + + TabletSchemaSPtr create_snii_schema(int64_t index_id = 1) { + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(DUP_KEYS); + schema_pb.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::SNII); + auto tablet_schema = std::make_shared(); + tablet_schema->init_from_pb(schema_pb); + + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + tablet_schema->append_column(column); + + TabletIndex index; + index._index_id = index_id; + index._index_type = IndexType::INVERTED; + index._col_unique_ids.push_back(1); + index._properties["parser"] = "standard"; + index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(index)); + return tablet_schema; + } + + Status write_legacy_v3_segment(const TabletSchemaSPtr& tablet_schema, + const std::string& segment_path) { + std::vector paths; + paths.emplace_back(test_dir_, 1024); + auto tmp_file_dirs = std::make_unique(paths); + RETURN_IF_ERROR(tmp_file_dirs->init()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + + const std::string index_path_prefix { + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix(segment_path)}; + io::FileWriterPtr compound_file; + io::FileWriterOptions options; + auto fs = io::global_local_filesystem(); + RETURN_IF_ERROR(fs->create_file( + segment_v2::InvertedIndexDescriptor::get_index_file_path_v2(index_path_prefix), + &compound_file, &options)); + + segment_v2::IndexFileWriter file_writer(fs, index_path_prefix, "legacy_v3", 0, + InvertedIndexStorageFormatPB::V3, + std::move(compound_file)); + const auto index_metas = tablet_schema->inverted_indexs(1); + DORIS_CHECK(index_metas.size() == 1); + std::unique_ptr column_writer; + RETURN_IF_ERROR(segment_v2::IndexColumnWriter::create( + &tablet_schema->column(0), &column_writer, &file_writer, index_metas[0])); + std::vector values {Slice("alpha beta")}; + RETURN_IF_ERROR(column_writer->add_values("content", values.data(), values.size())); + RETURN_IF_ERROR(column_writer->finish()); + RETURN_IF_ERROR(file_writer.begin_close()); + return file_writer.finish_close(); + } + + Status write_snii_common_grams_segment(const std::string& segment_path, + std::string base_analyzer_fingerprint, + uint64_t scoring_token_count = 3) { + const std::string index_path_prefix { + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix(segment_path)}; + io::FileWriterPtr file_writer; + io::FileWriterOptions options; + auto fs = io::global_local_filesystem(); + RETURN_IF_ERROR(fs->create_file( + segment_v2::InvertedIndexDescriptor::get_index_file_path_v2(index_path_prefix), + &file_writer, &options)); + + segment_v2::snii_doris::DorisSniiFileWriter adapter(file_writer.get()); + snii::writer::SniiCompoundWriter writer(&adapter); + auto metadata = segment_v2::inverted_index::make_common_grams_segment_metadata( + {.common_grams_dictionary_identity = "test-stopwords-v1", + .base_analyzer_fingerprint = std::move(base_analyzer_fingerprint), + .common_grams_fingerprint = "test-common-grams-v1"}); + metadata.scoring_doc_count = 2; + metadata.scoring_token_count = scoring_token_count; + + snii::writer::TermPostings alpha; + alpha.term = "alpha"; + alpha.docids = {0, 1}; + alpha.freqs = {1, 1}; + alpha.positions_flat = {0, 0}; + + snii::writer::TermPostings beta; + beta.term = "beta"; + beta.docids = {0}; + beta.freqs = {1}; + beta.positions_flat = {1}; + + snii::writer::TermPostings gram; + gram.term = DORIS_TRY(segment_v2::inverted_index::encode_common_gram("alpha", "beta")); + gram.docids = {0}; + gram.freqs = {1}; + gram.positions_flat = {0}; + + snii::writer::SniiIndexInput input; + input.index_id = 1; + input.config = snii::format::IndexConfig::kDocsPositionsScoring; + input.doc_count = 2; + input.encoded_norms = {snii::query::encode_norm(2), snii::query::encode_norm(1)}; + input.terms = {std::move(gram), std::move(alpha), std::move(beta)}; + std::ranges::sort(input.terms, {}, &snii::writer::TermPostings::term); + input.common_grams_metadata = std::move(metadata); + + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + return file_writer->close(false); + } + + Status write_plain_snii_scoring_segment(const std::string& segment_path) { + const std::string index_path_prefix { + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix(segment_path)}; + io::FileWriterPtr file_writer; + io::FileWriterOptions options; + auto fs = io::global_local_filesystem(); + RETURN_IF_ERROR(fs->create_file( + segment_v2::InvertedIndexDescriptor::get_index_file_path_v2(index_path_prefix), + &file_writer, &options)); + + segment_v2::snii_doris::DorisSniiFileWriter adapter(file_writer.get()); + snii::writer::SniiCompoundWriter writer(&adapter); + + snii::writer::TermPostings alpha; + alpha.term = "alpha"; + alpha.docids = {0, 1}; + alpha.freqs = {1, 1}; + alpha.positions_flat = {0, 0}; + + snii::writer::TermPostings beta; + beta.term = "beta"; + beta.docids = {0}; + beta.freqs = {1}; + beta.positions_flat = {1}; + + snii::writer::SniiIndexInput input; + input.index_id = 1; + input.config = snii::format::IndexConfig::kDocsPositionsScoring; + input.doc_count = 2; + input.encoded_norms = {snii::query::encode_norm(2), snii::query::encode_norm(1)}; + input.terms = {std::move(alpha), std::move(beta)}; + + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + return file_writer->close(false); + } + + VExprContextSPtrs create_search_contexts(const std::string& clause_type, + const std::string& value) { + TSearchClause clause; + clause.clause_type = clause_type; + clause.field_name = "content"; + clause.value = value; + clause.__isset.field_name = true; + clause.__isset.value = true; + + return create_search_contexts(std::move(clause)); + } + + VExprContextSPtrs create_search_contexts(TSearchClause root, + std::vector field_bindings = {}) { + TSearchParam search_param; + search_param.root = std::move(root); + search_param.field_bindings = std::move(field_bindings); + + TExprNode node; + node.node_type = TExprNodeType::SEARCH_EXPR; + TTypeNode type_node; + type_node.type = TTypeNodeType::SCALAR; + TScalarType scalar_type; + scalar_type.__set_type(TPrimitiveType::BOOLEAN); + type_node.__set_scalar_type(scalar_type); + TTypeDesc type_desc; + type_desc.types.push_back(type_node); + node.__set_type(type_desc); + node.search_param = std::move(search_param); + node.__isset.search_param = true; + + return {std::make_shared(VSearchExpr::create_shared(node))}; + } + + TabletSchemaSPtr create_tablet_schema_with_keyword_and_fulltext_indexes() { + auto tablet_schema = std::make_shared(); + + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + tablet_schema->append_column(column); + + TabletIndex keyword_index; + keyword_index._index_id = 10; + keyword_index._index_type = IndexType::INVERTED; + keyword_index._col_unique_ids.push_back(1); + tablet_schema->append_index(std::move(keyword_index)); + + TabletIndex fulltext_index; + fulltext_index._index_id = 20; + fulltext_index._index_type = IndexType::INVERTED; + fulltext_index._col_unique_ids.push_back(1); + fulltext_index._properties["parser"] = "standard"; + fulltext_index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(fulltext_index)); + + return tablet_schema; + } + + TabletSchemaSPtr create_array_tablet_schema_with_keyword_and_fulltext_indexes() { + auto tablet_schema = std::make_shared(); + + TabletColumn item; + item.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_ARRAY); + column.add_sub_column(item); + tablet_schema->append_column(column); + + TabletIndex keyword_index; + keyword_index._index_id = 10; + keyword_index._index_type = IndexType::INVERTED; + keyword_index._col_unique_ids.push_back(1); + tablet_schema->append_index(std::move(keyword_index)); + + TabletIndex fulltext_index; + fulltext_index._index_id = 20; + fulltext_index._index_type = IndexType::INVERTED; + fulltext_index._col_unique_ids.push_back(1); + fulltext_index._properties["parser"] = "standard"; + fulltext_index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(fulltext_index)); + + return tablet_schema; + } + + TabletSchemaSPtr create_tablet_schema_with_two_fulltext_indexes() { + auto tablet_schema = std::make_shared(); + + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + tablet_schema->append_column(column); + + for (const auto& [index_id, parser] : {std::pair {10, "standard"}, + std::pair {20, "english"}}) { + TabletIndex index; + index._index_id = index_id; + index._index_type = IndexType::INVERTED; + index._col_unique_ids.push_back(1); + index._properties["parser"] = parser; + index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(index)); + } + + return tablet_schema; + } + + VExprContextSPtrs create_reserved_exact_search_contexts() { + return create_search_contexts( + "EXACT", std::string(segment_v2::inverted_index::CG_V1_MARKER) + "user"); + } + + void expect_no_collected_tokens(const std::wstring& field_name) { + EXPECT_THROW(stats_->get_total_term_cnt_by_col(field_name), Exception); + } + + void expect_collected_tokens(const std::wstring& field_name, uint64_t token_count) { + EXPECT_EQ(stats_->get_total_term_cnt_by_col(field_name), token_count); + } + + void expect_no_collected_term(const std::wstring& field_name, const std::wstring& term) { + EXPECT_THROW(stats_->get_term_doc_freq_by_col(field_name, term), Exception); + } + + void expect_collected_term(const std::wstring& field_name, const std::wstring& term, + uint64_t doc_frequency) { + EXPECT_EQ(stats_->get_term_doc_freq_by_col(field_name, term), doc_frequency); + } + + struct SniiScoringFieldInput { + SniiScoringFieldInput( + std::wstring field_name, + std::optional metadata, + uint64_t index_doc_count, bool has_semantic_norms, + std::optional expected_base_analyzer_fingerprint = std::nullopt) + : field_name(std::move(field_name)), + metadata(std::move(metadata)), + index_doc_count(index_doc_count), + physical_sum_total_term_freq(this->metadata ? this->metadata->scoring_token_count + : 0), + has_semantic_norms(has_semantic_norms), + expected_base_analyzer_fingerprint( + std::move(expected_base_analyzer_fingerprint)) {} + + std::wstring field_name; + std::optional metadata; + uint64_t index_doc_count = 0; + uint64_t physical_sum_total_term_freq = 0; + bool has_scoring_tier = true; + bool has_positions = true; + bool has_semantic_norms = false; + std::optional expected_base_analyzer_fingerprint; + }; + + Status stage_snii_fields_for_test( + CollectionStatistics* statistics, const std::vector& fields, + CollectionStatistics::SniiScoringSegmentAccumulator* segment_accumulator) { + for (const auto& field : fields) { + segment_v2::inverted_index::PlainTermKeyVersion key_version; + const std::string_view expected_base_analyzer_fingerprint = + field.expected_base_analyzer_fingerprint.has_value() + ? *field.expected_base_analyzer_fingerprint + : field.metadata.has_value() + ? std::string_view(field.metadata->base_analyzer_fingerprint) + : std::string_view(); + RETURN_IF_ERROR(statistics->admit_snii_scoring_segment( + field.field_name, field.metadata, expected_base_analyzer_fingerprint, + field.index_doc_count, field.physical_sum_total_term_freq, + field.has_scoring_tier, field.has_positions, field.has_semantic_norms, + &key_version, segment_accumulator)); + } + return Status::OK(); + } + + Status admit_snii_fields_for_test(CollectionStatistics* statistics, + const std::vector& fields) { + CollectionStatistics::SniiScoringSegmentAccumulator segment_accumulator; + RETURN_IF_ERROR(stage_snii_fields_for_test(statistics, fields, &segment_accumulator)); + statistics->commit_snii_scoring_segment(std::move(segment_accumulator)); + return Status::OK(); + } + + Status admit_snii_segment_for_test( + CollectionStatistics* statistics, const std::wstring& field_name, + const std::optional& metadata, + uint64_t index_doc_count, bool has_semantic_norms) { + return admit_snii_fields_for_test( + statistics, {{field_name, metadata, index_doc_count, has_semantic_norms}}); + } + + Status stage_snii_fields_then_file_not_found_for_test( + CollectionStatistics* statistics, + const std::vector& fields_before_failure) { + CollectionStatistics::SniiScoringSegmentAccumulator segment_accumulator; + RETURN_IF_ERROR(stage_snii_fields_for_test(statistics, fields_before_failure, + &segment_accumulator)); + for (const auto& field : fields_before_failure) { + add_term_doc_frequency(&segment_accumulator.term_doc_freqs, field.field_name, L"staged", + 1); + } + return Status::Error( + "simulated later field failure"); + } + + void expect_collected_stats(const std::wstring& field_name, uint64_t doc_count, + uint64_t token_count) { + EXPECT_EQ(stats_->get_doc_num(), doc_count); + expect_collected_tokens(field_name, token_count); + } + std::unique_ptr stats_; std::shared_ptr runtime_state_; std::string test_dir_; @@ -428,7 +881,8 @@ TEST_F(CollectionStatisticsTest, CollectWithMockRowsetSplits) { auto status = stats_->collect(runtime_state_.get(), splits, tablet_schema, expr_contexts, nullptr); - EXPECT_TRUE(status.ok()); + EXPECT_TRUE(status.ok()) << status; + expect_no_collected_tokens(L"1"); } TEST_F(CollectionStatisticsTest, CollectWithEmptySegments) { @@ -442,6 +896,281 @@ TEST_F(CollectionStatisticsTest, CollectWithEmptySegments) { EXPECT_TRUE(status.ok()) << status.msg(); } +TEST_F(CollectionStatisticsTest, LegacyReservedTermUsesRawV3Namespace) { + auto tablet_schema = create_legacy_v3_schema(); + const std::string segment_path = test_dir_ + "/legacy_v3_0.dat"; + ASSERT_TRUE(write_legacy_v3_segment(tablet_schema, segment_path).ok()); + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(1); + rowset->set_segment_path(0, segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = stats_->collect(runtime_state_.get(), splits, tablet_schema, + create_reserved_exact_search_contexts(), nullptr); + + ASSERT_TRUE(status.ok()) << status; + expect_collected_stats(L"1", 1, 2); + expect_collected_term(L"1", + segment_v2::inverted_index::StringHelper::to_wstring( + std::string(segment_v2::inverted_index::CG_V1_MARKER) + "user"), + 0); +} + +TEST_F(CollectionStatisticsTest, LegacyV3SegmentsUsePhysicalScoringStatistics) { + auto tablet_schema = create_legacy_v3_schema(); + const std::string first_segment_path = test_dir_ + "/legacy_v3_0.dat"; + const std::string second_segment_path = test_dir_ + "/legacy_v3_1.dat"; + ASSERT_TRUE(write_legacy_v3_segment(tablet_schema, first_segment_path).ok()); + ASSERT_TRUE(write_legacy_v3_segment(tablet_schema, second_segment_path).ok()); + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(2); + rowset->set_segment_path(0, first_segment_path); + rowset->set_segment_path(1, second_segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = stats_->collect(runtime_state_.get(), splits, tablet_schema, + create_match_expr_contexts("alpha"), nullptr); + + ASSERT_TRUE(status.ok()) << status; + expect_collected_stats(L"1", 2, 4); + expect_collected_term(L"1", L"alpha", 2); +} + +TEST_F(CollectionStatisticsTest, LegacyV3SkipsMissingSegmentAfterCollectingAvailableStatistics) { + auto tablet_schema = create_legacy_v3_schema(); + const std::string first_segment_path = test_dir_ + "/legacy_v3_0.dat"; + ASSERT_TRUE(write_legacy_v3_segment(tablet_schema, first_segment_path).ok()); + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(2); + rowset->set_segment_path(0, first_segment_path); + rowset->set_segment_path(1, test_dir_ + "/missing_v3_1.dat"); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = stats_->collect(runtime_state_.get(), splits, tablet_schema, + create_match_expr_contexts("alpha"), nullptr); + + ASSERT_TRUE(status.ok()) << status; + expect_collected_stats(L"1", 1, 2); + expect_collected_term(L"1", L"alpha", 1); +} + +TEST_F(CollectionStatisticsTest, LegacyV3SkipsEmptySegmentAfterCollectingAvailableStatistics) { + auto tablet_schema = create_legacy_v3_schema(); + const std::string first_segment_path = test_dir_ + "/legacy_v3_0.dat"; + const std::string empty_segment_path = test_dir_ + "/legacy_v3_empty_1.dat"; + ASSERT_TRUE(write_legacy_v3_segment(tablet_schema, first_segment_path).ok()); + + const std::string empty_index_path = + segment_v2::InvertedIndexDescriptor::get_index_file_path_v2( + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix( + empty_segment_path)); + io::FileWriterPtr empty_file; + io::FileWriterOptions options; + ASSERT_TRUE(io::global_local_filesystem() + ->create_file(empty_index_path, &empty_file, &options) + .ok()); + ASSERT_TRUE(empty_file->close(false).ok()); + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(2); + rowset->set_segment_path(0, first_segment_path); + rowset->set_segment_path(1, empty_segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + const Status status = stats_->collect(runtime_state_.get(), splits, tablet_schema, + create_match_expr_contexts("alpha"), nullptr); + + ASSERT_TRUE(status.ok()) << status; + expect_collected_stats(L"1", 1, 2); + expect_collected_term(L"1", L"alpha", 1); +} + +TEST_F(CollectionStatisticsTest, SniiCommonGramsUsesSemanticScoringStatistics) { + auto tablet_schema = create_snii_schema(); + auto expr_contexts = create_match_expr_contexts("alpha", "test-base-v1"); + const auto* analyzer_ctx = expr_contexts.front()->root()->query_analyzer_ctx(); + ASSERT_NE(analyzer_ctx, nullptr); + ASSERT_NE(analyzer_ctx->analyzer_provider, nullptr); + + const std::string segment_path = test_dir_ + "/snii_common_grams_0.dat"; + auto write_status = write_snii_common_grams_segment( + segment_path, + std::string(analyzer_ctx->analyzer_provider->base_analyzer_fingerprint())); + ASSERT_TRUE(write_status.ok()) << write_status; + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(1); + rowset->set_segment_path(0, segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = + stats_->collect(runtime_state_.get(), splits, tablet_schema, expr_contexts, nullptr); + + ASSERT_TRUE(status.ok()) << status; + expect_collected_stats(L"1", 2, 3); + expect_collected_term(L"1", L"alpha", 2); +} + +TEST_F(CollectionStatisticsTest, SniiScoringLookupUsesCallerIoContext) { + snii::snii_test::ScopedEnv force_nonresident_dict("SNII_DICT_RESIDENT_MAX", "0"); + auto tablet_schema = create_snii_schema(); + auto expr_contexts = create_match_expr_contexts("alpha", "test-base-v1"); + const auto* analyzer_ctx = expr_contexts.front()->root()->query_analyzer_ctx(); + ASSERT_NE(analyzer_ctx, nullptr); + ASSERT_NE(analyzer_ctx->analyzer_provider, nullptr); + + const std::string segment_path = test_dir_ + "/snii_io_context_0.dat"; + ASSERT_TRUE(write_snii_common_grams_segment( + segment_path, + std::string(analyzer_ctx->analyzer_provider->base_analyzer_fingerprint())) + .ok()); + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(1); + rowset->set_segment_path(0, segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + const std::string index_path_prefix { + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix(segment_path)}; + io::FileCacheStatistics open_stats; + io::IOContext open_io_ctx; + open_io_ctx.file_cache_stats = &open_stats; + segment_v2::IndexFileReader file_reader(io::global_local_filesystem(), index_path_prefix, + InvertedIndexStorageFormatPB::SNII); + ASSERT_TRUE(file_reader.init(config::inverted_index_read_buffer_size, &open_io_ctx).ok()); + const auto index_metas = tablet_schema->inverted_indexs(1); + ASSERT_EQ(index_metas.size(), 1); + auto logical_reader = file_reader.open_snii_index(index_metas.front(), &open_io_ctx); + ASSERT_TRUE(logical_reader.has_value()) << logical_reader.error(); + ASSERT_GT(open_stats.inverted_index_range_read_count, 0); + + io::FileCacheStatistics collect_stats; + io::IOContext collect_io_ctx; + collect_io_ctx.file_cache_stats = &collect_stats; + const auto status = stats_->collect(runtime_state_.get(), splits, tablet_schema, expr_contexts, + &collect_io_ctx); + + ASSERT_TRUE(status.ok()) << status; + EXPECT_GT(collect_stats.inverted_index_range_read_count, + open_stats.inverted_index_range_read_count); +} + +TEST_F(CollectionStatisticsTest, SniiStatsProviderUsesSemanticCommonGramsTokenCount) { + auto tablet_schema = create_snii_schema(); + auto expr_contexts = create_match_expr_contexts("alpha", "test-base-v1"); + const auto* analyzer_ctx = expr_contexts.front()->root()->query_analyzer_ctx(); + ASSERT_NE(analyzer_ctx, nullptr); + ASSERT_NE(analyzer_ctx->analyzer_provider, nullptr); + + const std::string segment_path = test_dir_ + "/snii_semantic_stats_0.dat"; + auto write_status = write_snii_common_grams_segment( + segment_path, + std::string(analyzer_ctx->analyzer_provider->base_analyzer_fingerprint())); + ASSERT_TRUE(write_status.ok()) << write_status; + + const std::string index_path_prefix { + segment_v2::InvertedIndexDescriptor::get_index_file_path_prefix(segment_path)}; + segment_v2::IndexFileReader file_reader(io::global_local_filesystem(), index_path_prefix, + InvertedIndexStorageFormatPB::SNII); + ASSERT_TRUE(file_reader.init().ok()); + const auto index_metas = tablet_schema->inverted_indexs(1); + ASSERT_EQ(index_metas.size(), 1); + auto logical_reader = file_reader.open_snii_index(index_metas.front()); + ASSERT_TRUE(logical_reader.has_value()) << logical_reader.error(); + + snii::stats::SniiStatsProvider provider; + ASSERT_TRUE(snii::stats::SniiStatsProvider::open(logical_reader->get(), &provider).ok()); + EXPECT_EQ(provider.doc_count(), 2); + EXPECT_EQ(provider.indexed_doc_count(), 2); + EXPECT_EQ(provider.sum_total_term_freq(), 3); + EXPECT_DOUBLE_EQ(provider.avgdl(), 1.5); + EXPECT_TRUE(provider.has_norms()); +} + +TEST_F(CollectionStatisticsTest, SniiWriterRejectsMissingSemanticScoringMetadata) { + const std::string segment_path = test_dir_ + "/snii_missing_scoring_metadata_0.dat"; + auto write_status = write_plain_snii_scoring_segment(segment_path); + EXPECT_EQ(write_status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST_F(CollectionStatisticsTest, SniiScoringRejectsMissingSegmentForWholeCollection) { + auto tablet_schema = create_snii_schema(); + auto expr_contexts = create_match_expr_contexts("alpha", "test-base-v1"); + const auto* analyzer_ctx = expr_contexts.front()->root()->query_analyzer_ctx(); + ASSERT_NE(analyzer_ctx, nullptr); + ASSERT_NE(analyzer_ctx->analyzer_provider, nullptr); + + const std::string first_segment_path = test_dir_ + "/snii_complete_0.dat"; + auto write_status = write_snii_common_grams_segment( + first_segment_path, + std::string(analyzer_ctx->analyzer_provider->base_analyzer_fingerprint())); + ASSERT_TRUE(write_status.ok()) << write_status; + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(2); + rowset->set_segment_path(0, first_segment_path); + rowset->set_segment_path(1, test_dir_ + "/missing_snii_1.dat"); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = + stats_->collect(runtime_state_.get(), splits, tablet_schema, expr_contexts, nullptr); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST_F(CollectionStatisticsTest, SniiScoringRejectsMissingLogicalIndexForWholeCollection) { + auto tablet_schema = create_snii_schema(/*index_id=*/2); + auto expr_contexts = create_match_expr_contexts("alpha", "test-base-v1"); + const auto* analyzer_ctx = expr_contexts.front()->root()->query_analyzer_ctx(); + ASSERT_NE(analyzer_ctx, nullptr); + ASSERT_NE(analyzer_ctx->analyzer_provider, nullptr); + + const std::string segment_path = test_dir_ + "/snii_missing_logical_index_0.dat"; + auto write_status = write_snii_common_grams_segment( + segment_path, + std::string(analyzer_ctx->analyzer_provider->base_analyzer_fingerprint())); + ASSERT_TRUE(write_status.ok()) << write_status; + + auto rowset_meta = std::make_shared(); + auto rowset = std::make_shared(tablet_schema, rowset_meta); + rowset->set_num_segments(1); + rowset->set_segment_path(0, segment_path); + auto reader = std::make_shared(rowset); + std::vector splits {RowSetSplits(reader)}; + + auto status = + stats_->collect(runtime_state_.get(), splits, tablet_schema, expr_contexts, nullptr); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); +} + +TEST_F(CollectionStatisticsTest, SniiWriterRejectsZeroSemanticTokensForNonemptyPostings) { + auto status = write_snii_common_grams_segment(test_dir_ + "/snii_zero_semantic_tokens_0.dat", + "test-base-v1", + /*scoring_token_count=*/0); + + EXPECT_EQ(status.code(), ErrorCode::INVALID_ARGUMENT); + EXPECT_THAT(status.msg(), ::testing::HasSubstr("zero semantic scoring tokens")); +} + TEST_F(CollectionStatisticsTest, CollectWithMultipleRowsetSplits) { auto tablet_schema = create_tablet_schema_with_inverted_index(); auto expr_contexts = create_match_expr_contexts(); @@ -488,6 +1217,382 @@ class CollectionStatisticsDetailedTest : public ::testing::Test { std::unique_ptr stats_; }; +segment_v2::inverted_index::CommonGramsSegmentMetadata complete_snii_scoring_metadata( + std::string base_analyzer_fingerprint = "base-v1", uint64_t doc_count = 3, + uint64_t token_count = 7) { + using namespace segment_v2::inverted_index; + segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = CommonGramsCoverage::kComplete; + metadata.common_grams_semantics_version = COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = "builtin-stopwords:v1"; + metadata.base_analyzer_fingerprint = std::move(base_analyzer_fingerprint); + metadata.common_grams_fingerprint = "common-grams-v1"; + metadata.scoring_coverage = ScoringCoverage::kComplete; + metadata.scoring_stats_version = COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + metadata.scoring_doc_count = doc_count; + metadata.scoring_token_count = token_count; + return metadata; +} + +Result resolve_snii_scoring_segment_for_test( + std::optional metadata, + uint64_t physical_doc_count, bool has_norms) { + const uint64_t physical_sum_total_term_freq = metadata ? metadata->scoring_token_count : 0; + return resolve_snii_scoring_segment(metadata, physical_doc_count, physical_sum_total_term_freq, + /*has_scoring_tier=*/true, + /*has_positions=*/true, has_norms); +} + +TEST(CollectionStatisticsCommonGramsTest, MissingMetadataWithoutPersistedProofRejectsScoring) { + auto result = resolve_snii_scoring_segment_for_test(std::nullopt, 3, true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST(CollectionStatisticsCommonGramsTest, RawNoInternalMetadataWithoutScoringProofIsRejected) { + segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + metadata.common_grams_coverage = segment_v2::inverted_index::CommonGramsCoverage::kNone; + metadata.base_analyzer_fingerprint = "base-v1"; + + auto result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST(CollectionStatisticsCommonGramsTest, ScoringCoverageNoneRejectsScoringAdmission) { + auto metadata = complete_snii_scoring_metadata(); + metadata.scoring_coverage = segment_v2::inverted_index::ScoringCoverage::kNone; + metadata.scoring_stats_version = 0; + metadata.norm_semantics_version = 0; + + auto result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST(CollectionStatisticsCommonGramsTest, IncompatibleScoringVersionsRejectScoringAdmission) { + auto metadata = complete_snii_scoring_metadata(); + + ++metadata.scoring_stats_version; + auto scoring_version_result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + ASSERT_FALSE(scoring_version_result.has_value()); + EXPECT_EQ(scoring_version_result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + + --metadata.scoring_stats_version; + ++metadata.norm_semantics_version; + auto norm_version_result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + ASSERT_FALSE(norm_version_result.has_value()); + EXPECT_EQ(norm_version_result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST(CollectionStatisticsCommonGramsTest, ScoringDocCountMismatchRejectsScoringAdmission) { + auto result = resolve_snii_scoring_segment_for_test(complete_snii_scoring_metadata(), 4, true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST(CollectionStatisticsCommonGramsTest, MissingSemanticNormsRejectScoringAdmission) { + auto result = resolve_snii_scoring_segment_for_test(complete_snii_scoring_metadata(), 3, false); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST(CollectionStatisticsCommonGramsTest, CompleteMetadataOnNonScoringTierIsCorruption) { + auto metadata = complete_snii_scoring_metadata(); + auto result = resolve_snii_scoring_segment(metadata, 3, 7, + /*has_scoring_tier=*/false, + /*has_positions=*/true, + /*has_semantic_norms=*/true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST(CollectionStatisticsCommonGramsTest, ZeroSemanticTokensWithPhysicalTermsIsCorruption) { + auto metadata = complete_snii_scoring_metadata("base-v1", 3, 0); + auto result = resolve_snii_scoring_segment(metadata, 3, 1, + /*has_scoring_tier=*/true, + /*has_positions=*/true, + /*has_semantic_norms=*/true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST(CollectionStatisticsCommonGramsTest, EmptyPhysicalAndSemanticTokenCountsAreValid) { + auto metadata = complete_snii_scoring_metadata("base-v1", 3, 0); + auto result = resolve_snii_scoring_segment(metadata, 3, 0, + /*has_scoring_tier=*/true, + /*has_positions=*/true, + /*has_semantic_norms=*/true); + + ASSERT_TRUE(result.has_value()) << result.error(); + EXPECT_EQ(result->token_count, 0); +} + +TEST(CollectionStatisticsCommonGramsTest, LegacyPhysicalScoringRequiresDocumentLengthNorms) { + auto result = resolve_snii_scoring_segment_for_test(std::nullopt, 3, false); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); +} + +TEST(CollectionStatisticsCommonGramsTest, CompleteMetadataUsesSemanticDocAndTokenCounts) { + auto result = resolve_snii_scoring_segment_for_test(complete_snii_scoring_metadata(), 3, true); + + ASSERT_TRUE(result.has_value()) << result.error(); + EXPECT_EQ(result->doc_count, 3); + EXPECT_EQ(result->token_count, 7); + EXPECT_EQ(result->plain_term_key_version, + segment_v2::inverted_index::PlainTermKeyVersion::kEscapedV1); + EXPECT_EQ(result->base_analyzer_fingerprint, "base-v1"); +} + +TEST(CollectionStatisticsCommonGramsTest, CompletePlainMetadataIsExplicitSemanticProof) { + auto metadata = complete_snii_scoring_metadata(); + metadata.plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + metadata.common_grams_coverage = segment_v2::inverted_index::CommonGramsCoverage::kNone; + metadata.common_grams_semantics_version = 0; + metadata.common_grams_key_version = 0; + metadata.common_grams_dictionary_identity.clear(); + metadata.common_grams_fingerprint.clear(); + + auto result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + + ASSERT_TRUE(result.has_value()) << result.error(); + EXPECT_EQ(result->doc_count, 3); + EXPECT_EQ(result->token_count, 7); +} + +TEST(CollectionStatisticsCommonGramsTest, PlainSemanticTokenCountMustEqualPhysicalCount) { + auto metadata = complete_snii_scoring_metadata(); + metadata.plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + metadata.common_grams_coverage = segment_v2::inverted_index::CommonGramsCoverage::kNone; + metadata.common_grams_semantics_version = 0; + metadata.common_grams_key_version = 0; + metadata.common_grams_dictionary_identity.clear(); + metadata.common_grams_fingerprint.clear(); + + auto result = resolve_snii_scoring_segment(metadata, 3, 8, + /*has_scoring_tier=*/true, + /*has_positions=*/true, + /*has_semantic_norms=*/true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST(CollectionStatisticsCommonGramsTest, EmptySemanticFingerprintIsNotScoringProof) { + auto metadata = complete_snii_scoring_metadata(); + metadata.base_analyzer_fingerprint.clear(); + + auto result = resolve_snii_scoring_segment_for_test(metadata, 3, true); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), ErrorCode::INVERTED_INDEX_FILE_CORRUPTED); +} + +TEST_F(CollectionStatisticsTest, MixedBaseFingerprintRejectsAndClearsWholeCollection) { + auto first = complete_snii_scoring_metadata("base-v1", 3, 7); + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", first, 3, true).ok()); + EXPECT_FLOAT_EQ(stats_->get_or_calculate_avg_dl(L"1"), 7.0F / 3.0F); + + auto second = complete_snii_scoring_metadata("base-v2", 3, 5); + auto status = admit_snii_segment_for_test(stats_.get(), L"1", second, 3, true); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); + EXPECT_THROW(stats_->get_or_calculate_avg_dl(L"1"), Exception); +} + +TEST_F(CollectionStatisticsTest, PersistedFingerprintMustMatchRequestAnalyzer) { + auto metadata = complete_snii_scoring_metadata("persisted-base", 3, 7); + auto status = admit_snii_fields_for_test( + stats_.get(), {{L"1", metadata, 3, true, std::string("request-base")}}); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); + EXPECT_THROW(stats_->get_doc_num(), Exception); +} + +TEST_F(CollectionStatisticsTest, CollectionStatisticsInstancesKeepAdmissionStateIsolated) { + CollectionStatistics first; + CollectionStatistics second; + + ASSERT_TRUE(admit_snii_segment_for_test(&first, L"1", + complete_snii_scoring_metadata("base-a", 2, 6), 2, true) + .ok()); + ASSERT_TRUE(admit_snii_segment_for_test( + &second, L"1", complete_snii_scoring_metadata("base-b", 5, 25), 5, true) + .ok()); + + EXPECT_FLOAT_EQ(first.get_or_calculate_avg_dl(L"1"), 3.0F); + EXPECT_FLOAT_EQ(second.get_or_calculate_avg_dl(L"1"), 5.0F); +} + +TEST_F(CollectionStatisticsTest, LegacyAndUnprovedRawNoInternalRejectWholeCollection) { + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("base-v1", 2, 6), 2, + true) + .ok()); + + segment_v2::inverted_index::CommonGramsSegmentMetadata plain_metadata; + plain_metadata.plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + plain_metadata.common_grams_coverage = segment_v2::inverted_index::CommonGramsCoverage::kNone; + plain_metadata.base_analyzer_fingerprint = "base-v1"; + auto status = admit_snii_segment_for_test(stats_.get(), L"1", plain_metadata, 3, true); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); + EXPECT_THROW(stats_->get_doc_num(), Exception); +} + +TEST_F(CollectionStatisticsTest, LegacyAndCommonGramsMixRejectsWholeCollection) { + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("base-v1", 2, 6), 2, + true) + .ok()); + + auto status = admit_snii_segment_for_test(stats_.get(), L"1", std::nullopt, 3, true); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); + EXPECT_THROW(stats_->get_or_calculate_avg_dl(L"1"), Exception); +} + +TEST_F(CollectionStatisticsTest, ExplicitSemanticPlainAndCommonGramsSegmentsAccumulate) { + auto plain_metadata = complete_snii_scoring_metadata("base-v1", 2, 6); + plain_metadata.plain_term_key_version = + segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + plain_metadata.common_grams_coverage = segment_v2::inverted_index::CommonGramsCoverage::kNone; + plain_metadata.common_grams_semantics_version = 0; + plain_metadata.common_grams_key_version = 0; + plain_metadata.common_grams_dictionary_identity.clear(); + plain_metadata.common_grams_fingerprint.clear(); + + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", plain_metadata, 2, true).ok()); + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("base-v1", 3, 7), 3, + true) + .ok()); + + expect_collected_stats(L"1", 5, 13); + EXPECT_FLOAT_EQ(stats_->get_or_calculate_avg_dl(L"1"), 13.0F / 5.0F); +} + +TEST_F(CollectionStatisticsTest, CommonGramsSegmentsAccumulateSemanticStatistics) { + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("base-v1", 2, 6), 2, + true) + .ok()); + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("base-v1", 3, 9), 3, + true) + .ok()); + + expect_collected_stats(L"1", 5, 15); + EXPECT_FLOAT_EQ(stats_->get_or_calculate_avg_dl(L"1"), 3.0F); +} + +TEST_F(CollectionStatisticsTest, MultiFieldSegmentsCommitAndAccumulateAtomically) { + ASSERT_TRUE(admit_snii_fields_for_test( + stats_.get(), + {{L"1", complete_snii_scoring_metadata("field-1", 3, 7), 3, true}, + {L"2", complete_snii_scoring_metadata("field-2", 3, 12), 3, true}}) + .ok()); + ASSERT_TRUE(admit_snii_fields_for_test( + stats_.get(), + {{L"1", complete_snii_scoring_metadata("field-1", 2, 5), 2, true}, + {L"2", complete_snii_scoring_metadata("field-2", 2, 8), 2, true}}) + .ok()); + + expect_collected_stats(L"1", 5, 12); + expect_collected_tokens(L"2", 20); + EXPECT_FLOAT_EQ(stats_->get_or_calculate_avg_dl(L"1"), 12.0F / 5.0F); + EXPECT_FLOAT_EQ(stats_->get_or_calculate_avg_dl(L"2"), 4.0F); +} + +TEST_F(CollectionStatisticsTest, MultiFieldSegmentDocCountsMustAgree) { + auto status = admit_snii_fields_for_test( + stats_.get(), {{L"1", complete_snii_scoring_metadata("field-1", 3, 7), 3, true}, + {L"2", complete_snii_scoring_metadata("field-2", 4, 12), 4, true}}); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + expect_no_collected_tokens(L"1"); + expect_no_collected_tokens(L"2"); + EXPECT_THROW(stats_->get_doc_num(), Exception); +} + +TEST_F(CollectionStatisticsTest, LaterFieldFileNotFoundDoesNotPublishPartialSegment) { + ASSERT_TRUE(admit_snii_segment_for_test(stats_.get(), L"1", + complete_snii_scoring_metadata("field-1", 2, 6), 2, + true) + .ok()); + + auto status = stage_snii_fields_then_file_not_found_for_test( + stats_.get(), + {{L"2", complete_snii_scoring_metadata("staged-field-2", 3, 12), 3, true}}); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_FILE_NOT_FOUND); + expect_collected_stats(L"1", 2, 6); + expect_no_collected_tokens(L"2"); + expect_no_collected_term(L"2", L"staged"); + + ASSERT_TRUE(admit_snii_segment_for_test( + stats_.get(), L"2", + complete_snii_scoring_metadata("different-field-2", 1, 4), 1, true) + .ok()); + expect_collected_tokens(L"2", 4); +} + +TEST(CollectionStatisticsCommonGramsTest, DocFrequencySupportsLogicalAndPhysicalTermKeys) { + std::unordered_map> + logical_frequencies; + + add_term_doc_frequency(&logical_frequencies, L"field", L"logical", 3); + EXPECT_EQ(logical_frequencies[L"field"][L"logical"], 3); + + add_term_doc_frequency(&logical_frequencies, L"field", L"same", 5); + EXPECT_EQ(logical_frequencies[L"field"][L"same"], 5); +} + +TEST(CollectionStatisticsCommonGramsTest, PhysicalAliasesCannotMergeLogicalDocFrequencies) { + using Frequencies = + std::unordered_map>; + Frequencies logical_frequencies; + + const std::wstring alias = std::wstring(1, wchar_t {0x1e}) + L"G00000001:"; + add_term_doc_frequency(&logical_frequencies, L"field", L"\x1f", 3); + add_term_doc_frequency(&logical_frequencies, L"field", alias, 5); + + EXPECT_EQ(logical_frequencies[L"field"][alias], 5); +} + +TEST(CollectionStatisticsCommonGramsTest, UnrepresentablePlainTermRegistersZeroDocFrequency) { + using Frequencies = + std::unordered_map>; + Frequencies logical_frequencies; + + add_term_doc_frequency(&logical_frequencies, L"field", L"unrepresentable", 0); + + ASSERT_TRUE(logical_frequencies.contains(L"field")); + EXPECT_TRUE(logical_frequencies.at(L"field").contains(L"unrepresentable")); + EXPECT_EQ(logical_frequencies.at(L"field").at(L"unrepresentable"), 0); +} + TEST_F(CollectionStatisticsDetailedTest, GetStatisticsWithValidData) { std::wstring field_name = L"test_field"; std::wstring term = L"test_term"; @@ -694,13 +1799,7 @@ TEST_F(CollectionStatisticsTest, ExtractCollectInfoForVariantSubcolumnIndex) { EXPECT_EQ(it->second.index_meta->index_name(), "variant_subcolumn_idx"); } -// Regression for score on a dynamic variant sub-column inherited from a plain -// parent variant inverted index (no field_pattern template). Matches the -// scan-time schema shape: _init_variant_columns materializes the accessed -// path as an extracted VARIANT placeholder, so neither inverted_indexs(column) -// nor generate_sub_column_info resolves the parent index. Collector clones -// the parent's non-field-pattern indexes with the variant path as suffix. -TEST_F(CollectionStatisticsTest, ExtractCollectInfoForVariantParentIndexWithoutTemplate) { +TEST_F(CollectionStatisticsTest, MatchScoringUsesTextSemanticsForVariantParentIndexFallback) { auto tablet_schema = std::make_shared(); constexpr int32_t kVariantUid = 9004; @@ -761,13 +1860,14 @@ TEST_F(CollectionStatisticsTest, ExtractCollectInfoForVariantParentIndexWithoutT std::unordered_map collect_infos; auto status = stats_->extract_collect_info(runtime_state_.get(), contexts, tablet_schema, &collect_infos); - ASSERT_TRUE(status.ok()) << status.msg(); - ASSERT_EQ(collect_infos.size(), 1u); + ASSERT_TRUE(status.ok()) << status; + ASSERT_EQ(collect_infos.size(), 1U); auto it = collect_infos.find(StringHelper::to_wstring(std::to_string(kVariantUid) + ".v.key")); ASSERT_NE(it, collect_infos.end()); ASSERT_NE(it->second.index_meta, nullptr); ASSERT_NE(it->second.owned_index_meta, nullptr); - EXPECT_EQ(it->second.index_meta->index_name(), "variant_parent_idx"); + EXPECT_EQ(it->second.index_meta->index_id(), 2004); + EXPECT_EQ(it->second.unique_terms, std::vector({"abc"})); } namespace { @@ -1059,7 +2159,7 @@ TEST_F(CollectionStatisticsTest, CollectDirectIndexHitFromSchema) { ASSERT_NE(it, collect_infos.end()); EXPECT_NE(it->second.index_meta, nullptr); EXPECT_EQ(it->second.owned_index_meta, nullptr); // O1: schema-direct meta is not owned - EXPECT_FALSE(it->second.term_infos.empty()); + EXPECT_FALSE(it->second.unique_terms.empty()); } // I2: Plain string column with no index and not an extracted variant @@ -1176,10 +2276,7 @@ TEST_F(CollectionStatisticsTest, CollectSkipsIndexWithoutSimilarityScore) { EXPECT_TRUE(collect_infos.empty()); } -// L4: Two MATCH predicates on the same column produce CollectInfo entries -// keyed on the same field_name; the second insertion merges term_infos -// into the first entry. -TEST_F(CollectionStatisticsTest, CollectMergesTermsForSameFieldName) { +TEST_F(CollectionStatisticsTest, CollectPreservesLogicalClauseShapesForSameFieldName) { auto tablet_schema = std::make_shared(); constexpr int32_t kColUid = 1400; @@ -1215,16 +2312,539 @@ TEST_F(CollectionStatisticsTest, CollectMergesTermsForSameFieldName) { MatchPredicateCollector collector; std::unordered_map collect_infos; - auto first = collector.collect(runtime_state_.get(), tablet_schema, build_match("alpha"), + auto first = collector.collect(runtime_state_.get(), tablet_schema, build_match("alpha beta"), &collect_infos); ASSERT_TRUE(first.ok()) << first.msg(); - auto second = collector.collect(runtime_state_.get(), tablet_schema, build_match("beta"), - &collect_infos); + auto second = collector.collect(runtime_state_.get(), tablet_schema, + build_match("alpha alpha beta"), &collect_infos); ASSERT_TRUE(second.ok()) << second.msg(); ASSERT_EQ(collect_infos.size(), 1u); auto it = collect_infos.find(StringHelper::to_wstring(std::to_string(kColUid))); ASSERT_NE(it, collect_infos.end()); - EXPECT_GE(it->second.term_infos.size(), 2u); // both "alpha" and "beta" present + ASSERT_EQ(it->second.unique_terms, std::vector({"alpha", "beta"})); + ASSERT_EQ(it->second.unique_term_slots.size(), 2u); + EXPECT_EQ(it->second.unique_term_slots.at("alpha"), 0u); + EXPECT_EQ(it->second.unique_term_slots.at("beta"), 1u); + ASSERT_EQ(it->second.logical_scoring_leaves.size(), 2u); + ASSERT_EQ(it->second.logical_scoring_leaves[0].clauses.size(), 2u); + EXPECT_EQ(it->second.logical_scoring_leaves[0].clauses[0].df_slot, 0u); + EXPECT_EQ(it->second.logical_scoring_leaves[0].clauses[0].position, 1); + EXPECT_EQ(it->second.logical_scoring_leaves[0].clauses[1].df_slot, 1u); + EXPECT_EQ(it->second.logical_scoring_leaves[0].clauses[1].position, 2); + ASSERT_EQ(it->second.logical_scoring_leaves[1].clauses.size(), 3u); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[0].df_slot, 0u); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[0].position, 1); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[1].df_slot, 0u); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[1].position, 2); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[2].df_slot, 1u); + EXPECT_EQ(it->second.logical_scoring_leaves[1].clauses[2].position, 3); +} + +TEST_F(CollectionStatisticsTest, CollectUsesMatchRequestAnalyzerProviderAndFingerprint) { + auto tablet_schema = create_tablet_schema_with_inverted_index(); + + auto analyzer = segment_v2::inverted_index::InvertedIndexAnalyzer::create_builtin_analyzer( + InvertedIndexParserType::PARSER_ENGLISH, "", INVERTED_INDEX_PARSER_FALSE, "none"); + auto provider = std::make_shared( + std::move(analyzer), "request-base-v1"); + auto analyzer_ctx = std::make_shared(); + analyzer_ctx->analyzer_provider = provider; + + auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->set_analyzer_ctx(std::move(analyzer_ctx)); + auto slot_ref = std::make_shared("content", SlotId(1)); + auto literal = std::make_shared("Alpha ALPHA"); + match_expr->_children.push_back(slot_ref); + match_expr->_children.push_back(literal); + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto status = + collector.collect(runtime_state_.get(), tablet_schema, match_expr, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + EXPECT_EQ(collect_info.expected_base_analyzer_fingerprint, "request-base-v1"); + ASSERT_EQ(collect_info.unique_terms, std::vector({"Alpha", "ALPHA"})); + ASSERT_EQ(collect_info.logical_scoring_leaves.size(), 1u); + ASSERT_EQ(collect_info.logical_scoring_leaves[0].clauses.size(), 2u); + EXPECT_EQ(collect_info.logical_scoring_leaves[0].clauses[0].df_slot, 0u); + EXPECT_EQ(collect_info.logical_scoring_leaves[0].clauses[1].df_slot, 1u); +} + +TEST_F(CollectionStatisticsTest, CollectPhrasePrefixExcludesScoringTail) { + auto tablet_schema = create_tablet_schema_with_inverted_index(); + auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->set_opcode(TExprOpcode::MATCH_PHRASE_PREFIX); + match_expr->_children.push_back( + std::make_shared("content", SlotId(1))); + match_expr->_children.push_back( + std::make_shared("alpha beta gamma")); + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto status = + collector.collect(runtime_state_.get(), tablet_schema, match_expr, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + EXPECT_EQ(collect_info.unique_terms, std::vector({"alpha", "beta"})); + ASSERT_EQ(collect_info.logical_scoring_leaves.size(), 1u); + ASSERT_EQ(collect_info.logical_scoring_leaves[0].clauses.size(), 2u); + EXPECT_EQ(collect_info.logical_scoring_leaves[0].clauses[0].position, 1); + EXPECT_EQ(collect_info.logical_scoring_leaves[0].clauses[1].position, 2); +} + +TEST_F(CollectionStatisticsTest, CollectSingleTermPhrasePrefixHasEmptyScoringLeaf) { + auto tablet_schema = create_tablet_schema_with_inverted_index(); + auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->set_opcode(TExprOpcode::MATCH_PHRASE_PREFIX); + match_expr->_children.push_back( + std::make_shared("content", SlotId(1))); + match_expr->_children.push_back(std::make_shared("alpha")); + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto status = + collector.collect(runtime_state_.get(), tablet_schema, match_expr, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + EXPECT_TRUE(collect_info.unique_terms.empty()); + ASSERT_EQ(collect_info.logical_scoring_leaves.size(), 1u); + EXPECT_TRUE(collect_info.logical_scoring_leaves[0].clauses.empty()); +} + +TEST_F(CollectionStatisticsTest, SearchMatchCollectsRawExecutionTerm) { + auto tablet_schema = create_tablet_schema_with_inverted_index(); + auto contexts = create_search_contexts("MATCH", "alpha beta"); + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), contexts, tablet_schema, + &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + EXPECT_EQ(collect_info.unique_terms, std::vector({"alpha beta"})); + ASSERT_EQ(collect_info.logical_scoring_leaves.size(), 1u); + ASSERT_EQ(collect_info.logical_scoring_leaves[0].clauses.size(), 1u); + EXPECT_EQ(collect_info.logical_scoring_leaves[0].clauses[0].df_slot, 0u); +} + +TEST_F(CollectionStatisticsTest, MatchSelectsOnlyTheRuntimeAnalyzerIndex) { + auto tablet_schema = create_tablet_schema_with_two_fulltext_indexes(); + auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->_children.push_back( + std::make_shared("content", SlotId(1))); + match_expr->_children.push_back( + std::make_shared("running quickly")); + + InvertedIndexAnalyzerConfig config; + config.analyzer_name = "english"; + config.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + config.stop_words = "none"; + auto analyzer_ctx = std::make_shared(); + analyzer_ctx->analyzer_name = "english"; + analyzer_ctx->analyzer_provider = + segment_v2::inverted_index::InvertedIndexAnalyzer::create_analyzer_provider(&config); + match_expr->set_analyzer_ctx(std::move(analyzer_ctx)); + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto status = + collector.collect(runtime_state_.get(), tablet_schema, match_expr, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + ASSERT_NE(collect_info.index_meta, nullptr); + EXPECT_EQ(collect_info.index_meta->index_id(), 20); + EXPECT_EQ(collect_info.logical_scoring_leaves.size(), 1u); +} + +TEST_F(CollectionStatisticsTest, MatchArrayStringSelectsFulltextLeafIndex) { + auto tablet_schema = create_array_tablet_schema_with_keyword_and_fulltext_indexes(); + auto match_expr = std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->_children.push_back( + std::make_shared("content", SlotId(1))); + match_expr->_children.push_back( + std::make_shared("alpha beta")); + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto status = + collector.collect(runtime_state_.get(), tablet_schema, match_expr, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + ASSERT_NE(collect_infos.begin()->second.index_meta, nullptr); + EXPECT_EQ(collect_infos.begin()->second.index_meta->index_id(), 20); +} + +TEST_F(CollectionStatisticsTest, SearchTermSelectsOnlyTheRuntimeFullTextIndex) { + auto tablet_schema = create_tablet_schema_with_keyword_and_fulltext_indexes(); + auto contexts = create_search_contexts("TERM", "alpha beta"); + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), contexts, tablet_schema, + &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + ASSERT_NE(collect_info.index_meta, nullptr); + EXPECT_EQ(collect_info.index_meta->index_id(), 20); + EXPECT_EQ(collect_info.logical_scoring_leaves.size(), 1u); +} + +TEST_F(CollectionStatisticsTest, SearchArrayStringSelectsFulltextLeafIndex) { + auto tablet_schema = create_array_tablet_schema_with_keyword_and_fulltext_indexes(); + auto contexts = create_search_contexts("TERM", "alpha beta"); + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), contexts, tablet_schema, + &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + ASSERT_NE(collect_infos.begin()->second.index_meta, nullptr); + EXPECT_EQ(collect_infos.begin()->second.index_meta->index_id(), 20); +} + +TEST_F(CollectionStatisticsTest, SearchExactIgnoresAnalyzedBindingHint) { + auto tablet_schema = create_tablet_schema_with_keyword_and_fulltext_indexes(); + TSearchClause clause; + clause.clause_type = "EXACT"; + clause.field_name = "content"; + clause.value = "running quickly"; + clause.__isset.field_name = true; + clause.__isset.value = true; + TSearchFieldBinding binding; + binding.field_name = "content"; + binding.slot_index = 0; + binding.index_properties["parser"] = "english"; + binding.index_properties["support_phrase"] = "true"; + binding.__isset.index_properties = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info( + runtime_state_.get(), create_search_contexts(std::move(clause), {std::move(binding)}), + tablet_schema, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + ASSERT_NE(collect_infos.begin()->second.index_meta, nullptr); + EXPECT_EQ(collect_infos.begin()->second.index_meta->index_id(), 10); + EXPECT_EQ(collect_infos.begin()->second.unique_terms, + std::vector({"running quickly"})); +} + +TEST_F(CollectionStatisticsTest, SearchScoringRejectsMissingField) { + auto tablet_schema = create_tablet_schema_with_inverted_index(); + TSearchClause clause; + clause.clause_type = "TERM"; + clause.field_name = "missing"; + clause.value = "alpha"; + clause.__isset.field_name = true; + clause.__isset.value = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), + create_search_contexts(std::move(clause)), + tablet_schema, &collect_infos); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + EXPECT_TRUE(collect_infos.empty()); +} + +TEST_F(CollectionStatisticsTest, SearchScoringRejectsMissingIndex) { + auto tablet_schema = std::make_shared(); + TabletColumn column; + column.set_unique_id(1); + column.set_name("content"); + column.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + tablet_schema->append_column(column); + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), + create_search_contexts("TERM", "alpha"), + tablet_schema, &collect_infos); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + EXPECT_TRUE(collect_infos.empty()); +} + +TEST_F(CollectionStatisticsTest, SearchTypedVariantBindingSelectsItsAnalyzerIndex) { + auto tablet_schema = std::make_shared(); + constexpr int32_t kVariantUid = 9010; + + TabletColumn variant_column; + variant_column.set_unique_id(kVariantUid); + variant_column.set_name("v"); + variant_column.set_type(FieldType::OLAP_FIELD_TYPE_VARIANT); + tablet_schema->append_column(variant_column); + + TabletColumn subcolumn; + subcolumn.set_unique_id(-1); + subcolumn.set_name("v.host"); + subcolumn.set_type(FieldType::OLAP_FIELD_TYPE_STRING); + subcolumn.set_parent_unique_id(kVariantUid); + subcolumn.set_path_info(PathInData("v.host", true)); + tablet_schema->append_column(subcolumn); + + TabletSchema::PathsSetInfo path_set_info; + TabletSchema::SubColumnInfo typed_path_info; + typed_path_info.column = subcolumn; + for (const auto& [index_id, parser] : {std::pair {3010, "standard"}, + std::pair {3020, "english"}}) { + auto index = std::make_shared(); + TabletIndexPB index_pb; + index_pb.set_index_id(index_id); + index_pb.set_index_name(parser + "_variant_idx"); + index_pb.set_index_type(IndexType::INVERTED); + index_pb.add_col_unique_id(kVariantUid); + (*index_pb.mutable_properties())["parser"] = parser; + (*index_pb.mutable_properties())["support_phrase"] = "true"; + index->init_from_pb(index_pb); + typed_path_info.indexes.push_back(std::move(index)); + } + path_set_info.typed_path_set.emplace("host", std::move(typed_path_info)); + std::unordered_map path_set_info_map; + path_set_info_map.emplace(kVariantUid, std::move(path_set_info)); + tablet_schema->set_path_set_info(std::move(path_set_info_map)); + + TSearchClause clause; + clause.clause_type = "TERM"; + clause.field_name = "v.host"; + clause.value = "running"; + clause.__isset.field_name = true; + clause.__isset.value = true; + TSearchFieldBinding binding; + binding.field_name = "v.host"; + binding.slot_index = 0; + binding.is_variant_subcolumn = true; + binding.__isset.is_variant_subcolumn = true; + binding.parent_field_name = "v"; + binding.__isset.parent_field_name = true; + binding.subcolumn_path = "host"; + binding.__isset.subcolumn_path = true; + binding.index_properties["parser"] = "english"; + binding.index_properties["support_phrase"] = "true"; + binding.__isset.index_properties = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info( + runtime_state_.get(), create_search_contexts(std::move(clause), {std::move(binding)}), + tablet_schema, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + const auto& collect_info = collect_infos.begin()->second; + ASSERT_NE(collect_info.index_meta, nullptr); + EXPECT_EQ(collect_info.index_meta->index_id(), 3020); + EXPECT_EQ(collect_info.unique_terms, std::vector({"running"})); +} + +TEST_F(CollectionStatisticsTest, SearchScoringUsesTextSemanticsForVariantParentIndexFallback) { + auto tablet_schema = std::make_shared(); + constexpr int32_t kVariantUid = 9015; + + TabletColumn variant_column; + variant_column.set_unique_id(kVariantUid); + variant_column.set_name("v"); + variant_column.set_type(FieldType::OLAP_FIELD_TYPE_VARIANT); + tablet_schema->append_column(variant_column); + + TabletIndex parent_index; + parent_index._index_id = 3025; + parent_index._index_type = IndexType::INVERTED; + parent_index._col_unique_ids.push_back(kVariantUid); + parent_index._properties["parser"] = "standard"; + parent_index._properties["support_phrase"] = "true"; + tablet_schema->append_index(std::move(parent_index)); + + TSearchClause clause; + clause.clause_type = "PHRASE"; + clause.field_name = "v.dynamic"; + clause.value = "alpha beta"; + clause.__isset.field_name = true; + clause.__isset.value = true; + TSearchFieldBinding binding; + binding.field_name = "v.dynamic"; + binding.slot_index = 0; + binding.is_variant_subcolumn = true; + binding.__isset.is_variant_subcolumn = true; + binding.parent_field_name = "v"; + binding.__isset.parent_field_name = true; + binding.subcolumn_path = "dynamic"; + binding.__isset.subcolumn_path = true; + binding.index_properties["parser"] = "standard"; + binding.index_properties["support_phrase"] = "true"; + binding.__isset.index_properties = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info( + runtime_state_.get(), create_search_contexts(std::move(clause), {std::move(binding)}), + tablet_schema, &collect_infos); + + ASSERT_TRUE(status.ok()) << status; + ASSERT_EQ(collect_infos.size(), 1U); + auto it = collect_infos.find( + StringHelper::to_wstring(std::to_string(kVariantUid) + ".v.dynamic")); + ASSERT_NE(it, collect_infos.end()); + ASSERT_NE(it->second.index_meta, nullptr); + ASSERT_NE(it->second.owned_index_meta, nullptr); + EXPECT_EQ(it->second.index_meta->index_id(), 3025); + EXPECT_EQ(it->second.unique_terms, std::vector({"alpha", "beta"})); +} + +TEST_F(CollectionStatisticsTest, SearchVariantFieldPatternKeepsSelectedMetadataAlive) { + auto tablet_schema = std::make_shared(); + constexpr int32_t kVariantUid = 9020; + + TabletColumn variant_column; + variant_column.set_unique_id(kVariantUid); + variant_column.set_name("meta"); + variant_column.set_type(FieldType::OLAP_FIELD_TYPE_VARIANT); + auto subcolumn_template = make_subcolumn_template("user.*", PatternTypePB::MATCH_NAME_GLOB); + variant_column.add_sub_column(subcolumn_template); + tablet_schema->append_column(variant_column); + + TabletIndexPB index_pb; + index_pb.set_index_id(3030); + index_pb.set_index_name("variant_search_field_pattern_idx"); + index_pb.set_index_type(IndexType::INVERTED); + index_pb.add_col_unique_id(kVariantUid); + (*index_pb.mutable_properties())["parser"] = "standard"; + (*index_pb.mutable_properties())["support_phrase"] = "true"; + (*index_pb.mutable_properties())["field_pattern"] = "user.*"; + TabletIndex index; + index.init_from_pb(index_pb); + tablet_schema->append_index(std::move(index)); + + TSearchClause clause; + clause.clause_type = "PHRASE"; + clause.field_name = "meta.user.name"; + clause.value = "alice smith"; + clause.__isset.field_name = true; + clause.__isset.value = true; + TSearchFieldBinding binding; + binding.field_name = "meta.user.name"; + binding.slot_index = 0; + binding.is_variant_subcolumn = true; + binding.__isset.is_variant_subcolumn = true; + binding.parent_field_name = "meta"; + binding.__isset.parent_field_name = true; + binding.subcolumn_path = "user.name"; + binding.__isset.subcolumn_path = true; + binding.index_properties["parser"] = "standard"; + binding.index_properties["support_phrase"] = "true"; + binding.__isset.index_properties = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info( + runtime_state_.get(), create_search_contexts(std::move(clause), {std::move(binding)}), + tablet_schema, &collect_infos); + + ASSERT_TRUE(status.ok()) << status.msg(); + ASSERT_EQ(collect_infos.size(), 1u); + auto iter = collect_infos.find( + StringHelper::to_wstring(std::to_string(kVariantUid) + ".meta.user.name")); + ASSERT_NE(iter, collect_infos.end()); + ASSERT_NE(iter->second.index_meta, nullptr); + ASSERT_NE(iter->second.owned_index_meta, nullptr); + EXPECT_EQ(iter->second.index_meta->index_id(), 3030); +} + +TEST_F(CollectionStatisticsTest, SearchScoringRejectsNumericBkdLeaf) { + auto tablet_schema = std::make_shared(); + TabletColumn column; + column.set_unique_id(2); + column.set_name("number"); + column.set_type(FieldType::OLAP_FIELD_TYPE_INT); + tablet_schema->append_column(column); + TabletIndex index; + index._index_id = 3040; + index._index_type = IndexType::INVERTED; + index._col_unique_ids.push_back(2); + tablet_schema->append_index(std::move(index)); + TSearchClause clause; + clause.clause_type = "TERM"; + clause.field_name = "number"; + clause.value = "42"; + clause.__isset.field_name = true; + clause.__isset.value = true; + CollectInfoMap collect_infos; + + auto status = stats_->extract_collect_info(runtime_state_.get(), + create_search_contexts(std::move(clause)), + tablet_schema, &collect_infos); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + EXPECT_TRUE(collect_infos.empty()); +} + +TEST_F(CollectionStatisticsTest, NestedSearchScoringIsRejected) { + TSearchClause phrase; + phrase.clause_type = "PHRASE"; + phrase.field_name = "content"; + phrase.value = "alpha beta"; + phrase.__isset.field_name = true; + phrase.__isset.value = true; + + TSearchClause nested; + nested.clause_type = "NESTED"; + nested.children.push_back(std::move(phrase)); + nested.__isset.children = true; + + CollectInfoMap collect_infos; + auto status = stats_->extract_collect_info( + runtime_state_.get(), create_search_contexts(std::move(nested)), + create_tablet_schema_with_inverted_index(), &collect_infos); + + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); + EXPECT_TRUE(collect_infos.empty()); +} + +TEST_F(CollectionStatisticsTest, OneScoringFieldCannotSelectDifferentPhysicalIndexes) { + auto tablet_schema = create_tablet_schema_with_two_fulltext_indexes(); + auto build_match = [](const std::string& analyzer_name, InvertedIndexParserType parser_type) { + auto match_expr = + std::make_shared(TExprNodeType::MATCH_PRED); + match_expr->_children.push_back( + std::make_shared("content", SlotId(1))); + match_expr->_children.push_back( + std::make_shared("alpha beta")); + + InvertedIndexAnalyzerConfig config; + config.analyzer_name = analyzer_name; + config.parser_type = parser_type; + config.stop_words = "none"; + auto analyzer_ctx = std::make_shared(); + analyzer_ctx->analyzer_name = analyzer_name; + analyzer_ctx->analyzer_provider = + segment_v2::inverted_index::InvertedIndexAnalyzer::create_analyzer_provider( + &config); + match_expr->set_analyzer_ctx(std::move(analyzer_ctx)); + return match_expr; + }; + + MatchPredicateCollector collector; + CollectInfoMap collect_infos; + auto first = collector.collect( + runtime_state_.get(), tablet_schema, + build_match("standard", InvertedIndexParserType::PARSER_STANDARD), &collect_infos); + ASSERT_TRUE(first.ok()) << first.msg(); + + auto second = collector.collect(runtime_state_.get(), tablet_schema, + build_match("english", InvertedIndexParserType::PARSER_ENGLISH), + &collect_infos); + + EXPECT_EQ(second.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED); } // Test-only subclass that exposes the protected helpers of PredicateCollector. @@ -1270,41 +2890,4 @@ TEST_F(CollectionStatisticsTest, BuildFieldNameWithoutSuffix) { EXPECT_EQ(collector.build_field_name(42, ""), "42"); } -TEST(TermInfoComparerTest, OrdersByTermAndDedups) { - using doris::TermInfoComparer; - using doris::segment_v2::TermInfo; - - std::set terms; - - TermInfo t1; - t1.term = std::string("banana"); - t1.position = 2; - - TermInfo t2; - t2.term = std::string("apple"); - t2.position = 10; - - TermInfo t3; - t3.term = std::string("cherry"); - t3.position = 1; - - TermInfo dup; - dup.term = std::string("banana"); - dup.position = 100; - - terms.insert(t1); - terms.insert(t2); - terms.insert(t3); - terms.insert(dup); - - std::vector ordered; - ordered.reserve(terms.size()); - for (const auto& t : terms) { - ordered.push_back(t.get_single_term()); - } - - EXPECT_EQ(terms.size(), 3u); - EXPECT_THAT(ordered, ::testing::ElementsAre("apple", "banana", "cherry")); -} - } // namespace doris diff --git a/be/test/storage/compaction/segcompaction_test.cpp b/be/test/storage/compaction/segcompaction_test.cpp index d3b843c050da2e..2c6697bc592c59 100644 --- a/be/test/storage/compaction/segcompaction_test.cpp +++ b/be/test/storage/compaction/segcompaction_test.cpp @@ -19,12 +19,15 @@ #include #include +#include +#include #include #include #include #include #include "common/config.h" +#include "cpp/sync_point.h" #include "io/fs/local_file_system.h" #include "runtime/exec_env.h" #include "runtime/memory/mem_tracker.h" @@ -114,6 +117,10 @@ class SegCompactionTest : public testing::Test { } void TearDown() { + auto* sync_point = SyncPoint::get_instance(); + sync_point->disable_processing(); + sync_point->clear_all_call_backs(); + sync_point->clear_trace(); config::enable_segcompaction = false; ExecEnv* exec_env = doris::ExecEnv::GetInstance(); l_engine = nullptr; @@ -288,6 +295,55 @@ class SegCompactionTest : public testing::Test { std::unique_ptr _inverted_index_searcher_cache; }; +TEST_F(SegCompactionTest, FatalStatusPreservesOriginalError) { + config::segcompaction_candidate_max_rows = 10; + config::segcompaction_batch_size = 2; + + auto tablet_schema = std::make_shared(); + create_tablet_schema(tablet_schema, DUP_KEYS); + + RowsetWriterContext writer_context; + create_rowset_writer_context(10046, tablet_schema, &writer_context); + auto writer_result = RowsetFactory::create_rowset_writer(*l_engine, writer_context, false); + ASSERT_TRUE(writer_result.has_value()) << writer_result.error(); + auto rowset_writer = std::move(writer_result).value(); + + const auto expected_status = Status::Error( + "CommonGrams analysis failed; set enable_common_grams_index_build=false and retry " + "the load as a new transaction"); + std::promise worker_started; + auto worker_started_future = worker_started.get_future(); + auto* sync_point = SyncPoint::get_instance(); + SyncPoint::CallbackGuard callback_guard; + sync_point->set_call_back( + "SegcompactionWorker::_do_compact_segments", + [&worker_started, expected_status](auto&& args) { + auto* result = try_any_cast_ret(args); + result->first = expected_status; + result->second = true; + worker_started.set_value(); + }, + &callback_guard); + sync_point->enable_processing(); + + for (int segment_id = 0; segment_id < 3; ++segment_id) { + Block block = tablet_schema->create_block(); + auto columns = std::move(block).mutate_columns(); + for (uint32_t column_id = 0; column_id < columns.size(); ++column_id) { + const uint32_t value = segment_id + column_id; + columns[column_id]->insert_data(reinterpret_cast(&value), sizeof(value)); + } + ASSERT_TRUE(add_block_with_columns(rowset_writer.get(), &block, &columns).ok()); + ASSERT_TRUE(rowset_writer->flush().ok()); + } + ASSERT_EQ(worker_started_future.wait_for(std::chrono::seconds(10)), std::future_status::ready); + + RowsetSharedPtr rowset; + const auto status = rowset_writer->build(rowset); + EXPECT_EQ(status.code(), INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_EQ(status.to_string(), expected_status.to_string()); +} + TEST_F(SegCompactionTest, SegCompactionThenRead) { config::enable_segcompaction = true; Status s; diff --git a/be/test/storage/index/index_builder_test.cpp b/be/test/storage/index/index_builder_test.cpp index dd36ba3ab33159..c6b3ebb9dedf9c 100644 --- a/be/test/storage/index/index_builder_test.cpp +++ b/be/test/storage/index/index_builder_test.cpp @@ -20,6 +20,7 @@ #include #include +#include "common/config.h" #include "storage/olap_common.h" #include "storage/rowset/beta_rowset.h" #include "storage/rowset/rowset_factory.h" @@ -27,10 +28,29 @@ #include "storage/storage_engine.h" #include "storage/tablet/tablet_fwd.h" #include "storage/tablet/tablet_schema.h" +#include "util/debug_points.h" namespace doris { using namespace testing; +class ScopedIndexBuilderDebugPoints { +public: + ScopedIndexBuilderDebugPoints() : _debug_points_enabled(config::enable_debug_points) { + config::enable_debug_points = true; + DebugPoints::instance()->clear(); + } + + ~ScopedIndexBuilderDebugPoints() { + DebugPoints::instance()->clear(); + config::enable_debug_points = _debug_points_enabled; + } + + void enable(const std::string& name) { DebugPoints::instance()->add(name); } + +private: + bool _debug_points_enabled; +}; + class IndexBuilderTest : public ::testing::Test { protected: void SetUp() override { @@ -169,6 +189,58 @@ class IndexBuilderTest : public ::testing::Test { rs_meta->set_tablet_schema(tablet_schema); } + void prepare_single_index_build(int64_t rowset_id) { + auto tablet_path = _absolute_dir + "/" + std::to_string(rowset_id); + _tablet->_tablet_path = tablet_path; + ASSERT_TRUE(io::global_local_filesystem()->delete_directory(tablet_path).ok()); + ASSERT_TRUE(io::global_local_filesystem()->create_directory(tablet_path).ok()); + + RowsetWriterContext writer_context; + writer_context.rowset_id.init(rowset_id); + writer_context.tablet_id = _tablet->tablet_id(); + writer_context.tablet_schema_hash = _tablet_meta->schema_hash(); + writer_context.partition_id = 10; + writer_context.rowset_type = BETA_ROWSET; + writer_context.tablet_path = tablet_path; + writer_context.rowset_state = VISIBLE; + writer_context.tablet_schema = _tablet_schema; + writer_context.version = Version(10, 10); + + auto result = RowsetFactory::create_rowset_writer(*_engine_ref, writer_context, false); + ASSERT_TRUE(result.has_value()) << result.error(); + auto rowset_writer = std::move(result).value(); + + Block block = _tablet_schema->create_block(); + auto columns = std::move(block).mutate_columns(); + for (int i = 0; i < 8; ++i) { + int32_t k1 = i * 10; + int32_t k2 = i; + columns[0]->insert_data(reinterpret_cast(&k1), sizeof(k1)); + columns[1]->insert_data(reinterpret_cast(&k2), sizeof(k2)); + } + block.set_columns(std::move(columns)); + ASSERT_TRUE(rowset_writer->add_block(&block).ok()); + ASSERT_TRUE(rowset_writer->flush().ok()); + + RowsetSharedPtr rowset; + ASSERT_TRUE(rowset_writer->build(rowset).ok()); + ASSERT_TRUE(_tablet->add_rowset(rowset).ok()); + + TOlapTableIndex index; + index.index_id = 101; + index.index_name = "k1_index"; + index.columns.emplace_back("k1"); + index.column_unique_ids.push_back(1); + index.index_type = TIndexType::INVERTED; + _alter_indexes.push_back(std::move(index)); + } + + Status build_single_index() { + IndexBuilder builder(*_engine_ref, _tablet, _columns, _alter_indexes, false); + RETURN_IF_ERROR(builder.init()); + return builder.do_build_inverted_index(); + } + StorageEngine* _engine_ref = nullptr; TabletSharedPtr _tablet; TabletMetaSharedPtr _tablet_meta; @@ -202,6 +274,17 @@ TEST_F(IndexBuilderTest, BasicBuildTest) { EXPECT_EQ(builder._alter_index_ids.size(), 1); } +TEST_F(IndexBuilderTest, HandleSingleRowsetPreservesOrdinaryAppendFailure) { + prepare_single_index_build(16604); + ScopedIndexBuilderDebugPoints debug_points; + debug_points.enable("IndexBuilder::handle_single_rowset_write_inverted_index_data_error"); + + auto status = build_single_index(); + + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(status.msg(), "debug point: handle_single_rowset_write_inverted_index_data_error"); +} + TEST_F(IndexBuilderTest, DropInvertedIndexTest) { // 0. prepare tablet path auto tablet_path = _absolute_dir + "/" + std::to_string(15676); diff --git a/be/test/storage/index/inverted/ananlyzer/analyzer_test.cpp b/be/test/storage/index/inverted/ananlyzer/analyzer_test.cpp index 4a51609a68806d..6d7d6e02221f6d 100644 --- a/be/test/storage/index/inverted/ananlyzer/analyzer_test.cpp +++ b/be/test/storage/index/inverted/ananlyzer/analyzer_test.cpp @@ -345,4 +345,23 @@ TEST_F(AnalyzerTest, TestAnalyzerFunctionality) { } } +TEST_F(AnalyzerTest, CommonGramsQueryPurposeSelection) { + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_ANY_QUERY, 0, false), + AnalysisPurpose::kPlainQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_ALL_QUERY, 0, false), + AnalysisPurpose::kPlainQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_PHRASE_QUERY, 0, false), + AnalysisPurpose::kExactPhraseQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_PHRASE_QUERY, 2, false), + AnalysisPurpose::kPlainQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_PHRASE_QUERY, 0, true), + AnalysisPurpose::kPlainQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, 0, false), + AnalysisPurpose::kPhrasePrefixQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, 0, true), + AnalysisPurpose::kPlainQuery); + EXPECT_EQ(select_analysis_purpose(InvertedIndexQueryType::MATCH_REGEXP_QUERY, 0, false), + AnalysisPurpose::kPlainQuery); +} + } // namespace doris::segment_v2::inverted_index \ No newline at end of file diff --git a/be/test/storage/index/inverted/ananlyzer/common_grams_filter_test.cpp b/be/test/storage/index/inverted/ananlyzer/common_grams_filter_test.cpp new file mode 100644 index 00000000000000..9d4feb731bb150 --- /dev/null +++ b/be/test/storage/index/inverted/ananlyzer/common_grams_filter_test.cpp @@ -0,0 +1,644 @@ +// 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 "storage/index/inverted/token_filter/common_grams_filter.h" + +#include + +#include +#include +#include +#include +#include + +#include "common/exception.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_word_set.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +struct ScriptedToken { + std::string term; + int32_t position_increment = 1; + int32_t start_offset = 0; + int32_t end_offset = 0; + const TCHAR* type = Token::getDefaultType(); +}; + +class ScriptedTokenStream final : public TokenStream { +public: + explicit ScriptedTokenStream(std::vector tokens) : _tokens(std::move(tokens)) {} + + Token* next(Token* token) override { + if (_next == _tokens.size()) { + return nullptr; + } + const auto& scripted = _tokens[_next++]; + _scratch = scripted.term; + token->clear(); + token->setTextNoCopy(_scratch.data(), static_cast(_scratch.size())); + token->positionIncrement = scripted.position_increment; + token->setStartOffset(scripted.start_offset); + token->setEndOffset(scripted.end_offset); + token->setType(scripted.type); + return token; + } + + void close() override {} + void reset() override { _next = 0; } + + void set_tokens(std::vector tokens) { + _tokens = std::move(tokens); + _next = 0; + } + + size_t consumed() const { return _next; } + +private: + std::vector _tokens; + size_t _next = 0; + std::string _scratch; +}; + +struct ActualToken { + std::string term; + int32_t position_increment; + int32_t start_offset; + int32_t end_offset; + std::wstring type; + + bool operator==(const ActualToken&) const = default; +}; + +struct TokenSemantics { + std::string term; + int32_t position; + std::wstring type; + + bool operator==(const TokenSemantics&) const = default; +}; + +template +concept HasAnalyzerCommonGramClassification = requires(Event event) { + event.has_preceding_gram; + event.preceding_gram_both_common; +}; + +static_assert(!HasAnalyzerCommonGramClassification); + +std::string gram(std::string_view left, std::string_view right) { + auto encoded = encode_common_gram(left, right); + EXPECT_TRUE(encoded.has_value()) << encoded.error(); + return encoded.value(); +} + +std::vector collect(const TokenStreamPtr& stream) { + std::vector result; + Token token; + while (stream->next(&token) != nullptr) { + result.push_back({std::string(token.termBuffer(), token.termLength()), + token.getPositionIncrement(), token.startOffset(), token.endOffset(), + token.type()}); + } + return result; +} + +std::vector collect_semantics(const TokenStreamPtr& stream) { + std::vector result; + Token token; + int32_t position = 0; + while (stream->next(&token) != nullptr) { + position += token.getPositionIncrement(); + result.push_back({std::string(token.termBuffer(), token.termLength()), position, + token.type()}); + } + return result; +} + +std::vector expand_snii_index_events(CommonGramsFilter* stream, size_t* event_count, + size_t* both_common_gram_count) { + std::vector result; + std::optional previous_logical_term; + bool previous_is_common = false; + int32_t position = 0; + SniiCommonGramsIndexEvent event; + while (stream->next_snii_index_event(&event)) { + ++*event_count; + const std::string physical_plain_term(event.plain_term); + const std::string logical_term = + decode_plain_term(physical_plain_term, PlainTermKeyVersion::kEscapedV1).value(); + EXPECT_EQ(event.logical_term, logical_term); + const bool current_is_common = + CommonWordSet::builtin_english_stop_words_v1().contains(logical_term); + const bool has_preceding_gram = previous_logical_term.has_value() && + (previous_is_common || current_is_common) && + common_gram_component_sizes_encodable( + previous_logical_term->size(), logical_term.size()); + if (has_preceding_gram) { + EXPECT_TRUE(previous_logical_term.has_value()); + *both_common_gram_count += previous_is_common && current_is_common; + result.push_back({gram(previous_logical_term.value(), logical_term), position, + COMMON_GRAM_TOKEN_TYPE}); + } + ++position; + result.push_back({physical_plain_term, position, Token::getDefaultType()}); + previous_logical_term = logical_term; + previous_is_common = current_is_common; + } + return result; +} + +std::vector words(std::initializer_list terms) { + std::vector result; + int32_t offset = 0; + for (auto term : terms) { + result.push_back( + {std::string(term), 1, offset, offset + static_cast(term.size())}); + offset += static_cast(term.size()) + 1; + } + return result; +} + +std::vector terms(const std::vector& tokens) { + std::vector result; + for (const auto& token : tokens) { + EXPECT_EQ(token.position_increment, 1); + result.push_back(token.term); + } + return result; +} + +std::shared_ptr builtin_common_words() { + static const auto common_words = std::shared_ptr( + &CommonWordSet::builtin_english_stop_words_v1(), [](const CommonWordSet*) {}); + return common_words; +} + +TokenStreamPtr index_stream(const std::shared_ptr& input) { + return std::make_shared(input, builtin_common_words()); +} + +TokenStreamPtr escaped_index_stream(const std::shared_ptr& input) { + return std::make_shared(input, builtin_common_words(), + CommonGramsOutputMode::kEscapedV1Index); +} + +TokenStreamPtr spimi_index_stream(const std::shared_ptr& input) { + return std::make_shared(input, builtin_common_words(), + CommonGramsOutputMode::kEscapedV1SpimiIndex); +} + +TokenStreamPtr plain_stream(const std::shared_ptr& input) { + return std::make_shared(input); +} + +TokenStreamPtr exact_stream(const std::shared_ptr& input) { + return std::make_shared(index_stream(input), builtin_common_words()); +} + +TokenStreamPtr prefix_stream(const std::shared_ptr& input) { + return std::make_shared(index_stream(input), + builtin_common_words()); +} + +TEST(CommonGramsFilterTest, IndexPreservesUnigramsAndAddsEligibleOwnedGrams) { + auto input = std::make_shared(words({"man", "of", "the", "year"})); + auto stream = index_stream(input); + + EXPECT_EQ(collect(stream), (std::vector { + {"man", 1, 0, 3, Token::getDefaultType()}, + {gram("man", "of"), 0, 0, 6, COMMON_GRAM_TOKEN_TYPE}, + {"of", 1, 4, 6, Token::getDefaultType()}, + {gram("of", "the"), 0, 4, 10, COMMON_GRAM_TOKEN_TYPE}, + {"the", 1, 7, 10, Token::getDefaultType()}, + {gram("the", "year"), 0, 7, 15, COMMON_GRAM_TOKEN_TYPE}, + {"year", 1, 11, 15, Token::getDefaultType()}, + })); +} + +TEST(CommonGramsFilterTest, QueryModesNormalizeBothCommonGramType) { + auto input = std::make_shared(words({"of", "the"})); + EXPECT_EQ(collect(exact_stream(input)), + (std::vector { + {gram("of", "the"), 1, 0, 6, COMMON_GRAM_TOKEN_TYPE}, + })); + + input = std::make_shared(words({"of", "the"})); + EXPECT_EQ(collect(prefix_stream(input)), + (std::vector { + {gram("of", "the"), 1, 0, 6, COMMON_GRAM_TOKEN_TYPE}, + })); +} + +TEST(CommonGramsFilterTest, QueryGramEligibilityMatchesPurposeSpecificFilters) { + const std::vector> cases { + {}, + {"alpha"}, + {"alpha", "beta"}, + {"alpha", "the"}, + {"the", "alpha"}, + {"alpha", "beta", "the"}, + {"alpha", "the", "beta"}, + {"alpha", "beta", "gamma", "delta"}, + }; + for (const auto& query_terms : cases) { + SCOPED_TRACE(testing::PrintToString(query_terms)); + std::vector scripted; + scripted.reserve(query_terms.size()); + for (const auto& term : query_terms) { + scripted.push_back({.term = term}); + } + + for (const auto mode : + {CommonGramsQueryMode::kExact, CommonGramsQueryMode::kPhrasePrefix}) { + auto input = std::make_shared(scripted); + const auto output = + collect(mode == CommonGramsQueryMode::kExact ? exact_stream(input) + : prefix_stream(input)); + bool filter_used_gram = false; + for (const auto& token : output) { + filter_used_gram = + filter_used_gram || token.type == std::wstring(COMMON_GRAM_TOKEN_TYPE); + } + EXPECT_EQ(common_grams_query_may_use_gram(query_terms, mode, *builtin_common_words()), + filter_used_gram); + } + } +} + +TEST(CommonGramsFilterTest, PhysicalIndexModeEscapesPlainTermsButGramsUseLogicalBytes) { + const std::string internal_plain = std::string(1, '\x1f') + "literal"; + auto input = std::make_shared(words({internal_plain, "of"})); + auto stream = escaped_index_stream(input); + + EXPECT_EQ(collect(stream), + (std::vector { + {std::string(1, PLAIN_ESCAPE_PREFIX) + "Gliteral", 1, 0, + static_cast(internal_plain.size()), Token::getDefaultType()}, + {gram(internal_plain, "of"), 0, 0, + static_cast(internal_plain.size() + 3), COMMON_GRAM_TOKEN_TYPE}, + {"of", 1, static_cast(internal_plain.size() + 1), + static_cast(internal_plain.size() + 3), Token::getDefaultType()}, + })); + + const std::string escape_plain = std::string(1, PLAIN_ESCAPE_PREFIX) + "literal"; + input->set_tokens(words({escape_plain})); + stream->reset(); + EXPECT_EQ(terms(collect(stream)), + (std::vector {std::string(1, PLAIN_ESCAPE_PREFIX) + "Eliteral"})); +} + +TEST(CommonGramsFilterTest, SpimiIndexModeEmitsPhysicalGramsAndEscapedPlainTerms) { + const std::string internal_plain = std::string(1, '\x1f') + "literal"; + auto input = std::make_shared(words({internal_plain, "of"})); + + const std::string physical = gram(internal_plain, "of"); + EXPECT_EQ(collect(spimi_index_stream(input)), + (std::vector { + {std::string(1, PLAIN_ESCAPE_PREFIX) + "Gliteral", 1, 0, + static_cast(internal_plain.size()), Token::getDefaultType()}, + {physical, 0, 0, static_cast(internal_plain.size() + 3), + COMMON_GRAM_TOKEN_TYPE}, + {"of", 1, static_cast(internal_plain.size() + 1), + static_cast(internal_plain.size() + 3), Token::getDefaultType()}, + })); +} + +TEST(CommonGramsFilterTest, SniiIndexEventsExpandToExistingSpimiTokenSemantics) { + const std::string internal_plain = std::string(1, '\x1f') + "literal"; + const std::string escaped_plain = std::string(1, PLAIN_ESCAPE_PREFIX) + "literal"; + const std::vector input_tokens = + words({internal_plain, "of", "the", escaped_plain, "中文词"}); + + auto legacy_input = std::make_shared(input_tokens); + const std::vector expected = + collect_semantics(spimi_index_stream(legacy_input)); + + auto event_input = std::make_shared(input_tokens); + CommonGramsFilter event_stream(event_input, builtin_common_words(), + CommonGramsOutputMode::kEscapedV1SpimiIndex); + size_t event_count = 0; + size_t both_common_gram_count = 0; + const std::vector actual = + expand_snii_index_events(&event_stream, &event_count, &both_common_gram_count); + + EXPECT_EQ(event_count, input_tokens.size()); + EXPECT_EQ(both_common_gram_count, 1U); + EXPECT_EQ(actual, expected); +} + +TEST(CommonGramsFilterTest, SniiIndexEventsDeferCommonWordClassificationToWriter) { + auto input = std::make_shared( + words({"the", "database", "of", "the", "world", "and", "search"})); + CommonGramsFilter stream(input, builtin_common_words(), + CommonGramsOutputMode::kEscapedV1SpimiIndex); + + common_grams_testing::reset_common_word_membership_lookup_count(); + SniiCommonGramsIndexEvent event; + size_t event_count = 0; + while (stream.next_snii_index_event(&event)) { + EXPECT_FALSE(event.logical_term.empty()); + EXPECT_FALSE(event.plain_term.empty()); + ++event_count; + } + + EXPECT_EQ(event_count, 7U); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 0U); +} + +TEST(CommonGramsFilterTest, IndexOmitsNonCommonPair) { + auto input = std::make_shared(words({"man", "year"})); + EXPECT_EQ(terms(collect(index_stream(input))), (std::vector {"man", "year"})); +} + +TEST(CommonGramsFilterTest, CachesCommonWordMembershipWithoutDefeatingShortCircuit) { + auto input = std::make_shared(words({"single"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + collect(index_stream(input)); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 0); + + input = std::make_shared(words({"of", "dog"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + collect(index_stream(input)); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 1); + + input = std::make_shared(words({"of", "the", "and", "to"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + collect(exact_stream(input)); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 6); + + input = std::make_shared(words({"man", "dog"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + collect(prefix_stream(input)); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 3); + + input = std::make_shared(words({"man", "dog", "year", "thing"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + EXPECT_EQ(terms(collect(index_stream(input))), + (std::vector {"man", "dog", "year", "thing"})); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 4); + + input = std::make_shared(words({"man", "dog", "year", "thing"})); + common_grams_testing::reset_common_word_membership_lookup_count(); + EXPECT_EQ(terms(collect(exact_stream(input))), + (std::vector {"man", "dog", "year", "thing"})); + EXPECT_EQ(common_grams_testing::common_word_membership_lookup_count(), 8); +} + +TEST(CommonGramsFilterTest, ExactRewriteTruthTable) { + struct Case { + std::vector input; + std::vector expected; + }; + const std::vector cases = { + {words({"man", "dog", "year"}), {"man", "dog", "year"}}, + {words({"man", "dog", "the"}), {"man", gram("dog", "the")}}, + {words({"man", "of", "year"}), {gram("man", "of"), gram("of", "year")}}, + {words({"man", "of", "the"}), {gram("man", "of"), gram("of", "the")}}, + {words({"of", "dog", "year"}), {gram("of", "dog"), "dog", "year"}}, + {words({"of", "dog", "the"}), {gram("of", "dog"), gram("dog", "the")}}, + {words({"of", "the", "year"}), {gram("of", "the"), gram("the", "year")}}, + {words({"of", "the", "and"}), {gram("of", "the"), gram("the", "and")}}, + {words({"the", "the", "the"}), {gram("the", "the"), gram("the", "the")}}, + }; + + for (const auto& test_case : cases) { + auto input = std::make_shared(test_case.input); + EXPECT_EQ(terms(collect(exact_stream(input))), test_case.expected); + } +} + +TEST(CommonGramsFilterTest, PhrasePrefixRewritesOnlyCommonLeftBoundary) { + struct Case { + std::vector input; + std::vector expected; + }; + const std::vector cases = { + {words({"the", "wo"}), {gram("the", "wo")}}, + {words({"foo", "the"}), {"foo", "the"}}, + {words({"foo", "of", "th"}), {gram("foo", "of"), gram("of", "th")}}, + {words({"the", "bar", "ba"}), {gram("the", "bar"), "bar", "ba"}}, + {words({"the"}), {"the"}}, + }; + + for (const auto& test_case : cases) { + auto input = std::make_shared(test_case.input); + EXPECT_EQ(terms(collect(prefix_stream(input))), test_case.expected); + } +} + +TEST(CommonGramsFilterTest, ResetAndReuseRestoresUnigramMetadata) { + auto input = std::make_shared(words({"man", "of"})); + auto stream = index_stream(input); + ASSERT_EQ(collect(stream).size(), 3); + + input->set_tokens({{"plain", 1, 17, 22, Token::getDefaultType()}}); + stream->reset(); + EXPECT_EQ(collect(stream), (std::vector { + {"plain", 1, 17, 22, Token::getDefaultType()}, + })); +} + +TEST(CommonGramsFilterTest, PreservesUpstreamUnigramTypes) { + auto input = std::make_shared(std::vector { + {"man", 1, 3, 6, L"left_type"}, {"of", 1, 7, 9, L"right_type"}}); + EXPECT_EQ(collect(index_stream(input)), + (std::vector { + {"man", 1, 3, 6, L"left_type"}, + {gram("man", "of"), 0, 3, 9, COMMON_GRAM_TOKEN_TYPE}, + {"of", 1, 7, 9, L"right_type"}, + })); + + input->set_tokens({{"man", 1, 3, 6, L"left_type"}, {"year", 1, 7, 11, L"right_type"}}); + auto exact = exact_stream(input); + EXPECT_EQ(collect(exact), (std::vector { + {"man", 1, 3, 6, L"left_type"}, + {"year", 1, 7, 11, L"right_type"}, + })); +} + +TEST(CommonGramsFilterTest, ExactAndPrefixResetReuseIndependentInput) { + auto exact_input = std::make_shared(words({"man", "of"})); + auto exact = exact_stream(exact_input); + EXPECT_EQ(terms(collect(exact)), (std::vector {gram("man", "of")})); + exact_input->set_tokens(words({"plain", "terms"})); + exact->reset(); + EXPECT_EQ(terms(collect(exact)), (std::vector {"plain", "terms"})); + + auto prefix_input = std::make_shared(words({"the", "term"})); + auto prefix = prefix_stream(prefix_input); + EXPECT_EQ(terms(collect(prefix)), (std::vector {gram("the", "term")})); + prefix_input->set_tokens(words({"plain", "terms"})); + prefix->reset(); + EXPECT_EQ(terms(collect(prefix)), (std::vector {"plain", "terms"})); +} + +TEST(CommonGramsFilterTest, EmptyAndSingleTokenStreamsAreStable) { + auto input = std::make_shared(std::vector {}); + auto stream = index_stream(input); + EXPECT_TRUE(collect(stream).empty()); + + input->set_tokens(words({"of"})); + stream->reset(); + EXPECT_EQ(terms(collect(stream)), (std::vector {"of"})); +} + +TEST(CommonGramsFilterTest, RejectsNonUnitInputForEveryStreamPurpose) { + using Factory = TokenStreamPtr (*)(const std::shared_ptr&); + for (Factory factory : {index_stream, plain_stream, exact_stream, prefix_stream}) { + for (int32_t bad_increment : {-1, 0, 2}) { + auto input = std::make_shared( + std::vector {{"of", bad_increment, 0, 2}}); + try { + collect(factory(input)); + FAIL() << "expected analyzer error for increment " << bad_increment; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } + + for (int32_t bad_increment : {-1, 0, 2}) { + auto empty_input = std::make_shared( + std::vector {{"", bad_increment, 0, 0}}); + try { + collect(factory(empty_input)); + FAIL() << "expected analyzer error for empty token increment " << bad_increment; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } + + auto input = std::make_shared( + std::vector {{"man", 1, 0, 3}, {"of", 2, 4, 6}}); + auto stream = factory(input); + try { + collect(stream); + FAIL() << "expected analyzer error after a valid token"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_EQ(input->consumed(), 2); + } + } +} + +TEST(CommonGramsFilterTest, RejectsEmptyInputForEveryStreamPurpose) { + using Factory = TokenStreamPtr (*)(const std::shared_ptr&); + for (Factory factory : {index_stream, plain_stream, exact_stream, prefix_stream}) { + auto input = + std::make_shared(std::vector {{"", 1, 0, 0}}); + try { + collect(factory(input)); + FAIL() << "expected analyzer error for an empty token"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } +} + +TEST(CommonGramsFilterTest, QueryPreparationLatchesFailureUntilReset) { + for (auto factory : {exact_stream, prefix_stream}) { + auto input = std::make_shared( + std::vector {{"the", 1, 0, 3}, {"term", 2, 4, 8}}); + auto stream = factory(input); + for (int attempt = 0; attempt < 2; ++attempt) { + try { + collect(stream); + FAIL() << "expected latched analyzer error on attempt " << attempt; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } + + input->set_tokens(words({"the", "term"})); + stream->reset(); + EXPECT_EQ(terms(collect(stream)), (std::vector {gram("the", "term")})); + } +} + +TEST(CommonGramsFilterTest, UnencodableRequiredGramFallsBackToWholePlainSequence) { + const std::string huge(COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + + for (auto factory : {exact_stream, prefix_stream}) { + auto input = std::make_shared(words({"the", huge})); + EXPECT_EQ(terms(collect(factory(input))), (std::vector {"the", huge})); + + input = std::make_shared(words({"the", "of", huge})); + EXPECT_EQ(terms(collect(factory(input))), (std::vector {"the", "of", huge})); + } + + auto input = std::make_shared(words({"the", huge})); + EXPECT_EQ(terms(collect(index_stream(input))), (std::vector {"the", huge})); +} + +TEST(CommonGramsFilterTest, MaximumMarkerLeadingUnigramFailsWithBuildSwitchRecovery) { + for (const char marker : {PLAIN_ESCAPE_PREFIX, '\x1f'}) { + std::string term(COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + term.front() = marker; + auto input = std::make_shared( + std::vector {{term, 1, 0, static_cast(term.size())}}); + try { + static_cast(collect(escaped_index_stream(input))); + FAIL() << "expected CommonGrams escaped-term overflow"; + } catch (const Exception& e) { + EXPECT_EQ(e.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_NE(std::string(e.what()).find("enable_common_grams_index_build=false"), + std::string::npos); + EXPECT_NE(std::string(e.what()).find("new transaction"), std::string::npos); + } + } +} + +TEST(CommonGramsFilterTest, LargestEscapableMarkerLeadingUnigramUsesPhysicalKeyLimit) { + for (const char marker : {PLAIN_ESCAPE_PREFIX, '\x1f'}) { + std::string term(COMMON_GRAM_MAX_ENCODED_BYTES - 1, 'x'); + term.front() = marker; + auto input = std::make_shared( + std::vector {{term, 1, 0, static_cast(term.size())}}); + const auto tokens = collect(escaped_index_stream(input)); + ASSERT_EQ(tokens.size(), 1); + EXPECT_EQ(tokens[0].term.size(), COMMON_GRAM_MAX_ENCODED_BYTES); + } +} + +TEST(CommonGramsFilterTest, InvalidLogicalTermsRemainHardAnalyzerErrors) { + for (const std::string& term : {std::string("bad\0term", 8), std::string("\xc3", 1)}) { + auto input = std::make_shared( + std::vector {{term, 1, 0, static_cast(term.size())}}); + try { + collect(escaped_index_stream(input)); + FAIL() << "expected analyzer error"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } +} + +TEST(CommonGramsFilterTest, RejectsOverlongLogicalToken) { + auto input = std::make_shared( + std::vector {{std::string(COMMON_GRAM_MAX_ENCODED_BYTES + 1, 'x'), 1, 0, + static_cast(COMMON_GRAM_MAX_ENCODED_BYTES + 1)}}); + try { + collect(index_stream(input)); + FAIL() << "expected analyzer error"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } +} + +} // namespace +} // namespace doris::segment_v2::inverted_index diff --git a/be/test/storage/index/inverted/ananlyzer/common_grams_key_codec_test.cpp b/be/test/storage/index/inverted/ananlyzer/common_grams_key_codec_test.cpp new file mode 100644 index 00000000000000..05ef2bb4e56453 --- /dev/null +++ b/be/test/storage/index/inverted/ananlyzer/common_grams_key_codec_test.cpp @@ -0,0 +1,553 @@ +// 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 "storage/index/inverted/common_grams/common_grams_key_codec.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "util/utf8_check.h" + +namespace doris::segment_v2::inverted_index { +namespace { + +TEST(CommonGramsKeyCodecTest, GramGoldenBytesAndMarkerRange) { + const std::string expected = + "\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "00000001:ab"; + auto encoded = encode_common_gram("a", "b"); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_EQ(encoded.value(), expected); + + EXPECT_EQ(CG_V1_MARKER, std::string_view("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f", + 22)); + EXPECT_EQ(CG_V1_MARKER_END, std::string_view("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x20", + 22)); + EXPECT_GE(encoded.value(), CG_V1_MARKER); + EXPECT_LT(encoded.value(), CG_V1_MARKER_END); + EXPECT_EQ(encoded->find('\0'), std::string::npos); + EXPECT_TRUE(validate_utf8(encoded->data(), encoded->size())); + + auto hexadecimal_length = encode_common_gram("0123456789", "x"); + ASSERT_TRUE(hexadecimal_length.has_value()) << hexadecimal_length.error(); + EXPECT_EQ(hexadecimal_length.value(), + "\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "0000000a:0123456789x"); + + auto different_boundary = encode_common_gram("ab", "c"); + ASSERT_TRUE(different_boundary.has_value()) << different_boundary.error(); + auto same_text = encode_common_gram("a", "bc"); + ASSERT_TRUE(same_text.has_value()) << same_text.error(); + EXPECT_NE(different_boundary.value(), same_text.value()); +} + +TEST(CommonGramsKeyCodecTest, ValidatesLogicalTermsWithoutEncoding) { + EXPECT_TRUE(validate_common_grams_logical_term("valid", "test term").ok()); + EXPECT_TRUE(validate_common_grams_logical_term("", "test term").ok()); + + const std::string nul_term("a\0b", 3); + const std::string invalid_utf8("\xc3\x28", 2); + const std::string overlong(COMMON_GRAM_MAX_ENCODED_BYTES + 1, 'x'); + for (const auto& term : {nul_term, invalid_utf8, overlong}) { + auto status = validate_common_grams_logical_term(term, "test term"); + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } +} + +TEST(CommonGramsKeyCodecTest, TryEncodeDistinguishesTooLongAndReusesOutputCapacity) { + std::string output; + output.reserve(COMMON_GRAM_MAX_ENCODED_BYTES); + const char* reserved_data = output.data(); + + auto encoded = try_encode_common_gram("a", "b", &output); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_TRUE(encoded.value()); + EXPECT_EQ(output, encode_common_gram("a", "b").value()); + EXPECT_EQ(output.data(), reserved_data); + + const std::string huge(COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + auto too_long = try_encode_common_gram("the", huge, &output); + ASSERT_TRUE(too_long.has_value()) << too_long.error(); + EXPECT_FALSE(too_long.value()); + EXPECT_TRUE(output.empty()); + EXPECT_EQ(output.data(), reserved_data); + + const std::string invalid_utf8("\xc3\x28", 2); + auto invalid = try_encode_common_gram(invalid_utf8, "valid", &output); + EXPECT_FALSE(invalid.has_value()); + EXPECT_EQ(invalid.error().code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_TRUE(output.empty()); +} + +TEST(CommonGramsKeyCodecTest, PrevalidatedEncoderMatchesCheckedEncoder) { + for (const auto& [left, right] : std::vector> { + {"a", "b"}, {"the", "term"}, {"", "right"}}) { + std::string checked; + auto checked_result = try_encode_common_gram(left, right, &checked); + ASSERT_TRUE(checked_result.has_value()) << checked_result.error(); + + std::string prevalidated; + EXPECT_EQ(try_encode_common_gram_prevalidated(left, right, prevalidated), + checked_result.value()); + EXPECT_EQ(prevalidated, checked); + } + + std::string too_long(COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + std::string output = "stale"; + EXPECT_FALSE(try_encode_common_gram_prevalidated("left", too_long, output)); + EXPECT_TRUE(output.empty()); +} + +TEST(CommonGramsKeyCodecTest, PrevalidatedPlainEncoderMatchesCheckedEncoderAndBoundaries) { + std::string output; + output.reserve(COMMON_GRAM_MAX_ENCODED_BYTES); + const char* reserved_data = output.data(); + + for (const char marker : {PLAIN_ESCAPE_PREFIX, '\x1f'}) { + const std::string logical = std::string(1, marker) + "plain"; + std::string checked; + auto checked_result = + try_encode_plain_term(logical, PlainTermKeyVersion::kEscapedV1, &checked); + ASSERT_TRUE(checked_result.has_value()) << checked_result.error(); + ASSERT_TRUE(checked_result.value()); + + EXPECT_TRUE(try_encode_escaped_plain_term_prevalidated(logical, output)); + EXPECT_EQ(output, checked); + EXPECT_EQ(output.data(), reserved_data); + } + + const std::string maximum_encodable = std::string(1, PLAIN_ESCAPE_PREFIX) + + std::string(COMMON_GRAM_MAX_ENCODED_BYTES - 2, 'x'); + EXPECT_TRUE(try_encode_escaped_plain_term_prevalidated(maximum_encodable, output)); + EXPECT_EQ(output.size(), COMMON_GRAM_MAX_ENCODED_BYTES); + EXPECT_EQ(output.data(), reserved_data); + + const std::string too_long = std::string(1, PLAIN_ESCAPE_PREFIX) + + std::string(COMMON_GRAM_MAX_ENCODED_BYTES - 1, 'x'); + EXPECT_FALSE(try_encode_escaped_plain_term_prevalidated(too_long, output)); + EXPECT_TRUE(output.empty()); + EXPECT_EQ(output.data(), reserved_data); +} + +TEST(CommonGramsKeyCodecTest, PlainTermViewBorrowsOrdinaryKeysAndUsesScratchOnlyForEscapes) { + const std::string ordinary = "ordinary"; + std::string scratch = "stale"; + auto borrowed = try_encode_plain_term_view(ordinary, PlainTermKeyVersion::kEscapedV1, &scratch); + ASSERT_TRUE(borrowed.has_value()) << borrowed.error(); + ASSERT_TRUE(borrowed->has_value()); + EXPECT_EQ(**borrowed, ordinary); + EXPECT_EQ(borrowed->value().data(), ordinary.data()); + EXPECT_TRUE(scratch.empty()); + + const std::string escaped = std::string(1, '\x1f') + "literal"; + auto encoded = try_encode_plain_term_view(escaped, PlainTermKeyVersion::kEscapedV1, &scratch); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + ASSERT_TRUE(encoded->has_value()); + EXPECT_EQ(**encoded, *encode_plain_term(escaped, PlainTermKeyVersion::kEscapedV1)); + EXPECT_EQ(encoded->value().data(), scratch.data()); + + std::string unrepresentable(COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + unrepresentable.front() = PLAIN_ESCAPE_PREFIX; + auto absent = + try_encode_plain_term_view(unrepresentable, PlainTermKeyVersion::kEscapedV1, &scratch); + ASSERT_TRUE(absent.has_value()) << absent.error(); + EXPECT_FALSE(absent->has_value()); + EXPECT_TRUE(scratch.empty()); +} + +TEST(CommonGramsKeyCodecTest, GramComponentsUseUnescapedLogicalBytes) { + const std::array cases { + std::tuple {std::string(1, '\x1e'), std::string(1, '\x1f'), + std::string("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "00000001:\x1e\x1f")}, + std::tuple {std::string(1, '\x1f'), std::string(1, '\x1e'), + std::string("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "00000001:\x1f\x1e")}, + std::tuple {std::string("\x1e" + "L"), + std::string("\x1f" + "R"), + std::string("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "00000002:\x1e" + "L\x1f" + "R")}, + std::tuple {std::string("\x1f" + "L"), + std::string("\x1e" + "R"), + std::string("\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f" + "00000002:\x1f" + "L\x1e" + "R")}, + }; + for (const auto& [left, right, expected] : cases) { + auto encoded = encode_common_gram(left, right); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_EQ(encoded.value(), expected); + } +} + +TEST(CommonGramsKeyCodecTest, PlainKeyVersionsHaveReversibleGoldenBytes) { + const std::array raw_versions {PlainTermKeyVersion::kLegacyRaw, + PlainTermKeyVersion::kRawNoInternal}; + const std::string escape_leading = + "\x1e" + "alpha"; + const std::string gram_leading = + "\x1f" + "alpha"; + + for (PlainTermKeyVersion version : raw_versions) { + auto escape_encoded = encode_plain_term(escape_leading, version); + ASSERT_TRUE(escape_encoded.has_value()) << escape_encoded.error(); + EXPECT_EQ(escape_encoded.value(), escape_leading); + auto escape_decoded = decode_plain_term(escape_encoded.value(), version); + ASSERT_TRUE(escape_decoded.has_value()) << escape_decoded.error(); + EXPECT_EQ(escape_decoded.value(), escape_leading); + + auto gram_encoded = encode_plain_term(gram_leading, version); + ASSERT_TRUE(gram_encoded.has_value()) << gram_encoded.error(); + EXPECT_EQ(gram_encoded.value(), gram_leading); + auto gram_decoded = decode_plain_term(gram_encoded.value(), version); + ASSERT_TRUE(gram_decoded.has_value()) << gram_decoded.error(); + EXPECT_EQ(gram_decoded.value(), gram_leading); + } + + auto escape_encoded = encode_plain_term(escape_leading, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(escape_encoded.has_value()) << escape_encoded.error(); + EXPECT_EQ(escape_encoded.value(), + "\x1e" + "Ealpha"); + auto escape_decoded = + decode_plain_term(escape_encoded.value(), PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(escape_decoded.has_value()) << escape_decoded.error(); + EXPECT_EQ(escape_decoded.value(), escape_leading); + + auto gram_encoded = encode_plain_term(gram_leading, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(gram_encoded.has_value()) << gram_encoded.error(); + EXPECT_EQ(gram_encoded.value(), + "\x1e" + "Galpha"); + auto gram_decoded = decode_plain_term(gram_encoded.value(), PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(gram_decoded.has_value()) << gram_decoded.error(); + EXPECT_EQ(gram_decoded.value(), gram_leading); + + EXPECT_FALSE(decode_plain_term("\x1e", PlainTermKeyVersion::kEscapedV1).has_value()); + EXPECT_FALSE(decode_plain_term("\x1e" + "Xalpha", + PlainTermKeyVersion::kEscapedV1) + .has_value()); +} + +TEST(CommonGramsKeyCodecTest, DecodePlainTermViewBorrowsOrdinaryInput) { + const std::array cases { + std::pair {PlainTermKeyVersion::kLegacyRaw, std::string("legacy")}, + std::pair {PlainTermKeyVersion::kEscapedV1, std::string("ordinary")}, + }; + for (const auto& [version, physical_term] : cases) { + std::string scratch; + auto decoded = decode_plain_term_view(physical_term, version, &scratch); + ASSERT_TRUE(decoded.has_value()) << decoded.error(); + EXPECT_EQ(decoded.value(), physical_term); + EXPECT_EQ(decoded->data(), physical_term.data()); + EXPECT_TRUE(scratch.empty()); + } +} + +TEST(CommonGramsKeyCodecTest, DecodePlainTermViewUsesScratchForEscapes) { + const std::array cases { + std::pair {std::string("\x1e" + "Ealpha"), + std::string("\x1e" + "alpha")}, + std::pair {std::string("\x1e" + "Galpha"), + std::string("\x1f" + "alpha")}, + }; + for (const auto& [physical_term, logical_term] : cases) { + std::string scratch = "stale"; + auto decoded = + decode_plain_term_view(physical_term, PlainTermKeyVersion::kEscapedV1, &scratch); + ASSERT_TRUE(decoded.has_value()) << decoded.error(); + EXPECT_EQ(decoded.value(), logical_term); + EXPECT_EQ(scratch, logical_term); + EXPECT_EQ(decoded->data(), scratch.data()); + } +} + +TEST(CommonGramsKeyCodecTest, EscapedPlainKeysCannotEnterGramMarkerRange) { + const std::array logical_terms {std::string("plain"), + std::string("\x1e" + "plain"), + std::string("\x1f" + "plain")}; + for (const auto& term : logical_terms) { + auto encoded = encode_plain_term(term, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_FALSE(encoded.value() >= CG_V1_MARKER && encoded.value() < CG_V1_MARKER_END); + auto decoded = decode_plain_term(encoded.value(), PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(decoded.has_value()) << decoded.error(); + EXPECT_EQ(decoded.value(), term); + } +} + +TEST(CommonGramsKeyCodecTest, InternalNamespaceAndLegacyBypassAreSeparated) { + const std::string legacy_marker = + "\x1f" + "SNII_PHRASE_BIGRAM" + "\x1f"; + EXPECT_EQ(INTERNAL_TERM_NAMESPACE_BEGIN, std::string_view("\x1f", 1)); + EXPECT_EQ(INTERNAL_TERM_NAMESPACE_END, std::string_view("\x20", 1)); + + EXPECT_TRUE(is_internal_term_key(encode_common_gram("a", "b").value())); + EXPECT_TRUE(is_internal_term_key(legacy_marker + "payload")); + EXPECT_TRUE( + is_internal_term_key("\x1f" + "FUTURE_INTERNAL_TERM")); + EXPECT_FALSE( + is_internal_term_key("\x1e" + "Gescaped")); + EXPECT_FALSE(is_internal_term_key("plain")); + + EXPECT_TRUE(legacy_raw_exact_requires_bypass(CG_V1_MARKER)); + EXPECT_TRUE(legacy_raw_exact_requires_bypass(legacy_marker + "payload")); + EXPECT_FALSE(legacy_raw_exact_requires_bypass(std::string(1, '\x1f'))); + EXPECT_FALSE( + legacy_raw_exact_requires_bypass("\x1f" + "literal")); + EXPECT_FALSE(legacy_raw_exact_requires_bypass("plain")); + + EXPECT_TRUE(legacy_raw_prefix_requires_bypass("")); + EXPECT_TRUE(legacy_raw_prefix_requires_bypass(std::string(1, '\x1f'))); + EXPECT_TRUE( + legacy_raw_prefix_requires_bypass("\x1f" + "DORIS_COMMON")); + EXPECT_TRUE(legacy_raw_prefix_requires_bypass(CG_V1_MARKER)); + EXPECT_TRUE(legacy_raw_prefix_requires_bypass(legacy_marker + "payload")); + EXPECT_FALSE( + legacy_raw_prefix_requires_bypass("\x1f" + "literal")); + EXPECT_FALSE(legacy_raw_prefix_requires_bypass("plain")); +} + +TEST(CommonGramsKeyCodecTest, PlainEncodingPreservesLogicalPrefixRanges) { + const std::array logical_prefixes { + std::string("plain"), + std::string("\x1e" + "escape"), + std::string("\x1f" + "marker"), + }; + const std::array logical_terms { + logical_prefixes[0] + "-suffix", + logical_prefixes[1] + "-suffix", + logical_prefixes[2] + "-suffix", + }; + for (const std::string& logical_prefix : logical_prefixes) { + auto physical_prefix = encode_plain_term(logical_prefix, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(physical_prefix.has_value()) << physical_prefix.error(); + for (const std::string& logical_term : logical_terms) { + auto physical_term = encode_plain_term(logical_term, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(physical_term.has_value()) << physical_term.error(); + EXPECT_EQ(logical_term.starts_with(logical_prefix), + physical_term->starts_with(physical_prefix.value())); + } + } +} + +TEST(CommonGramsKeyCodecTest, EncodedLengthBoundariesAreShared) { + EXPECT_EQ(COMMON_GRAM_MAX_ENCODED_BYTES, 16383); + + const std::array all_versions {PlainTermKeyVersion::kLegacyRaw, PlainTermKeyVersion::kEscapedV1, + PlainTermKeyVersion::kRawNoInternal}; + for (PlainTermKeyVersion version : all_versions) { + for (size_t size : {size_t {16382}, size_t {16383}}) { + std::string term(size, 'p'); + auto encoded = encode_plain_term(term, version); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_EQ(encoded->size(), size); + auto decoded = decode_plain_term(encoded.value(), version); + ASSERT_TRUE(decoded.has_value()) << decoded.error(); + EXPECT_EQ(decoded.value(), term); + } + const std::string overlong(16384, 'p'); + EXPECT_FALSE(encode_plain_term(overlong, version).has_value()); + EXPECT_FALSE(decode_plain_term(overlong, version).has_value()); + } + + for (const auto& [logical_size, encoded_size] : + std::array {std::pair {size_t {16381}, size_t {16382}}, + std::pair {size_t {16382}, size_t {16383}}}) { + std::string term(logical_size, 'p'); + term.front() = '\x1e'; + auto encoded = encode_plain_term(term, PlainTermKeyVersion::kEscapedV1); + ASSERT_TRUE(encoded.has_value()) << encoded.error(); + EXPECT_EQ(encoded->size(), encoded_size); + } + + std::string escape_overflow(16383, 'p'); + escape_overflow.front() = '\x1f'; + EXPECT_FALSE(encode_plain_term(escape_overflow, PlainTermKeyVersion::kEscapedV1).has_value()); + std::string try_output = "stale"; + auto try_overflow = + try_encode_plain_term(escape_overflow, PlainTermKeyVersion::kEscapedV1, &try_output); + ASSERT_TRUE(try_overflow.has_value()) << try_overflow.error(); + EXPECT_FALSE(try_overflow.value()); + EXPECT_TRUE(try_output.empty()); + + const std::string invalid_utf8("\xc3\x28", 2); + auto try_invalid = + try_encode_plain_term(invalid_utf8, PlainTermKeyVersion::kEscapedV1, &try_output); + EXPECT_FALSE(try_invalid.has_value()); + EXPECT_EQ(try_invalid.error().code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_TRUE(try_output.empty()); + + auto raw_fallback = encode_plain_term(escape_overflow, PlainTermKeyVersion::kRawNoInternal); + ASSERT_TRUE(raw_fallback.has_value()) << raw_fallback.error(); + EXPECT_EQ(raw_fallback->size(), 16383); + EXPECT_EQ(raw_fallback.value(), escape_overflow); + + for (const auto& [encoded_size, right_size, expected_encodable] : + std::array {std::tuple {size_t {16382}, size_t {16350}, true}, + std::tuple {size_t {16383}, size_t {16351}, true}, + std::tuple {size_t {16384}, size_t {16352}, false}}) { + std::string right(right_size, 'r'); + EXPECT_EQ(is_common_gram_encodable("a", right), expected_encodable); + auto encoded = encode_common_gram("a", right); + EXPECT_EQ(encoded.has_value(), expected_encodable); + if (encoded.has_value()) { + EXPECT_EQ(encoded->size(), encoded_size); + } + } +} + +TEST(CommonGramsKeyCodecTest, GramLengthBoundariesCanBeConcentratedInLeft) { + for (const auto& [encoded_size, left_size, expected_encodable] : + std::array {std::tuple {size_t {16382}, size_t {16350}, true}, + std::tuple {size_t {16383}, size_t {16351}, true}, + std::tuple {size_t {16384}, size_t {16352}, false}}) { + std::string left(left_size, 'l'); + EXPECT_EQ(is_common_gram_encodable(left, "r"), expected_encodable); + auto encoded = encode_common_gram(left, "r"); + EXPECT_EQ(encoded.has_value(), expected_encodable); + if (encoded.has_value()) { + EXPECT_EQ(encoded->size(), encoded_size); + } + } + + const std::string e_acute = "\xc3\xa9"; + std::string multibyte_left; + multibyte_left.reserve(16352); + for (size_t i = 0; i < 8175; ++i) { + multibyte_left.append(e_acute); + } + auto at_16382 = encode_common_gram(multibyte_left, "r"); + ASSERT_TRUE(at_16382.has_value()) << at_16382.error(); + EXPECT_EQ(at_16382->size(), 16382); + + multibyte_left.push_back('x'); + auto at_16383 = encode_common_gram(multibyte_left, "r"); + ASSERT_TRUE(at_16383.has_value()) << at_16383.error(); + EXPECT_EQ(at_16383->size(), 16383); + + multibyte_left.push_back('x'); + EXPECT_FALSE(is_common_gram_encodable(multibyte_left, "r")); + EXPECT_FALSE(encode_common_gram(multibyte_left, "r").has_value()); +} + +TEST(CommonGramsKeyCodecTest, LengthUsesUtf8BytesRatherThanCodePoints) { + const std::string chinese = "\xe4\xbd\xa0"; + const std::string e_acute = "\xc3\xa9"; + std::string right; + right.reserve(16350); + for (size_t i = 0; i < 8174; ++i) { + right.append(e_acute); + } + + auto at_16382 = encode_common_gram(chinese, right); + ASSERT_TRUE(at_16382.has_value()) << at_16382.error(); + EXPECT_EQ(at_16382->size(), 16382); + EXPECT_NE(at_16382->find("00000003:"), std::string::npos); + + right.push_back('x'); + auto at_16383 = encode_common_gram(chinese, right); + ASSERT_TRUE(at_16383.has_value()) << at_16383.error(); + EXPECT_EQ(at_16383->size(), 16383); + + right.push_back('x'); + EXPECT_FALSE(is_common_gram_encodable(chinese, right)); + EXPECT_FALSE(encode_common_gram(chinese, right).has_value()); +} + +TEST(CommonGramsKeyCodecTest, RejectsNulAndInvalidUtf8) { + const std::string nul_term("a\0b", 3); + const std::string invalid_utf8("\xc3\x28", 2); + for (const auto& term : {nul_term, invalid_utf8}) { + EXPECT_FALSE(encode_plain_term(term, PlainTermKeyVersion::kLegacyRaw).has_value()); + EXPECT_FALSE(encode_plain_term(term, PlainTermKeyVersion::kEscapedV1).has_value()); + EXPECT_FALSE(encode_plain_term(term, PlainTermKeyVersion::kRawNoInternal).has_value()); + EXPECT_FALSE(decode_plain_term(term, PlainTermKeyVersion::kLegacyRaw).has_value()); + EXPECT_FALSE(is_common_gram_encodable(term, "valid")); + EXPECT_FALSE(is_common_gram_encodable("valid", term)); + EXPECT_FALSE(encode_common_gram(term, "valid").has_value()); + EXPECT_FALSE(encode_common_gram("valid", term).has_value()); + } + + const std::string escaped_nul( + "\x1e" + "Ea\0b", + 5); + const std::string escaped_invalid( + "\x1e" + "G\xc3\x28", + 4); + EXPECT_FALSE(decode_plain_term(escaped_nul, PlainTermKeyVersion::kEscapedV1).has_value()); + EXPECT_FALSE(decode_plain_term(escaped_invalid, PlainTermKeyVersion::kEscapedV1).has_value()); + for (PlainTermKeyVersion version : + {PlainTermKeyVersion::kLegacyRaw, PlainTermKeyVersion::kRawNoInternal}) { + EXPECT_FALSE(decode_plain_term(nul_term, version).has_value()); + EXPECT_FALSE(decode_plain_term(invalid_utf8, version).has_value()); + } +} + +} // namespace +} // namespace doris::segment_v2::inverted_index diff --git a/be/test/storage/index/inverted/ananlyzer/common_word_set_test.cpp b/be/test/storage/index/inverted/ananlyzer/common_word_set_test.cpp new file mode 100644 index 00000000000000..2f1ce9fca4fe00 --- /dev/null +++ b/be/test/storage/index/inverted/ananlyzer/common_word_set_test.cpp @@ -0,0 +1,186 @@ +// 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 "storage/index/inverted/common_grams/common_word_set.h" + +#include + +#include +#include +#include + +#include "common/config.h" + +namespace doris::segment_v2::inverted_index { + +namespace common_grams_testing { +uint64_t common_word_hash_lookup_count(); +void reset_common_word_hash_lookup_count(); +} // namespace common_grams_testing + +namespace { + +TEST(CommonWordSetTest, DefaultWordSetLivesUnderTheSharedDictionaryRoot) { + // The word list follows the same layout as the icu, ik and pinyin dictionaries, so relocating + // inverted_index_dict_path moves all of them together instead of leaving this one behind. + const std::string saved = config::inverted_index_dict_path; + config::inverted_index_dict_path = "/somewhere/else"; + EXPECT_EQ(CommonWordSet::default_word_set_path(), + "/somewhere/else/common_grams/default_words.txt"); + config::inverted_index_dict_path = saved; +} + +TEST(CommonWordSetTest, BuiltinEnglishStopV1IsTheExactFrozen33WordSet) { + EXPECT_EQ(BUILTIN_COMMON_WORDS_RESOURCE, "builtin:lucene_english_stop:v1"); + const auto& words = CommonWordSet::builtin_english_stop_words_v1(); + constexpr std::array expected = { + "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", + "in", "into", "is", "it", "no", "not", "of", "on", "or", "such", "that", + "the", "their", "then", "there", "these", "they", "this", "to", "was", "will", "with"}; + + EXPECT_EQ(words.size(), expected.size()); + for (std::string_view word : expected) { + EXPECT_TRUE(words.contains(word)) << word; + } + EXPECT_FALSE(words.contains("english")); + EXPECT_FALSE(words.contains("The")); +} + +TEST(CommonWordSetTest, RejectsImpossibleShapesBeforeHashLookup) { + const auto& words = CommonWordSet::builtin_english_stop_words_v1(); + common_grams_testing::reset_common_word_hash_lookup_count(); + + EXPECT_FALSE(words.contains("encyclopedia")); // longer than every builtin word + EXPECT_FALSE(words.contains("zoo")); // impossible first byte + EXPECT_FALSE(words.contains("tea")); // possible shape, absent from the set + EXPECT_TRUE(words.contains("the")); + + EXPECT_EQ(common_grams_testing::common_word_hash_lookup_count(), 2); +} + +TEST(CommonWordSetTest, WordsetV1IgnoresEmptyAndCommentLinesAndDeduplicates) { + EXPECT_EQ(WORDSET_FORMAT_V1, "wordset:v1"); + const std::string content = + "# comment\n" + "alpha\n" + "\n" + "beta\n" + "alpha\n" + "# trailing comment\n" + "\xe4\xbd\xa0\xe5\xa5\xbd\n"; + auto parsed = CommonWordSet::parse_words(content); + ASSERT_TRUE(parsed.has_value()) << parsed.error(); + EXPECT_EQ(parsed->size(), 3); + EXPECT_TRUE(parsed->contains("alpha")); + EXPECT_TRUE(parsed->contains("beta")); + EXPECT_TRUE(parsed->contains("\xe4\xbd\xa0\xe5\xa5\xbd")); + EXPECT_FALSE(parsed->contains("# comment")); + EXPECT_FALSE(parsed->contains("")); +} + +TEST(CommonWordSetTest, WordsetV1PreservesTermBytesAndAcceptsCrLfAndFinalLine) { + const std::string content = + "alpha\r\n" + "beta\r\n" + "inline#hash\n" + " leading\n" + "trailing \n" + " #not-comment\n" + "final"; + auto parsed = CommonWordSet::parse_words(content); + ASSERT_TRUE(parsed.has_value()) << parsed.error(); + EXPECT_EQ(parsed->size(), 7); + for (std::string_view term : + {"alpha", "beta", "inline#hash", " leading", "trailing ", " #not-comment", "final"}) { + EXPECT_TRUE(parsed->contains(term)) << term; + } + EXPECT_FALSE(parsed->contains("alpha\r")); + EXPECT_FALSE(parsed->contains("beta\r")); +} + +TEST(CommonWordSetTest, UnterminatedFinalCarriageReturnIsTermData) { + auto final_with_carriage_return = CommonWordSet::parse_words("final\r"); + ASSERT_TRUE(final_with_carriage_return.has_value()) << final_with_carriage_return.error(); + EXPECT_EQ(final_with_carriage_return->size(), 1); + EXPECT_TRUE(final_with_carriage_return->contains("final\r")); + EXPECT_FALSE(final_with_carriage_return->contains("final")); + + auto lone_carriage_return = CommonWordSet::parse_words("\r"); + ASSERT_TRUE(lone_carriage_return.has_value()) << lone_carriage_return.error(); + EXPECT_EQ(lone_carriage_return->size(), 1); + EXPECT_TRUE(lone_carriage_return->contains("\r")); + EXPECT_FALSE(lone_carriage_return->contains("")); +} + +TEST(CommonWordSetTest, RejectsNulAndInvalidUtf8Terms) { + const std::string nul_content("alpha\na\0b\n", 10); + EXPECT_FALSE(CommonWordSet::parse_words(nul_content).has_value()); + + const std::string invalid_utf8("alpha\n\xc3\x28\n", 9); + EXPECT_FALSE(CommonWordSet::parse_words(invalid_utf8).has_value()); + + const std::string nul_comment("# bad\0comment\nalpha\n", 20); + EXPECT_FALSE(CommonWordSet::parse_words(nul_comment).has_value()); + const std::string invalid_comment("# bad\xc3\x28\nalpha\n", 14); + EXPECT_FALSE(CommonWordSet::parse_words(invalid_comment).has_value()); +} + +TEST(CommonWordSetTest, WordsetV1IsCaseSensitive) { + auto parsed = CommonWordSet::parse_words("Alpha\nalpha\n"); + ASSERT_TRUE(parsed.has_value()) << parsed.error(); + EXPECT_EQ(parsed->size(), 2); + EXPECT_TRUE(parsed->contains("Alpha")); + EXPECT_TRUE(parsed->contains("alpha")); +} + +TEST(CommonWordSetTest, BuiltinIdentityIsTheFrozenResourceName) { + EXPECT_EQ(CommonWordSet::builtin_english_stop_words_v1().identity(), + BUILTIN_COMMON_WORDS_RESOURCE); +} + +// A segment stamps this identity and the querying analyzer compares against it, so two BEs +// reading different word lists must not agree. Deriving it from the file bytes is what makes +// the BE-local word list safe; a fixed constant would let mismatched grams pass unnoticed. +TEST(CommonWordSetTest, ParsedIdentityIsDerivedFromContent) { + auto first = CommonWordSet::parse_words("alpha\nbeta\n"); + auto same = CommonWordSet::parse_words("alpha\nbeta\n"); + auto different = CommonWordSet::parse_words("alpha\ngamma\n"); + ASSERT_TRUE(first.has_value()) << first.error(); + ASSERT_TRUE(same.has_value()) << same.error(); + ASSERT_TRUE(different.has_value()) << different.error(); + + EXPECT_EQ(first->identity(), same->identity()); + EXPECT_NE(first->identity(), different->identity()); + EXPECT_TRUE(first->identity().starts_with("wordset:md5:")); + EXPECT_NE(first->identity(), BUILTIN_COMMON_WORDS_RESOURCE); +} + +// Comments and blank lines do not change the parsed word set, but they do change the identity. +// That direction is the safe one: a superfluous re-plan costs a cost comparison, whereas reusing +// grams across an edit that DID change the list would be wrong. +TEST(CommonWordSetTest, IdentityTracksRawBytesNotJustTheParsedWords) { + auto plain = CommonWordSet::parse_words("alpha\nbeta\n"); + auto commented = CommonWordSet::parse_words("# a note\nalpha\nbeta\n"); + ASSERT_TRUE(plain.has_value()) << plain.error(); + ASSERT_TRUE(commented.has_value()) << commented.error(); + + EXPECT_EQ(plain->size(), commented->size()); + EXPECT_NE(plain->identity(), commented->identity()); +} + +} // namespace +} // namespace doris::segment_v2::inverted_index diff --git a/be/test/storage/index/inverted/ananlyzer/custom_analyzer_test.cpp b/be/test/storage/index/inverted/ananlyzer/custom_analyzer_test.cpp index a96edc1f38b346..e6d0d73f5cf2a9 100644 --- a/be/test/storage/index/inverted/ananlyzer/custom_analyzer_test.cpp +++ b/be/test/storage/index/inverted/ananlyzer/custom_analyzer_test.cpp @@ -21,12 +21,16 @@ #include #include +#include #include "CLucene/store/Directory.h" #include "CLucene/store/FSDirectory.h" #include "roaring/roaring.hh" #include "runtime/exec_env.h" #include "storage/index/inverted/analysis_factory_mgr.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_word_set.h" #include "storage/index/inverted/query/phrase_prefix_query.h" #include "storage/index/inverted/query/phrase_query.h" #include "storage/index/inverted/setting.h" @@ -255,6 +259,437 @@ TEST_F(CustomAnalyzerTest, TokenStreamWithReaderPtr) { delete token_stream; } +CustomAnalyzerConfigPtr common_grams_config( + const std::string& tokenizer, const Settings& tokenizer_settings, + const std::vector>& filters = {}, + const Settings& common_grams_settings = {}) { + CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config(tokenizer, tokenizer_settings); + for (const auto& [name, settings] : filters) { + builder.add_token_filter_config(name, settings); + } + builder.add_token_filter_config("common_grams", common_grams_settings); + return builder.build(); +} + +Settings whitespace_tokenizer_settings() { + Settings settings; + settings.set("tokenize_on_chars", "[whitespace]"); + return settings; +} + +std::string expected_gram(std::string_view left, std::string_view right) { + auto result = encode_common_gram(left, right); + EXPECT_TRUE(result.has_value()) << result.error(); + return result.value(); +} + +void expect_common_grams_analyzer_error(const CustomAnalyzerConfigPtr& config, + AnalysisPurpose purpose) { + try { + CustomAnalyzer::build_custom_analyzer(config, purpose); + FAIL() << "expected CommonGrams analyzer error"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } +} + +TEST_F(CustomAnalyzerTest, CommonGramsBuildsIndependentPurposeStreams) { + auto config = + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}}); + auto index = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kIndex); + auto snii_index = + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kSniiTransientIndex); + auto plain = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPlainQuery); + auto exact = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kExactPhraseQuery); + auto prefix = + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPhrasePrefixQuery); + + EXPECT_EQ(tokenize1(index, "Man of the Year"), + (std::vector {{"man", 1}, + {expected_gram("man", "of"), 0}, + {"of", 1}, + {expected_gram("of", "the"), 0}, + {"the", 1}, + {expected_gram("the", "year"), 0}, + {"year", 1}})); + EXPECT_EQ(tokenize1(snii_index, "Man of the Year"), + (std::vector {{"man", 1}, + {expected_gram("man", "of"), 0}, + {"of", 1}, + {expected_gram("of", "the"), 0}, + {"the", 1}, + {expected_gram("the", "year"), 0}, + {"year", 1}})); + EXPECT_EQ(tokenize1(plain, "Man of the Year"), + (std::vector {{"man", 1}, {"of", 1}, {"the", 1}, {"year", 1}})); + EXPECT_EQ(tokenize1(exact, "Man of the Year"), + (std::vector {{expected_gram("man", "of"), 1}, + {expected_gram("of", "the"), 1}, + {expected_gram("the", "year"), 1}})); + EXPECT_EQ(tokenize1(prefix, "the wo"), + (std::vector {{expected_gram("the", "wo"), 1}})); + + EXPECT_EQ(tokenize1(exact, "plain terms"), + (std::vector {{"plain", 1}, {"terms", 1}})); + EXPECT_EQ(tokenize1(prefix, "of term"), + (std::vector {{expected_gram("of", "term"), 1}})); +} + +TEST_F(CustomAnalyzerTest, CommonGramsReusableIndexStreamDoesNotBridgeRows) { + auto analyzer = CustomAnalyzer::build_custom_analyzer( + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}}), + AnalysisPurpose::kIndex); + auto reader = std::make_shared>(); + + auto tokenize_row = [&](std::string_view row) { + reader->init(row.data(), static_cast(row.size()), false); + auto* stream = analyzer->reusableTokenStream(L"", reader); + stream->reset(); + std::vector tokens; + Token token; + while (stream->next(&token)) { + tokens.emplace_back(std::string(token.termBuffer(), token.termLength()), + token.getPositionIncrement()); + } + return std::pair {stream, std::move(tokens)}; + }; + + auto [first_stream, first] = tokenize_row("foo of"); + auto [second_stream, second] = tokenize_row("the bar"); + + EXPECT_EQ(first_stream, second_stream); + EXPECT_EQ(first, (std::vector { + {"foo", 1}, {expected_gram("foo", "of"), 0}, {"of", 1}})); + EXPECT_EQ(second, (std::vector { + {"the", 1}, {expected_gram("the", "bar"), 0}, {"bar", 1}})); +} + +TEST_F(CustomAnalyzerTest, CommonGramsEscapesPlainKeysOnlyForIndexPurpose) { + Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[\\u0020]"); + auto config = common_grams_config("char_group", tokenizer_settings); + auto index = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kIndex); + auto plain = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPlainQuery); + auto exact = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kExactPhraseQuery); + auto prefix = + CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPhrasePrefixQuery); + + const std::string logical = std::string(1, '\x1f') + "literal"; + const std::string input = "the " + logical; + const std::string physical = std::string(1, PLAIN_ESCAPE_PREFIX) + "Gliteral"; + const std::string common_gram = expected_gram("the", logical); + + EXPECT_EQ(tokenize1(index, input), + (std::vector {{"the", 1}, {common_gram, 0}, {physical, 1}})); + EXPECT_EQ(tokenize1(plain, input), (std::vector {{"the", 1}, {logical, 1}})); + EXPECT_EQ(tokenize1(exact, input), (std::vector {{common_gram, 1}})); + EXPECT_EQ(tokenize1(prefix, input), (std::vector {{common_gram, 1}})); +} + +TEST_F(CustomAnalyzerTest, AnalyseResultPreservesCommonGramTermKind) { + auto config = + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}}); + auto exact = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kExactPhraseQuery); + auto reader = InvertedIndexAnalyzer::create_reader({}); + const std::string input = "Man of year"; + reader->init(input.data(), static_cast(input.size()), true); + + const auto terms = InvertedIndexAnalyzer::get_analyse_result(reader, exact.get()); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].get_single_term(), expected_gram("man", "of")); + EXPECT_EQ(terms[0].key_kind, TermKeyKind::kCommonGram); + EXPECT_EQ(terms[1].get_single_term(), expected_gram("of", "year")); + EXPECT_EQ(terms[1].key_kind, TermKeyKind::kCommonGram); + + auto plain = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kPlainQuery); + reader = InvertedIndexAnalyzer::create_reader({}); + reader->init(input.data(), static_cast(input.size()), true); + const auto plain_terms = InvertedIndexAnalyzer::get_analyse_result(reader, plain.get()); + ASSERT_EQ(plain_terms.size(), 3U); + for (const auto& term : plain_terms) { + EXPECT_EQ(term.key_kind, TermKeyKind::kPlain); + } +} + +TEST_F(CustomAnalyzerTest, AnalyseResultPreservesBothCommonGramTermKind) { + auto config = + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}}); + auto index = CustomAnalyzer::build_custom_analyzer(config, AnalysisPurpose::kIndex); + auto reader = InvertedIndexAnalyzer::create_reader({}); + const std::string input = "of the"; + reader->init(input.data(), static_cast(input.size()), true); + + const auto terms = InvertedIndexAnalyzer::get_analyse_result(reader, index.get()); + ASSERT_EQ(terms.size(), 3U); + EXPECT_EQ(terms[1].get_single_term(), expected_gram("of", "the")); + EXPECT_EQ(terms[1].key_kind, TermKeyKind::kCommonGram); +} + +TEST_F(CustomAnalyzerTest, CommonGramsAllowsAbsentOrOneTerminalFilter) { + CustomAnalyzerConfig::Builder missing; + missing.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + auto missing_config = missing.build(); + for (AnalysisPurpose purpose : + {AnalysisPurpose::kIndex, AnalysisPurpose::kPlainQuery, AnalysisPurpose::kExactPhraseQuery, + AnalysisPurpose::kPhrasePrefixQuery}) { + EXPECT_NO_THROW(CustomAnalyzer::build_custom_analyzer(missing_config, purpose)); + } + + CustomAnalyzerConfig::Builder duplicate; + duplicate.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + duplicate.add_token_filter_config("common_grams", {}); + duplicate.add_token_filter_config("common_grams", {}); + expect_common_grams_analyzer_error(duplicate.build(), AnalysisPurpose::kIndex); + + CustomAnalyzerConfig::Builder not_last; + not_last.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + not_last.add_token_filter_config("common_grams", {}); + not_last.add_token_filter_config("lowercase", {}); + expect_common_grams_analyzer_error(not_last.build(), AnalysisPurpose::kIndex); +} + +TEST_F(CustomAnalyzerTest, CommonGramsValidatesPlainConfigurationButOmitsGrams) { + Settings unknown; + unknown.set("unknown_setting", "true"); + auto invalid = common_grams_config("char_group", whitespace_tokenizer_settings(), {}, unknown); + expect_common_grams_analyzer_error(invalid, AnalysisPurpose::kPlainQuery); + + auto valid = common_grams_config("char_group", whitespace_tokenizer_settings()); + auto plain = CustomAnalyzer::build_custom_analyzer(valid, AnalysisPurpose::kPlainQuery); + EXPECT_EQ(tokenize1(plain, "the term"), (std::vector {{"the", 1}, {"term", 1}})); +} + +TEST_F(CustomAnalyzerTest, CommonGramsIndexChainRejectsInvalidUtf8AfterAValidToken) { + auto analyzer = CustomAnalyzer::build_custom_analyzer( + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}}), + AnalysisPurpose::kIndex); + const std::string input = std::string("valid b") + static_cast(0xFF) + std::string("ad"); + auto reader = std::make_shared>(); + reader->init(input.data(), static_cast(input.size()), false); + auto* stream = analyzer->reusableTokenStream(L"", reader); + stream->reset(); + + Token token; + ASSERT_NE(stream->next(&token), nullptr); + EXPECT_EQ(std::string(token.termBuffer(), token.termLength()), "valid"); + try { + stream->next(&token); + FAIL() << "expected malformed UTF-8 to fail the analyzer chain"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } +} + +TEST_F(CustomAnalyzerTest, CommonGramsDeniesUnsafePositionFactories) { + for (const std::string& tokenizer : {"standard", "pinyin", "basic", "icu", "keyword"}) { + expect_common_grams_analyzer_error(common_grams_config(tokenizer, {}), + AnalysisPurpose::kIndex); + } + + Settings preserve_original; + preserve_original.set("preserve_original", "true"); + expect_common_grams_analyzer_error( + common_grams_config("char_group", whitespace_tokenizer_settings(), + {{"asciifolding", preserve_original}}), + AnalysisPurpose::kIndex); + expect_common_grams_analyzer_error( + common_grams_config("char_group", whitespace_tokenizer_settings(), + {{"word_delimiter", {}}}), + AnalysisPurpose::kIndex); + expect_common_grams_analyzer_error( + common_grams_config("char_group", whitespace_tokenizer_settings(), {{"pinyin", {}}}), + AnalysisPurpose::kIndex); +} + +TEST_F(CustomAnalyzerTest, CommonGramsAcceptsReviewedUnitPositionFactories) { + for (const std::string& tokenizer : {"empty", "char_group", "ngram", "edge_ngram"}) { + EXPECT_NO_THROW(CustomAnalyzer::build_custom_analyzer( + common_grams_config(tokenizer, tokenizer == "char_group" + ? whitespace_tokenizer_settings() + : Settings {}), + AnalysisPurpose::kIndex)); + } + + for (const std::string& filter : {"empty", "lowercase", "icu_normalizer", "asciifolding"}) { + EXPECT_NO_THROW(CustomAnalyzer::build_custom_analyzer( + common_grams_config("char_group", whitespace_tokenizer_settings(), {{filter, {}}}), + AnalysisPurpose::kIndex)); + } +} + +TEST_F(CustomAnalyzerTest, CommonGramsRejectsTokensNormalizedToEmpty) { + const auto config = common_grams_config("char_group", whitespace_tokenizer_settings(), + {{"icu_normalizer", {}}}); + const std::string input = std::string("\xC2\xAD") + " the"; + for (AnalysisPurpose purpose : + {AnalysisPurpose::kIndex, AnalysisPurpose::kPlainQuery, AnalysisPurpose::kExactPhraseQuery, + AnalysisPurpose::kPhrasePrefixQuery}) { + auto analyzer = CustomAnalyzer::build_custom_analyzer(config, purpose); + try { + tokenize1(analyzer, input); + ADD_FAILURE() << "expected analyzer error for purpose " << static_cast(purpose); + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } +} + +TEST_F(CustomAnalyzerTest, CommonGramsInternalQueryFiltersAreNotRegistered) { + for (const std::string& internal : {"common_grams_query", "common_grams_phrase_prefix"}) { + CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + builder.add_token_filter_config(internal, {}); + builder.add_token_filter_config("common_grams", {}); + EXPECT_THROW( + CustomAnalyzer::build_custom_analyzer(builder.build(), AnalysisPurpose::kIndex), + Exception); + } +} + +// "the" is a member of the built-in stop-word list, which is what default_word_set() resolves to +// when no wordset file is installed -- the provider can no longer be handed a word list of its own. +TEST_F(CustomAnalyzerTest, CommonGramsProviderCachesPurposeAnalyzersWithOneWordSetSnapshot) { + auto provider = std::make_shared(common_grams_config( + "char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}})); + + auto index = provider->get_analyzer(AnalysisPurpose::kIndex); + auto snii_index = provider->get_analyzer(AnalysisPurpose::kSniiTransientIndex); + auto plain = provider->get_analyzer(AnalysisPurpose::kPlainQuery); + auto exact = provider->get_analyzer(AnalysisPurpose::kExactPhraseQuery); + auto prefix = provider->get_analyzer(AnalysisPurpose::kPhrasePrefixQuery); + + EXPECT_EQ(tokenize1(std::dynamic_pointer_cast(index), "The year"), + (std::vector { + {"the", 1}, {expected_gram("the", "year"), 0}, {"year", 1}})); + EXPECT_EQ(tokenize1(std::dynamic_pointer_cast(snii_index), "The year"), + (std::vector { + {"the", 1}, {expected_gram("the", "year"), 0}, {"year", 1}})); + EXPECT_EQ(tokenize1(std::dynamic_pointer_cast(plain), "The year"), + (std::vector {{"the", 1}, {"year", 1}})); + EXPECT_EQ(tokenize1(std::dynamic_pointer_cast(exact), "The year"), + (std::vector {{expected_gram("the", "year"), 1}})); + EXPECT_EQ(tokenize1(std::dynamic_pointer_cast(prefix), "The ye"), + (std::vector {{expected_gram("the", "ye"), 1}})); + + // Every purpose analyzer shares the one process-wide word list. + EXPECT_EQ(provider->common_words(), CommonWordSet::default_word_set()); + EXPECT_EQ(provider->get_analyzer(AnalysisPurpose::kPlainQuery), plain); + EXPECT_NE(index, plain); + EXPECT_NE(index, snii_index); + EXPECT_NE(plain, exact); + EXPECT_NE(exact, prefix); +} + +TEST_F(CustomAnalyzerTest, CommonGramsProviderOwnsDeterministicBuiltinIdentity) { + Settings first_settings; + first_settings.set("max_token_length", "16383"); + first_settings.set("tokenize_on_chars", "[whitespace]"); + Settings second_settings; + second_settings.set("tokenize_on_chars", "[whitespace]"); + second_settings.set("max_token_length", "16383"); + + CustomAnalyzerProvider first( + common_grams_config("char_group", first_settings, {{"lowercase", {}}})); + CustomAnalyzerProvider second( + common_grams_config("char_group", second_settings, {{"lowercase", {}}})); + + ASSERT_NE(first.common_grams_identity(), nullptr); + ASSERT_NE(second.common_grams_identity(), nullptr); + EXPECT_EQ(first.common_grams_identity()->common_grams_dictionary_identity, + BUILTIN_COMMON_WORDS_RESOURCE); + EXPECT_EQ(*first.common_grams_identity(), *second.common_grams_identity()); + EXPECT_EQ(first.common_grams_identity()->base_analyzer_fingerprint.size(), 64U); + EXPECT_EQ(first.common_grams_identity()->common_grams_fingerprint.size(), 64U); +} + +TEST_F(CustomAnalyzerTest, CommonGramsProviderSeparatesBaseAndDictionaryIdentity) { + Settings first_settings = whitespace_tokenizer_settings(); + first_settings.set("max_token_length", "100"); + Settings second_settings = whitespace_tokenizer_settings(); + second_settings.set("max_token_length", "101"); + CustomAnalyzerProvider first(common_grams_config("char_group", first_settings)); + CustomAnalyzerProvider second(common_grams_config("char_group", second_settings)); + + ASSERT_NE(first.common_grams_identity(), nullptr); + ASSERT_NE(second.common_grams_identity(), nullptr); + EXPECT_NE(first.common_grams_identity()->base_analyzer_fingerprint, + second.common_grams_identity()->base_analyzer_fingerprint); + EXPECT_EQ(first.common_grams_identity()->common_grams_fingerprint, + second.common_grams_identity()->common_grams_fingerprint); + + // The dictionary half of the identity is the word list's own content identity. Neither the + // provider nor an index policy can supply one, so every provider in this process agrees on it. + EXPECT_EQ(first.common_grams_identity()->common_grams_dictionary_identity, + CommonWordSet::default_word_set()->identity()); + EXPECT_EQ(second.common_grams_identity()->common_grams_dictionary_identity, + first.common_grams_identity()->common_grams_dictionary_identity); +} + +TEST_F(CustomAnalyzerTest, CommonGramsIdentityIncludesOuterCharFilter) { + CustomAnalyzerProvider no_outer_filter( + common_grams_config("char_group", whitespace_tokenizer_settings())); + ASSERT_NE(no_outer_filter.common_grams_identity(), nullptr); + + const std::map slash_to_space = { + {INVERTED_INDEX_PARSER_CHAR_FILTER_TYPE, INVERTED_INDEX_CHAR_FILTER_CHAR_REPLACE}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_PATTERN, "/"}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_REPLACEMENT, " "}}; + const std::map dash_to_space = { + {INVERTED_INDEX_PARSER_CHAR_FILTER_TYPE, INVERTED_INDEX_CHAR_FILTER_CHAR_REPLACE}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_PATTERN, "-"}, + {INVERTED_INDEX_PARSER_CHAR_FILTER_REPLACEMENT, " "}}; + CustomAnalyzerProvider slash_filter( + common_grams_config("char_group", whitespace_tokenizer_settings()), slash_to_space); + CustomAnalyzerProvider repeated_slash_filter( + common_grams_config("char_group", whitespace_tokenizer_settings()), slash_to_space); + CustomAnalyzerProvider dash_filter( + common_grams_config("char_group", whitespace_tokenizer_settings()), dash_to_space); + ASSERT_NE(slash_filter.common_grams_identity(), nullptr); + ASSERT_NE(repeated_slash_filter.common_grams_identity(), nullptr); + ASSERT_NE(dash_filter.common_grams_identity(), nullptr); + + EXPECT_EQ(*slash_filter.common_grams_identity(), + *repeated_slash_filter.common_grams_identity()); + EXPECT_EQ(slash_filter.common_grams_identity()->common_grams_dictionary_identity, + no_outer_filter.common_grams_identity()->common_grams_dictionary_identity); + EXPECT_EQ(slash_filter.common_grams_identity()->common_grams_fingerprint, + no_outer_filter.common_grams_identity()->common_grams_fingerprint); + EXPECT_NE(slash_filter.common_grams_identity()->base_analyzer_fingerprint, + no_outer_filter.common_grams_identity()->base_analyzer_fingerprint); + EXPECT_NE(slash_filter.common_grams_identity()->base_analyzer_fingerprint, + dash_filter.common_grams_identity()->base_analyzer_fingerprint); +} + +TEST_F(CustomAnalyzerTest, ProviderSharesOneLegacyAnalyzerWhenCommonGramsIsAbsent) { + CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + builder.add_token_filter_config("lowercase", {}); + auto provider = std::make_shared(builder.build()); + + auto analyzer = provider->get_analyzer(AnalysisPurpose::kIndex); + EXPECT_EQ(provider->get_analyzer(AnalysisPurpose::kPlainQuery), analyzer); + EXPECT_EQ(provider->get_analyzer(AnalysisPurpose::kExactPhraseQuery), analyzer); + EXPECT_EQ(provider->get_analyzer(AnalysisPurpose::kPhrasePrefixQuery), analyzer); + EXPECT_EQ(provider->common_grams_identity(), nullptr); +} + +TEST_F(CustomAnalyzerTest, ProviderExposesBaseFingerprintWithoutCommonGrams) { + CustomAnalyzerConfig::Builder plain_builder; + plain_builder.with_tokenizer_config("char_group", whitespace_tokenizer_settings()); + plain_builder.add_token_filter_config("lowercase", {}); + CustomAnalyzerProvider plain(plain_builder.build()); + + CustomAnalyzerProvider common_grams(common_grams_config( + "char_group", whitespace_tokenizer_settings(), {{"lowercase", {}}})); + + EXPECT_EQ(plain.base_analyzer_fingerprint().size(), 64U); + EXPECT_EQ(plain.base_analyzer_fingerprint(), common_grams.base_analyzer_fingerprint()); + ASSERT_NE(common_grams.common_grams_identity(), nullptr); + EXPECT_EQ(common_grams.base_analyzer_fingerprint(), + common_grams.common_grams_identity()->base_analyzer_fingerprint); +} + // TEST_F(CustomAnalyzerTest, test) { // std::string name = "name"; // std::string path = "/mnt/disk3/yangsiyu/clucene"; diff --git a/be/test/storage/index/inverted/common/single_flight_test.cpp b/be/test/storage/index/inverted/common/single_flight_test.cpp new file mode 100644 index 00000000000000..5f6527230552ae --- /dev/null +++ b/be/test/storage/index/inverted/common/single_flight_test.cpp @@ -0,0 +1,220 @@ +// 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 "storage/index/inverted/common/single_flight.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" + +namespace doris::segment_v2::inverted_index { + +struct BlockingMoveState { + std::mutex mutex; + std::condition_variable condition; + bool move_started = false; + bool release_move = false; +}; + +struct BlockingMoveResult { + explicit BlockingMoveResult(std::shared_ptr state_) + : state(std::move(state_)) {} + + BlockingMoveResult(const BlockingMoveResult&) = default; + BlockingMoveResult& operator=(const BlockingMoveResult&) = default; + + BlockingMoveResult(BlockingMoveResult&& other) noexcept : state(std::move(other.state)) { + std::unique_lock lock(state->mutex); + state->move_started = true; + state->condition.notify_all(); + state->condition.wait(lock, [&] { return state->release_move; }); + } + + BlockingMoveResult& operator=(BlockingMoveResult&&) = delete; + + std::shared_ptr state; +}; + +// A lone caller leads, a second concurrent caller follows, and the follower reuses the +// leader's published result. The in-flight entry is cleared on publish. +TEST(SingleFlight, LeadFollowReuse) { + SingleFlight sf; + + auto leader = sf.join_or_lead("k"); + EXPECT_FALSE(leader.has_value()); // first caller leads + EXPECT_EQ(sf.inflight_size(), 1U); + + auto follower = sf.join_or_lead("k"); + ASSERT_TRUE(follower.has_value()); // second caller follows the in-flight leader + EXPECT_EQ(sf.inflight_size(), 1U); // still one in-flight key + + sf.publish("k", 7); + EXPECT_EQ(follower->get(), 7); // follower reuses the leader's result + EXPECT_EQ(sf.inflight_size(), 0U); +} + +// Different keys execute independently -- each first caller leads. +TEST(SingleFlight, DistinctKeysLeadIndependently) { + SingleFlight sf; + EXPECT_FALSE(sf.join_or_lead("a").has_value()); + EXPECT_FALSE(sf.join_or_lead("b").has_value()); + EXPECT_EQ(sf.inflight_size(), 2U); + sf.publish("a", 1); + sf.publish("b", 2); + EXPECT_EQ(sf.inflight_size(), 0U); +} + +// After a key is published, the next caller of the same key leads a fresh execution. +TEST(SingleFlight, ReLeadAfterPublish) { + SingleFlight sf; + EXPECT_FALSE(sf.join_or_lead("k").has_value()); + sf.publish("k", 1); + EXPECT_EQ(sf.inflight_size(), 0U); + EXPECT_FALSE(sf.join_or_lead("k").has_value()); // leads again + EXPECT_EQ(sf.inflight_size(), 1U); + sf.publish("k", 2); +} + +TEST(SingleFlight, JoinDuringPublicationStillFollowsPublishedFlight) { + SingleFlight sf; + ASSERT_FALSE(sf.join_or_lead("k").has_value()); + + auto state = std::make_shared(); + BlockingMoveResult result(state); + std::thread publisher([&] { sf.publish("k", result); }); + { + std::unique_lock lock(state->mutex); + state->condition.wait(lock, [&] { return state->move_started; }); + } + + auto follower = sf.join_or_lead("k"); + EXPECT_TRUE(follower.has_value()); + { + std::lock_guard lock(state->mutex); + state->release_move = true; + } + state->condition.notify_all(); + publisher.join(); + + ASSERT_TRUE(follower.has_value()); + EXPECT_EQ(follower->get().state, state); + EXPECT_EQ(sf.inflight_size(), 0U); +} + +// The motivating scenario: many threads issue the same key concurrently. Exactly one leads; +// all followers receive the single shared result (one shared_ptr payload, shared read-only). +TEST(SingleFlight, ConcurrentCollapsesToOneLeader) { + constexpr int kThreads = 16; + SingleFlight> sf; + + std::atomic leader_count {0}; + std::latch joined(kThreads); + std::latch release_leader(1); + std::vector> results(kThreads); + auto payload = std::make_shared(42); + + std::vector threads; + threads.reserve(kThreads); + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + auto follow = sf.join_or_lead("hot-segment"); + const bool is_leader = !follow.has_value(); + joined.count_down(); + if (is_leader) { + leader_count.fetch_add(1, std::memory_order_relaxed); + joined.wait(); // ensure every thread has joined before we publish + release_leader.wait(); + sf.publish("hot-segment", payload); + } else { + results[i] = follow->get(); + } + }); + } + joined.wait(); + EXPECT_EQ(sf.inflight_size(), 1U); + release_leader.count_down(); + for (auto& t : threads) { + t.join(); + } + + EXPECT_EQ(leader_count.load(), 1); // only one open/execute happened + EXPECT_EQ(sf.inflight_size(), 0U); + for (int i = 0; i < kThreads; ++i) { + if (results[i] != nullptr) { + EXPECT_EQ(results[i], payload); // same shared object, not a recomputed copy + EXPECT_EQ(*results[i], 42); + } + } +} + +TEST(SingleFlight, ConcurrentDifferentRawQueriesNeverJoin) { + constexpr int kThreads = 16; + SingleFlight sf; + std::atomic leader_count {0}; + std::latch all_registered(kThreads); + std::latch release_publishers(1); + std::vector threads; + threads.reserve(kThreads); + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + const std::string key = "raw-query-" + std::to_string(i); + auto follower = sf.join_or_lead(key); + if (!follower.has_value()) { + leader_count.fetch_add(1, std::memory_order_relaxed); + } + all_registered.count_down(); + release_publishers.wait(); + DORIS_CHECK(!follower.has_value()); + sf.publish(key, i); + }); + } + all_registered.wait(); + EXPECT_EQ(sf.inflight_size(), kThreads); + release_publishers.count_down(); + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(leader_count.load(std::memory_order_relaxed), kThreads); + EXPECT_EQ(sf.inflight_size(), 0U); +} + +// A leader that fails publishes its failure; followers observe it (and would then fall back +// to computing independently in the production path). +TEST(SingleFlight, ErrorResultPropagates) { + SingleFlight> sf; + auto leader = sf.join_or_lead("k"); + ASSERT_FALSE(leader.has_value()); + auto follower = sf.join_or_lead("k"); + ASSERT_TRUE(follower.has_value()); + + sf.publish("k", std::make_pair(false, 0)); // leader failed + auto [ok, value] = follower->get(); + EXPECT_FALSE(ok); + EXPECT_EQ(value, 0); +} + +} // namespace doris::segment_v2::inverted_index diff --git a/be/test/storage/index/inverted/compaction/index_compaction_test.cpp b/be/test/storage/index/inverted/compaction/index_compaction_test.cpp index 0f96372a5b6058..acff45e65fea32 100644 --- a/be/test/storage/index/inverted/compaction/index_compaction_test.cpp +++ b/be/test/storage/index/inverted/compaction/index_compaction_test.cpp @@ -17,8 +17,12 @@ #include +#include + #include "storage/index/inverted/compaction/util/index_compaction_utils.cpp" #include "storage/utils.h" +#include "util/debug_points.h" +#include "util/defer_op.h" namespace doris { @@ -731,6 +735,62 @@ class IndexCompactionTest : public ::testing::Test { } } + void _build_snii_multi_index_tablet(bool second_supports_phrase = true) { + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(KeysType::DUP_KEYS); + schema_pb.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::SNII); + + IndexCompactionUtils::construct_column(schema_pb.add_column(), 0, "INT", "key"); + IndexCompactionUtils::construct_column(schema_pb.add_column(), 1, "STRING", "v1"); + IndexCompactionUtils::construct_column(schema_pb.add_column(), 2, "STRING", "v2"); + IndexCompactionUtils::construct_column(schema_pb.add_column(), 3, "INT", "v3"); + + auto add_index = [&schema_pb](int64_t index_id, std::string_view name, + bool supports_phrase) { + TabletIndexPB* index = schema_pb.add_index(); + index->set_index_id(index_id); + index->set_index_name(std::string(name)); + index->set_index_type(IndexType::INVERTED); + index->add_col_unique_id(1); + auto* properties = index->mutable_properties(); + (*properties)[INVERTED_INDEX_PARSER_KEY] = INVERTED_INDEX_PARSER_UNICODE; + (*properties)[INVERTED_INDEX_PARSER_PHRASE_SUPPORT_KEY] = + supports_phrase ? INVERTED_INDEX_PARSER_PHRASE_SUPPORT_YES + : INVERTED_INDEX_PARSER_PHRASE_SUPPORT_NO; + (*properties)[INVERTED_INDEX_PARSER_LOWERCASE_KEY] = INVERTED_INDEX_PARSER_TRUE; + }; + add_index(11001, "v1_phrase_a", true); + add_index(11002, "v1_phrase_b", second_supports_phrase); + + _tablet_schema = std::make_shared(); + _tablet_schema->init_from_pb(schema_pb); + TabletMetaSharedPtr tablet_meta(new TabletMeta(_tablet_schema)); + _tablet = std::make_shared(*_engine_ref, tablet_meta, _data_dir.get()); + EXPECT_TRUE(_tablet->init().ok()); + } + + std::vector _build_snii_source_rowsets() { + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(_tablet->tablet_path()).ok()); + EXPECT_TRUE(io::global_local_filesystem()->create_directory(_tablet->tablet_path()).ok()); + const std::vector data_files = { + _current_dir + "/be/test/storage/index/inverted/data/data1.csv", + _current_dir + "/be/test/storage/index/inverted/data/data2.csv"}; + std::vector rowsets(data_files.size()); + auto check_indexes = [](const int32_t& size) { EXPECT_EQ(size, 2); }; + IndexCompactionUtils::build_rowsets( + _data_dir, _tablet_schema, _tablet, _engine_ref, rowsets, data_files, _inc_id, + check_indexes, false, 1000); + return rowsets; + } + + static std::string _read_index_file_bytes(const RowsetSharedPtr& rowset, uint32_t segment_id) { + const std::string path = fmt::format("{}/{}_{}.idx", rowset->tablet_path(), + rowset->rowset_id().to_string(), segment_id); + std::ifstream input(path, std::ios::binary); + EXPECT_TRUE(input.is_open()) << path; + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; + } + private: TabletSchemaSPtr _tablet_schema = nullptr; StorageEngine* _engine_ref = nullptr; @@ -1687,4 +1747,208 @@ TEST_F(IndexCompactionTest, test_inverted_index_ram_dir_disable_with_debug_point config::inverted_index_ram_dir_enable = original_ram_dir_enable; config::enable_debug_points = original_enable_debug_points; } + +TEST_F(IndexCompactionTest, snii_native_merge_validates_rowids_once_and_matches_raw_rebuild) { + const bool old_common_grams = config::enable_common_grams_index_build; + const bool old_debug_points = config::enable_debug_points; + const bool old_write_freq = config::snii_positions_index_write_freq; + config::enable_common_grams_index_build = false; + config::enable_debug_points = true; + config::snii_positions_index_write_freq = false; + constexpr std::string_view kValidationPoint = + "Compaction::snii_validated_rowid_conversion_created"; + constexpr std::string_view kReaderInitPoint = "Compaction::snii_eligibility_reader_initialized"; + DEFER({ + DebugPoints::instance()->remove(std::string(kValidationPoint)); + DebugPoints::instance()->remove(std::string(kReaderInitPoint)); + config::enable_common_grams_index_build = old_common_grams; + config::enable_debug_points = old_debug_points; + config::snii_positions_index_write_freq = old_write_freq; + }); + + _build_snii_multi_index_tablet(); + const std::vector rowsets = _build_snii_source_rowsets(); + size_t validation_count = 0; + size_t reader_init_count = 0; + std::function count_validation = [&validation_count]() { ++validation_count; }; + std::function count_reader_init = [&reader_init_count]() { ++reader_init_count; }; + DebugPoints::instance()->add_with_callback(std::string(kValidationPoint), count_validation); + DebugPoints::instance()->add_with_callback(std::string(kReaderInitPoint), count_reader_init); + + auto check_native_merge = [](const BaseCompaction& compaction, const RowsetWriterContext& ctx) { + EXPECT_EQ(ctx.columns_to_do_index_compaction, (std::set {1})); + EXPECT_EQ(compaction._output_rowset->num_segments(), 2); + }; + RowsetSharedPtr native_merge; + Status status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, true, + native_merge, check_native_merge, 1000); + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(native_merge, nullptr); + EXPECT_EQ(validation_count, 1); + EXPECT_EQ(reader_init_count, rowsets.size()); + DebugPoints::instance()->remove(std::string(kValidationPoint)); + DebugPoints::instance()->remove(std::string(kReaderInitPoint)); + + auto check_raw_rebuild = [](const BaseCompaction& compaction, const RowsetWriterContext& ctx) { + EXPECT_TRUE(ctx.columns_to_do_index_compaction.empty()); + EXPECT_EQ(compaction._output_rowset->num_segments(), 2); + }; + RowsetSharedPtr raw_rebuild; + status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, false, raw_rebuild, + check_raw_rebuild, 1000); + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(raw_rebuild, nullptr); + ASSERT_EQ(native_merge->num_segments(), raw_rebuild->num_segments()); + for (uint32_t segment_id = 0; segment_id < native_merge->num_segments(); ++segment_id) { + EXPECT_EQ(_read_index_file_bytes(native_merge, segment_id), + _read_index_file_bytes(raw_rebuild, segment_id)); + } +} + +TEST_F(IndexCompactionTest, snii_native_merge_falls_back_for_entire_multi_index_column) { + const bool old_common_grams = config::enable_common_grams_index_build; + const bool old_write_freq = config::snii_positions_index_write_freq; + config::enable_common_grams_index_build = false; + config::snii_positions_index_write_freq = false; + DEFER({ + config::enable_common_grams_index_build = old_common_grams; + config::snii_positions_index_write_freq = old_write_freq; + }); + + _build_snii_multi_index_tablet(/*second_supports_phrase=*/false); + const std::vector rowsets = _build_snii_source_rowsets(); + auto check_fallback = [](const BaseCompaction& compaction, const RowsetWriterContext& ctx) { + EXPECT_TRUE(ctx.columns_to_do_index_compaction.empty()); + EXPECT_EQ(compaction._output_rowset->num_segments(), 2); + }; + RowsetSharedPtr output; + const Status status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, true, + output, check_fallback, 1000); + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(output, nullptr); + for (uint32_t segment_id = 0; segment_id < output->num_segments(); ++segment_id) { + const auto segment_path = output->segment_path(segment_id); + ASSERT_TRUE(segment_path.has_value()) << segment_path.error(); + auto file_reader = IndexCompactionUtils::init_index_file_reader( + output, segment_path.value(), InvertedIndexStorageFormatPB::SNII); + for (const TabletIndex* index : _tablet_schema->inverted_indexes()) { + const auto logical_index = file_reader->open_snii_index(index); + EXPECT_TRUE(logical_index.has_value()) << logical_index.error(); + } + } +} + +TEST_F(IndexCompactionTest, snii_native_merge_aborts_after_partial_destination_creation) { + const bool old_common_grams = config::enable_common_grams_index_build; + const bool old_debug_points = config::enable_debug_points; + const bool old_write_freq = config::snii_positions_index_write_freq; + config::enable_common_grams_index_build = false; + config::enable_debug_points = true; + config::snii_positions_index_write_freq = false; + constexpr std::string_view kFailurePoint = "Compaction::before_add_snii_destination_session"; + constexpr std::string_view kAbortPoint = "Compaction::snii_destination_session_aborted"; + DEFER({ + DebugPoints::instance()->remove(std::string(kFailurePoint)); + DebugPoints::instance()->remove(std::string(kAbortPoint)); + config::enable_common_grams_index_build = old_common_grams; + config::enable_debug_points = old_debug_points; + config::snii_positions_index_write_freq = old_write_freq; + }); + + _build_snii_multi_index_tablet(); + const std::vector rowsets = _build_snii_source_rowsets(); + std::function fail_second_destination = [](size_t destination_ordinal, + Status* status) { + if (destination_ordinal == 1) { + *status = Status::Error( + "injected SNII destination session failure"); + } + }; + DebugPoints::instance()->add_with_callback(std::string(kFailurePoint), fail_second_destination); + size_t abort_count = 0; + std::function count_abort = [&abort_count](size_t) { ++abort_count; }; + DebugPoints::instance()->add_with_callback(std::string(kAbortPoint), count_abort); + + RowsetSharedPtr output; + const Status status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, true, + output, nullptr, 1000); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_THAT(status.to_string(), + testing::HasSubstr("injected SNII destination session failure")); + EXPECT_EQ(output, nullptr); + EXPECT_EQ(abort_count, 1); + for (const RowsetSharedPtr& rowset : rowsets) { + EXPECT_FALSE(rowset->is_skip_index_compaction(1)); + } +} + +TEST_F(IndexCompactionTest, snii_native_merge_mem_limit_arms_raw_rebuild_without_wrapping) { + const bool old_common_grams = config::enable_common_grams_index_build; + const bool old_debug_points = config::enable_debug_points; + const bool old_write_freq = config::snii_positions_index_write_freq; + config::enable_common_grams_index_build = false; + config::enable_debug_points = true; + config::snii_positions_index_write_freq = false; + constexpr std::string_view kFailurePoint = "Compaction::before_execute_snii_merge"; + DEFER({ + DebugPoints::instance()->remove(std::string(kFailurePoint)); + config::enable_common_grams_index_build = old_common_grams; + config::enable_debug_points = old_debug_points; + config::snii_positions_index_write_freq = old_write_freq; + }); + + _build_snii_multi_index_tablet(); + const std::vector rowsets = _build_snii_source_rowsets(); + std::function inject_mem_limit = [](Status* status) { + *status = Status::Error("injected SNII merge memory limit"); + }; + DebugPoints::instance()->add_with_callback(std::string(kFailurePoint), inject_mem_limit); + + RowsetSharedPtr output; + const Status status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, true, + output, nullptr, 1000); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_THAT(status.to_string(), testing::HasSubstr("injected SNII merge memory limit")); + EXPECT_EQ(output, nullptr); + for (const RowsetSharedPtr& rowset : rowsets) { + EXPECT_TRUE(rowset->is_skip_index_compaction(1)); + } +} + +TEST_F(IndexCompactionTest, snii_native_merge_corruption_arms_raw_rebuild_fallback) { + const bool old_common_grams = config::enable_common_grams_index_build; + const bool old_debug_points = config::enable_debug_points; + const bool old_write_freq = config::snii_positions_index_write_freq; + config::enable_common_grams_index_build = false; + config::enable_debug_points = true; + config::snii_positions_index_write_freq = false; + constexpr std::string_view kFailurePoint = "Compaction::before_execute_snii_merge"; + DEFER({ + DebugPoints::instance()->remove(std::string(kFailurePoint)); + config::enable_common_grams_index_build = old_common_grams; + config::enable_debug_points = old_debug_points; + config::snii_positions_index_write_freq = old_write_freq; + }); + + _build_snii_multi_index_tablet(); + const std::vector rowsets = _build_snii_source_rowsets(); + std::function inject_corruption = [](Status* status) { + *status = Status::Error( + "injected SNII merge corruption"); + }; + DebugPoints::instance()->add_with_callback(std::string(kFailurePoint), inject_corruption); + + RowsetSharedPtr output; + const Status status = IndexCompactionUtils::do_compaction(rowsets, _engine_ref, _tablet, true, + output, nullptr, 1000); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()) << status; + EXPECT_THAT(status.to_string(), testing::HasSubstr("injected SNII merge corruption")); + EXPECT_EQ(output, nullptr); + for (const RowsetSharedPtr& rowset : rowsets) { + EXPECT_TRUE(rowset->is_skip_index_compaction(1)); + } +} } // namespace doris diff --git a/be/test/storage/index/inverted/compaction/index_compaction_write_index_test.cpp b/be/test/storage/index/inverted/compaction/index_compaction_write_index_test.cpp index a530be96da3775..48803988d876f6 100644 --- a/be/test/storage/index/inverted/compaction/index_compaction_write_index_test.cpp +++ b/be/test/storage/index/inverted/compaction/index_compaction_write_index_test.cpp @@ -28,7 +28,6 @@ #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wshadow-field" #include // IWYU pragma: keep -#include #include #include "CLucene/analysis/Analyzers.h" diff --git a/be/test/storage/index/inverted/compaction/util/index_compaction_utils.cpp b/be/test/storage/index/inverted/compaction/util/index_compaction_utils.cpp index 6ca726770cfb5d..7feb99bf16d2ec 100644 --- a/be/test/storage/index/inverted/compaction/util/index_compaction_utils.cpp +++ b/be/test/storage/index/inverted/compaction/util/index_compaction_utils.cpp @@ -427,7 +427,8 @@ class IndexCompactionUtils { const TabletSharedPtr& tablet, bool is_index_compaction, RowsetSharedPtr& rowset_ptr, const std::function custom_check = nullptr, - int64_t max_rows_per_segment = 100000) { + int64_t max_rows_per_segment = 100000, + const std::optional& output_storage_resource = std::nullopt) { config::inverted_index_compaction_enable = is_index_compaction; // control max rows in one block config::compaction_batch_size = max_rows_per_segment; @@ -441,6 +442,8 @@ class IndexCompactionUtils { RowsetWriterContext ctx; ctx.max_rows_per_segment = max_rows_per_segment; + // Empty for a local output rowset, which is what is_local_rowset() keys off. + ctx.storage_resource = output_storage_resource; RETURN_IF_ERROR(compaction.construct_output_rowset_writer(ctx)); compaction._stats.rowid_conversion = compaction._rowid_conversion.get(); @@ -585,7 +588,8 @@ class IndexCompactionUtils { InvertedIndexDescriptor::get_index_file_path_prefix(seg_path.value()); auto index_file_reader = std::make_shared( fs, std::string(index_file_path_prefix), - tablet_schema->get_inverted_index_storage_format(), index_info); + tablet_schema->get_inverted_index_storage_format(), index_info, + output_rowset->rowset_meta()->tablet_id()); EXPECT_TRUE(index_file_reader->init().ok()); const auto& dirs = index_file_reader->get_all_directories(); EXPECT_TRUE(dirs.has_value()); @@ -616,11 +620,11 @@ class IndexCompactionUtils { } } - static RowsetWriterContext rowset_writer_context(const std::unique_ptr& data_dir, - const TabletSchemaSPtr& schema, - const std::string& tablet_path, - int64_t& inc_id, - int64_t max_rows_per_segment = 200) { + static RowsetWriterContext rowset_writer_context( + const std::unique_ptr& data_dir, const TabletSchemaSPtr& schema, + const std::string& tablet_path, int64_t& inc_id, int64_t max_rows_per_segment = 200, + const std::optional& storage_resource = std::nullopt, + int64_t tablet_id = 0, bool write_file_cache = false) { RowsetWriterContext context; RowsetId rowset_id; rowset_id.init(inc_id); @@ -630,8 +634,16 @@ class IndexCompactionUtils { context.rowset_state = VISIBLE; context.tablet_schema = schema; context.tablet_path = tablet_path; + // Remote rowsets need this: CachedRemoteFileReader asserts tablet_id > 0 for Doris tables. + context.tablet_id = tablet_id; context.version = Version(inc_id, inc_id); context.max_rows_per_segment = max_rows_per_segment; + // Set only by callers that want the rowset to live on remote storage; a local rowset + // leaves it empty, which is what is_local_rowset() keys off. + context.storage_resource = storage_resource; + // Cloud sets this from the load request (cloud_rowset_builder.cpp), which is what leaves + // the block cache warm for compaction. Off by default, matching non-cloud writes. + context.write_file_cache = write_file_cache; inc_id++; return context; } @@ -643,7 +655,9 @@ class IndexCompactionUtils { const std::vector& data_files, int64_t& inc_id, const std::function custom_check = nullptr, const bool& is_performance = false, - int64_t max_rows_per_segment = 200) { + int64_t max_rows_per_segment = 200, + const std::optional& storage_resource = std::nullopt, + bool write_file_cache = false) { std::vector> data; for (const auto& file : data_files) { data.emplace_back(read_data(file)); @@ -652,7 +666,8 @@ class IndexCompactionUtils { const auto& res = RowsetFactory::create_rowset_writer( *engine_ref, rowset_writer_context(data_dir, schema, tablet->tablet_path(), inc_id, - max_rows_per_segment), + max_rows_per_segment, storage_resource, + tablet->tablet_id(), write_file_cache), false); EXPECT_TRUE(res.has_value()) << res.error(); const auto& rowset_writer = res.value(); @@ -715,8 +730,10 @@ class IndexCompactionUtils { (rowsets[i]->num_rows() / max_rows_per_segment)) << rowsets[i]->num_segments(); - // check rowset meta and file - for (int seg_id = 0; seg_id < rowsets[i]->num_segments(); seg_id++) { + // check rowset meta and file -- local only: the paths below are local, so a remote + // rowset would always report size 0 here. + for (int seg_id = 0; rowsets[i]->is_local() && seg_id < rowsets[i]->num_segments(); + seg_id++) { const auto& index_info = rowsets[i]->_rowset_meta->inverted_index_file_info(seg_id); EXPECT_TRUE(index_info.has_index_size()); const auto& fs = rowsets[i]->_rowset_meta->fs(); @@ -733,13 +750,26 @@ class IndexCompactionUtils { InvertedIndexDescriptor::get_index_file_path_prefix(seg_path.value()); auto index_file_reader = std::make_shared( fs, std::string(index_file_path_prefix), - schema->get_inverted_index_storage_format(), index_info); + schema->get_inverted_index_storage_format(), index_info, + rowsets[i]->rowset_meta()->tablet_id()); st = index_file_reader->init(); EXPECT_TRUE(st.ok()) << st.to_string(); - const auto& dirs = index_file_reader->get_all_directories(); - EXPECT_TRUE(dirs.has_value()); - if (custom_check) { - custom_check(dirs.value().size()); + if (schema->get_inverted_index_storage_format() == + InvertedIndexStorageFormatPB::SNII) { + const auto& indexes = schema->inverted_indexes(); + for (const TabletIndex* index : indexes) { + const auto logical_index = index_file_reader->open_snii_index(index); + EXPECT_TRUE(logical_index.has_value()) << logical_index.error(); + } + if (custom_check) { + custom_check(indexes.size()); + } + } else { + const auto& dirs = index_file_reader->get_all_directories(); + EXPECT_TRUE(dirs.has_value()); + if (custom_check) { + custom_check(dirs.value().size()); + } } } } diff --git a/be/test/storage/index/inverted/empty_index_file_test.cpp b/be/test/storage/index/inverted/empty_index_file_test.cpp index 94f1a6493209bb..cfe67dd5d9411d 100644 --- a/be/test/storage/index/inverted/empty_index_file_test.cpp +++ b/be/test/storage/index/inverted/empty_index_file_test.cpp @@ -31,11 +31,18 @@ constexpr int64_t LOAD_ID_HI = 2; constexpr int64_t NUM_STREAM = 3; constexpr static std::string_view tmp_dir = "./ut_dir/tmp"; class EmptyIndexFileTest : public testing::Test { + struct WriteState { + size_t bytes_appended = 0; + int data_calls = 0; + int eos_calls = 0; + }; + class MockStreamStub : public LoadStreamStub { public: - MockStreamStub(PUniqueId load_id, int64_t src_id) + MockStreamStub(PUniqueId load_id, int64_t src_id, std::shared_ptr state) : LoadStreamStub(load_id, src_id, std::make_shared(), - std::make_shared()) {}; + std::make_shared()), + _state(std::move(state)) {}; virtual ~MockStreamStub() = default; @@ -44,9 +51,21 @@ class EmptyIndexFileTest : public testing::Test { int32_t segment_id, uint64_t offset, std::span data, bool segment_eos = false, FileType file_type = FileType::SEGMENT_FILE) override { - EXPECT_TRUE(segment_eos); + EXPECT_EQ(offset, _state->bytes_appended); + if (segment_eos) { + ++_state->eos_calls; + EXPECT_TRUE(data.empty()); + return Status::OK(); + } + ++_state->data_calls; + for (const auto& slice : data) { + _state->bytes_appended += slice.size; + } return Status::OK(); } + + private: + std::shared_ptr _state; }; public: @@ -58,7 +77,9 @@ class EmptyIndexFileTest : public testing::Test { _load_id.set_hi(LOAD_ID_HI); _load_id.set_lo(LOAD_ID_LO); for (int src_id = 0; src_id < NUM_STREAM; src_id++) { - _streams.emplace_back(new MockStreamStub(_load_id, src_id)); + auto state = std::make_shared(); + _write_states.push_back(state); + _streams.emplace_back(new MockStreamStub(_load_id, src_id, std::move(state))); } EXPECT_TRUE(io::global_local_filesystem()->delete_directory(tmp_dir).ok()); EXPECT_TRUE(io::global_local_filesystem()->create_directory(tmp_dir).ok()); @@ -75,9 +96,10 @@ class EmptyIndexFileTest : public testing::Test { PUniqueId _load_id; std::vector> _streams; + std::vector> _write_states; }; -TEST_F(EmptyIndexFileTest, test_empty_index_file) { +TEST_F(EmptyIndexFileTest, WritesValidV2HeaderForNoLogicalIndexes) { io::FileWriterPtr file_writer = std::make_unique(_streams); auto fs = io::global_local_filesystem(); std::string index_path = "/tmp/empty_index_file_test"; @@ -88,6 +110,11 @@ TEST_F(EmptyIndexFileTest, test_empty_index_file) { std::move(file_writer), false); EXPECT_TRUE(index_file_writer->begin_close().ok()); EXPECT_TRUE(index_file_writer->finish_close().ok()); + for (const auto& state : _write_states) { + EXPECT_EQ(state->bytes_appended, sizeof(int32_t) * 2); + EXPECT_GT(state->data_calls, 0); + EXPECT_EQ(state->eos_calls, 1); + } } } // namespace doris diff --git a/be/test/storage/index/inverted/inverted_index_reader_analysis_purpose_test.cpp b/be/test/storage/index/inverted/inverted_index_reader_analysis_purpose_test.cpp new file mode 100644 index 00000000000000..b21d2c8790a374 --- /dev/null +++ b/be/test/storage/index/inverted/inverted_index_reader_analysis_purpose_test.cpp @@ -0,0 +1,722 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include + +#include "common/exception.h" +#include "io/fs/local_file_system.h" +#include "runtime/exec_env.h" +#include "runtime/runtime_state.h" +#include "storage/compaction/collection_similarity.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/analyzer/analyzer_provider.h" +#include "storage/index/inverted/analyzer/custom_analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/inverted_index_cache.h" +#include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/inverted/token_filter/common_grams_filter.h" +#include "storage/index/snii/snii_index_reader.h" +#include "storage/tablet/tablet_schema.h" +#include "util/defer_op.h" +#include "util/time.h" + +namespace doris::segment_v2 { +namespace { + +using inverted_index::AnalysisPurpose; +using inverted_index::AnalyzerProvider; + +class RecordingFailingAnalyzerProvider final : public AnalyzerProvider { +public: + std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const override { + purposes.push_back(purpose); + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "forced analyzer provider failure"); + } + + mutable std::vector purposes; +}; + +class IdentityFailingAnalyzerProvider final : public AnalyzerProvider { +public: + explicit IdentityFailingAnalyzerProvider(inverted_index::CommonGramsQueryIdentity identity) + : _identity(std::move(identity)) {} + + std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const override { + purposes.push_back(purpose); + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "forced analyzer provider failure"); + } + + bool uses_common_grams() const override { return true; } + + const inverted_index::CommonGramsQueryIdentity* common_grams_identity() const override { + return &_identity; + } + + mutable std::vector purposes; + +private: + inverted_index::CommonGramsQueryIdentity _identity; +}; + +class PartialFailureTokenStream final : public lucene::analysis::TokenStream { +public: + explicit PartialFailureTokenStream(std::shared_ptr> emitted_tokens) + : _emitted_tokens(std::move(emitted_tokens)) {} + + lucene::analysis::Token* next(lucene::analysis::Token* token) override { + if (!_emitted) { + _emitted = true; + _term = "partial"; + token->clear(); + token->setTextNoCopy(_term.data(), static_cast(_term.size())); + token->positionIncrement = 1; + _emitted_tokens->fetch_add(1, std::memory_order_relaxed); + return token; + } + throw Exception(ErrorCode::INVERTED_INDEX_ANALYZER_ERROR, + "forced failure after first token"); + } + + void close() override {} + void reset() override { _emitted = false; } + +private: + std::shared_ptr> _emitted_tokens; + bool _emitted = false; + std::string _term; +}; + +class PartialFailureAnalyzer final : public lucene::analysis::Analyzer { +public: + explicit PartialFailureAnalyzer(std::shared_ptr> emitted_tokens) + : _emitted_tokens(std::move(emitted_tokens)) {} + + bool isSDocOpt() override { return true; } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, lucene::util::Reader*) override { + return new PartialFailureTokenStream(_emitted_tokens); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + lucene::util::Reader*) override { + _reusable = std::make_unique(_emitted_tokens); + return _reusable.get(); + } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + return new PartialFailureTokenStream(_emitted_tokens); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + _reusable = std::make_unique(_emitted_tokens); + return _reusable.get(); + } + +private: + std::shared_ptr> _emitted_tokens; + std::unique_ptr _reusable; +}; + +class RecordingPartialFailureAnalyzerProvider final : public AnalyzerProvider { +public: + std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const override { + purposes.push_back(purpose); + return std::make_shared(emitted_tokens); + } + + mutable std::vector purposes; + std::shared_ptr> emitted_tokens = + std::make_shared>(0); +}; + +class GeneratedTokenStream final : public lucene::analysis::TokenStream { +public: + GeneratedTokenStream(std::string term, bool mark_common_gram) + : _term(std::move(term)), _mark_common_gram(mark_common_gram) {} + + lucene::analysis::Token* next(lucene::analysis::Token* token) override { + if (_emitted) { + return nullptr; + } + _emitted = true; + token->clear(); + token->setTextNoCopy(_term.data(), static_cast(_term.size())); + token->setPositionIncrement(1); + if (_mark_common_gram) { + token->setType(inverted_index::COMMON_GRAM_TOKEN_TYPE); + } + return token; + } + + void close() override {} + void reset() override { _emitted = false; } + +private: + std::string _term; + bool _mark_common_gram = false; + bool _emitted = false; +}; + +class GeneratedTokenAnalyzer final : public lucene::analysis::Analyzer { +public: + GeneratedTokenAnalyzer(std::string term, bool mark_common_gram) + : _term(std::move(term)), _mark_common_gram(mark_common_gram) {} + + bool isSDocOpt() override { return true; } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, lucene::util::Reader*) override { + return new GeneratedTokenStream(_term, _mark_common_gram); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + lucene::util::Reader*) override { + _reusable = std::make_unique(_term, _mark_common_gram); + return _reusable.get(); + } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + return new GeneratedTokenStream(_term, _mark_common_gram); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + _reusable = std::make_unique(_term, _mark_common_gram); + return _reusable.get(); + } + +private: + std::string _term; + bool _mark_common_gram = false; + std::unique_ptr _reusable; +}; + +class GeneratedGramAnalyzerProvider final : public AnalyzerProvider { +public: + GeneratedGramAnalyzerProvider() + : _gram(*inverted_index::encode_common_gram("the", "history")), + _identity {.common_grams_dictionary_identity = "builtin-stopwords:v1", + .base_analyzer_fingerprint = "base:v1", + .common_grams_fingerprint = "grams:v1"} {} + + std::shared_ptr get_analyzer( + AnalysisPurpose purpose) const override { + purposes.push_back(purpose); + return std::make_shared(_gram, true); + } + + bool uses_common_grams() const override { return true; } + const inverted_index::CommonGramsQueryIdentity* common_grams_identity() const override { + return &_identity; + } + + mutable std::vector purposes; + +private: + std::string _gram; + inverted_index::CommonGramsQueryIdentity _identity; +}; + +struct QueryExecutionContext { + explicit QueryExecutionContext(bool scoring) { + TQueryOptions query_options; + query_options.enable_inverted_index_query_cache = true; + query_options.enable_inverted_index_searcher_cache = true; + runtime_state.set_query_options(query_options); + context->io_ctx = &io_ctx; + context->stats = &stats; + context->runtime_state = &runtime_state; + if (scoring) { + context->collection_similarity = std::make_shared(); + } + } + + OlapReaderStatistics stats; + io::IOContext io_ctx; + RuntimeState runtime_state; + IndexQueryContextPtr context = std::make_shared(); +}; + +class InvertedIndexReaderAnalysisPurposeTest : public testing::Test { +protected: + void SetUp() override { + _previous_searcher_cache = ExecEnv::GetInstance()->get_inverted_index_searcher_cache(); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); + _searcher_cache.reset(InvertedIndexSearcherCache::create_global_instance(1024 * 1024, 1)); + _query_cache.reset(InvertedIndexQueryCache::create_global_cache(1024 * 1024, 1)); + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_searcher_cache.get()); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_query_cache.get()); + + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(73); + pb.set_index_name("analysis_purpose_idx"); + pb.add_col_unique_id(0); + pb.mutable_properties()->insert({"parser", "english"}); + pb.mutable_properties()->insert({"lower_case", "true"}); + pb.mutable_properties()->insert({"support_phrase", "true"}); + _meta.init_from_pb(pb); + + _snii_file_reader = std::make_shared( + io::global_local_filesystem(), "./ut_dir/missing_snii_analysis_purpose", + InvertedIndexStorageFormatPB::SNII); + _snii_reader = SniiIndexReader::create_shared(&_meta, _snii_file_reader, + InvertedIndexReaderType::FULLTEXT); + } + + void TearDown() override { + _snii_reader.reset(); + _snii_file_reader.reset(); + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_previous_searcher_cache); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); + _searcher_cache.reset(); + _query_cache.reset(); + } + + // This entry is admission-only: tests must exit during analysis before dereferencing the + // SNII postings reader. + void preload_legacy_searcher_cache_entries() { + const InvertedIndexSearcherCache::CacheKey snii_key( + _snii_file_reader->get_index_file_cache_key(&_meta)); + _searcher_cache->insert( + snii_key, new InvertedIndexSearcherCache::CacheValue( + std::make_unique(), 1, + UnixMillis(), _snii_file_reader)); + } + + template + void expect_analysis_failure_after_segment_admission(const std::shared_ptr& reader, + InvertedIndexQueryType query_type, + std::string query, bool scoring, + AnalysisPurpose expected_purpose, + const std::shared_ptr& provider, + int64_t expected_query_cache_lookups = 0, + int64_t expected_searcher_cache_hits = 1) { + QueryExecutionContext execution(scoring); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + + auto original_bitmap = std::make_shared(); + original_bitmap->add(999); + std::shared_ptr bitmap = original_bitmap; + const Field query_value = Field::create_field(std::move(query)); + + Status status; + EXPECT_NO_THROW(status = reader->query(execution.context, "content", query_value, + query_type, bitmap, &analyzer_ctx)); + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR) << status; + EXPECT_EQ(provider->purposes, (std::vector {expected_purpose})); + EXPECT_EQ(bitmap, original_bitmap); + EXPECT_EQ(bitmap->cardinality(), 1); + EXPECT_TRUE(bitmap->contains(999)); + EXPECT_EQ(execution.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, expected_query_cache_lookups); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, expected_query_cache_lookups); + EXPECT_EQ(execution.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_hit, expected_searcher_cache_hits); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_miss, 0); + } + + template + void expect_provider_failure_after_segment_admission(const std::shared_ptr& reader, + InvertedIndexQueryType query_type, + std::string query, bool scoring, + AnalysisPurpose expected_purpose, + int64_t expected_query_cache_lookups = 0, + int64_t expected_searcher_cache_hits = 1) { + expect_analysis_failure_after_segment_admission( + reader, query_type, std::move(query), scoring, expected_purpose, + std::make_shared(), expected_query_cache_lookups, + expected_searcher_cache_hits); + } + + template + void expect_generated_gram_bypass_after_segment_admission( + const std::shared_ptr& reader) { + QueryExecutionContext execution(/*scoring=*/false); + auto provider = std::make_shared(); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + auto original_bitmap = std::make_shared(); + original_bitmap->add(999); + std::shared_ptr bitmap = original_bitmap; + const Field query_value = Field::create_field("the history"); + + const Status status = + reader->query(execution.context, "content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap, &analyzer_ctx); + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_BYPASS) << status; + EXPECT_EQ(provider->purposes, + (std::vector {AnalysisPurpose::kPlainQuery})); + EXPECT_EQ(bitmap, original_bitmap); + EXPECT_EQ(execution.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_hit, 1); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_miss, 0); + } + + void expect_snii_sloppy_phrase_bypass_before_segment_admission() { + QueryExecutionContext execution(/*scoring=*/false); + auto provider = std::make_shared(); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + auto original_bitmap = std::make_shared(); + original_bitmap->add(999); + std::shared_ptr bitmap = original_bitmap; + const Field query_value = Field::create_field("the history ~2"); + + const Status status = _snii_reader->query(execution.context, "content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + bitmap, &analyzer_ctx); + EXPECT_EQ(status.code(), ErrorCode::INVERTED_INDEX_BYPASS) << status; + EXPECT_TRUE(provider->purposes.empty()); + EXPECT_EQ(bitmap, original_bitmap); + EXPECT_EQ(bitmap->cardinality(), 1); + EXPECT_TRUE(bitmap->contains(999)); + EXPECT_EQ(execution.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_miss, 0); + } + + template + void expect_raw_query_bypasses_analyzer(const std::shared_ptr& reader, + InvertedIndexQueryType query_type, std::string query) { + QueryExecutionContext execution(/*scoring=*/false); + auto provider = std::make_shared(); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + std::shared_ptr bitmap; + const Field query_value = Field::create_field(std::move(query)); + + const Status status = reader->query(execution.context, "content", query_value, query_type, + bitmap, &analyzer_ctx); + EXPECT_NE(status.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR) << status; + EXPECT_TRUE(provider->purposes.empty()); + } + + template + void expect_raw_cache_hit_before_analysis(const std::shared_ptr& reader, + const std::shared_ptr& file_reader, + const std::shared_ptr& provider, + uint64_t common_grams_cache_generation) { + QueryExecutionContext execution(/*scoring=*/false); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + + const std::string raw_query = "the history"; + const InvertedIndexRawQuerySemantic semantic { + .raw_query_bytes = raw_query, + .query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY, + .slop = 0, + .ordered = false, + .max_expansions = + execution.runtime_state.query_options().inverted_index_max_expansions, + .common_grams_cache_generation = common_grams_cache_generation}; + const InvertedIndexQueryCache::CacheKey key { + file_reader->get_index_file_cache_key(&_meta), "content", + InvertedIndexQueryType::MATCH_PHRASE_QUERY, semantic.encode()}; + auto cached = std::make_shared(); + cached->add(7); + InvertedIndexQueryCacheHandle insert_handle; + _query_cache->insert(key, cached, &insert_handle); + + std::shared_ptr bitmap; + const Field query_value = Field::create_field(raw_query); + const Status status = + reader->query(execution.context, "content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap, &analyzer_ctx); + EXPECT_TRUE(status.ok()) << status; + EXPECT_TRUE(provider->purposes.empty()); + ASSERT_NE(bitmap, nullptr); + EXPECT_EQ(bitmap->cardinality(), 1); + EXPECT_TRUE(bitmap->contains(7)); + EXPECT_EQ(execution.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_searcher_cache_miss, 0); + } + + InvertedIndexSearcherCache* _previous_searcher_cache = nullptr; + InvertedIndexQueryCache* _previous_query_cache = nullptr; + std::unique_ptr _searcher_cache; + std::unique_ptr _query_cache; + TabletIndex _meta; + std::shared_ptr _snii_file_reader; + std::shared_ptr _snii_reader; +}; + +TEST(InvertedIndexRawQuerySemanticTest, EncodesOnlyRawSemanticDimensionsWithoutDelimiters) { + const std::string raw_query("a/b\0c", 5); + InvertedIndexRawQuerySemantic base {.raw_query_bytes = raw_query, + .query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY, + .slop = 2, + .ordered = true, + .max_expansions = 50, + .cache_semantics_version = 3, + .common_grams_cache_generation = 7}; + const std::string encoded = base.encode(); + constexpr size_t kFixedEncodedBytes = sizeof(uint32_t) + sizeof(uint64_t) + sizeof(uint32_t) + + sizeof(uint32_t) + sizeof(uint8_t) + sizeof(uint32_t) + + sizeof(uint64_t); + EXPECT_EQ(encoded.size(), kFixedEncodedBytes + raw_query.size()); + + auto changed = base; + changed.raw_query_bytes = std::string_view(raw_query).substr(0, 3); + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.query_type = InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY; + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.slop = 3; + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.ordered = false; + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.max_expansions = 51; + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.cache_semantics_version = 4; + EXPECT_NE(changed.encode(), encoded); + changed = base; + changed.common_grams_cache_generation = 8; + EXPECT_NE(changed.encode(), encoded); +} + +TEST(InvertedIndexRawQuerySemanticTest, CacheEnvelopeSeparatesSlashAndNulBoundaries) { + const InvertedIndexQueryCache::CacheKey slash_left { + io::Path("a/b"), "c", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "d"}; + const InvertedIndexQueryCache::CacheKey slash_right { + io::Path("a"), "b/c", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "d"}; + EXPECT_NE(slash_left.encode(), slash_right.encode()); + + const InvertedIndexQueryCache::CacheKey nul_left { + io::Path("a"), std::string("b\0c", 3), InvertedIndexQueryType::MATCH_PHRASE_QUERY, "d"}; + const InvertedIndexQueryCache::CacheKey nul_right { + io::Path(std::string("a\0b", 3)), "c", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "d"}; + EXPECT_NE(nul_left.encode(), nul_right.encode()); +} + +TEST(InvertedIndexRawQuerySemanticTest, KillSwitchTransitionsAdvanceCacheGeneration) { + const bool original = config::enable_common_grams_query_plan; + const auto before = config::common_grams_query_plan_config_snapshot(); + + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", original ? "false" : "true", + /*need_persist=*/false) + .ok()); + const auto disabled_or_enabled = config::common_grams_query_plan_config_snapshot(); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", original ? "true" : "false", + /*need_persist=*/false) + .ok()); + const auto restored = config::common_grams_query_plan_config_snapshot(); + + EXPECT_EQ(before.enabled, original); + EXPECT_EQ(disabled_or_enabled.enabled, !original); + EXPECT_GT(disabled_or_enabled.cache_generation, before.cache_generation); + EXPECT_EQ(restored.enabled, original); + EXPECT_GT(restored.cache_generation, disabled_or_enabled.cache_generation); +} + +TEST(InvertedIndexRawQuerySemanticTest, CommonGramsQueryPlanIsDisabledByDefault) { + EXPECT_FALSE(config::enable_common_grams_query_plan); + EXPECT_FALSE(config::common_grams_query_plan_config_snapshot().enabled); +} + +TEST(InvertedIndexRawQuerySemanticTest, CostModelSnapshotTracksDynamicConfigAsOneVersion) { + const int32_t original_ratio = config::common_grams_plan_cost_ratio_percent; + const int32_t original_factor = config::common_grams_position_verify_factor; + const int32_t changed_ratio = original_ratio == 84 ? 85 : 84; + const int32_t changed_factor = original_factor == 7 ? 8 : 7; + + ASSERT_TRUE(config::set_config("common_grams_plan_cost_ratio_percent", + std::to_string(changed_ratio)) + .ok()); + ASSERT_TRUE(config::set_config("common_grams_position_verify_factor", + std::to_string(changed_factor)) + .ok()); + const auto changed = config::common_grams_query_plan_config_snapshot(); + EXPECT_EQ(changed.plan_cost_ratio_percent, changed_ratio); + EXPECT_EQ(changed.position_verify_factor, changed_factor); + EXPECT_EQ(changed.cost_model_generation, + (static_cast(static_cast(changed_factor)) << 32) | + static_cast(changed_ratio)); + + ASSERT_TRUE(config::set_config("common_grams_plan_cost_ratio_percent", + std::to_string(original_ratio)) + .ok()); + ASSERT_TRUE(config::set_config("common_grams_position_verify_factor", + std::to_string(original_factor)) + .ok()); +} + +TEST(InvertedIndexRawQuerySemanticTest, InvalidCostModelConfigDoesNotMutatePublishedState) { + const auto before = config::common_grams_query_plan_config_snapshot(); + + for (const auto& [field, value] : std::vector> { + {"common_grams_plan_cost_ratio_percent", "-1"}, + {"common_grams_plan_cost_ratio_percent", "101"}, + {"common_grams_position_verify_factor", "-1"}}) { + EXPECT_FALSE(config::set_config(field, value).ok()); + const auto after = config::common_grams_query_plan_config_snapshot(); + EXPECT_EQ(after.plan_cost_ratio_percent, before.plan_cost_ratio_percent); + EXPECT_EQ(after.position_verify_factor, before.position_verify_factor); + EXPECT_EQ(after.cost_model_generation, before.cost_model_generation); + } +} + +TEST(InvertedIndexAnalyzerCtxTest, UsesProviderCommonGramsIdentityWithoutCopyingIt) { + auto provider = std::make_shared(); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.analyzer_provider = provider; + EXPECT_EQ(analyzer_ctx.get_common_grams_identity(), provider->common_grams_identity()); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, + SniiRawCacheLookupIsIndependentOfRequestAnalyzerIdentity) { + const bool original = config::enable_common_grams_query_plan; + Defer restore([original] { + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original ? "true" : "false", /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + const auto safety = config::common_grams_query_plan_config_snapshot(); + const inverted_index::CommonGramsQueryIdentity complete_identity { + .common_grams_dictionary_identity = "dictionary:complete", + .base_analyzer_fingerprint = "base:complete", + .common_grams_fingerprint = "grams:complete"}; + const inverted_index::CommonGramsQueryIdentity empty_identity; + for (const auto& identity : {complete_identity, empty_identity}) { + expect_raw_cache_hit_before_analysis( + _snii_reader, _snii_file_reader, + std::make_shared(identity), + safety.cache_generation); + } + expect_raw_cache_hit_before_analysis(_snii_reader, _snii_file_reader, + std::make_shared(), + safety.cache_generation); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, DisabledResultCacheDoesNotLookupCountOrInsert) { + QueryExecutionContext execution(/*scoring=*/false); + TQueryOptions disabled_options; + disabled_options.enable_inverted_index_query_cache = false; + disabled_options.enable_inverted_index_searcher_cache = true; + execution.runtime_state.set_query_options(disabled_options); + + const InvertedIndexQueryCache::CacheKey key {io::Path("disabled-cache"), "content", + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + "raw-semantic"}; + auto bitmap = std::make_shared(); + bitmap->add(3); + InvertedIndexQueryCacheHandle handle; + _snii_reader->insert_query_cache(execution.context, _query_cache.get(), key, bitmap, &handle); + std::shared_ptr lookup_bitmap; + EXPECT_FALSE(_snii_reader->handle_query_cache(execution.context, _query_cache.get(), key, + &handle, lookup_bitmap)); + EXPECT_EQ(execution.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, 0); + EXPECT_EQ(execution.stats.inverted_index_query_cache_insert, 0); + + TQueryOptions enabled_options; + enabled_options.enable_inverted_index_query_cache = true; + enabled_options.enable_inverted_index_searcher_cache = true; + execution.runtime_state.set_query_options(enabled_options); + EXPECT_FALSE(_snii_reader->handle_query_cache(execution.context, _query_cache.get(), key, + &handle, lookup_bitmap)); + EXPECT_EQ(execution.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(execution.stats.inverted_index_query_cache_lookup, 1); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, SniiSelectsPurposeAfterSegmentAdmission) { + preload_legacy_searcher_cache_entries(); + expect_provider_failure_after_segment_admission( + _snii_reader, InvertedIndexQueryType::MATCH_PHRASE_QUERY, "the history", false, + AnalysisPurpose::kPlainQuery, /*expected_query_cache_lookups=*/1); + expect_snii_sloppy_phrase_bypass_before_segment_admission(); + expect_provider_failure_after_segment_admission( + _snii_reader, InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "the hist", false, + AnalysisPurpose::kPlainQuery, /*expected_query_cache_lookups=*/1); + expect_provider_failure_after_segment_admission( + _snii_reader, InvertedIndexQueryType::MATCH_PHRASE_QUERY, "the history", true, + AnalysisPurpose::kPlainQuery); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, PartialAnalysisFailureDoesNotPublishState) { + preload_legacy_searcher_cache_entries(); + auto snii_provider = std::make_shared(); + expect_analysis_failure_after_segment_admission( + _snii_reader, InvertedIndexQueryType::MATCH_PHRASE_QUERY, "the history", false, + AnalysisPurpose::kPlainQuery, snii_provider, /*expected_query_cache_lookups=*/1); + EXPECT_EQ(snii_provider->emitted_tokens->load(std::memory_order_relaxed), 1); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, + SniiGeneratedGramBypassesAfterSegmentAdmissionBeforeSearch) { + preload_legacy_searcher_cache_entries(); + const bool original = config::enable_common_grams_query_plan; + Defer restore([original] { + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original ? "true" : "false", /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + expect_generated_gram_bypass_after_segment_admission(_snii_reader); +} + +TEST_F(InvertedIndexReaderAnalysisPurposeTest, RegexpAndWildcardBypassAnalyzer) { + for (const auto query_type : + {InvertedIndexQueryType::MATCH_REGEXP_QUERY, InvertedIndexQueryType::WILDCARD_QUERY}) { + expect_raw_query_bypasses_analyzer(_snii_reader, query_type, "hist.*"); + } +} + +} // namespace +} // namespace doris::segment_v2 diff --git a/be/test/storage/index/inverted/query/phrase_edge_query_test.cpp b/be/test/storage/index/inverted/query/phrase_edge_query_test.cpp index 2da72758470a27..975ae80051678f 100644 --- a/be/test/storage/index/inverted/query/phrase_edge_query_test.cpp +++ b/be/test/storage/index/inverted/query/phrase_edge_query_test.cpp @@ -62,12 +62,19 @@ class PhraseEdgeQueryTest : public testing::Test { _inverted_index_query_cache = std::unique_ptr( InvertedIndexQueryCache::create_global_cache(inverted_index_cache_limit, 1)); + // Both caches are owned by this fixture, so the previous globals must come back in + // TearDown -- otherwise ExecEnv keeps pointing at them after the fixture is destroyed and + // the next test that reaches InvertedIndexQueryCache::instance() reads freed memory. + _previous_searcher_cache = ExecEnv::GetInstance()->get_inverted_index_searcher_cache(); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); ExecEnv::GetInstance()->set_inverted_index_searcher_cache( _inverted_index_searcher_cache.get()); - ExecEnv::GetInstance()->_inverted_index_query_cache = _inverted_index_query_cache.get(); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_inverted_index_query_cache.get()); } void TearDown() override { + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_previous_searcher_cache); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); EXPECT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); } @@ -180,6 +187,8 @@ class PhraseEdgeQueryTest : public testing::Test { ~PhraseEdgeQueryTest() override = default; private: + InvertedIndexSearcherCache* _previous_searcher_cache = nullptr; + InvertedIndexQueryCache* _previous_query_cache = nullptr; std::unique_ptr _inverted_index_searcher_cache; std::unique_ptr _inverted_index_query_cache; }; diff --git a/be/test/storage/index/inverted/similarity/bm25_similarity_test.cpp b/be/test/storage/index/inverted/similarity/bm25_similarity_test.cpp index 24fdfbd2e62709..9ceba92421e87f 100644 --- a/be/test/storage/index/inverted/similarity/bm25_similarity_test.cpp +++ b/be/test/storage/index/inverted/similarity/bm25_similarity_test.cpp @@ -22,8 +22,8 @@ #include #include "common/be_mock_util.h" -#include "storage/compaction/collection_statistics.h" #include "storage/index/index_query_context.h" +#include "storage/index/inverted/similarity/collection_statistics.h" using namespace doris; using namespace doris::segment_v2; diff --git a/be/test/storage/index/inverted/token_filter/lower_case_filter_factory_test.cpp b/be/test/storage/index/inverted/token_filter/lower_case_filter_factory_test.cpp index cd8c72755e16ce..6c19d437c9da59 100644 --- a/be/test/storage/index/inverted/token_filter/lower_case_filter_factory_test.cpp +++ b/be/test/storage/index/inverted/token_filter/lower_case_filter_factory_test.cpp @@ -20,10 +20,17 @@ #include +#include + #include "storage/index/inverted/tokenizer/keyword/keyword_tokenizer_factory.h" namespace doris::segment_v2::inverted_index { +namespace lower_case_testing { +uint64_t unicode_path_count(); +void reset_unicode_path_count(); +} // namespace lower_case_testing + TokenStreamPtr create_lowercase_filter(const std::string& text, Settings settings = Settings()) { ReaderPtr reader = std::make_shared>(); reader->init(text.data(), text.size(), false); @@ -45,6 +52,32 @@ struct ExpectedToken { int pos_inc; }; +class ScriptedLowercaseInput final : public TokenStream { +public: + explicit ScriptedLowercaseInput(std::vector terms) : _terms(std::move(terms)) {} + + Token* next(Token* token) override { + if (_next == _terms.size()) { + return nullptr; + } + const auto& term = _terms[_next++]; + token->clear(); + token->setTextNoCopy(term.data(), static_cast(term.size())); + token->setPositionIncrement(3); + token->setStartOffset(7); + token->setEndOffset(19); + token->setType(_T("scripted")); + return token; + } + + void close() override {} + void reset() override { _next = 0; } + +private: + std::vector _terms; + size_t _next = 0; +}; + class LowerCaseFilterTest : public ::testing::Test { protected: void assert_filter_output(const std::string& text, const std::vector& expected) { @@ -75,10 +108,73 @@ TEST_F(LowerCaseFilterTest, HandlesMixedCase) { assert_filter_output("HeLLo WoRLd", {{"hello world", 1}}); } +TEST_F(LowerCaseFilterTest, ASCIIBypassesUnicodeConversion) { + lower_case_testing::reset_unicode_path_count(); + assert_filter_output("already lowercase", {{"already lowercase", 1}}); + assert_filter_output("ASCII UPPER", {{"ascii upper", 1}}); + EXPECT_EQ(lower_case_testing::unicode_path_count(), 0); + + assert_filter_output( + "\xC3\x9C" + "BER", + {{"\xC3\xBC" + "ber", + 1}}); + EXPECT_EQ(lower_case_testing::unicode_path_count(), 1); +} + +TEST_F(LowerCaseFilterTest, ASCIIPreservesMetadataAndEmbeddedNul) { + auto input = std::make_shared( + std::vector {std::string("A\0B", 3), "already lower"}); + LowerCaseFilterFactory factory; + factory.initialize({}); + auto filter = factory.create(input); + filter->reset(); + + Token token; + ASSERT_NE(filter->next(&token), nullptr); + EXPECT_EQ(std::string(token.termBuffer(), token.termLength()), + std::string("a\0b", 3)); + EXPECT_EQ(token.getPositionIncrement(), 3); + EXPECT_EQ(token.startOffset(), 7); + EXPECT_EQ(token.endOffset(), 19); + EXPECT_EQ(std::wstring(token.type()), L"scripted"); + + ASSERT_NE(filter->next(&token), nullptr); + EXPECT_EQ(std::string(token.termBuffer(), token.termLength()), "already lower"); + EXPECT_EQ(token.getPositionIncrement(), 3); + EXPECT_EQ(token.startOffset(), 7); + EXPECT_EQ(token.endOffset(), 19); + EXPECT_EQ(std::wstring(token.type()), L"scripted"); +} + TEST_F(LowerCaseFilterTest, ConvertsUnicodeCharacters) { assert_filter_output("ÜBER ΜΈΓΑ", {{"über μέγα", 1}}); } +TEST_F(LowerCaseFilterTest, RetriesUnicodeExpansionWithRequiredBufferSize) { + assert_filter_output("\xC4\xB0", {{"i\xCC\x87", 1}}); +} + +TEST_F(LowerCaseFilterTest, RejectsInvalidUtf8WithAnalyzerError) { + auto input = std::make_shared( + std::vector {"VALID", std::string(1, static_cast(0xFF))}); + LowerCaseFilterFactory factory; + factory.initialize({}); + auto filter = factory.create(input); + filter->reset(); + + Token token; + ASSERT_NE(filter->next(&token), nullptr); + EXPECT_EQ(std::string(token.termBuffer(), token.termLength()), "valid"); + try { + filter->next(&token); + FAIL() << "expected malformed UTF-8 to fail analysis"; + } catch (const Exception& error) { + EXPECT_EQ(error.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } +} + TEST_F(LowerCaseFilterTest, HandlesNumbersAndSymbols) { assert_filter_output("123!@# ABC", {{"123!@# abc", 1}}); } diff --git a/be/test/storage/index/inverted/tokenizer/char_group_tokenizer_factory_test.cpp b/be/test/storage/index/inverted/tokenizer/char_group_tokenizer_factory_test.cpp index 986074b052ce7e..a124c1793eb2ba 100644 --- a/be/test/storage/index/inverted/tokenizer/char_group_tokenizer_factory_test.cpp +++ b/be/test/storage/index/inverted/tokenizer/char_group_tokenizer_factory_test.cpp @@ -23,6 +23,11 @@ namespace doris::segment_v2::inverted_index { +namespace char_tokenizer_testing { +uint64_t non_ascii_decode_count(); +void reset_non_ascii_decode_count(); +} // namespace char_tokenizer_testing + class CharGroupTokenizerTest : public ::testing::Test { protected: std::vector tokenize(CharGroupTokenizerFactory& factory, const std::string& text) { @@ -65,6 +70,22 @@ TEST_F(CharGroupTokenizerTest, TokenizeOnSpace) { ASSERT_EQ(tokens, expected); } +TEST_F(CharGroupTokenizerTest, ASCIIUsesPrecomputedClassification) { + CharGroupTokenizerFactory factory; + Settings settings; + settings.set("tokenize_on_chars", "[whitespace], [punctuation]"); + factory.initialize(settings); + + char_tokenizer_testing::reset_non_ascii_decode_count(); + EXPECT_EQ(tokenize(factory, "Hello, ASCII world!"), + (std::vector {"Hello", "ASCII", "world"})); + EXPECT_EQ(char_tokenizer_testing::non_ascii_decode_count(), 0); + + EXPECT_EQ(tokenize(factory, "Hello \xE4\xB8\x96\xE7\x95\x8C"), + (std::vector {"Hello", "\xE4\xB8\x96\xE7\x95\x8C"})); + EXPECT_EQ(char_tokenizer_testing::non_ascii_decode_count(), 2); +} + TEST_F(CharGroupTokenizerTest, TokenizeOnLetter) { CharGroupTokenizerFactory factory; Settings settings; diff --git a/be/test/storage/index/inverted_index_parser_test.cpp b/be/test/storage/index/inverted_index_parser_test.cpp index 719f8fb170f21e..00174207fec417 100644 --- a/be/test/storage/index/inverted_index_parser_test.cpp +++ b/be/test/storage/index/inverted_index_parser_test.cpp @@ -418,6 +418,23 @@ TEST_F(InvertedIndexParserTest, BuildAnalyzerKeyFromProperties_CustomOverridesPa EXPECT_EQ(build_analyzer_key_from_properties(properties), "my_custom"); } +TEST_F(InvertedIndexParserTest, BuildAnalyzerKeyFromProperties_ParserOverridesNormalizer) { + std::map properties; + properties[INVERTED_INDEX_NORMALIZER_NAME_KEY] = "MY_NORMALIZER"; + properties[INVERTED_INDEX_PARSER_KEY] = "chinese"; + + EXPECT_EQ(build_analyzer_key_from_properties(properties), "chinese"); +} + +TEST_F(InvertedIndexParserTest, BuildAnalyzerKeyFromProperties_AnalyzerOverridesNormalizer) { + std::map properties; + properties[INVERTED_INDEX_ANALYZER_NAME_KEY] = "MY_ANALYZER"; + properties[INVERTED_INDEX_NORMALIZER_NAME_KEY] = "MY_NORMALIZER"; + properties[INVERTED_INDEX_PARSER_KEY] = "chinese"; + + EXPECT_EQ(build_analyzer_key_from_properties(properties), "my_analyzer"); +} + // ============================================================================ // AnalyzerConfigParser Tests // ============================================================================ diff --git a/be/test/storage/index/snii/bench/snii_vs_v3_benchmark_test.cpp b/be/test/storage/index/snii/bench/snii_vs_v3_benchmark_test.cpp new file mode 100644 index 00000000000000..96357dc5719d34 --- /dev/null +++ b/be/test/storage/index/snii/bench/snii_vs_v3_benchmark_test.cpp @@ -0,0 +1,1222 @@ +// 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. +// +// SNII vs V3 baseline over the same wikipedia corpus, covering the three phases that matter for the +// format: building the index while loading, compacting it, and querying it across MATCH_ANY, +// MATCH_ALL, MATCH_PHRASE and MATCH_PHRASE_PREFIX. +// +// Why this lives in a unit test rather than a cluster benchmark: a UT iterates in minutes instead +// of a deploy cycle, and it measures one process we control end to end. +// +// Why CPU time is the headline number: this machine is shared, so wall clock moves with whatever +// else is running. Process CPU time barely does. Wall time is still reported -- a large wall/CPU +// gap means the run was IO bound or descheduled and the comparison should be rerun -- but the +// SNII/V3 verdict is taken from CPU. +// +// Corpus is not committed (~41 MB). Point SNII_BENCH_CORPUS_DIR at a directory of wikipedia_*.json +// with {"title","content"} per line: +// +// SNII_BENCH_CORPUS_DIR=/path/to/corpus \ +// SNII_BENCH_QUERY_ITERATIONS=30 \ +// ./run-be-ut.sh --run --filter='*SniiVsV3Benchmark*' -j +// +// SNII_BENCH_QUERY_ITERATIONS defaults to 30. The benchmark reports nearest-rank +// p50/p99 query CPU and wall time from the sorted per-iteration samples. +// +// The fixture is DISABLED_ so CI never runs it; --gtest_also_run_disabled_tests or the filter +// above opts in. It is still compiled into doris_be_test on every build. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/data_type/data_type_string.h" +#include "io/cache/block_file_cache.h" +#include "io/cache/block_file_cache_factory.h" +#include "io/cache/fs_file_cache_storage.h" +#include "io/fs/local_file_system.h" +#include "io/fs/s3_file_system.h" +#include "io/io_common.h" +#include "runtime/exec_env.h" +#include "runtime/memory/cache_manager.h" +#include "runtime/runtime_state.h" +#include "storage/index/index_iterator.h" +#include "storage/index/index_query_context.h" +#include "storage/index/inverted/compaction/util/index_compaction_utils.cpp" +#include "storage/index/inverted/inverted_index_cache.h" +#include "storage/index/inverted/inverted_index_desc.h" +#include "storage/index/inverted/inverted_index_iterator.h" +#include "storage/index/inverted/inverted_index_parser.h" +#include "storage/olap_common.h" +#include "storage/rowset/rowset_factory.h" +#include "storage/segment/segment.h" +#include "storage/storage_engine.h" +#include "storage/storage_policy.h" +#include "storage/tablet/tablet.h" +#include "util/defer_op.h" +#include "util/threadpool.h" +#include "util/time.h" + +namespace doris { + +constexpr static uint32_t MAX_PATH_LEN = 1024; +constexpr static std::string_view kDestDir = "./ut_dir/snii_bench"; +constexpr static std::string_view kTmpDir = "./ut_dir/snii_bench_tmp"; + +struct QueryCase { + const char* label; + InvertedIndexQueryType type; + const char* text; +}; + +// Spread across the shapes the two formats plan differently: single common term, single rare term, +// two/three-word phrases, stopword-led phrases (where CommonGrams matters), and prefixes of varying +// length. A single query shape would only exercise one code path. +const std::vector kQueryCases = { + {"any_common", InvertedIndexQueryType::MATCH_ANY_QUERY, "anarchism"}, + {"any_rare", InvertedIndexQueryType::MATCH_ANY_QUERY, "syndicalism"}, + {"any_multi", InvertedIndexQueryType::MATCH_ANY_QUERY, "philosophy movement state"}, + {"all_multi", InvertedIndexQueryType::MATCH_ALL_QUERY, "political philosophy"}, + {"phrase_2", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "anarchism is"}, + {"phrase_3", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "political philosophy and"}, + {"phrase_stop", InvertedIndexQueryType::MATCH_PHRASE_QUERY, "of the"}, + {"prefix_1", InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "anarch"}, + {"prefix_2", InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "anarchism is"}, + {"prefix_3", InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "philosophy and mov"}, + {"prefix_stop", InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "the united sta"}, + {"prefix_long", InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, + "in the early twentieth cen"}, +}; + +// Which storage the rowset lives on. Local is the fast iteration loop; RemoteS3 is the shape a +// real deployment has -- S3FileSystem -> CachedRemoteFileReader -> BlockFileCache -> query -- and is +// the only mode where the read-amplification counters mean anything. +enum class IoMode { kLocal, kRemoteS3 }; + +// How the run treats the block file cache. Applied uniformly to load, compaction and query so the +// three phases describe one consistent deployment instead of contradicting each other. +// +// kDirect - file cache off end to end. Load writes straight to S3, compaction re-reads from +// S3, the query reads from S3. This is the read-amplification measurement: every +// byte and every GET is real. +// kWriteBack - file cache on, load populates it (cloud sets RowsetWriterContext::write_file_cache +// from the load request). Compaction then reads local SSD, which is CPU bound and is +// what cloud E2E actually does. Two query numbers come out of this mode: a cold one +// (cache emptied first, so the query fetches from S3 and repopulates) and a hot one +// (same query again against the cache it just filled). +// +// Mixing them is what made earlier runs incoherent: a warm-cache compaction paired with a +// cold-cache query described no real deployment. +enum class BenchCachePolicy { kDirect, kWriteBack }; + +struct Measurement { + double wall_s = 0; + double cpu_s = 0; +}; + +double nearest_rank_percentile(const std::vector& sorted_samples, size_t percentile) { + DORIS_CHECK(!sorted_samples.empty()); + DORIS_CHECK(percentile > 0); + DORIS_CHECK(percentile <= 100); + const size_t whole_hundreds = sorted_samples.size() / 100; + const size_t remainder = sorted_samples.size() % 100; + const size_t percentile_rank = + whole_hundreds * percentile + (remainder * percentile + 99) / 100; + const size_t percentile_index = percentile_rank - 1; + return sorted_samples[percentile_index]; +} + +int parse_query_iterations(std::string_view value) { + int query_iterations = 0; + const auto [end, error] = + std::from_chars(value.data(), value.data() + value.size(), query_iterations); + DORIS_CHECK(error == std::errc()); + DORIS_CHECK(end == value.data() + value.size()); + DORIS_CHECK(query_iterations > 0); + return query_iterations; +} + +// Prevent a future timing-report change from silently using the upper median for +// even-sized samples or from treating p99 as a maximum by convention. +TEST(SniiBenchmarkPercentile, SelectsNearestRankForOddAndEvenSamples) { + const std::vector odd {1.0, 2.0, 3.0, 4.0, 5.0}; + EXPECT_DOUBLE_EQ(nearest_rank_percentile(odd, 50), 3.0); + EXPECT_DOUBLE_EQ(nearest_rank_percentile(odd, 99), 5.0); + + const std::vector even {1.0, 2.0, 3.0, 4.0}; + EXPECT_DOUBLE_EQ(nearest_rank_percentile(even, 50), 2.0); + EXPECT_DOUBLE_EQ(nearest_rank_percentile(even, 99), 4.0); + + std::vector more_than_one_hundred_samples; + for (double sample = 0; sample <= 100; ++sample) { + more_than_one_hundred_samples.push_back(sample); + } + EXPECT_DOUBLE_EQ(nearest_rank_percentile(more_than_one_hundred_samples, 99), 99.0); + EXPECT_NE(nearest_rank_percentile(more_than_one_hundred_samples, 99), + more_than_one_hundred_samples.back()); +} + +TEST(SniiBenchmarkPercentile, RejectsInvalidPercentiles) { + const std::vector samples {1.0}; + EXPECT_DEATH({ static_cast(nearest_rank_percentile(samples, 0)); }, ""); + EXPECT_DEATH({ static_cast(nearest_rank_percentile(samples, 101)); }, ""); +} + +TEST(SniiBenchmarkConfig, ParsesPositiveQueryIterations) { + EXPECT_EQ(parse_query_iterations("30"), 30); +} + +TEST(SniiBenchmarkConfig, RejectsInvalidQueryIterations) { + EXPECT_DEATH({ static_cast(parse_query_iterations("0")); }, ""); + EXPECT_DEATH({ static_cast(parse_query_iterations("-1")); }, ""); + EXPECT_DEATH({ static_cast(parse_query_iterations("30junk")); }, ""); + EXPECT_DEATH({ static_cast(parse_query_iterations("abc")); }, ""); + + const std::string overflow = + std::to_string(static_cast(std::numeric_limits::max()) + 1); + EXPECT_DEATH({ static_cast(parse_query_iterations(overflow)); }, ""); +} + +struct PhaseResult { + Measurement import; + Measurement compaction; + Measurement cold_query; // p50 over query iterations + Measurement cold_query_p99; + // Only filled for BenchCachePolicy::kWriteBack: the same queries run again without emptying the + // cache, so they are served by the blocks the cold pass just wrote back. + Measurement hot_query; + Measurement hot_query_p99; + bool has_hot = false; + int64_t hot_remote_read_bytes = 0; + int64_t hot_range_read_count = 0; + int64_t index_bytes = -1; // -1 = not measured (remote rowsets expose S3 keys, not paths) + // Deterministic remote-IO comparables. Unlike latency these do not move with machine or network + // load, so they carry the SNII-vs-V3 verdict in remote mode. + int64_t remote_physical_read_bytes = 0; + int64_t range_read_count = 0; + int64_t serial_read_rounds = 0; + // How many columns actually went through index compaction. Zero means the index was rebuilt + // from raw data instead, which is a completely different cost profile -- comparing a format + // that compacted against one that rebuilt would be meaningless. + int64_t index_compaction_columns = -1; + int64_t matched_docs = 0; + double cold_wall_min = 0; + double cold_wall_max = 0; + std::vector cold_cpu_samples; + std::vector cold_wall_samples; + std::vector hot_cpu_samples; + std::vector hot_wall_samples; + std::vector per_case; +}; + +class DISABLED_SniiVsV3BenchmarkTest : public ::testing::Test { +protected: + static constexpr int kDefaultQueryIterations = 30; + // build_rowsets() asserts num_segments == num_rows / max_rows_per_segment, so this must divide + // the per-file document count exactly. The prepared corpus is 200 documents per file. + static constexpr int64_t kRowsPerSegment = 200; + static constexpr int64_t kBenchTabletId = 990001; + // What the cache directory still weighs once it is empty of data: the RocksDB meta store and + // the LRU dump files, which are never removed. ~37 KB on a run that caches everything, but it + // grows with the number of blocks ever inserted, so a small-capacity run that churns 120 MB + // through a 3 MB queue leaves ~330 KB of metadata behind. 2 MB separates that from any + // meaningful amount of retained data (a cold pass fetches 12-16 MB). + static constexpr int64_t kColdCacheResidueBytes = 2 * 1024 * 1024; + // 400 x 25 ms = 10 s for the GC thread to finish removing what clear_file_caches() marked. + static constexpr int kCacheDrainSpins = 400; + static constexpr int kCacheDrainSleepMs = 25; + // 600 x 50 ms = 30 s for BlockFileCache to finish opening asynchronously. + static constexpr int kCacheOpenSpins = 600; + static constexpr int kCacheOpenSleepMs = 50; + static constexpr const char* kCacheDirPrefix = "snii_bench_file_cache_"; + static constexpr const char* kS3Prefix = "snii_bench_"; + + // Drops every block from the file cache and returns what the directory still weighs. + // clear_file_caches() only marks referenced blocks for recycling, so it can return with the + // cache still full (num_cells_wait_recycle); BlockFileCache's background GC thread does the + // actual removal. Poll until it has, otherwise a "cold" query is served from the cache it was + // supposed to have dropped -- which showed up as a wall spread of [0.010, 0.234] across + // iterations that were all supposed to be cold. + int64_t _empty_file_cache() { + int64_t remaining = _cache_dir_bytes(); + for (int spin = 0; spin < kCacheDrainSpins && remaining > kColdCacheResidueBytes; ++spin) { + static_cast(io::FileCacheFactory::instance()->clear_file_caches(true)); + std::this_thread::sleep_for(std::chrono::milliseconds(kCacheDrainSleepMs)); + remaining = _cache_dir_bytes(); + } + return remaining; + } + + // Bytes the block file cache is actually holding, read off the filesystem. The queue metrics + // are not usable for this -- they report zero while the block files are still on disk. + int64_t _cache_dir_bytes() const { + if (_file_cache_dir.empty()) { + return 0; + } + int64_t total = 0; + std::error_code ec; + for (auto it = std::filesystem::recursive_directory_iterator(_file_cache_dir, ec); + !ec && it != std::filesystem::recursive_directory_iterator(); it.increment(ec)) { + if (it->is_regular_file(ec)) { + total += static_cast(it->file_size(ec)); + } + } + return total; + } + + void SetUp() override { + char buffer[MAX_PATH_LEN]; + ASSERT_NE(getcwd(buffer, MAX_PATH_LEN), nullptr); + _current_dir = std::string(buffer); + _absolute_dir = _current_dir + std::string(kDestDir); + ASSERT_TRUE(io::global_local_filesystem()->delete_directory(_absolute_dir).ok()); + ASSERT_TRUE(io::global_local_filesystem()->create_directory(_absolute_dir).ok()); + + ASSERT_TRUE(io::global_local_filesystem()->delete_directory(kTmpDir).ok()); + ASSERT_TRUE(io::global_local_filesystem()->create_directory(kTmpDir).ok()); + std::vector paths; + paths.emplace_back(std::string(kTmpDir), 1024000000); + auto tmp_file_dirs = std::make_unique(paths); + ASSERT_TRUE(tmp_file_dirs->init().ok()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + + doris::EngineOptions options; + auto engine = std::make_unique(options); + _engine_ref = engine.get(); + _data_dir = std::make_unique(*_engine_ref, _absolute_dir); + static_cast(_data_dir->update_capacity()); + ExecEnv::GetInstance()->set_storage_engine(std::move(engine)); + + // Match the existing index-compaction perf test so the two are comparable. + config::enable_segcompaction = false; + config::enable_ordered_data_compaction = false; + config::total_permits_for_compaction_score = 200000; + config::inverted_index_ram_dir_enable = true; + config::string_type_length_soft_limit_bytes = 10485760; + + // The query path goes through InvertedIndexQueryCache::instance() unconditionally, and + // ExecEnv hands out a null pointer unless something installs one. Own them here and put the + // previous globals back in TearDown so a later test never sees a dangling pointer. + constexpr int64_t kIndexCacheBytes = 1024L * 1024L * 1024L; + _searcher_cache.reset(segment_v2::InvertedIndexSearcherCache::create_global_instance( + kIndexCacheBytes, 1)); + _query_cache.reset( + segment_v2::InvertedIndexQueryCache::create_global_cache(kIndexCacheBytes, 1)); + _previous_searcher_cache = ExecEnv::GetInstance()->get_inverted_index_searcher_cache(); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_searcher_cache.get()); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_query_cache.get()); + + _corpus_files = _discover_corpus(); + } + + void TearDown() override { + if (_tablet != nullptr) { + static_cast( + io::global_local_filesystem()->delete_directory(_tablet->tablet_path())); + } + static_cast(io::global_local_filesystem()->delete_directory(_absolute_dir)); + static_cast(io::global_local_filesystem()->delete_directory(kTmpDir)); + _engine_ref = nullptr; + ExecEnv::GetInstance()->set_storage_engine(nullptr); + + config::enable_segcompaction = true; + config::enable_ordered_data_compaction = true; + config::total_permits_for_compaction_score = 1000000; + config::string_type_length_soft_limit_bytes = 1048576; + // do_compaction() raises this; the sibling DISABLED_IndexCompactionPerformanceTest + // restores it the same way. Leaving it set changes later tests in the same binary. + config::compaction_batch_size = _origin_compaction_batch_size; + config::inverted_index_compaction_enable = _origin_index_compaction_enable; + config::inverted_index_ram_dir_enable = _origin_ram_dir_enable; + + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_previous_searcher_cache); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); + + // Unconditionally, not inside _teardown_remote(): _setup_remote() installs the fixture's + // factory before any of its five later failure paths, and each of those makes the caller + // GTEST_SKIP(), which skips _teardown_remote() entirely. ExecEnv would then keep a pointer + // to a factory this fixture is about to destroy, and every later test in doris_be_test that + // touches the file cache would read freed memory. + _restore_file_cache_globals(); + } + + // Puts back everything _setup_remote() replaced. Safe to call when setup never ran. + void _restore_file_cache_globals() { + if (_origin_file_cache_factory_saved) { + ExecEnv::GetInstance()->_file_cache_factory = _origin_file_cache_factory; + _origin_file_cache_factory_saved = false; + // Deliberately leaked, never destroyed. FileWriter::init_cache_builder hands the S3 + // write path a raw BlockFileCache* (file_writer.h), and S3FileWriter fills the cache + // from its upload threads -- its own comment notes the writer may already be gone by + // then. The cache's own GC/monitor threads are likewise only stopped by ~BlockFileCache. + // In a real BE the factory is a process-lifetime singleton so none of that can dangle; + // destroying it here frees the caches out from under those threads. + static_cast(_owned_file_cache_factory.release()); + } + config::enable_file_cache = _origin_enable_file_cache; + } + + // Process CPU time, which is what the verdict is based on; wall time is kept only to detect a + // run that was descheduled or IO bound. + static Measurement measure(const std::function& body) { + timespec cpu_start {}; + timespec cpu_end {}; + clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &cpu_start); + const int64_t wall_start = MonotonicNanos(); + body(); + const int64_t wall_end = MonotonicNanos(); + clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &cpu_end); + + Measurement m; + m.wall_s = static_cast(wall_end - wall_start) / 1e9; + m.cpu_s = static_cast(cpu_end.tv_sec - cpu_start.tv_sec) + + static_cast(cpu_end.tv_nsec - cpu_start.tv_nsec) / 1e9; + return m; + } + + static int _query_iterations() { + const char* const env_iterations = std::getenv("SNII_BENCH_QUERY_ITERATIONS"); + return env_iterations != nullptr ? parse_query_iterations(env_iterations) + : kDefaultQueryIterations; + } + + std::vector _discover_corpus() const { + const char* env_dir = std::getenv("SNII_BENCH_CORPUS_DIR"); + const std::string dir = + env_dir != nullptr + ? std::string(env_dir) + : _current_dir + "/be/test/storage/index/inverted/data/performance"; + std::vector files; + if (!std::filesystem::exists(dir)) { + return files; + } + for (const auto& entry : std::filesystem::directory_iterator(dir)) { + const std::string name = entry.path().filename().string(); + if (entry.is_regular_file() && name.starts_with("wikipedia") && + name.ends_with(".json")) { + files.push_back(entry.path().string()); + } + } + // Deterministic load order: the segment layout must not depend on readdir order, or the two + // formats would not be indexing the same thing in the same sequence. + std::sort(files.begin(), files.end()); + return files; + } + + void _build_tablet(InvertedIndexStorageFormatPB storage_format) { + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(KeysType::DUP_KEYS); + schema_pb.set_inverted_index_storage_format(storage_format); + + std::map properties; + properties.emplace(INVERTED_INDEX_PARSER_KEY, INVERTED_INDEX_PARSER_ENGLISH); + properties.emplace(INVERTED_INDEX_PARSER_PHRASE_SUPPORT_KEY, + INVERTED_INDEX_PARSER_PHRASE_SUPPORT_YES); + properties.emplace(INVERTED_INDEX_PARSER_LOWERCASE_KEY, INVERTED_INDEX_PARSER_TRUE); + + IndexCompactionUtils::construct_column(schema_pb.add_column(), 0, "STRING", "title"); + IndexCompactionUtils::construct_column(schema_pb.add_column(), schema_pb.add_index(), 10001, + "idx_content", 1, "STRING", "content", properties); + IndexCompactionUtils::construct_column(schema_pb.add_column(), 2, "STRING", "redirect"); + IndexCompactionUtils::construct_column(schema_pb.add_column(), 3, "STRING", "namespace"); + + _tablet_schema = std::make_shared(); + _tablet_schema->init_from_pb(schema_pb); + + // Default-construct rather than TabletMeta(schema): only the default ctor initialises + // _delete_bitmap, and init_from_pb writes through it. Going via the PB is the only way to + // set tablet_id, which CachedRemoteFileReader requires to be > 0 for Doris tables. + TabletMetaPB meta_pb; + meta_pb.set_tablet_id(kBenchTabletId); + meta_pb.set_schema_hash(1); + meta_pb.set_table_id(1); + meta_pb.set_partition_id(1); + meta_pb.set_replica_id(1); + *meta_pb.mutable_schema() = schema_pb; + TabletMetaSharedPtr tablet_meta(new TabletMeta()); + tablet_meta->init_from_pb(meta_pb); + _tablet = std::make_shared(*_engine_ref, tablet_meta, _data_dir.get()); + ASSERT_TRUE(_tablet->init().ok()); + std::cout << " tablet_id=" << _tablet->tablet_id() << std::endl; + + static_cast(io::global_local_filesystem()->delete_directory(_tablet->tablet_path())); + ASSERT_TRUE(io::global_local_filesystem()->create_directory(_tablet->tablet_path()).ok()); + } + + // Stand up the real deployment chain: S3FileSystem -> CachedRemoteFileReader -> BlockFileCache. + // Credentials come from the environment so they never live in the source tree; the endpoint is + // a real S3 service, so this measures genuine remote behaviour rather than a local stand-in. + // Returns false (and skips) when the environment is not configured. + bool _setup_remote() { + auto env = [](const char* name) -> std::string { + const char* v = std::getenv(name); + return v != nullptr ? std::string(v) : std::string(); + }; + const std::string ak = env("SNII_BENCH_S3_AK"); + const std::string sk = env("SNII_BENCH_S3_SK"); + const std::string endpoint = env("SNII_BENCH_S3_ENDPOINT"); + const std::string region = env("SNII_BENCH_S3_REGION"); + const std::string bucket = env("SNII_BENCH_S3_BUCKET"); + if (ak.empty() || sk.empty() || endpoint.empty() || bucket.empty()) { + return false; + } + + _origin_enable_file_cache = config::enable_file_cache; + config::enable_file_cache = true; + + // Disk-backed, not memory-backed: production caches to disk, and a cold query has to be + // able to evict the cache's own backing pages as well. + // Must live on a disk with real free space. BlockFileCache watches the filesystem holding + // this path and, above config::file_cache_enter_disk_resource_limit_mode_percent (90%), + // enters disk resource limit mode: its background thread then evicts everything as fast as + // the write path caches it, so the cache stays empty no matter what write_file_cache says. + // The repo checkout sits on a 98%-full disk, which is what silently disabled caching here. + const char* cache_root = std::getenv("SNII_BENCH_CACHE_DIR"); + _file_cache_dir = (cache_root != nullptr && *cache_root != '\0') + ? std::string(cache_root) + "/" + kCacheDirPrefix + + std::to_string(::getpid()) + : _absolute_dir + "/file_cache"; + // The cache dir is never deleted on teardown (the factory is leaked, see _teardown_remote), + // so sweep leftovers from earlier runs here instead. + if (cache_root != nullptr && *cache_root != '\0') { + std::error_code sec; + for (auto it = std::filesystem::directory_iterator(cache_root, sec); + !sec && it != std::filesystem::directory_iterator(); it.increment(sec)) { + const std::string name = it->path().filename().string(); + if (!name.starts_with(kCacheDirPrefix)) { + continue; + } + // Only reap directories whose owning process is gone. A second benchmark running + // concurrently has a live pid, and deleting its cache out from under it would + // corrupt that run rather than this one. + const std::string pid_part = name.substr(std::strlen(kCacheDirPrefix)); + const int64_t owner = std::atoll(pid_part.c_str()); + if (owner > 0 && ::kill(static_cast(owner), 0) == 0) { + continue; + } + std::error_code rec; + std::filesystem::remove_all(it->path(), rec); + } + } + static_cast(io::global_local_filesystem()->delete_directory(_file_cache_dir)); + static_cast(io::global_local_filesystem()->create_directory(_file_cache_dir)); + // Reading a block back out of the cache goes FSFileCacheStorage::get_or_open_file_reader -> + // FDCache::instance(), which is just ExecEnv::file_cache_open_fd_cache(). Only + // exec_env_init.cpp:318 ever creates it, and this fixture does not run that, so the pointer + // is null and the first cache hit dereferences it. Never seen until the cache actually + // retained data. + if (ExecEnv::GetInstance()->file_cache_open_fd_cache() == nullptr) { + ExecEnv::GetInstance()->set_file_cache_open_fd_cache(std::make_unique()); + } + _origin_file_cache_factory = ExecEnv::GetInstance()->_file_cache_factory; + _origin_file_cache_factory_saved = true; + _owned_file_cache_factory = std::make_unique(); + ExecEnv::GetInstance()->_file_cache_factory = _owned_file_cache_factory.get(); + // Built the way a real BE builds it. Hand-filling FileCacheSettings leaves ttl_queue_* + // at zero and `storage` empty, and the queue sizes then do not add up to capacity, so + // reservations for the queue the write path uses can never be satisfied. + // Shrinking this is how the benchmark models a working set larger than the cache: index + // blocks land in the normal queue, which get_file_cache_settings sizes at 40% of capacity, + // so a 16 MB index stops fitting once capacity drops below ~40 MB and the LRU starts + // evicting between queries. Default is large enough that nothing is ever evicted. + int64_t capacity_mb = 4096; + if (const char* cap = std::getenv("SNII_BENCH_CACHE_CAPACITY_MB"); + cap != nullptr && *cap != '\0') { + capacity_mb = std::max(1, std::atoll(cap)); + } + io::FileCacheSettings settings = + io::get_file_cache_settings(/*capacity=*/capacity_mb * 1024L * 1024L, + /*max_query_cache_size=*/0); + settings.max_file_block_size = 1024 * 1024; + std::cout << " file cache capacity " << capacity_mb << " MB (normal queue " + << settings.query_queue_size / (1024 * 1024) << " MB)" << std::endl; + if (!io::FileCacheFactory::instance()->create_file_cache(_file_cache_dir, settings).ok()) { + return false; + } + // create_file_cache() returns before the cache finishes opening. Until then try_reserve() + // diverts to try_reserve_during_async_load(), which throttles reservations, so a load that + // starts immediately writes almost nothing into the cache -- the whole point of + // write_file_cache. A real BE has been up long enough for this to be done. + { + io::BlockFileCache* cache = + io::FileCacheFactory::instance()->get_by_path(_file_cache_dir); + if (cache == nullptr) { + return false; + } + for (int i = 0; i < kCacheOpenSpins && !cache->get_async_open_success(); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(kCacheOpenSleepMs)); + } + if (!cache->get_async_open_success()) { + std::cerr << "file cache did not finish async open" << std::endl; + return false; + } + } + + if (ExecEnv::GetInstance()->s3_file_upload_thread_pool() == nullptr) { + std::unique_ptr pool; + if (!ThreadPoolBuilder("snii_bench_s3_upload") + .set_min_threads(1) + .set_max_threads(8) + .build(&pool) + .ok()) { + return false; + } + ExecEnv::GetInstance()->_s3_file_upload_thread_pool = std::move(pool); + } + + S3Conf s3_conf; + s3_conf.client_conf.ak = ak; + s3_conf.client_conf.sk = sk; + s3_conf.client_conf.endpoint = endpoint; + s3_conf.client_conf.region = region; + s3_conf.bucket = bucket; + s3_conf.prefix = kS3Prefix + std::to_string(::getpid()); + auto fs = io::S3FileSystem::create(std::move(s3_conf), "snii-bench-s3-fs"); + if (!fs.has_value()) { + std::cout << " S3FileSystem::create failed: " << fs.error() << std::endl; + return false; + } + _remote_fs = fs.value(); + std::cout << " remote: s3://" << bucket << "/" << kS3Prefix << ::getpid() << " via " + << endpoint << std::endl; + return true; + } + + void _teardown_remote() { + if (_remote_fs == nullptr) { + return; + } + // Otherwise every run leaves a full corpus (~40 MB per format, plus compaction output) + // in the user's bucket forever. + const Status st = _remote_fs->delete_directory(""); + if (!st.ok()) { + std::cout << " warning: could not remove s3://.../" << kS3Prefix << ::getpid() << ": " + << st.to_string() << std::endl; + } + _remote_fs.reset(); + _restore_file_cache_globals(); + } + + // Pages of `path` still resident in the OS page cache. Used to prove the eviction below + // actually worked instead of assuming it did. + static std::pair _resident_pages(const std::string& path) { + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + return {0, 0}; + } + struct stat st {}; + if (::fstat(fd, &st) != 0 || st.st_size == 0) { + ::close(fd); + return {0, 0}; + } + void* addr = ::mmap(nullptr, static_cast(st.st_size), PROT_READ, MAP_SHARED, fd, 0); + if (addr == MAP_FAILED) { + ::close(fd); + return {0, 0}; + } + const size_t page_size = static_cast(::sysconf(_SC_PAGESIZE)); + const size_t pages = (static_cast(st.st_size) + page_size - 1) / page_size; + std::vector vec(pages, 0); + size_t resident = 0; + if (::mincore(addr, static_cast(st.st_size), vec.data()) == 0) { + for (unsigned char v : vec) { + resident += (v & 1u); + } + } + ::munmap(addr, static_cast(st.st_size)); + ::close(fd); + return {resident, pages}; + } + + // A real cold query means three things, and the first two are not enough on their own: + // 1. Doris' own caches (query/searcher/page/segment) emptied; + // 2. the OS page cache emptied for the data and index files, otherwise every "cold" read is + // served from RAM and measures nothing; + // 3. no Doris file cache in the way -- enable_file_cache is false here and the UT reads local + // files through LocalFileReader. In remote write-back mode BlockFileCache does + // participate, which is why the caller also empties it via _empty_file_cache(). + // fsync first: POSIX_FADV_DONTNEED only evicts clean pages, and compaction has just written + // these files. + static void _drop_caches(const std::string& tablet_dir) { + // Not a null check: if the manager were missing, nothing would be pruned and every + // "cold" number below would silently be a warm one. Fail instead of measuring garbage. + auto* manager = ExecEnv::GetInstance()->get_cache_manager(); + DORIS_CHECK(manager != nullptr); + static_cast(manager->for_each_cache_prune_all(nullptr, /*force=*/true)); + std::error_code ec; + for (const auto& entry : std::filesystem::directory_iterator(tablet_dir, ec)) { + if (!entry.is_regular_file(ec)) { + continue; + } + const int fd = ::open(entry.path().c_str(), O_RDONLY); + if (fd < 0) { + continue; + } + ::fsync(fd); + ::posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED); + ::close(fd); + } + } + + static int64_t _index_bytes(const RowsetSharedPtr& rowset) { + // Only this rowset's own index files. The tablet directory also holds the source rowsets + // that compaction consumed, and summing the whole directory silently reported + // sources + output as if it were the compacted index. + int64_t total = 0; + const auto& seg_path = rowset->segment_path(0); + if (!seg_path.has_value()) { + return total; + } + const std::string prefix = rowset->rowset_id().to_string(); + const auto dir = std::filesystem::path(seg_path.value()).parent_path(); + std::error_code ec; + for (const auto& entry : std::filesystem::directory_iterator(dir, ec)) { + const std::string name = entry.path().filename().string(); + if (name.starts_with(prefix) && (name.ends_with(".idx") || name.ends_with(".snii"))) { + total += static_cast(entry.file_size(ec)); + } + } + return total; + } + + // Runs both formats and returns {v3, snii}. Order matters and is not neutral: whichever runs + // first pays Aws::InitAPI, the TLS handshake and connection pool, the CLucene/analyzer and + // codec-pool lazy init, and faults in the tcmalloc arena the second run then reuses warm. + // SNII_BENCH_REVERSE_ORDER=1 runs SNII first; comparing the two orders quantifies that bias + // instead of assuming it away. + std::pair _run_both(IoMode io_mode, BenchCachePolicy policy) { + const bool reverse = std::getenv("SNII_BENCH_REVERSE_ORDER") != nullptr; + std::cout << " format order: " << (reverse ? "SNII first" : "V3 first") << std::endl; + if (reverse) { + PhaseResult snii = _run_format(InvertedIndexStorageFormatPB::SNII, io_mode, policy); + PhaseResult v3 = _run_format(InvertedIndexStorageFormatPB::V3, io_mode, policy); + return {std::move(v3), std::move(snii)}; + } + PhaseResult v3 = _run_format(InvertedIndexStorageFormatPB::V3, io_mode, policy); + PhaseResult snii = _run_format(InvertedIndexStorageFormatPB::SNII, io_mode, policy); + return {std::move(v3), std::move(snii)}; + } + + PhaseResult _run_format(InvertedIndexStorageFormatPB storage_format, IoMode io_mode, + BenchCachePolicy policy) { + PhaseResult result; + const int query_iterations = _query_iterations(); + const bool write_back = + io_mode == IoMode::kRemoteS3 && policy == BenchCachePolicy::kWriteBack; + const bool local_warm = io_mode == IoMode::kLocal; + // Held for the whole run so load, compaction and query all see the same policy. With the + // cache off, S3FileSystem hands out a plain reader instead of a CachedRemoteFileReader, so + // every read is a real GET. + const bool saved_enable_file_cache = config::enable_file_cache; + config::enable_file_cache = write_back; + Defer restore_cache_cfg {[&] { config::enable_file_cache = saved_enable_file_cache; }}; + _build_tablet(storage_format); + + // Both formats must start from the same cache state. _setup_remote() runs once per test and + // _run_format() is then called for V3 and SNII against the same BlockFileCache, so without + // this the second format would load into an LRU the first one left full -- exactly the + // regime SNII_BENCH_CACHE_CAPACITY_MB exists to probe, where residency decides the result. + int64_t cache_bytes_before_load = 0; + if (write_back) { + _empty_file_cache(); + cache_bytes_before_load = _cache_dir_bytes(); + } + + // --- Phase 1: load + build the index --- + std::optional storage_resource; + if (io_mode == IoMode::kRemoteS3) { + storage_resource = StorageResource(_remote_fs); + } + std::vector rowsets(_corpus_files.size()); + result.import = measure([&] { + IndexCompactionUtils::build_rowsets( + _data_dir, _tablet_schema, _tablet, _engine_ref, rowsets, _corpus_files, + _inc_id, nullptr, /*is_performance=*/true, kRowsPerSegment, storage_resource, + /*write_file_cache=*/write_back); + }); + + // --- Phase 2: index compaction --- + // Cloud load leaves the file cache warm for the data it just wrote (write_file_cache above), + // so compaction there reads local SSD rather than S3. The S3 upload path fills the cache + // asynchronously, so drain it before timing compaction, otherwise the first compaction read + // races the upload and still goes to S3. + if (write_back) { + const int64_t cached = _cache_dir_bytes(); + std::cout << " load populated file cache with " << cached - cache_bytes_before_load + << " B (dir now " << cached << " B)" << std::endl; + // An empty cache dir still weighs kColdCacheResidueBytes of RocksDB metadata, so + // "> 0" would pass even if nothing was cached. If the load did not fill the cache, + // compaction below reads S3 and the timing is not comparable to cloud. + EXPECT_GT(cached - cache_bytes_before_load, kColdCacheResidueBytes) + << "write_file_cache did not populate " << _file_cache_dir; + } + + RowsetSharedPtr output_rowset; + Status compaction_status; + result.compaction = measure([&] { + compaction_status = IndexCompactionUtils::do_compaction( + rowsets, _engine_ref, _tablet, /*is_index_compaction=*/true, output_rowset, + [&result](const BaseCompaction&, const RowsetWriterContext& cctx) { + result.index_compaction_columns = + static_cast(cctx.columns_to_do_index_compaction.size()); + }, + 10000000, storage_resource); + }); + EXPECT_TRUE(compaction_status.ok()) << compaction_status.to_string(); + // Only meaningful for a local rowset: for a remote one segment_path() returns an S3 key, + // so the directory walk in _index_bytes() would find nothing and report 0 -- which the + // report would print as if the index were empty. Leave it unset (-1 = not measured). + result.index_bytes = io_mode == IoMode::kLocal ? _index_bytes(output_rowset) : -1; + + // --- Phase 3: cold match_phrase_prefix --- + std::vector cpu_samples; + std::vector wall_samples; + cpu_samples.reserve(query_iterations); + wall_samples.reserve(query_iterations); + const std::string tablet_dir = _tablet->tablet_path(); + // Same reason: an S3 key is not a path mincore() can open, and "0/0 pages resident" + // would be indistinguishable from a successful eviction. + std::string probe_file; + if (io_mode == IoMode::kLocal) { + const auto& sp = output_rowset->segment_path(0); + DORIS_CHECK(sp.has_value()); + probe_file = std::filesystem::path(sp.value()).replace_extension(".idx").string(); + } + for (int i = 0; i < query_iterations; ++i) { + _drop_caches(tablet_dir); + if (_remote_fs != nullptr) { + // Evicting the cache directory's OS pages is not enough: BlockFileCache still + // believes it holds the blocks and serves them from local disk, so the query never + // goes back to S3 and the read counters stay at zero. Invalidate the cache itself + // so a cold query is a genuine remote fetch. The rowset's own segments pin blocks, + // so release them first or nothing becomes evictable. + output_rowset->clear_cache(); + const int64_t before = _cache_dir_bytes(); + const int64_t after = _empty_file_cache(); + if (i == 0) { + std::cout << " cache before/after clear: " << before << " -> " << after << " B" + << std::endl; + } + // A "cold" query served from a cache that never emptied is not a measurement. + EXPECT_LE(after, kColdCacheResidueBytes) + << "file cache still holds " << after << " B; cold query is not cold"; + } + if (!_file_cache_dir.empty()) { + // The block file cache is the layer a real deployment actually hits; leaving it + // warm would make every "cold" query a cache hit and measure nothing remote. + _drop_caches(_file_cache_dir); + } + if (i == 0 && !probe_file.empty()) { + const auto [resident, pages] = _resident_pages(probe_file); + std::cout << " cold check: " << resident << "/" << pages + << " index pages resident after evict" << std::endl; + } + int64_t matched = 0; + io::FileCacheStatistics io_stats; + const Measurement m = measure([&] { + matched = _run_query_cases(output_rowset, &result.per_case, &io_stats); + }); + // These are deterministic per iteration; the report pairs them with a median + // timing, so assert they really are identical rather than quietly publishing + // iteration N's IO next to iteration N/2's time. + if (i == 0) { + result.remote_physical_read_bytes = + io_stats.inverted_index_remote_physical_read_bytes; + result.range_read_count = io_stats.inverted_index_range_read_count; + result.serial_read_rounds = io_stats.inverted_index_serial_read_rounds; + } else { + EXPECT_EQ(result.range_read_count, io_stats.inverted_index_range_read_count) + << "range_reads moved between iterations; the IO counters are not " + "comparable to a median timing"; + } + cpu_samples.push_back(m.cpu_s); + wall_samples.push_back(m.wall_s); + result.matched_docs = matched; + } + std::sort(cpu_samples.begin(), cpu_samples.end()); + std::sort(wall_samples.begin(), wall_samples.end()); + result.cold_cpu_samples = std::move(cpu_samples); + result.cold_wall_samples = std::move(wall_samples); + result.cold_query.cpu_s = nearest_rank_percentile(result.cold_cpu_samples, 50); + result.cold_query.wall_s = nearest_rank_percentile(result.cold_wall_samples, 50); + result.cold_query_p99.cpu_s = nearest_rank_percentile(result.cold_cpu_samples, 99); + result.cold_query_p99.wall_s = nearest_rank_percentile(result.cold_wall_samples, 99); + // A cold query is dominated by IO, so wall spread is the honest signal-to-noise indicator. + result.cold_wall_min = result.cold_wall_samples.front(); + result.cold_wall_max = result.cold_wall_samples.back(); + + // --- Phase 3b: hot query --- + // The cold pass above ended by pulling everything it touched into the active cache layer. + // Repeating the same queries without dropping anything measures the steady state: the + // block cache for remote write-back, or the OS page cache for a local run. + if (write_back || local_warm) { + std::vector hot_cpu; + std::vector hot_wall; + hot_cpu.reserve(query_iterations); + hot_wall.reserve(query_iterations); + for (int i = 0; i < query_iterations; ++i) { + io::FileCacheStatistics hot_stats; + std::vector hot_hits; + const Measurement m = measure([&] { + static_cast(_run_query_cases(output_rowset, &hot_hits, &hot_stats)); + }); + result.hot_remote_read_bytes = hot_stats.inverted_index_remote_physical_read_bytes; + result.hot_range_read_count = hot_stats.inverted_index_range_read_count; + hot_cpu.push_back(m.cpu_s); + hot_wall.push_back(m.wall_s); + } + std::sort(hot_cpu.begin(), hot_cpu.end()); + std::sort(hot_wall.begin(), hot_wall.end()); + result.hot_cpu_samples = std::move(hot_cpu); + result.hot_wall_samples = std::move(hot_wall); + result.hot_query.cpu_s = nearest_rank_percentile(result.hot_cpu_samples, 50); + result.hot_query.wall_s = nearest_rank_percentile(result.hot_wall_samples, 50); + result.hot_query_p99.cpu_s = nearest_rank_percentile(result.hot_cpu_samples, 99); + result.hot_query_p99.wall_s = nearest_rank_percentile(result.hot_wall_samples, 99); + result.has_hot = true; + } + return result; + } + + int64_t _run_query_cases(const RowsetSharedPtr& rowset, + std::vector* per_case = nullptr, + io::FileCacheStatistics* io_out = nullptr) { + // Both are reachable only when an earlier phase failed. Returning 0 quietly would let + // the run print a full report in which every query "matched" nothing. + if (rowset == nullptr) { + ADD_FAILURE() << "no output rowset; compaction must have failed"; + return 0; + } + SegmentCacheHandle segment_cache; + const Status load_st = SegmentLoader::instance()->load_segments( + std::static_pointer_cast(rowset), &segment_cache); + if (!load_st.ok()) { + ADD_FAILURE() << "load_segments failed: " << load_st.to_string(); + return 0; + } + int64_t matched = 0; + io::FileCacheStatistics collected_io; + if (per_case != nullptr) { + per_case->assign(kQueryCases.size(), 0); + } + // enable_inverted_index_query_cache defaults to true (PaloInternalService.thrift:370), and + // InvertedIndexReader::handle_query_cache then answers a repeated query straight from the + // cached bitmap without touching the searcher or any file reader. Every phase here reruns + // the same 12 queries against the same segments, so leaving it on would measure a bitmap + // lookup instead of the index read path this benchmark exists to compare -- the hot phase + // in particular would never reach the block cache it claims to be measuring. + TQueryOptions query_options; + query_options.enable_inverted_index_query_cache = false; + RuntimeState runtime_state(query_options, TQueryGlobals()); + for (const auto& segment : segment_cache.get_segments()) { + const auto& indexes = _tablet_schema->inverted_indexes(); + if (indexes.empty()) { + continue; + } + OlapReaderStatistics stats; + StorageReadOptions read_options; + read_options.stats = &stats; + std::unique_ptr iter; + const Status iter_st = segment->new_index_iterator(_tablet_schema->column(1), + indexes[0], read_options, &iter); + if (!iter_st.ok()) { + // Every segment must have the index; skipping one would compare a format that read + // 4 segments against one that read 3. + ADD_FAILURE() << "new_index_iterator failed: " << iter_st.to_string(); + continue; + } + DORIS_CHECK(iter != nullptr); + // read_from_index() dereferences the context unconditionally; SegmentIterator normally + // supplies it, so a hand-built iterator has to as well. + // The index stores the field under the column's unique id, not its name + // (index_writer.cpp: field_name = std::to_string(column->unique_id())). Querying by + // "content" silently matched no field and returned OK with an empty bitmap. + const std::string index_field_name = + std::to_string(_tablet_schema->column(1).unique_id()); + auto query_context = std::make_shared(); + query_context->stats = &stats; + // Without this the inverted-index IO counters stay at zero and the remote comparison + // has nothing to compare. + read_options.io_ctx.file_cache_stats = &stats.file_cache_stats; + query_context->io_ctx = &read_options.io_ctx; + // FullTextIndexReader reads query_options() off it unconditionally (max_expansions for + // phrase prefix), so a default-constructed state is required, not optional. + query_context->runtime_state = &runtime_state; + iter->set_context(query_context); + // new_index_iterator() alone leaves the index file reader uninitialised, and every + // query then returns OK with an empty bitmap -- which is indistinguishable from "no + // matches". SegmentIterator does this as part of its own setup. + const Status reader_init = segment->_index_file_reader->init( + config::inverted_index_read_buffer_size, &read_options.io_ctx); + if (!reader_init.ok()) { + ADD_FAILURE() << "index file reader init failed: " << reader_init.to_string(); + continue; + } + for (size_t pi = 0; pi < kQueryCases.size(); ++pi) { + const auto& qc = kQueryCases[pi]; + InvertedIndexParam param; + param.column_name = index_field_name; + // Required: the reader DCHECKs on it. The indexed column is STRING. + param.column_type = std::make_shared(); + param.query_value = Field::create_field(std::string(qc.text)); + param.query_type = qc.type; + param.num_rows = segment->num_rows(); + param.roaring = std::make_shared(); + const Status qs = iter->read_from_index(¶m); + if (!qs.ok()) { + // Never swallow this: a format that errors out looks identical to a format that + // is simply fast, and the timings would be comparing work against no work. + ADD_FAILURE() << qc.label << " ('" << qc.text + << "') failed: " << qs.to_string(); + continue; + } + { + const auto hits = static_cast(param.roaring->cardinality()); + matched += hits; + if (per_case != nullptr) { + (*per_case)[pi] += hits; + } + } + } + collected_io.merge_from(stats.file_cache_stats); + } + if (io_out != nullptr) { + *io_out = collected_io; + } + return matched; + } + + static void _report(const PhaseResult& v3, const PhaseResult& snii) { + auto ratio = [](double snii_v, double v3_v) { return v3_v > 0 ? snii_v / v3_v : 0.0; }; + auto line = [&](const char* phase, const Measurement& a, const Measurement& b) { + std::cout << std::left << std::setw(18) << phase << std::right << std::fixed + << std::setprecision(6) << std::setw(12) << a.cpu_s << std::setw(12) + << b.cpu_s << std::setw(9) << ratio(b.cpu_s, a.cpu_s) << "x" + << " " << std::setw(12) << a.wall_s << std::setw(12) << b.wall_s + << std::endl; + }; + + std::cout << "\n=== SNII vs V3 (CPU seconds; ratio <1 means SNII is cheaper) ===\n" + << std::left << std::setw(18) << "phase" << std::right << std::setw(12) + << "V3 cpu" << std::setw(12) << "SNII cpu" << std::setw(10) << "ratio" + << " " << std::setw(12) << "V3 wall" << std::setw(12) << "SNII wall" + << std::endl; + line("import", v3.import, snii.import); + line("compaction", v3.compaction, snii.compaction); + line("cold_query_p50", v3.cold_query, snii.cold_query); + line("cold_query_p99", v3.cold_query_p99, snii.cold_query_p99); + const double n = static_cast(kQueryCases.size()); + auto per_query = [&](const char* label, const Measurement& a, const Measurement& b) { + std::cout << std::left << std::setw(18) << label << std::right << std::fixed + << std::setprecision(4) << std::setw(10) << a.wall_s / n << std::setw(10) + << b.wall_s / n << std::setw(9) << (a.wall_s > 0 ? b.wall_s / a.wall_s : 0.0) + << "x (" << kQueryCases.size() << " queries)" << std::endl; + }; + per_query(" per-query wall", v3.cold_query, snii.cold_query); + std::cout << " cold wall spread V3 [" << std::fixed << std::setprecision(6) + << v3.cold_wall_min << ", " << v3.cold_wall_max << "] SNII [" + << snii.cold_wall_min << ", " << snii.cold_wall_max << "]" << std::endl; + if (v3.has_hot && snii.has_hot) { + line("hot_query_p50", v3.hot_query, snii.hot_query); + line("hot_query_p99", v3.hot_query_p99, snii.hot_query_p99); + per_query(" per-query wall", v3.hot_query, snii.hot_query); + } + std::cout << "\nper-query hits (V3 / SNII):" << std::endl; + for (size_t i = 0; i < kQueryCases.size(); ++i) { + const int64_t a = i < v3.per_case.size() ? v3.per_case[i] : -1; + const int64_t b = i < snii.per_case.size() ? snii.per_case[i] : -1; + std::cout << " " << std::left << std::setw(24) << (std::string(kQueryCases[i].label)) + << std::right << std::setw(8) << a << std::setw(8) << b + << (a == b ? "" : " <-- differs") << std::endl; + } + auto io_line = [&](const char* name, int64_t a, int64_t b) { + std::cout << std::left << std::setw(18) << name << std::right << std::setw(10) << a + << std::setw(10) << b << std::setw(9); + // A zero denominator means "nothing was measured"; printing 0.00x would read as + // "SNII is infinitely cheaper". + if (a > 0) { + std::cout << static_cast(b) / static_cast(a) << "x"; + } else { + std::cout << "n/a"; + } + std::cout << std::endl; + }; + io_line("idx_compact_cols", v3.index_compaction_columns, snii.index_compaction_columns); + io_line("remote_read_B", v3.remote_physical_read_bytes, snii.remote_physical_read_bytes); + io_line("range_reads", v3.range_read_count, snii.range_read_count); + io_line("serial_rounds", v3.serial_read_rounds, snii.serial_read_rounds); + if (v3.has_hot && snii.has_hot) { + // Should be ~0 remote bytes: if the hot pass still fetches from S3 the cold pass failed + // to write back and the hot number is not a cache-hit measurement. + io_line("hot_remote_read_B", v3.hot_remote_read_bytes, snii.hot_remote_read_bytes); + io_line("hot_range_reads", v3.hot_range_read_count, snii.hot_range_read_count); + } + if (v3.index_bytes < 0 || snii.index_bytes < 0) { + std::cout << std::left << std::setw(18) << "index_bytes" << std::right << std::setw(10) + << "n/a" + << " (remote rowsets expose S3 keys, not local paths)" << std::endl; + } else { + std::cout << std::left << std::setw(18) << "index_bytes" << std::right << std::setw(10) + << v3.index_bytes << std::setw(10) << snii.index_bytes << std::setw(9) + << ratio(static_cast(snii.index_bytes), + static_cast(v3.index_bytes)) + << "x" << std::endl; + } + + // A run where wall clock ran far ahead of CPU was competing for the machine; the ratios + // above are still CPU-based and usable, but say so rather than let it pass silently. + auto flag = [](const char* phase, const Measurement& m) { + if (m.cpu_s > 0.05 && m.wall_s > m.cpu_s * 2.0) { + std::cout << " [warn] " << phase << " wall " << m.wall_s << "s vs cpu " << m.cpu_s + << "s -- machine was busy or the phase was IO bound" << std::endl; + } + }; + flag("V3 import", v3.import); + flag("SNII import", snii.import); + flag("V3 compaction", v3.compaction); + flag("SNII compaction", snii.compaction); + } + + TabletSchemaSPtr _tablet_schema; + StorageEngine* _engine_ref = nullptr; + std::unique_ptr _data_dir; + TabletSharedPtr _tablet; + std::string _absolute_dir; + std::string _current_dir; + std::unique_ptr _searcher_cache; + std::unique_ptr _query_cache; + segment_v2::InvertedIndexSearcherCache* _previous_searcher_cache = nullptr; + segment_v2::InvertedIndexQueryCache* _previous_query_cache = nullptr; + std::shared_ptr _remote_fs; + std::unique_ptr _owned_file_cache_factory; + io::FileCacheFactory* _origin_file_cache_factory = nullptr; + bool _origin_file_cache_factory_saved = false; + int64_t _origin_compaction_batch_size = config::compaction_batch_size; + bool _origin_index_compaction_enable = config::inverted_index_compaction_enable; + bool _origin_ram_dir_enable = config::inverted_index_ram_dir_enable; + std::string _file_cache_dir; + bool _origin_enable_file_cache = false; + std::vector _corpus_files; + int64_t _inc_id = 1000; +}; + +TEST_F(DISABLED_SniiVsV3BenchmarkTest, wikipedia_english_local) { + ASSERT_FALSE(_corpus_files.empty()) + << "no wikipedia_*.json found; set SNII_BENCH_CORPUS_DIR to a corpus directory"; + std::cout << "corpus: " << _corpus_files.size() << " files, io=local" << std::endl; + + // Local rowsets never touch the block file cache, so the policy is irrelevant here. + const auto [v3, snii] = _run_both(IoMode::kLocal, BenchCachePolicy::kDirect); + EXPECT_EQ(v3.matched_docs, snii.matched_docs) + << "V3 and SNII matched a different number of documents, timings are not comparable"; + // A run where compaction or segment loading failed silently produces zeros everywhere and + // still satisfies the equality above as 0 == 0, printing a clean but meaningless report. + EXPECT_GT(v3.matched_docs, 0) << "no documents matched; the run produced no measurement"; + // Zero means the index was rebuilt from raw data instead of compacted (compaction.cpp:1520 + // rejects ineligible SNII postings and falls back), which is a different cost profile + // entirely -- comparing a format that compacted against one that rebuilt measures nothing. + EXPECT_GT(v3.index_compaction_columns, 0) << "V3 did not run index compaction"; + EXPECT_EQ(v3.index_compaction_columns, snii.index_compaction_columns) + << "formats compacted a different number of index columns"; + _report(v3, snii); +} + +// File cache off end to end: load writes straight to S3, compaction re-reads from S3, the query +// reads from S3. Every byte in remote_read_B and every GET in range_reads is physical, which makes +// this the read-amplification comparison. +TEST_F(DISABLED_SniiVsV3BenchmarkTest, wikipedia_english_remote_s3_direct) { + ASSERT_FALSE(_corpus_files.empty()) + << "no wikipedia_*.json found; set SNII_BENCH_CORPUS_DIR to a corpus directory"; + if (!_setup_remote()) { + GTEST_SKIP() << "remote S3 not configured; set SNII_BENCH_S3_{AK,SK,ENDPOINT,REGION,BUCKET}" + << " and make sure HTTP(S)_PROXY is unset"; + } + std::cout << "corpus: " << _corpus_files.size() << " files, io=remote-s3, cache=direct" + << std::endl; + + const auto [v3, snii] = _run_both(IoMode::kRemoteS3, BenchCachePolicy::kDirect); + EXPECT_EQ(v3.matched_docs, snii.matched_docs) + << "V3 and SNII matched a different number of documents, timings are not comparable"; + // A run where compaction or segment loading failed silently produces zeros everywhere and + // still satisfies the equality above as 0 == 0, printing a clean but meaningless report. + EXPECT_GT(v3.matched_docs, 0) << "no documents matched; the run produced no measurement"; + // Zero means the index was rebuilt from raw data instead of compacted (compaction.cpp:1520 + // rejects ineligible SNII postings and falls back), which is a different cost profile + // entirely -- comparing a format that compacted against one that rebuilt measures nothing. + EXPECT_GT(v3.index_compaction_columns, 0) << "V3 did not run index compaction"; + EXPECT_EQ(v3.index_compaction_columns, snii.index_compaction_columns) + << "formats compacted a different number of index columns"; + // The whole point of this mode: if nothing was fetched from S3 the cache was not actually off. + EXPECT_GT(v3.remote_physical_read_bytes, 0) << "direct mode read nothing from S3"; + EXPECT_GT(snii.remote_physical_read_bytes, 0) << "direct mode read nothing from S3"; + _report(v3, snii); + _teardown_remote(); +} + +// File cache on and populated by the load, the way cloud does it (write_file_cache comes from the +// load request). Compaction then reads local SSD, and two query numbers come out: cold (cache +// emptied, query refetches and writes back) and hot (same query against what it just cached). +TEST_F(DISABLED_SniiVsV3BenchmarkTest, wikipedia_english_remote_s3_writeback) { + ASSERT_FALSE(_corpus_files.empty()) + << "no wikipedia_*.json found; set SNII_BENCH_CORPUS_DIR to a corpus directory"; + if (!_setup_remote()) { + GTEST_SKIP() << "remote S3 not configured; set SNII_BENCH_S3_{AK,SK,ENDPOINT,REGION,BUCKET}" + << " and make sure HTTP(S)_PROXY is unset"; + } + std::cout << "corpus: " << _corpus_files.size() << " files, io=remote-s3, cache=write-back" + << std::endl; + + const auto [v3, snii] = _run_both(IoMode::kRemoteS3, BenchCachePolicy::kWriteBack); + EXPECT_EQ(v3.matched_docs, snii.matched_docs) + << "V3 and SNII matched a different number of documents, timings are not comparable"; + // A run where compaction or segment loading failed silently produces zeros everywhere and + // still satisfies the equality above as 0 == 0, printing a clean but meaningless report. + EXPECT_GT(v3.matched_docs, 0) << "no documents matched; the run produced no measurement"; + // Zero means the index was rebuilt from raw data instead of compacted (compaction.cpp:1520 + // rejects ineligible SNII postings and falls back), which is a different cost profile + // entirely -- comparing a format that compacted against one that rebuilt measures nothing. + EXPECT_GT(v3.index_compaction_columns, 0) << "V3 did not run index compaction"; + EXPECT_EQ(v3.index_compaction_columns, snii.index_compaction_columns) + << "formats compacted a different number of index columns"; + _report(v3, snii); + _teardown_remote(); +} + +} // namespace doris diff --git a/be/test/storage/index/snii/common/slice_test.cpp b/be/test/storage/index/snii/common/slice_test.cpp new file mode 100644 index 00000000000000..a86f0a9d7b9aa1 --- /dev/null +++ b/be/test/storage/index/snii/common/slice_test.cpp @@ -0,0 +1,42 @@ +// 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 "storage/index/snii/common/slice.h" + +#include + +#include + +#include "common/status.h" + +using doris::snii::Slice; + +TEST(SniiSlice, BasicAccess) { + const uint8_t buf[] = {1, 2, 3, 4}; + Slice s(buf, 4); + EXPECT_EQ(s.size(), 4U); + EXPECT_EQ(s[2], 3U); + Slice sub = s.subslice(1, 2); + EXPECT_EQ(sub.size(), 2U); + EXPECT_EQ(sub[0], 2U); +} + +TEST(SniiSlice, EmptyDefault) { + Slice s; + EXPECT_TRUE(s.empty()); + EXPECT_EQ(s.size(), 0U); +} diff --git a/be/test/storage/index/snii/common/uninitialized_buffer_test.cpp b/be/test/storage/index/snii/common/uninitialized_buffer_test.cpp new file mode 100644 index 00000000000000..bc498c83e64860 --- /dev/null +++ b/be/test/storage/index/snii/common/uninitialized_buffer_test.cpp @@ -0,0 +1,154 @@ +// 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 "storage/index/snii/common/uninitialized_buffer.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/prx_pod.h" + +using doris::snii::ByteSink; +using doris::snii::ByteSource; +using doris::snii::Slice; +using doris::snii::zstd_compress; +using doris::snii::zstd_decompress; +using doris::snii::format::build_prx_window; +using doris::snii::format::read_prx_window_csr; + +namespace { + +using PerDoc = std::vector>; + +// The uninitialized_vector overload must NOT value-initialize the regrown tail: +// after a sentinel fill + clear + regrow into the SAME storage, the object +// representation still holds the sentinel bytes (proves no zero-fill pass). +TEST(SniiUninitializedBuffer, UninitVectorGrowSkipsZeroFill) { + constexpr size_t kN = 64; + doris::snii::uninitialized_vector v; + doris::snii::resize_uninitialized(v, kN); + for (size_t i = 0; i < kN; ++i) { + v[i] = 0xAAAAAAAAU; + } + const uint32_t* old_data = v.data(); + v.clear(); // trivial element destruction is a no-op; storage/capacity retained + doris::snii::resize_uninitialized(v, kN); + ASSERT_EQ(v.data(), old_data) << "regrow within capacity must not reallocate"; + // Examining the object representation through unsigned char* is well-defined. + const auto* bytes = reinterpret_cast(v.data()); + bool all_sentinel = true; + for (size_t i = 0; i < kN * sizeof(uint32_t); ++i) { + if (bytes[i] != 0xAA) { + all_sentinel = false; + break; + } + } + EXPECT_TRUE(all_sentinel) << "uninitialized_vector regrow must default-init (no zero-fill)"; +} + +// Contrast: a plain std::vector DOES value-initialize the regrown tail to zero. +TEST(SniiUninitializedBuffer, StdVectorGrowZeroFillsForContrast) { + constexpr size_t kN = 64; + std::vector v; + doris::snii::resize_uninitialized(v, kN); // std::vector overload == plain resize + for (size_t i = 0; i < kN; ++i) { + v[i] = 0xAAAAAAAAU; + } + v.clear(); + doris::snii::resize_uninitialized(v, kN); + for (size_t i = 0; i < kN; ++i) { + EXPECT_EQ(v[i], 0U) << "std::vector regrow value-initializes to 0"; + } +} + +// Shrinking a warm buffer must reuse storage (no realloc). +TEST(SniiUninitializedBuffer, ResizeUninitializedShrinkKeepsCapacity) { + std::vector v; + doris::snii::resize_uninitialized(v, 1000); + const size_t cap = v.capacity(); + doris::snii::resize_uninitialized(v, 10); + EXPECT_EQ(v.capacity(), cap) << "shrink must not reallocate"; +} + +// Warm-reused CSR decode (large window then small window into the SAME buffers) +// must equal a fresh decode of the small window -- no stale tail leaks past size(). +TEST(SniiUninitializedBuffer, CsrWarmReuseLargeThenSmallNoStaleTail) { + PerDoc large(64); + for (auto& doc : large) { + uint32_t p = 0; + for (int i = 0; i < 256; ++i) { + doc.push_back(p += 1); + } + } + const PerDoc small = {{1, 2, 3}, {10}, {7, 9}}; + + ByteSink large_sink; + ByteSink small_sink; + ASSERT_TRUE(build_prx_window(large, -1, &large_sink).ok()); + ASSERT_TRUE(build_prx_window(small, -1, &small_sink).ok()); + + std::vector pos_flat; + std::vector pos_off; + { + ByteSource src(large_sink.view()); + ASSERT_TRUE(read_prx_window_csr(&src, &pos_flat, &pos_off).ok()); + } + { + ByteSource src(small_sink.view()); + ASSERT_TRUE(read_prx_window_csr(&src, &pos_flat, &pos_off).ok()); + } + std::vector exp_flat; + std::vector exp_off; + { + ByteSource src(small_sink.view()); + ASSERT_TRUE(read_prx_window_csr(&src, &exp_flat, &exp_off).ok()); + } + EXPECT_EQ(pos_flat, exp_flat) << "warm-reused decode must equal a fresh decode"; + EXPECT_EQ(pos_off, exp_off); + EXPECT_EQ(pos_flat.size(), 6U) << "small window total_pos = 3 + 1 + 2"; +} + +// zstd decode into a reused buffer must not reallocate and stays byte-identical. +TEST(SniiUninitializedBuffer, ZstdReusedBufferNoReallocAndByteIdentical) { + std::string payload(50000, '\0'); + for (size_t i = 0; i < payload.size(); ++i) { + payload[i] = static_cast(i * 7 + 3); + } + std::vector comp; + ASSERT_TRUE(zstd_compress(Slice(payload), 3, &comp).ok()); + + std::vector out; + ASSERT_TRUE(zstd_decompress(Slice(comp), payload.size(), &out).ok()); + ASSERT_EQ(out.size(), payload.size()); + EXPECT_EQ(0, std::memcmp(out.data(), payload.data(), payload.size())); + const size_t cap = out.capacity(); + + ASSERT_TRUE(zstd_decompress(Slice(comp), payload.size(), &out).ok()); + EXPECT_EQ(out.capacity(), cap) << "warm-reused decompress must not reallocate"; + EXPECT_EQ(0, std::memcmp(out.data(), payload.data(), payload.size())); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_compaction_eligibility_test.cpp b/be/test/storage/index/snii/compaction/snii_compaction_eligibility_test.cpp new file mode 100644 index 00000000000000..02acbfa1227dd1 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_compaction_eligibility_test.cpp @@ -0,0 +1,553 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "storage/index/inverted/analyzer/analyzer_provider.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/inverted_index_parser.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/compaction/eligibility.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" +#include "storage/tablet/tablet_schema.h" + +namespace { + +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using doris::TabletIndex; +using doris::TabletIndexPB; +using namespace doris::snii; // NOLINT +namespace inverted_index = doris::segment_v2::inverted_index; + +class BufferFileReader final : public io::FileReader { +public: + explicit BufferFileReader(std::vector bytes) : bytes_(std::move(bytes)) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (offset > bytes_.size() || len > bytes_.size() - offset) { + return Status::Error( + "eligibility fixture: read past EOF"); + } + out->resize(len); + if (len != 0) { + std::memcpy(out->data(), bytes_.data() + offset, len); + } + return Status::OK(); + } + + uint64_t size() const override { return bytes_.size(); } + +private: + std::vector bytes_; +}; + +struct IndexShape { + format::IndexTier tier = format::IndexTier::kT2; + bool has_norms = false; + bool has_common_grams_metadata = false; + std::optional common_grams_metadata = std::nullopt; + format::CommonGramsPostingPolicy common_grams_posting_policy = + format::CommonGramsPostingPolicy::kNone; + format::StatsBlock stats { + .doc_count = 8, + .indexed_doc_count = 6, + .term_count = 0, + .sum_total_term_freq = 0, + .null_count = 2, + }; +}; + +struct OpenedIndex { + std::unique_ptr file; + reader::LogicalIndexReader reader; +}; + +std::unique_ptr open_index(const IndexShape& shape) { + std::vector file_bytes; + format::SectionRefs refs; + if (shape.has_norms || shape.tier == format::IndexTier::kT3) { + format::NormsPodWriter norms; + for (uint64_t doc = 0; doc < shape.stats.doc_count; ++doc) { + norms.add(1); + } + ByteSink norms_frame; + norms.finish(&norms_frame); + file_bytes = norms_frame.buffer(); + refs.norms = {.offset = 0, .length = file_bytes.size()}; + } + + ByteSink sampled_frame; + format::SampledTermIndexBuilder sampled; + sampled.finish(&sampled_frame); + ByteSink directory_frame; + format::DictBlockDirectoryBuilder directory; + directory.finish(&directory_frame); + + format::CoreMetadata core; + core.index_config = shape.tier == format::IndexTier::kT1 ? format::IndexConfig::kDocsOnly + : shape.tier == format::IndexTier::kT2 + ? format::IndexConfig::kDocsPositions + : format::IndexConfig::kDocsPositionsScoring; + core.stats = shape.stats; + core.section_refs = refs; + if (shape.common_grams_metadata.has_value()) { + core.common_grams_metadata = *shape.common_grams_metadata; + } else if (shape.has_common_grams_metadata) { + inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = inverted_index::PlainTermKeyVersion::kRawNoInternal; + core.common_grams_metadata = metadata; + } + if (shape.common_grams_posting_policy != format::CommonGramsPostingPolicy::kNone) { + core.common_grams_posting_policy = shape.common_grams_posting_policy; + } + ByteSink core_frame; + const Status core_status = format::encode_core_metadata(core, &core_frame); + EXPECT_TRUE(core_status.ok()) << core_status.to_string(); + + auto opened = std::make_unique(); + opened->file = std::make_unique(std::move(file_bytes)); + const Status open_status = reader::LogicalIndexReader::open( + opened->file.get(), core_frame.view(), sampled_frame.view(), directory_frame.view(), + &opened->reader); + EXPECT_TRUE(open_status.ok()) << open_status.to_string(); + return opened; +} + +std::unique_ptr make_index(std::map properties, + doris::IndexType type = doris::IndexType::INVERTED, + int64_t index_id = 7, std::string index_suffix = "body") { + TabletIndexPB pb; + pb.set_index_id(index_id); + pb.set_index_name("body_idx"); + pb.set_index_type(type); + pb.set_index_suffix_name(std::move(index_suffix)); + for (auto& [key, value] : properties) { + (*pb.mutable_properties())[key] = std::move(value); + } + auto index = std::make_unique(); + index->init_from_pb(pb); + return index; +} + +std::map plain_properties() { + return {{"lower_case", "true"}, {"parser", "standard"}, {"support_phrase", "true"}}; +} + +std::map common_grams_properties() { + return {{"analyzer", "common_grams_analyzer"}, {"support_phrase", "true"}}; +} + +inverted_index::CommonGramsQueryIdentity common_grams_identity(std::string dictionary = "dict-a") { + return {.common_grams_dictionary_identity = std::move(dictionary), + .base_analyzer_fingerprint = "base-a", + .common_grams_fingerprint = "grams-a"}; +} + +inverted_index::CommonGramsSegmentMetadata complete_common_grams_metadata( + const inverted_index::CommonGramsQueryIdentity& identity = common_grams_identity()) { + auto metadata = inverted_index::make_common_grams_segment_metadata(identity); + metadata.scoring_doc_count = 8; + metadata.scoring_token_count = 0; + return metadata; +} + +inverted_index::CommonGramsSegmentMetadata hybrid_common_grams_metadata( + const inverted_index::CommonGramsQueryIdentity& identity = common_grams_identity()) { + auto metadata = complete_common_grams_metadata(identity); + metadata.common_grams_coverage = inverted_index::CommonGramsCoverage::kMixed; + return metadata; +} + +compaction::PlainT2CompactionSource source(const OpenedIndex& index, + const TabletIndex& index_meta) { + return {.reader = std::cref(index.reader), .index_meta = std::cref(index_meta)}; +} + +void expect_rejected(const Status& status, std::string_view reason) { + EXPECT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find(reason), std::string::npos) << status.to_string(); +} + +class StubAnalyzerProvider final : public inverted_index::AnalyzerProvider { +public: + explicit StubAnalyzerProvider( + bool uses_common_grams, + std::optional identity = std::nullopt) + : uses_common_grams_(uses_common_grams), identity_(std::move(identity)) {} + + std::shared_ptr get_analyzer( + inverted_index::AnalysisPurpose) const override { + return nullptr; + } + bool uses_common_grams() const override { return uses_common_grams_; } + const inverted_index::CommonGramsQueryIdentity* common_grams_identity() const override { + return identity_ ? &*identity_ : nullptr; + } + +private: + bool uses_common_grams_ = false; + std::optional identity_; +}; + +class CommonGramsBuildSwitchGuard { +public: + CommonGramsBuildSwitchGuard() : saved_(doris::config::enable_common_grams_index_build) {} + ~CommonGramsBuildSwitchGuard() { doris::config::enable_common_grams_index_build = saved_; } + +private: + bool saved_; +}; + +TEST(SniiCompactionEligibilityTest, AcceptsIdenticalPlainT2SourcesAndDestination) { + auto first = open_index({}); + auto second = open_index({}); + auto first_meta = make_index(plain_properties()); + auto second_meta = make_index(plain_properties()); + auto destination = make_index(plain_properties()); + std::vector sources {source(*first, *first_meta), source(*second, *second_meta)}; + + EXPECT_TRUE(compaction::validate_plain_t2_compaction_eligibility(sources, *destination).ok()); +} + +TEST(SniiCompactionEligibilityTest, AcceptsHomogeneousCompleteCommonGramsT3Sources) { + CommonGramsBuildSwitchGuard guard; + doris::config::enable_common_grams_index_build = true; + const auto identity = common_grams_identity(); + IndexShape shape {.tier = format::IndexTier::kT3, + .has_norms = true, + .common_grams_metadata = complete_common_grams_metadata(identity)}; + auto first = open_index(shape); + auto second = open_index(shape); + auto first_meta = make_index(common_grams_properties()); + auto second_meta = make_index(common_grams_properties()); + auto destination = make_index(common_grams_properties()); + std::vector sources {source(*first, *first_meta), source(*second, *second_meta)}; + compaction::AnalyzerProviderFactory factory = [identity](const auto&) { + return std::make_shared(true, identity); + }; + + compaction::SniiCompactionEligibility eligibility; + ASSERT_TRUE(compaction::validate_snii_compaction_eligibility(sources, *destination, + &eligibility, factory) + .ok()); + EXPECT_EQ(eligibility.kind, compaction::SniiStreamedMergeKind::kCommonGramsT3); + ASSERT_TRUE(eligibility.common_grams_metadata_seed.has_value()); + EXPECT_TRUE(inverted_index::common_grams_identity_matches( + *eligibility.common_grams_metadata_seed, identity)); + EXPECT_EQ(eligibility.common_grams_metadata_seed->scoring_doc_count, 0U); + EXPECT_EQ(eligibility.common_grams_metadata_seed->scoring_token_count, 0U); + EXPECT_EQ(eligibility.common_grams_posting_policy, format::CommonGramsPostingPolicy::kNone); +} + +TEST(SniiCompactionEligibilityTest, AcceptsHomogeneousHybridCommonGramsT3Sources) { + CommonGramsBuildSwitchGuard guard; + doris::config::enable_common_grams_index_build = true; + const auto identity = common_grams_identity(); + IndexShape shape {.tier = format::IndexTier::kT3, + .has_norms = true, + .common_grams_metadata = hybrid_common_grams_metadata(identity), + .common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1}; + auto first = open_index(shape); + auto second = open_index(shape); + auto first_meta = make_index(common_grams_properties()); + auto second_meta = make_index(common_grams_properties()); + auto destination = make_index(common_grams_properties()); + const std::vector sources {source(*first, *first_meta), source(*second, *second_meta)}; + compaction::AnalyzerProviderFactory factory = [identity](const auto&) { + return std::make_shared(true, identity); + }; + + compaction::SniiCompactionEligibility eligibility; + ASSERT_TRUE(compaction::validate_snii_compaction_eligibility(sources, *destination, + &eligibility, factory) + .ok()); + EXPECT_EQ(eligibility.kind, compaction::SniiStreamedMergeKind::kCommonGramsT3); + ASSERT_TRUE(eligibility.common_grams_metadata_seed.has_value()); + EXPECT_EQ(eligibility.common_grams_metadata_seed->common_grams_coverage, + inverted_index::CommonGramsCoverage::kMixed); + EXPECT_EQ(eligibility.common_grams_posting_policy, format::CommonGramsPostingPolicy::kHybridV1); + EXPECT_TRUE(compaction::validate_snii_source_eligibility(first->reader, + /*source_ordinal=*/0, eligibility) + .ok()); +} + +TEST(SniiCompactionEligibilityTest, RejectsMixedMismatchedAndBuildDisabledCommonGramsSources) { + CommonGramsBuildSwitchGuard guard; + doris::config::enable_common_grams_index_build = true; + const auto identity = common_grams_identity(); + IndexShape common_shape {.tier = format::IndexTier::kT3, + .has_norms = true, + .common_grams_metadata = complete_common_grams_metadata(identity)}; + auto common = open_index(common_shape); + auto plain = open_index({}); + auto first_meta = make_index(common_grams_properties()); + auto second_meta = make_index(common_grams_properties()); + auto destination = make_index(common_grams_properties()); + compaction::AnalyzerProviderFactory factory = [identity](const auto&) { + return std::make_shared(true, identity); + }; + compaction::SniiCompactionEligibility eligibility; + + std::vector mixed {source(*common, *first_meta), source(*plain, *second_meta)}; + expect_rejected(compaction::validate_snii_compaction_eligibility(mixed, *destination, + &eligibility, factory), + "homogeneous"); + + const auto other_identity = common_grams_identity("dict-b"); + common_shape.common_grams_metadata = complete_common_grams_metadata(other_identity); + auto mismatched = open_index(common_shape); + std::vector mismatch_sources {source(*common, *first_meta), source(*mismatched, *second_meta)}; + expect_rejected(compaction::validate_snii_compaction_eligibility(mismatch_sources, *destination, + &eligibility, factory), + "identity"); + + IndexShape hybrid_shape { + .tier = format::IndexTier::kT3, + .has_norms = true, + .common_grams_metadata = hybrid_common_grams_metadata(identity), + .common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1}; + auto hybrid = open_index(hybrid_shape); + std::vector policy_mismatch_sources {source(*common, *first_meta), + source(*hybrid, *second_meta)}; + expect_rejected(compaction::validate_snii_compaction_eligibility( + policy_mismatch_sources, *destination, &eligibility, factory), + "homogeneous"); + + doris::config::enable_common_grams_index_build = false; + std::vector disabled_sources {source(*common, *first_meta)}; + expect_rejected(compaction::validate_snii_compaction_eligibility(disabled_sources, *destination, + &eligibility, factory), + "disabled"); +} + +TEST(SniiCompactionEligibilityTest, ExposesReusableSourceShapeValidation) { + auto eligible = open_index({}); + EXPECT_TRUE(compaction::validate_plain_t2_source(eligible->reader, /*source_ordinal=*/3).ok()); + + IndexShape invalid_shape; + invalid_shape.stats.indexed_doc_count = 7; + invalid_shape.stats.null_count = 2; + auto invalid = open_index(invalid_shape); + expect_rejected(compaction::validate_plain_t2_source(invalid->reader, /*source_ordinal=*/3), + "source 3"); +} + +TEST(SniiCompactionEligibilityTest, RejectsPhysicalShapesOutsidePlainT2) { + struct Case { + IndexShape shape; + std::string_view reason; + }; + std::vector cases; + cases.push_back({IndexShape {.tier = format::IndexTier::kT1}, "T2"}); + cases.push_back({IndexShape {.has_norms = true}, "norms"}); + + for (const auto& test_case : cases) { + auto index = open_index(test_case.shape); + auto source_meta = make_index(plain_properties()); + auto destination = make_index(plain_properties()); + std::vector sources {source(*index, *source_meta)}; + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(sources, *destination), + test_case.reason); + } +} + +TEST(SniiCompactionEligibilityTest, RejectsCompanionMetadata) { + auto index = open_index(IndexShape {.has_common_grams_metadata = true}); + auto source_meta = make_index(plain_properties()); + auto destination = make_index(plain_properties()); + std::vector sources {source(*index, *source_meta)}; + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(sources, *destination), + "CommonGrams"); +} + +TEST(SniiCompactionEligibilityTest, RejectsLegacyBigramMarkerBeforeMergeExecution) { + snii_test::MemoryFile file; + writer::SniiIndexInput input; + input.index_id = 7; + input.index_suffix = "body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 1; + input.terms.push_back( + snii_test::make_term(std::string(format::kPhraseBigramTermMarker) + "left\x1Fright", + {{.docid = 0, .positions = {1}}})); + writer::SniiCompoundWriter compound_writer(&file); + ASSERT_TRUE(compound_writer.add_logical_index(input).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + reader::SniiSegmentReader segment; + ASSERT_TRUE(reader::SniiSegmentReader::open(&file, &segment).ok()); + reader::LogicalIndexReader index; + ASSERT_TRUE(segment.open_index(7, "body", &index).ok()); + + auto source_meta = make_index(plain_properties()); + auto destination = make_index(plain_properties()); + std::vector sources { + {.reader = std::cref(index), .index_meta = std::cref(*source_meta)}}; + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(sources, *destination), + "legacy phrase-bigram"); +} + +TEST(SniiCompactionEligibilityTest, RejectsBrokenDocumentStatistics) { + IndexShape shape; + shape.stats.indexed_doc_count = 7; + shape.stats.null_count = 2; + auto index = open_index(shape); + auto source_meta = make_index(plain_properties()); + auto destination = make_index(plain_properties()); + std::vector sources {source(*index, *source_meta)}; + + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(sources, *destination), + "statistics"); +} + +TEST(SniiCompactionEligibilityTest, RequiresExactSourceAndDestinationProperties) { + auto first = open_index({}); + auto second = open_index({}); + auto first_meta = make_index(plain_properties()); + + auto changed_properties = plain_properties(); + changed_properties["dict_compression"] = "true"; + auto changed_meta = make_index(changed_properties); + auto destination = make_index(plain_properties()); + std::vector mismatched_sources {source(*first, *first_meta), source(*second, *changed_meta)}; + expect_rejected( + compaction::validate_plain_t2_compaction_eligibility(mismatched_sources, *destination), + "properties"); + + auto matching_second_meta = make_index(plain_properties()); + std::vector matching_sources {source(*first, *first_meta), + source(*second, *matching_second_meta)}; + expect_rejected( + compaction::validate_plain_t2_compaction_eligibility(matching_sources, *changed_meta), + "destination properties"); +} + +TEST(SniiCompactionEligibilityTest, RequiresDestinationLogicalIndexIdentity) { + auto index = open_index({}); + auto destination = make_index(plain_properties()); + + auto wrong_id = make_index(plain_properties(), doris::IndexType::INVERTED, /*index_id=*/8); + std::vector wrong_id_sources {source(*index, *wrong_id)}; + expect_rejected( + compaction::validate_plain_t2_compaction_eligibility(wrong_id_sources, *destination), + "index id"); + + auto wrong_suffix = make_index(plain_properties(), doris::IndexType::INVERTED, + /*index_id=*/7, /*index_suffix=*/"other"); + std::vector wrong_suffix_sources {source(*index, *wrong_suffix)}; + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(wrong_suffix_sources, + *destination), + "index suffix"); +} + +TEST(SniiCompactionEligibilityTest, RequiresInvertedPhraseDestination) { + auto index = open_index({}); + auto no_phrase_properties = plain_properties(); + no_phrase_properties["support_phrase"] = "false"; + auto no_phrase_source_meta = make_index(no_phrase_properties); + auto no_phrase = make_index(no_phrase_properties); + std::vector no_phrase_sources {source(*index, *no_phrase_source_meta)}; + expect_rejected( + compaction::validate_plain_t2_compaction_eligibility(no_phrase_sources, *no_phrase), + "phrase positions"); + + auto source_meta = make_index(plain_properties()); + std::vector sources {source(*index, *source_meta)}; + auto non_inverted = make_index(plain_properties(), doris::IndexType::BITMAP); + expect_rejected(compaction::validate_plain_t2_compaction_eligibility(sources, *non_inverted), + "inverted index"); +} + +TEST(SniiCompactionEligibilityTest, RejectsDestinationThatWouldBuildCommonGrams) { + CommonGramsBuildSwitchGuard guard; + doris::config::enable_common_grams_index_build = true; + auto index = open_index({}); + auto properties = plain_properties(); + properties.erase("parser"); + properties["analyzer"] = "common_grams_analyzer"; + auto source_meta = make_index(properties); + auto destination = make_index(properties); + std::vector sources {source(*index, *source_meta)}; + + compaction::AnalyzerProviderFactory factory = [](const doris::InvertedIndexAnalyzerConfig&) { + return std::make_shared(true); + }; + expect_rejected( + compaction::validate_plain_t2_compaction_eligibility(sources, *destination, factory), + "CommonGrams"); + + doris::config::enable_common_grams_index_build = false; + EXPECT_TRUE(compaction::validate_plain_t2_compaction_eligibility(sources, *destination, factory) + .ok()); +} + +TEST(SniiCompactionEligibilityTest, ResolvesDestinationProviderFromWriterAnalyzerConfig) { + auto index = open_index({}); + auto properties = plain_properties(); + properties["analyzer"] = "plain_custom_analyzer"; + properties["parser_mode"] = "fine_grained"; + properties["stopwords"] = "none"; + properties["char_filter_type"] = "char_replace"; + properties["char_filter_pattern"] = "_"; + properties["char_filter_replacement"] = " "; + auto source_meta = make_index(properties); + auto destination = make_index(properties); + std::vector sources {source(*index, *source_meta)}; + std::optional captured; + + compaction::AnalyzerProviderFactory factory = + [&](const doris::InvertedIndexAnalyzerConfig& config) { + captured = config; + return std::make_shared(false); + }; + ASSERT_TRUE(compaction::validate_plain_t2_compaction_eligibility(sources, *destination, factory) + .ok()); + ASSERT_TRUE(captured.has_value()); + EXPECT_EQ(captured->analyzer_name, "plain_custom_analyzer"); + EXPECT_EQ(captured->parser_type, doris::InvertedIndexParserType::PARSER_STANDARD); + EXPECT_EQ(captured->parser_mode, "fine_grained"); + EXPECT_EQ(captured->lower_case, "true"); + EXPECT_EQ(captured->stop_words, "none"); + EXPECT_EQ(captured->char_filter_map.at("char_filter_pattern"), "_"); + EXPECT_EQ(captured->char_filter_map.at("char_filter_replacement"), " "); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_index_compaction_test.cpp b/be/test/storage/index/snii/compaction/snii_index_compaction_test.cpp new file mode 100644 index 00000000000000..bb95bb4d4440e6 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_index_compaction_test.cpp @@ -0,0 +1,1284 @@ +// 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 "storage/index/snii/compaction/snii_index_compaction.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/compaction/posting_run_merger.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::snii_test; // NOLINT +using compaction::RowIdConversionMap; +using compaction::SniiPlainT2MergePlan; +using compaction::ValidatedRowIdConversion; +using query::phrase_query; +using query::term_query; +using writer::SniiCompoundWriter; +using writer::SniiIndexInput; +using writer::SniiStreamedIndexSession; +using writer::TermPostings; +using writer::MemoryReporter; +namespace inverted_index = doris::segment_v2::inverted_index; + +constexpr uint64_t kIndexId = 41; +constexpr std::string_view kIndexSuffix = "body"; +constexpr std::pair kDeleted = {std::numeric_limits::max(), + std::numeric_limits::max()}; + +struct OpenedIndex { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; +}; + +SniiIndexInput make_input(uint32_t doc_count, std::vector null_docids, + std::vector terms) { + SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = kIndexSuffix; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = doc_count; + input.null_docids = std::move(null_docids); + input.terms = std::move(terms); + input.write_freq = false; + return input; +} + +inverted_index::CommonGramsSegmentMetadata common_grams_metadata(uint64_t doc_count, + uint64_t token_count) { + inverted_index::CommonGramsQueryIdentity identity {.common_grams_dictionary_identity = "dict-a", + .base_analyzer_fingerprint = "base-a", + .common_grams_fingerprint = "grams-a"}; + auto metadata = inverted_index::make_common_grams_segment_metadata(identity); + metadata.scoring_doc_count = doc_count; + metadata.scoring_token_count = token_count; + return metadata; +} + +std::string common_gram(std::string_view left, std::string_view right) { + auto encoded = inverted_index::encode_common_gram(left, right); + EXPECT_TRUE(encoded.has_value()); + return encoded.has_value() ? std::move(encoded.value()) : std::string(); +} + +SniiIndexInput make_common_grams_input(uint32_t doc_count, std::vector null_docids, + std::vector norms, std::vector terms, + uint64_t token_count) { + SniiIndexInput input = make_input(doc_count, std::move(null_docids), std::move(terms)); + input.config = format::IndexConfig::kDocsPositionsScoring; + input.write_freq = true; + input.encoded_norms = std::move(norms); + input.common_grams_metadata = common_grams_metadata(doc_count, token_count); + return input; +} + +TermPostings make_docs_only_gram(std::string term, std::vector docids) { + TermPostings postings; + postings.term = std::move(term); + postings.docids = std::move(docids); + postings.retain_positions = false; + return postings; +} + +SniiIndexInput make_hybrid_common_grams_input(uint32_t doc_count, std::vector null_docids, + std::vector norms, + std::vector terms, + uint64_t token_count) { + std::ranges::sort(terms, [](const auto& lhs, const auto& rhs) { return lhs.term < rhs.term; }); + SniiIndexInput input = make_common_grams_input(doc_count, std::move(null_docids), + std::move(norms), std::move(terms), token_count); + input.common_grams_metadata->common_grams_coverage = + inverted_index::CommonGramsCoverage::kMixed; + input.common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1; + return input; +} + +void build_index(SniiIndexInput input, OpenedIndex* out, + reader::LogicalIndexOpenMode open_mode = reader::LogicalIndexOpenMode::kQuery) { + SniiCompoundWriter compound(&out->file); + assert_ok(compound.add_logical_index(input)); + assert_ok(compound.finish()); + assert_ok(reader::SniiSegmentReader::open(&out->file, &out->segment)); + assert_ok(out->segment.open_index(kIndexId, kIndexSuffix, &out->index, open_mode)); +} + +std::vector copy_region(const MemoryFile& file, const format::RegionRef& region) { + EXPECT_LE(region.offset + region.length, file.data().size()); + if (region.offset + region.length > file.data().size()) { + return {}; + } + return {file.data().begin() + static_cast(region.offset), + file.data().begin() + static_cast(region.offset + region.length)}; +} + +void expect_identical_index_image(MemoryFile* actual, MemoryFile* expected) { + reader::SniiSegmentReader actual_segment; + reader::SniiSegmentReader expected_segment; + reader::LogicalIndexReader actual_index; + reader::LogicalIndexReader expected_index; + assert_ok(reader::SniiSegmentReader::open(actual, &actual_segment)); + assert_ok(reader::SniiSegmentReader::open(expected, &expected_segment)); + assert_ok(actual_segment.open_index(kIndexId, kIndexSuffix, &actual_index)); + assert_ok(expected_segment.open_index(kIndexId, kIndexSuffix, &expected_index)); + + const format::SectionRefs& actual_refs = actual_index.section_refs(); + const format::SectionRefs& expected_refs = expected_index.section_refs(); + auto expect_region = [&](std::string_view name, const format::RegionRef& actual_region, + const format::RegionRef& expected_region) { + SCOPED_TRACE(name); + EXPECT_EQ(actual_region.offset, expected_region.offset); + EXPECT_EQ(actual_region.length, expected_region.length); + EXPECT_EQ(copy_region(*actual, actual_region), copy_region(*expected, expected_region)); + }; + expect_region("posting", actual_refs.posting_region, expected_refs.posting_region); + expect_region("DICT", actual_refs.dict_region, expected_refs.dict_region); + expect_region("norms", actual_refs.norms, expected_refs.norms); + expect_region("null bitmap", actual_refs.null_bitmap, expected_refs.null_bitmap); + expect_region("BSBF", actual_refs.bsbf, expected_refs.bsbf); + + EXPECT_EQ(actual->data(), expected->data()) << "complete SNII image differs"; +} + +std::unique_ptr make_validated_conversion( + const RowIdConversionMap* conversion, const std::vector& source_rows, + const std::vector& destination_rows) { + std::unique_ptr validated; + const Status status = + ValidatedRowIdConversion::create(conversion, source_rows, destination_rows, &validated); + EXPECT_TRUE(status.ok()) << status; + return validated; +} + +struct MergedPostingCopy { + uint32_t destination = 0; + uint32_t docid = 0; + uint32_t frequency = 0; + std::vector positions; + + bool operator==(const MergedPostingCopy&) const = default; +}; + +struct MergedRunHarness { + std::vector> read_contexts; + std::vector> cursors; +}; + +Status make_merged_run_harness(const std::vector& sources, std::string_view term, + const ValidatedRowIdConversion* rowid_conversion, + MergedRunHarness* harness) { + for (size_t source_ordinal = 0; source_ordinal < sources.size(); ++source_ordinal) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR( + sources[source_ordinal]->index.lookup(term, &found, &entry, &frq_base, &prx_base)); + if (!found) { + continue; + } + auto read_context = std::make_unique( + &sources[source_ordinal]->index, /*total_read_ahead_budget_bytes=*/1U << 20); + RETURN_IF_ERROR(read_context->init()); + auto cursor = std::make_unique( + read_context.get(), std::move(entry), frq_base, prx_base, + static_cast(source_ordinal), rowid_conversion); + RETURN_IF_ERROR(cursor->init()); + harness->read_contexts.push_back(std::move(read_context)); + harness->cursors.push_back(std::move(cursor)); + } + return Status::OK(); +} + +Status drain_merged_runs(compaction::MergedPostingRuns* runs, uint32_t max_run_docs, + std::vector* output) { + while (!runs->empty()) { + const uint32_t destination = runs->next_destination(); + RETURN_IF_ERROR(runs->begin_destination(destination)); + while (true) { + writer::PostingRunView run; + bool has_run = false; + RETURN_IF_ERROR(runs->next_run(max_run_docs, &run, &has_run)); + if (!has_run) { + break; + } + EXPECT_FALSE(run.docids.empty()); + EXPECT_TRUE(run.freqs.empty() || run.freqs.size() == run.docids.size()); + EXPECT_TRUE(run.position_offsets.empty() || + run.position_offsets.size() == run.docids.size() + 1); + if (!run.position_offsets.empty()) { + EXPECT_EQ(run.position_offsets.back() - run.position_offsets.front(), + run.positions_flat.size()); + } + for (size_t ordinal = 0; ordinal < run.docids.size(); ++ordinal) { + std::vector positions; + if (!run.position_offsets.empty()) { + const uint32_t position_base = run.position_offsets.front(); + const uint32_t begin = run.position_offsets[ordinal] - position_base; + const uint32_t end = run.position_offsets[ordinal + 1] - position_base; + positions.assign(run.positions_flat.begin() + begin, + run.positions_flat.begin() + end); + } + output->push_back({.destination = destination, + .docid = run.docids[ordinal], + .frequency = run.freqs.empty() ? 0 : run.freqs[ordinal], + .positions = std::move(positions)}); + } + } + } + return Status::OK(); +} + +std::vector merge_posting_oracle( + const std::vector>& source_postings, + const RowIdConversionMap& conversion) { + std::vector output; + for (size_t source = 0; source < source_postings.size(); ++source) { + for (const PostingDoc& posting : source_postings[source]) { + const auto [destination, docid] = conversion[source][posting.docid]; + if (std::pair(destination, docid) == kDeleted) { + continue; + } + output.push_back({.destination = destination, + .docid = docid, + .frequency = static_cast(posting.positions.size()), + .positions = posting.positions}); + } + } + std::ranges::sort(output, [](const MergedPostingCopy& lhs, const MergedPostingCopy& rhs) { + return std::pair(lhs.destination, lhs.docid) < std::pair(rhs.destination, rhs.docid); + }); + return output; +} + +std::vector source_zero_terms() { + return { + make_term("alpha", {{.docid = 0, .positions = {0}}, {.docid = 3, .positions = {0, 2}}}), + make_term("beta", {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {3}}}), + make_term("gamma", {{.docid = 0, .positions = {2}}, {.docid = 2, .positions = {1}}}), + }; +} + +std::vector source_one_terms() { + return { + make_term("alpha", {{.docid = 0, .positions = {0}}}), + make_term("beta", {{.docid = 1, .positions = {0}}}), + make_term("gamma", {{.docid = 0, .positions = {1}}}), + }; +} + +std::vector destination_zero_terms() { + return { + make_term("alpha", {{.docid = 0, .positions = {0}}, {.docid = 1, .positions = {0}}}), + make_term("beta", {{.docid = 0, .positions = {1}}, {.docid = 3, .positions = {0}}}), + make_term("gamma", {{.docid = 0, .positions = {2}}, {.docid = 1, .positions = {1}}}), + }; +} + +std::vector destination_one_terms() { + return { + make_term("beta", {{.docid = 0, .positions = {0}}}), + make_term("gamma", {{.docid = 0, .positions = {1}}}), + }; +} + +void expect_query_semantics(const reader::LogicalIndexReader& index, + const std::vector& alpha, + const std::vector& alpha_beta, + const std::vector& alpha_beta_gamma) { + std::vector docs; + assert_ok(term_query(index, "alpha", &docs)); + EXPECT_EQ(docs, alpha); + assert_ok(phrase_query(index, {"alpha", "beta"}, &docs)); + EXPECT_EQ(docs, alpha_beta); + assert_ok(phrase_query(index, {"alpha", "beta", "gamma"}, &docs)); + EXPECT_EQ(docs, alpha_beta_gamma); +} + +TEST(SniiIndexCompactionTest, DirectPostingRunsMatchPositionedOracleAcrossSources) { + const std::vector source_zero_postings = { + {.docid = 0, .positions = {1, 4}}, {.docid = 2, .positions = {8}}, + {.docid = 3, .positions = {2}}, {.docid = 4, .positions = {3, 8, 9}}, + {.docid = 6, .positions = {6}}, + }; + const std::vector source_one_postings = { + {.docid = 0, .positions = {0}}, {.docid = 1, .positions = {1, 5}}, + {.docid = 2, .positions = {4}}, {.docid = 4, .positions = {11}}, + {.docid = 5, .positions = {7, 10}}, + }; + OpenedIndex source_zero; + OpenedIndex source_one; + build_index(make_input(/*doc_count=*/7, /*null_docids=*/ {}, + {make_term("shared", source_zero_postings)}), + &source_zero, reader::LogicalIndexOpenMode::kCompaction); + build_index(make_input(/*doc_count=*/6, /*null_docids=*/ {}, + {make_term("shared", source_one_postings)}), + &source_one, reader::LogicalIndexOpenMode::kCompaction); + + const RowIdConversionMap conversion = { + {{0, 0}, {0, 2}, kDeleted, {0, 4}, {1, 0}, {1, 2}, {1, 4}}, + {{0, 1}, {0, 3}, {0, 5}, {1, 1}, kDeleted, {1, 3}}, + }; + const std::vector destination_rows = {6, 5}; + auto validated = make_validated_conversion(&conversion, {7, 6}, destination_rows); + ASSERT_NE(validated, nullptr); + + MergedRunHarness harness; + assert_ok(make_merged_run_harness({&source_zero, &source_one}, "shared", validated.get(), + &harness)); + std::vector semantic_token_counts(destination_rows.size(), 0); + compaction::testing::reset_posting_run_merge_counters(); + compaction::MergedPostingRuns runs(std::move(harness.cursors), /*retain_positions=*/true, + /*counts_as_semantic_token=*/false, destination_rows, + semantic_token_counts); + assert_ok(runs.init()); + std::vector got; + assert_ok(drain_merged_runs(&runs, /*max_run_docs=*/2, &got)); + + const std::vector expected = + merge_posting_oracle({source_zero_postings, source_one_postings}, conversion); + EXPECT_EQ(got, expected); + EXPECT_EQ(compaction::testing::posting_run_documents(), expected.size()); + EXPECT_EQ(compaction::testing::posting_run_shape_scan_documents(), expected.size()); + EXPECT_GT(compaction::testing::posting_run_emitted_runs(), 1U); + EXPECT_EQ(compaction::testing::posting_run_legacy_fill_calls(), 0U); + EXPECT_EQ(compaction::testing::posting_run_copied_documents(), 0U); +} + +TEST(SniiIndexCompactionTest, DirectPostingRunsMatchSingleSourceDocsOnlyOracle) { + const std::string gram = common_gram("a", "term"); + OpenedIndex source; + build_index(make_hybrid_common_grams_input( + /*doc_count=*/5, /*null_docids=*/ {}, /*norms=*/ {1, 1, 1, 1, 1}, + {make_docs_only_gram(gram, {0, 2, 4}), + make_term("plain", {{.docid = 1, .positions = {0}}})}, + /*token_count=*/1), + &source, reader::LogicalIndexOpenMode::kCompaction); + + const RowIdConversionMap conversion = {{{0, 0}, {0, 1}, {0, 2}, {1, 0}, {1, 1}}}; + const std::vector destination_rows = {3, 2}; + auto validated = make_validated_conversion(&conversion, {5}, destination_rows); + ASSERT_NE(validated, nullptr); + + MergedRunHarness harness; + assert_ok(make_merged_run_harness({&source}, gram, validated.get(), &harness)); + std::vector semantic_token_counts(destination_rows.size(), 0); + compaction::testing::reset_posting_run_merge_counters(); + compaction::MergedPostingRuns runs(std::move(harness.cursors), /*retain_positions=*/false, + /*counts_as_semantic_token=*/false, destination_rows, + semantic_token_counts); + assert_ok(runs.init()); + std::vector got; + assert_ok(drain_merged_runs(&runs, /*max_run_docs=*/1, &got)); + + EXPECT_EQ(got, (std::vector { + {0, 0, 0, {}}, + {0, 2, 0, {}}, + {1, 1, 0, {}}, + })); + EXPECT_EQ(compaction::testing::posting_run_shape_scan_documents(), got.size()); + EXPECT_EQ(compaction::testing::posting_run_legacy_fill_calls(), 0U); + EXPECT_EQ(compaction::testing::posting_run_copied_documents(), 0U); +} + +TEST(SniiIndexCompactionTest, MergesDeletesNullsAndInterleavedSourcesByteIdenticallyToRebuild) { + OpenedIndex source_zero; + OpenedIndex source_one; + build_index(make_input(/*doc_count=*/4, /*null_docids=*/ {1}, source_zero_terms()), + &source_zero, reader::LogicalIndexOpenMode::kCompaction); + build_index(make_input(/*doc_count=*/3, /*null_docids=*/ {2}, source_one_terms()), &source_one); + + const RowIdConversionMap conversion = { + {{0, 0}, {0, 2}, {1, 0}, kDeleted}, + {{0, 1}, {0, 3}, {1, 1}}, + }; + const std::vector destination_rows = {4, 2}; + auto validated = make_validated_conversion(&conversion, {4, 3}, destination_rows); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source_zero.index, &source_one.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + ASSERT_NE(plan, nullptr); + EXPECT_EQ(plan->destination_null_docids(0), (std::vector {2})); + EXPECT_EQ(plan->destination_null_docids(1), (std::vector {1})); + + std::array merged_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t i = 0; i < compounds.size(); ++i) { + compounds[i] = std::make_unique(&merged_files[i]); + SniiIndexInput input = + make_input(destination_rows[i], plan->destination_null_docids(i), {}); + assert_ok(compounds[i]->begin_streamed_index(std::move(input), &sessions[i])); + } + assert_ok(plan->execute(sessions)); + for (auto& compound : compounds) { + assert_ok(compound->finish()); + } + + std::array expected; + build_index(make_input(/*doc_count=*/4, /*null_docids=*/ {2}, destination_zero_terms()), + &expected[0]); + build_index(make_input(/*doc_count=*/2, /*null_docids=*/ {1}, destination_one_terms()), + &expected[1]); + expect_identical_index_image(&merged_files[0], &expected[0].file); + expect_identical_index_image(&merged_files[1], &expected[1].file); + + std::array merged_segments; + std::array merged_indexes; + for (size_t i = 0; i < merged_files.size(); ++i) { + assert_ok(reader::SniiSegmentReader::open(&merged_files[i], &merged_segments[i])); + assert_ok(merged_segments[i].open_index(kIndexId, kIndexSuffix, &merged_indexes[i])); + } + expect_query_semantics(merged_indexes[0], /*alpha=*/ {0, 1}, /*alpha_beta=*/ {0}, + /*alpha_beta_gamma=*/ {0}); + expect_query_semantics(merged_indexes[1], /*alpha=*/ {}, /*alpha_beta=*/ {}, + /*alpha_beta_gamma=*/ {}); + EXPECT_EQ(merged_indexes[0].stats().doc_count, 4U); + EXPECT_EQ(merged_indexes[0].stats().indexed_doc_count, 3U); + EXPECT_EQ(merged_indexes[0].stats().null_count, 1U); + EXPECT_EQ(merged_indexes[1].stats().doc_count, 2U); + EXPECT_EQ(merged_indexes[1].stats().indexed_doc_count, 1U); + EXPECT_EQ(merged_indexes[1].stats().null_count, 1U); +} + +TEST(SniiIndexCompactionTest, MergesTwentyFourRunSourcesByteIdenticallyToReferenceRebuild) { + constexpr size_t kSourceCount = 24; + constexpr uint32_t kDocsPerSource = 400; + constexpr uint32_t kSourceRunDocs = 17; + + std::array sources; + std::array, kSourceCount> source_docs; + for (size_t source = 0; source < kSourceCount; ++source) { + source_docs[source].reserve(kDocsPerSource); + for (uint32_t docid = 0; docid < kDocsPerSource; ++docid) { + const uint32_t first_position = static_cast((source + docid) % 19); + source_docs[source].push_back( + {.docid = docid, .positions = {first_position, first_position + 3}}); + } + std::vector terms; + if (source == 0) { + terms.push_back(make_term("first-destination-only", {{.docid = 0, .positions = {1}}})); + } + terms.push_back(make_term("shared", source_docs[source])); + build_index(make_input(kDocsPerSource, {}, std::move(terms)), &sources[source], + reader::LogicalIndexOpenMode::kCompaction); + } + + RowIdConversionMap conversion(kSourceCount); + for (auto& source_mapping : conversion) { + source_mapping.assign(kDocsPerSource, kDeleted); + } + std::vector> reference_order; + reference_order.reserve(kSourceCount * kDocsPerSource); + for (uint32_t run_begin = 0; run_begin < kDocsPerSource; run_begin += kSourceRunDocs) { + const uint32_t run_end = std::min(kDocsPerSource, run_begin + kSourceRunDocs); + for (size_t source = 0; source < kSourceCount; ++source) { + for (uint32_t docid = run_begin; docid < run_end; ++docid) { + const bool deleted = source == kSourceCount - 1 || + (source * 31 + static_cast(docid) * 17) % 113 == 7; + if (!deleted) { + reference_order.emplace_back(source, docid); + } + } + } + } + ASSERT_GT(reference_order.size(), 8500U); + const uint32_t first_destination_rows = 8500; + const std::vector destination_rows = { + first_destination_rows, + static_cast(reference_order.size() - first_destination_rows)}; + std::array, 2> reference_docs; + for (size_t ordinal = 0; ordinal < reference_order.size(); ++ordinal) { + const auto [source, source_docid] = reference_order[ordinal]; + const uint32_t destination = ordinal < first_destination_rows ? 0 : 1; + const uint32_t destination_docid = + destination == 0 ? static_cast(ordinal) + : static_cast(ordinal - first_destination_rows); + conversion[source][source_docid] = {destination, destination_docid}; + reference_docs[destination].push_back( + {.docid = destination_docid, + .positions = source_docs[source][source_docid].positions}); + } + ASSERT_EQ(conversion[0][0].first, 0U); + + std::vector source_indexes; + source_indexes.reserve(kSourceCount); + for (const OpenedIndex& source : sources) { + source_indexes.push_back(&source.index); + } + auto validated = make_validated_conversion( + &conversion, std::vector(kSourceCount, kDocsPerSource), destination_rows); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare( + std::move(source_indexes), *validated, + kSourceCount * SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource, &plan)); + + std::array merged_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t destination = 0; destination < compounds.size(); ++destination) { + compounds[destination] = std::make_unique(&merged_files[destination]); + assert_ok(compounds[destination]->begin_streamed_index( + make_input(destination_rows[destination], {}, {}), &sessions[destination])); + } + + compaction::testing::reset_posting_run_merge_counters(); + assert_ok(plan->execute(sessions)); + EXPECT_GT(compaction::testing::posting_run_frontier_updates(), 0U); + EXPECT_LE(compaction::testing::posting_run_frontier_updates(), + compaction::testing::posting_run_emitted_runs()); + EXPECT_GT(compaction::testing::posting_run_frontier_comparisons(), 0U); + EXPECT_EQ(compaction::testing::posting_run_documents(), reference_order.size() + 1); + EXPECT_LE(compaction::testing::posting_run_boundary_searches(), + compaction::testing::posting_run_emitted_runs()); + for (auto& compound : compounds) { + assert_ok(compound->finish()); + } + + for (size_t destination = 0; destination < merged_files.size(); ++destination) { + std::vector expected_terms; + if (destination == 0) { + expected_terms.push_back( + make_term("first-destination-only", + {{.docid = conversion[0][0].second, .positions = {1}}})); + } + expected_terms.push_back(make_term("shared", reference_docs[destination])); + OpenedIndex expected; + build_index(make_input(destination_rows[destination], {}, std::move(expected_terms)), + &expected); + expect_identical_index_image(&merged_files[destination], &expected.file); + } +} + +TEST(SniiIndexCompactionTest, CommonGramsMergeMatchesRebuildAfterDeletesAndRemap) { + const std::string gram = common_gram("the", "cat"); + OpenedIndex source_zero; + OpenedIndex source_one; + build_index(make_common_grams_input( + /*doc_count=*/3, /*null_docids=*/ {2}, /*norms=*/ {10, 20, 30}, + {make_term(gram, {{.docid = 0, .positions = {1}}, + {.docid = 1, .positions = {1}}}), + make_term("alpha", {{.docid = 0, .positions = {0, 2}}, + {.docid = 1, .positions = {0}}}), + make_term("beta", {{.docid = 2, .positions = {0}}})}, + /*token_count=*/4), + &source_zero, reader::LogicalIndexOpenMode::kCompaction); + build_index(make_common_grams_input( + /*doc_count=*/2, /*null_docids=*/ {}, /*norms=*/ {40, 50}, + {make_term(gram, {{.docid = 1, .positions = {1}}}), + make_term("alpha", {{.docid = 0, .positions = {0}}}), + make_term("gamma", {{.docid = 1, .positions = {0, 2}}})}, + /*token_count=*/3), + &source_one, reader::LogicalIndexOpenMode::kCompaction); + + const RowIdConversionMap conversion = { + {{0, 0}, kDeleted, {1, 0}}, + {{0, 1}, {0, 2}}, + }; + const std::vector destination_rows = {3, 1}; + auto validated = make_validated_conversion(&conversion, {3, 2}, destination_rows); + ASSERT_NE(validated, nullptr); + compaction::SniiCompactionEligibility eligibility { + .kind = compaction::SniiStreamedMergeKind::kCommonGramsT3, + .common_grams_metadata_seed = common_grams_metadata(0, 0)}; + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source_zero.index, &source_one.index}, *validated, + eligibility, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + ASSERT_NE(plan, nullptr); + EXPECT_EQ(plan->destination_null_docids(0), (std::vector {})); + EXPECT_EQ(plan->destination_null_docids(1), (std::vector {0})); + EXPECT_EQ(plan->destination_encoded_norms(0), (std::vector {10, 40, 50})); + EXPECT_EQ(plan->destination_encoded_norms(1), (std::vector {30})); + + std::array merged_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t i = 0; i < compounds.size(); ++i) { + compounds[i] = std::make_unique(&merged_files[i]); + SniiIndexInput input = make_input(destination_rows[i], {}, {}); + input.config = plan->destination_index_config(); + input.write_freq = true; + input.common_grams_metadata = plan->destination_common_grams_metadata(i); + assert_ok(compounds[i]->begin_streamed_index( + std::move(input), plan->take_destination_null_docids(i), + plan->take_destination_encoded_norms(i), &sessions[i])); + } + assert_ok(plan->execute(sessions)); + for (auto& compound : compounds) { + assert_ok(compound->finish()); + } + + std::array rebuilt; + build_index(make_common_grams_input( + /*doc_count=*/3, /*null_docids=*/ {}, /*norms=*/ {10, 40, 50}, + {make_term(gram, {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {1}}}), + make_term("alpha", {{.docid = 0, .positions = {0, 2}}, + {.docid = 1, .positions = {0}}}), + make_term("gamma", {{.docid = 2, .positions = {0, 2}}})}, + /*token_count=*/5), + &rebuilt[0]); + build_index(make_common_grams_input( + /*doc_count=*/1, /*null_docids=*/ {0}, /*norms=*/ {30}, + {make_term("beta", {{.docid = 0, .positions = {0}}})}, + /*token_count=*/1), + &rebuilt[1]); + expect_identical_index_image(&merged_files[0], &rebuilt[0].file); + expect_identical_index_image(&merged_files[1], &rebuilt[1].file); + + std::array merged_segments; + std::array merged_indexes; + for (size_t i = 0; i < merged_files.size(); ++i) { + assert_ok(reader::SniiSegmentReader::open(&merged_files[i], &merged_segments[i])); + assert_ok(merged_segments[i].open_index(kIndexId, kIndexSuffix, &merged_indexes[i])); + ASSERT_NE(merged_indexes[i].common_grams_metadata(), nullptr); + EXPECT_EQ(*merged_indexes[i].common_grams_metadata(), + *rebuilt[i].index.common_grams_metadata()); + } + EXPECT_EQ(merged_indexes[0].common_grams_metadata()->scoring_token_count, 5U); + EXPECT_EQ(merged_indexes[1].common_grams_metadata()->scoring_token_count, 1U); + format::NormsPodReader first_norms; + format::NormsPodReader second_norms; + assert_ok(merged_indexes[0].open_norms(&first_norms)); + assert_ok(merged_indexes[1].open_norms(&second_norms)); + EXPECT_EQ(first_norms.encoded_norm(0), 10U); + EXPECT_EQ(first_norms.encoded_norm(1), 40U); + EXPECT_EQ(first_norms.encoded_norm(2), 50U); + EXPECT_EQ(second_norms.encoded_norm(0), 30U); + + std::vector merged_docs; + std::vector rebuilt_docs; + assert_ok(term_query(merged_indexes[0], "alpha", &merged_docs)); + assert_ok(term_query(rebuilt[0].index, "alpha", &rebuilt_docs)); + EXPECT_EQ(merged_docs, rebuilt_docs); + assert_ok(phrase_query(merged_indexes[0], {"alpha", "gamma"}, &merged_docs)); + assert_ok(phrase_query(rebuilt[0].index, {"alpha", "gamma"}, &rebuilt_docs)); + EXPECT_EQ(merged_docs, rebuilt_docs); +} + +TEST(SniiIndexCompactionTest, HybridCommonGramsMergePreservesPostingShapesAfterDeleteAndSplit) { + const std::string positioned_gram = common_gram("the", "of"); + const std::string docs_only_gram = common_gram("of", "wolf"); + const auto source_terms = [&] { + return std::vector { + make_term(positioned_gram, {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}}), + make_docs_only_gram(docs_only_gram, {0, 1, 2}), + make_term("of", {{.docid = 0, .positions = {1}}, + {.docid = 1, .positions = {1}}, + {.docid = 2, .positions = {1}}}), + make_term("the", {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}}), + make_term("wolf", {{.docid = 0, .positions = {2}}, + {.docid = 1, .positions = {2}}, + {.docid = 2, .positions = {2}}})}; + }; + OpenedIndex source; + build_index(make_hybrid_common_grams_input( + /*doc_count=*/4, /*null_docids=*/ {3}, /*norms=*/ {10, 20, 30, 40}, + source_terms(), /*token_count=*/9), + &source, reader::LogicalIndexOpenMode::kCompaction); + + const RowIdConversionMap conversion = {{{0, 0}, kDeleted, {1, 0}, {1, 1}}}; + const std::vector destination_rows = {1, 2}; + auto validated = make_validated_conversion(&conversion, {4}, destination_rows); + ASSERT_NE(validated, nullptr); + compaction::SniiCompactionEligibility eligibility { + .kind = compaction::SniiStreamedMergeKind::kCommonGramsT3, + .common_grams_metadata_seed = + [] { + auto metadata = common_grams_metadata(0, 0); + metadata.common_grams_coverage = + inverted_index::CommonGramsCoverage::kMixed; + return metadata; + }(), + .common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1}; + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, eligibility, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + + std::array merged_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t i = 0; i < compounds.size(); ++i) { + compounds[i] = std::make_unique(&merged_files[i]); + SniiIndexInput input = make_input(destination_rows[i], {}, {}); + input.config = plan->destination_index_config(); + input.write_freq = true; + input.common_grams_metadata = plan->destination_common_grams_metadata(i); + input.common_grams_posting_policy = plan->destination_common_grams_posting_policy(); + assert_ok(compounds[i]->begin_streamed_index( + std::move(input), plan->take_destination_null_docids(i), + plan->take_destination_encoded_norms(i), &sessions[i])); + } + assert_ok(plan->execute(sessions)); + for (auto& compound : compounds) { + assert_ok(compound->finish()); + } + + std::array merged_segments; + std::array merged_indexes; + for (size_t i = 0; i < merged_indexes.size(); ++i) { + assert_ok(reader::SniiSegmentReader::open(&merged_files[i], &merged_segments[i])); + assert_ok(merged_segments[i].open_index(kIndexId, kIndexSuffix, &merged_indexes[i])); + ASSERT_NE(merged_indexes[i].common_grams_metadata(), nullptr); + EXPECT_EQ(merged_indexes[i].common_grams_metadata()->common_grams_coverage, + inverted_index::CommonGramsCoverage::kMixed); + EXPECT_EQ(merged_indexes[i].common_grams_posting_policy(), + format::CommonGramsPostingPolicy::kHybridV1); + + bool found = false; + format::DictEntry positioned_entry; + format::DictEntry docs_only_entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(merged_indexes[i].lookup(positioned_gram, &found, &positioned_entry, &frq_base, + &prx_base)); + ASSERT_TRUE(found); + assert_ok(merged_indexes[i].lookup(docs_only_gram, &found, &docs_only_entry, &frq_base, + &prx_base)); + ASSERT_TRUE(found); + EXPECT_TRUE(compaction::posting_entry_has_positions(positioned_entry)); + EXPECT_FALSE(compaction::posting_entry_has_positions(docs_only_entry)); + + std::vector docs; + assert_ok(term_query(merged_indexes[i], docs_only_gram, &docs)); + EXPECT_EQ(docs, (std::vector {0})); + assert_ok(phrase_query(merged_indexes[i], {"the", "of", "wolf"}, &docs)); + EXPECT_EQ(docs, (std::vector {0})); + } + EXPECT_EQ(merged_indexes[0].common_grams_metadata()->scoring_token_count, 3U); + EXPECT_EQ(merged_indexes[1].common_grams_metadata()->scoring_token_count, 3U); +} + +TEST(SniiIndexCompactionTest, HybridCommonGramsRejectsSameTermPositionShapeMismatch) { + const std::string gram = common_gram("the", "of"); + OpenedIndex positioned; + OpenedIndex docs_only; + build_index(make_hybrid_common_grams_input(1, {}, {1}, + {make_term(gram, {{.docid = 0, .positions = {0}}}), + make_term("the", {{.docid = 0, .positions = {0}}})}, + 1), + &positioned, reader::LogicalIndexOpenMode::kCompaction); + build_index(make_hybrid_common_grams_input(1, {}, {1}, + {make_docs_only_gram(gram, {0}), + make_term("the", {{.docid = 0, .positions = {0}}})}, + 1), + &docs_only, reader::LogicalIndexOpenMode::kCompaction); + + const RowIdConversionMap conversion = {{{0, 0}}, {{0, 1}}}; + auto validated = make_validated_conversion(&conversion, {1, 1}, {2}); + ASSERT_NE(validated, nullptr); + compaction::SniiCompactionEligibility eligibility { + .kind = compaction::SniiStreamedMergeKind::kCommonGramsT3, + .common_grams_metadata_seed = + [] { + auto metadata = common_grams_metadata(0, 0); + metadata.common_grams_coverage = + inverted_index::CommonGramsCoverage::kMixed; + return metadata; + }(), + .common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1}; + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&positioned.index, &docs_only.index}, *validated, + eligibility, /*total_read_ahead_budget_bytes=*/1U << 20, + &plan)); + + MemoryFile merged_file; + SniiCompoundWriter compound(&merged_file); + SniiIndexInput input = make_input(2, {}, {}); + input.config = plan->destination_index_config(); + input.write_freq = true; + input.common_grams_metadata = plan->destination_common_grams_metadata(0); + input.common_grams_posting_policy = plan->destination_common_grams_posting_policy(); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), plan->take_destination_null_docids(0), + plan->take_destination_encoded_norms(0), &session)); + const Status status = plan->execute(std::span(&session, 1)); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiIndexCompactionTest, CommonGramsWindowedMergeKeepsGramDocsOnlyAndPlainScoring) { + constexpr uint32_t kDocCount = 600; + const std::string gram = common_gram("the", "cat"); + std::vector gram_docs; + std::vector plain_docs; + gram_docs.reserve(kDocCount); + plain_docs.reserve(kDocCount); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + gram_docs.push_back({.docid = docid, .positions = {1}}); + plain_docs.push_back({.docid = docid, .positions = {0, 2}}); + } + + OpenedIndex source; + build_index(make_common_grams_input(kDocCount, /*null_docids=*/ {}, + /*norms=*/std::vector(kDocCount, 1), + {make_term(gram, std::move(gram_docs)), + make_term("alpha", std::move(plain_docs))}, + /*token_count=*/2 * kDocCount), + &source, reader::LogicalIndexOpenMode::kCompaction); + + RowIdConversionMap conversion(1); + conversion[0].reserve(kDocCount); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + conversion[0].emplace_back(0, docid); + } + auto validated = make_validated_conversion(&conversion, {kDocCount}, {kDocCount}); + ASSERT_NE(validated, nullptr); + compaction::SniiCompactionEligibility eligibility { + .kind = compaction::SniiStreamedMergeKind::kCommonGramsT3, + .common_grams_metadata_seed = common_grams_metadata(0, 0)}; + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, eligibility, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + ASSERT_NE(plan, nullptr); + + MemoryFile merged_file; + SniiCompoundWriter compound(&merged_file); + SniiStreamedIndexSession* session = nullptr; + SniiIndexInput input = make_input(kDocCount, {}, {}); + input.config = plan->destination_index_config(); + input.write_freq = true; + input.common_grams_metadata = plan->destination_common_grams_metadata(0); + assert_ok(compound.begin_streamed_index(std::move(input), plan->take_destination_null_docids(0), + plan->take_destination_encoded_norms(0), &session)); + std::array sessions = {session}; + assert_ok(plan->execute(sessions)); + assert_ok(compound.finish()); + + reader::SniiSegmentReader merged_segment; + reader::LogicalIndexReader merged_index; + assert_ok(reader::SniiSegmentReader::open(&merged_file, &merged_segment)); + assert_ok(merged_segment.open_index(kIndexId, kIndexSuffix, &merged_index)); + + for (const reader::LogicalIndexReader* index : {&source.index, &merged_index}) { + bool found = false; + format::DictEntry gram_entry; + format::DictEntry plain_entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index->lookup(gram, &found, &gram_entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + assert_ok(index->lookup("alpha", &found, &plain_entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + + EXPECT_EQ(gram_entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(gram_entry.enc, format::DictEntryEnc::kWindowed); + EXPECT_FALSE(gram_entry.term_stats_present); + EXPECT_EQ(gram_entry.frq_len, gram_entry.frq_docs_len); + EXPECT_EQ(plain_entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(plain_entry.enc, format::DictEntryEnc::kWindowed); + EXPECT_TRUE(plain_entry.term_stats_present); + EXPECT_EQ(plain_entry.ttf_delta, 2 * kDocCount); + EXPECT_EQ(plain_entry.max_freq, 2U); + EXPECT_GT(plain_entry.frq_len, plain_entry.frq_docs_len); + } + + ASSERT_NE(merged_index.common_grams_metadata(), nullptr); + EXPECT_EQ(merged_index.common_grams_metadata()->common_grams_coverage, + inverted_index::CommonGramsCoverage::kComplete); + EXPECT_EQ(merged_index.common_grams_metadata()->scoring_token_count, 2 * kDocCount); +} + +TEST(SniiIndexCompactionTest, CommonGramsMergeReclaimsResidentDictBeforeLargePlainTerm) { + constexpr size_t kGramPositions = 4096; + constexpr size_t kPlainPositions = 16384; + auto make_positions = [](size_t count, uint32_t salt) { + std::vector positions; + positions.reserve(count); + uint32_t position = salt; + for (size_t i = 0; i < count; ++i) { + position += 17 + static_cast((i * 257 + salt) & 1023); + positions.push_back(position); + } + return positions; + }; + + std::vector source_terms; + for (uint32_t i = 0; i < 32; ++i) { + source_terms.push_back( + make_term(common_gram("the", std::string("gram-") + std::to_string(100 + i)), + {{.docid = 0, .positions = make_positions(kGramPositions, i + 1)}})); + } + source_terms.push_back(make_term( + "zeta", {{.docid = 0, .positions = make_positions(kPlainPositions, /*salt=*/99)}})); + + SniiIndexInput source_input = + make_common_grams_input(/*doc_count=*/1, /*null_docids=*/ {}, /*norms=*/ {1}, + std::move(source_terms), /*token_count=*/kPlainPositions); + source_input.target_dict_block_bytes = 1; + OpenedIndex source; + build_index(std::move(source_input), &source, reader::LogicalIndexOpenMode::kCompaction); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + + constexpr size_t kReadAhead = SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource; + constexpr size_t kHardCap = 256U << 10; + auto reporter = std::make_shared(nullptr, kHardCap); + compaction::SniiCompactionEligibility eligibility { + .kind = compaction::SniiStreamedMergeKind::kCommonGramsT3, + .common_grams_metadata_seed = common_grams_metadata(0, 0)}; + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, eligibility, kReadAhead, + reporter, &plan)); + ASSERT_NE(plan, nullptr); + + MemoryFile merged_file; + SniiCompoundWriter compound(&merged_file); + SniiIndexInput input = make_input(/*doc_count=*/1, {}, {}); + input.config = plan->destination_index_config(); + input.write_freq = true; + input.target_dict_block_bytes = 1; + input.dict_resident_cap_bytes = kHardCap / 8; + input.mem_reporter = reporter.get(); + input.common_grams_metadata = plan->destination_common_grams_metadata(0); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), plan->take_destination_null_docids(0), + plan->take_destination_encoded_norms(0), &session)); + std::array sessions = {session}; + assert_ok(plan->execute(sessions)); + assert_ok(compound.finish()); + + reader::SniiSegmentReader merged_segment; + reader::LogicalIndexReader merged_index; + assert_ok(reader::SniiSegmentReader::open(&merged_file, &merged_segment)); + assert_ok(merged_segment.open_index(kIndexId, kIndexSuffix, &merged_index)); + std::vector docs; + assert_ok(term_query(merged_index, "zeta", &docs)); + EXPECT_EQ(docs, (std::vector {0})); +} + +TEST(SniiIndexCompactionTest, DecodesWindowedSourceAndReencodesEachDestination) { + std::vector source_docs; + source_docs.reserve(600); + for (uint32_t docid = 0; docid < 600; ++docid) { + source_docs.push_back({.docid = docid, .positions = {1, 4, 9}}); + } + OpenedIndex source; + build_index(make_input(/*doc_count=*/600, /*null_docids=*/ {}, + {make_term("wide", std::move(source_docs))}), + &source); + bool found = false; + format::DictEntry source_entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(source.index.lookup("wide", &found, &source_entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + ASSERT_EQ(source_entry.enc, format::DictEntryEnc::kWindowed); + + RowIdConversionMap conversion(1); + conversion[0].reserve(600); + for (uint32_t docid = 0; docid < 600; ++docid) { + conversion[0].emplace_back(docid / 300, docid % 300); + } + const std::vector destination_rows = {300, 300}; + auto validated = make_validated_conversion(&conversion, {600}, destination_rows); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + + std::array merged_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t i = 0; i < compounds.size(); ++i) { + compounds[i] = std::make_unique(&merged_files[i]); + SniiIndexInput input = make_input(destination_rows[i], {}, {}); + assert_ok(compounds[i]->begin_streamed_index(std::move(input), &sessions[i])); + } + compaction::testing::reset_posting_run_merge_counters(); + assert_ok(plan->execute(sessions)); + EXPECT_EQ(compaction::testing::posting_run_frontier_comparisons(), 0U); + EXPECT_GT(compaction::testing::posting_run_documents(), 0U); + for (auto& compound : compounds) { + assert_ok(compound->finish()); + } + + std::vector expected_docs; + expected_docs.reserve(300); + for (uint32_t docid = 0; docid < 300; ++docid) { + expected_docs.push_back({.docid = docid, .positions = {1, 4, 9}}); + } + for (size_t i = 0; i < merged_files.size(); ++i) { + OpenedIndex expected; + build_index(make_input(/*doc_count=*/300, /*null_docids=*/ {}, + {make_term("wide", expected_docs)}), + &expected); + expect_identical_index_image(&merged_files[i], &expected.file); + } +} + +TEST(SniiIndexCompactionTest, RejectsInsufficientBudgetBeforeCreatingPlan) { + OpenedIndex source; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("alpha", {{.docid = 0, .positions = {0}}})}), + &source); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + const Status status = SniiPlainT2MergePlan::prepare({&source.index}, *validated, + /*total_read_ahead_budget_bytes=*/1, &plan); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(plan, nullptr); +} + +TEST(SniiIndexCompactionTest, ReadAheadHardReservationFailureUnwindsAllCharges) { + OpenedIndex source; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("alpha", {{.docid = 0, .positions = {0}}})}), + &source); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + + constexpr size_t kReadAhead = SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource; + auto reporter = std::make_shared(nullptr, kReadAhead - 1); + std::unique_ptr plan; + const Status status = + SniiPlainT2MergePlan::prepare({&source.index}, *validated, kReadAhead, reporter, &plan); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(plan, nullptr); + EXPECT_EQ(reporter->current_bytes(), 0); +} + +TEST(SniiIndexCompactionTest, NullRemapMemoryIsReservedBeforeReadAhead) { + constexpr uint32_t kDocCount = 128; + std::vector null_docids; + RowIdConversionMap conversion(1); + null_docids.reserve(kDocCount); + conversion[0].reserve(kDocCount); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + null_docids.push_back(docid); + conversion[0].emplace_back(0, docid); + } + + OpenedIndex source; + build_index(make_input(kDocCount, std::move(null_docids), {}), &source); + auto validated = make_validated_conversion(&conversion, {kDocCount}, {kDocCount}); + ASSERT_NE(validated, nullptr); + + constexpr size_t kReadAhead = SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource; + auto reporter = std::make_shared(nullptr, kReadAhead); + std::unique_ptr plan; + const Status status = + SniiPlainT2MergePlan::prepare({&source.index}, *validated, kReadAhead, reporter, &plan); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(plan, nullptr); + EXPECT_EQ(reporter->current_bytes(), 0); +} + +TEST(SniiIndexCompactionTest, NullRemapExactCapIsRetainedUntilPlanDestruction) { + constexpr uint32_t kDocCount = 128; + std::vector null_docids; + RowIdConversionMap conversion(1); + null_docids.reserve(kDocCount); + conversion[0].reserve(kDocCount); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + null_docids.push_back(docid); + conversion[0].emplace_back(0, docid); + } + + OpenedIndex source; + build_index(make_input(kDocCount, null_docids, {}), &source); + auto validated = make_validated_conversion(&conversion, {kDocCount}, {kDocCount}); + ASSERT_NE(validated, nullptr); + + constexpr size_t kReadAhead = SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource; + constexpr size_t kRetainedNullBytes = kDocCount * sizeof(uint32_t); + constexpr size_t kExactCap = kReadAhead + kRetainedNullBytes; + auto exact_reporter = std::make_shared(nullptr, kExactCap); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, kReadAhead, exact_reporter, + &plan)); + ASSERT_NE(plan, nullptr); + EXPECT_EQ(plan->destination_null_docids(0), null_docids); + EXPECT_EQ(exact_reporter->current_bytes(), static_cast(kExactCap)); + writer::TrackedNullDocids transferred = plan->take_destination_null_docids(0); + EXPECT_EQ(transferred.size(), null_docids.size()); + plan.reset(); + EXPECT_EQ(exact_reporter->current_bytes(), static_cast(kRetainedNullBytes)); + transferred.release(); + EXPECT_EQ(exact_reporter->current_bytes(), 0); + + auto insufficient_reporter = std::make_shared(nullptr, kExactCap - 1); + const Status status = SniiPlainT2MergePlan::prepare({&source.index}, *validated, kReadAhead, + insufficient_reporter, &plan); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(plan, nullptr); + EXPECT_EQ(insufficient_reporter->current_bytes(), 0); +} + +TEST(SniiIndexCompactionTest, ExecuteBudgetFailureAbortsSessionAndKeepsChargeBalanced) { + OpenedIndex source; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("alpha", {{.docid = 0, .positions = {0}}})}), + &source); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + + constexpr size_t kReadAhead = SniiPlainT2MergePlan::kMinReadAheadBudgetPerSource; + auto reporter = std::make_shared(nullptr, kReadAhead); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, kReadAhead, reporter, + &plan)); + ASSERT_NE(plan, nullptr); + EXPECT_EQ(reporter->current_bytes(), static_cast(kReadAhead)); + + MemoryFile destination_file; + SniiCompoundWriter compound(&destination_file); + SniiIndexInput destination_input = make_input(/*doc_count=*/1, /*null_docids=*/ {}, {}); + destination_input.mem_reporter = reporter.get(); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(destination_input), &session)); + ASSERT_NE(session, nullptr); + + const Status status = plan->execute(std::span(&session, 1)); + EXPECT_TRUE(status.is()) << status; + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_TRUE(plan->execute(std::span(&session, 1)).is()); + EXPECT_EQ(reporter->current_bytes(), static_cast(kReadAhead)); + + plan.reset(); + EXPECT_EQ(reporter->current_bytes(), 0); +} + +TEST(SniiIndexCompactionTest, ReusesValidatedConversionAndRejectsLogicalIndexDocCountMismatch) { + OpenedIndex first; + OpenedIndex second; + OpenedIndex mismatched; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("first", {{.docid = 0, .positions = {0}}})}), + &first); + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("second", {{.docid = 0, .positions = {0}}})}), + &second); + build_index(make_input(/*doc_count=*/2, /*null_docids=*/ {}, + {make_term("mismatch", {{.docid = 1, .positions = {0}}})}), + &mismatched); + + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + std::unique_ptr first_plan; + std::unique_ptr second_plan; + std::unique_ptr mismatched_plan; + assert_ok(SniiPlainT2MergePlan::prepare({&first.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, + &first_plan)); + assert_ok(SniiPlainT2MergePlan::prepare({&second.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, + &second_plan)); + const Status mismatch_status = SniiPlainT2MergePlan::prepare( + {&mismatched.index}, *validated, /*total_read_ahead_budget_bytes=*/1U << 20, + &mismatched_plan); + EXPECT_TRUE(mismatch_status.is()) << mismatch_status; + EXPECT_EQ(mismatched_plan, nullptr); +} + +TEST(SniiIndexCompactionTest, ExecuteFailureAbortsEveryDestinationSession) { + OpenedIndex source; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("alpha", {{.docid = 0, .positions = {0}}})}), + &source); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + + std::array destination_files; + std::array, 2> compounds; + std::array sessions = {nullptr, nullptr}; + for (size_t i = 0; i < compounds.size(); ++i) { + compounds[i] = std::make_unique(&destination_files[i]); + assert_ok(compounds[i]->begin_streamed_index(make_input(1, {}, {}), &sessions[i])); + } + + const Status status = plan->execute(sessions); + EXPECT_TRUE(status.is()) << status; + for (size_t i = 0; i < sessions.size(); ++i) { + EXPECT_TRUE(sessions[i]->finish().is()); + EXPECT_TRUE(compounds[i]->finish().is()); + EXPECT_FALSE(destination_files[i].finalized()); + } +} + +TEST(SniiIndexCompactionTest, StickyExecuteFailureAbortsNewDestinationSession) { + OpenedIndex source; + build_index(make_input(/*doc_count=*/1, /*null_docids=*/ {}, + {make_term("alpha", {{.docid = 0, .positions = {0}}})}), + &source); + const RowIdConversionMap conversion = {{{0, 0}}}; + auto validated = make_validated_conversion(&conversion, {1}, {1}); + ASSERT_NE(validated, nullptr); + std::unique_ptr plan; + assert_ok(SniiPlainT2MergePlan::prepare({&source.index}, *validated, + /*total_read_ahead_budget_bytes=*/1U << 20, &plan)); + + MemoryFile first_file; + SniiCompoundWriter first_compound(&first_file); + SniiStreamedIndexSession* first_session = nullptr; + assert_ok(first_compound.begin_streamed_index(make_input(1, {}, {}), &first_session)); + std::array mismatched_sessions = {first_session, first_session}; + const Status first_status = plan->execute(mismatched_sessions); + EXPECT_TRUE(first_status.is()) << first_status; + + MemoryFile retry_file; + SniiCompoundWriter retry_compound(&retry_file); + SniiStreamedIndexSession* retry_session = nullptr; + assert_ok(retry_compound.begin_streamed_index(make_input(1, {}, {}), &retry_session)); + const Status retry_status = plan->execute(std::span(&retry_session, 1)); + EXPECT_EQ(retry_status.to_string(), first_status.to_string()); + EXPECT_TRUE(retry_session->finish().is()); + EXPECT_TRUE(retry_compound.finish().is()); + EXPECT_FALSE(retry_file.finalized()); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_logical_null_reader_test.cpp b/be/test/storage/index/snii/compaction/snii_logical_null_reader_test.cpp new file mode 100644 index 00000000000000..459666668ef358 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_logical_null_reader_test.cpp @@ -0,0 +1,289 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" + +namespace { + +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using namespace doris::snii; // NOLINT + +class BufferFileReader final : public io::FileReader { +public: + explicit BufferFileReader(std::vector bytes) : bytes_(std::move(bytes)) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (offset > bytes_.size() || len > bytes_.size() - offset) { + return Status::Error( + "buffer reader: read past EOF"); + } + size_t actual = len; + if (short_read_ && actual != 0) { + --actual; + } + out->resize(actual); + if (actual != 0) { + std::memcpy(out->data(), bytes_.data() + offset, actual); + } + return Status::OK(); + } + + uint64_t size() const override { return bytes_.size(); } + void set_short_read(bool value) { short_read_ = value; } + +private: + std::vector bytes_; + bool short_read_ = false; +}; + +struct LogicalImage { + std::vector file; + std::vector core; + std::vector sampled_term_index; + std::vector dict_block_directory; +}; + +std::vector build_null_frame(const std::vector& null_docids, + uint32_t recorded_doc_count) { + format::NullBitmapWriter writer; + for (uint32_t docid : null_docids) { + writer.add_null(docid); + } + ByteSink sink; + EXPECT_TRUE(writer.finish(recorded_doc_count, &sink).ok()); + return sink.buffer(); +} + +std::vector reframe(Slice frame, uint8_t type, bool append_payload_byte) { + ByteSource source(frame); + FramedSection section; + if (!SectionFramer::read(source, §ion).ok()) { + return {}; + } + std::vector payload(section.payload.data(), + section.payload.data() + section.payload.size()); + if (append_payload_byte) { + payload.push_back(0); + } + ByteSink sink; + SectionFramer::write(sink, type, Slice(payload)); + return sink.buffer(); +} + +LogicalImage build_logical_image(uint64_t stats_doc_count, uint64_t stats_null_count, + std::optional> null_frame, + std::optional null_ref = std::nullopt) { + LogicalImage image; + format::SectionRefs refs; + if (null_frame.has_value()) { + image.file = std::move(*null_frame); + refs.null_bitmap = {0, image.file.size()}; + } + if (null_ref.has_value()) { + refs.null_bitmap = *null_ref; + } + + ByteSink sampled_frame; + format::SampledTermIndexBuilder sampled; + sampled.finish(&sampled_frame); + ByteSink directory_frame; + format::DictBlockDirectoryBuilder directory; + directory.finish(&directory_frame); + + format::StatsBlock stats; + stats.doc_count = stats_doc_count; + stats.indexed_doc_count = stats_doc_count - stats_null_count; + stats.null_count = stats_null_count; + + format::CoreMetadata core; + core.index_config = format::IndexConfig::kDocsPositions; + core.stats = stats; + core.section_refs = refs; + ByteSink core_frame; + EXPECT_TRUE(format::encode_core_metadata(core, &core_frame).ok()); + image.core = core_frame.buffer(); + image.sampled_term_index = sampled_frame.buffer(); + image.dict_block_directory = directory_frame.buffer(); + return image; +} + +Status open_logical(BufferFileReader* file, const LogicalImage& image, + reader::LogicalIndexReader* out) { + return reader::LogicalIndexReader::open(file, Slice(image.core), + Slice(image.sampled_term_index), + Slice(image.dict_block_directory), out); +} + +TEST(SniiLogicalNullReaderTest, ReturnsSparseNullDocidsInAscendingOrder) { + constexpr uint32_t kDocCount = 1000000000; + LogicalImage image = build_logical_image( + kDocCount, 4, build_null_frame({999999999, 7, 500000, 1}, kDocCount)); + BufferFileReader file(std::move(image.file)); + reader::LogicalIndexReader index; + ASSERT_TRUE(open_logical(&file, image, &index).ok()); + + std::vector null_docids; + ASSERT_TRUE(index.read_null_docids(&null_docids).ok()); + EXPECT_EQ(null_docids, (std::vector {1, 7, 500000, 999999999})); +} + +TEST(SniiLogicalNullReaderTest, MissingSectionIsEmptyOnlyWhenStatsAreEmpty) { + for (uint64_t null_count : {0, 1}) { + LogicalImage image = build_logical_image(/*stats_doc_count=*/8, null_count, std::nullopt); + BufferFileReader file(std::move(image.file)); + reader::LogicalIndexReader index; + ASSERT_TRUE(open_logical(&file, image, &index).ok()); + + std::vector null_docids {99}; + Status status = index.read_null_docids(&null_docids); + if (null_count == 0) { + EXPECT_TRUE(status.ok()) << status.to_string(); + EXPECT_TRUE(null_docids.empty()); + } else { + EXPECT_TRUE(status.is()); + EXPECT_TRUE(null_docids.empty()); + } + } +} + +TEST(SniiLogicalNullReaderTest, AccountsEncodedEmptyBitmapFrameWhenStatsAreEmpty) { + std::vector frame = build_null_frame({}, /*recorded_doc_count=*/8); + const uint64_t frame_bytes = frame.size(); + uint64_t expected_decode_bytes = 0; + ASSERT_TRUE(format::NullBitmapReader::decoded_memory_bytes(Slice(frame), &expected_decode_bytes) + .ok()); + LogicalImage image = + build_logical_image(/*stats_doc_count=*/8, /*stats_null_count=*/0, std::move(frame)); + BufferFileReader file(std::move(image.file)); + reader::LogicalIndexReader index; + ASSERT_TRUE(open_logical(&file, image, &index).ok()); + + reader::NullDocidsScanMemory memory; + ASSERT_TRUE(index.null_docids_scan_memory(&memory).ok()); + EXPECT_EQ(memory.frame_bytes, frame_bytes); + EXPECT_EQ(memory.output_bytes, 0); + + uint64_t reserved_decode_bytes = 0; + std::vector null_docids; + ASSERT_TRUE(index.read_null_docids(&null_docids, [&](uint64_t bytes) { + reserved_decode_bytes = bytes; + return Status::OK(); + }).ok()); + EXPECT_TRUE(null_docids.empty()); + EXPECT_EQ(reserved_decode_bytes, expected_decode_bytes); +} + +TEST(SniiLogicalNullReaderTest, RejectsRegionPastFileAndShortRead) { + std::vector frame = build_null_frame({2}, 8); + LogicalImage past_eof = build_logical_image( + /*stats_doc_count=*/8, /*stats_null_count=*/1, frame, + format::RegionRef {.offset = 0, .length = frame.size() + 1}); + BufferFileReader past_eof_file(std::move(past_eof.file)); + reader::LogicalIndexReader past_eof_index; + ASSERT_TRUE(open_logical(&past_eof_file, past_eof, &past_eof_index).ok()); + reader::NullDocidsScanMemory memory; + EXPECT_TRUE(past_eof_index.null_docids_scan_memory(&memory) + .is()); + EXPECT_EQ(memory.frame_bytes, 0); + EXPECT_EQ(memory.output_bytes, 0); + std::vector null_docids; + EXPECT_TRUE(past_eof_index.read_null_docids(&null_docids) + .is()); + + LogicalImage short_read = build_logical_image( + /*stats_doc_count=*/8, /*stats_null_count=*/1, build_null_frame({2}, 8)); + BufferFileReader short_read_file(std::move(short_read.file)); + reader::LogicalIndexReader short_read_index; + ASSERT_TRUE(open_logical(&short_read_file, short_read, &short_read_index).ok()); + short_read_file.set_short_read(true); + EXPECT_TRUE(short_read_index.read_null_docids(&null_docids) + .is()); +} + +TEST(SniiLogicalNullReaderTest, RejectsInvalidSectionEnvelope) { + std::vector corrupt_crc = build_null_frame({2}, 8); + corrupt_crc.back() ^= 0xff; + std::vector wrong_type = + reframe(Slice(build_null_frame({2}, 8)), format::kNullBitmapSectionType + 1, + /*append_payload_byte=*/false); + std::vector trailing_payload = + reframe(Slice(build_null_frame({2}, 8)), format::kNullBitmapSectionType, + /*append_payload_byte=*/true); + std::vector trailing_frame = build_null_frame({2}, 8); + trailing_frame.push_back(0); + + std::vector> frames {std::move(corrupt_crc), std::move(wrong_type), + std::move(trailing_payload), + std::move(trailing_frame)}; + for (std::vector& frame : frames) { + const uint64_t frame_size = frame.size(); + LogicalImage image = build_logical_image( + /*stats_doc_count=*/8, /*stats_null_count=*/1, std::move(frame), + format::RegionRef {.offset = 0, .length = frame_size}); + BufferFileReader file(std::move(image.file)); + reader::LogicalIndexReader index; + ASSERT_TRUE(open_logical(&file, image, &index).ok()); + std::vector null_docids; + EXPECT_TRUE(index.read_null_docids(&null_docids) + .is()); + } +} + +TEST(SniiLogicalNullReaderTest, RejectsStatsAndDocumentDomainMismatches) { + struct Case { + uint64_t stats_doc_count; + uint64_t stats_null_count; + std::vector frame; + }; + std::vector cases; + cases.push_back({8, 1, build_null_frame({2}, /*recorded_doc_count=*/9)}); + cases.push_back({8, 2, build_null_frame({2}, /*recorded_doc_count=*/8)}); + cases.push_back({8, 1, build_null_frame({8}, /*recorded_doc_count=*/8)}); + + for (Case& test_case : cases) { + LogicalImage image = build_logical_image( + test_case.stats_doc_count, test_case.stats_null_count, std::move(test_case.frame)); + BufferFileReader file(std::move(image.file)); + reader::LogicalIndexReader index; + ASSERT_TRUE(open_logical(&file, image, &index).ok()); + std::vector null_docids; + EXPECT_TRUE(index.read_null_docids(&null_docids) + .is()); + } +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_memory_reporter_test.cpp b/be/test/storage/index/snii/compaction/snii_memory_reporter_test.cpp new file mode 100644 index 00000000000000..0b017290bb7e65 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_memory_reporter_test.cpp @@ -0,0 +1,180 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/memory_reporter.h" + +namespace { + +namespace ErrorCode = doris::ErrorCode; +using doris::snii::writer::MemoryReporter; + +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(std::is_nothrow_move_constructible_v); +static_assert(std::is_nothrow_move_assignable_v); + +TEST(SniiMemoryReporterTest, ReservationGrowthFailureLeavesStateUnchanged) { + int64_t mirrored_bytes = 0; + MemoryReporter reporter([&](int64_t delta) { mirrored_bytes += delta; }, 100); + + { + auto first = reporter.make_reservation(); + ASSERT_TRUE(first.set_bytes(60).ok()); + EXPECT_EQ(first.bytes(), 60); + EXPECT_EQ(reporter.current_bytes(), 60); + EXPECT_EQ(mirrored_bytes, 60); + + const auto rejected = first.set_bytes(101); + EXPECT_TRUE(rejected.is()); + EXPECT_EQ(first.bytes(), 60); + EXPECT_EQ(reporter.current_bytes(), 60); + EXPECT_EQ(mirrored_bytes, 60); + + auto second = reporter.make_reservation(); + ASSERT_TRUE(second.set_bytes(40).ok()); + EXPECT_EQ(reporter.current_bytes(), 100); + EXPECT_EQ(mirrored_bytes, 100); + + ASSERT_TRUE(first.set_bytes(20).ok()); + EXPECT_EQ(reporter.current_bytes(), 60); + EXPECT_EQ(mirrored_bytes, 60); + second.reset(); + EXPECT_EQ(reporter.current_bytes(), 20); + EXPECT_EQ(mirrored_bytes, 20); + } + + EXPECT_EQ(reporter.current_bytes(), 0); + EXPECT_EQ(mirrored_bytes, 0); +} + +TEST(SniiMemoryReporterTest, ReservationIncludesLegacyReportedBytesAtAcquireTime) { + MemoryReporter reporter(nullptr, 100); + reporter.report(40); + + auto reservation = reporter.make_reservation(); + const auto rejected = reservation.set_bytes(61); + EXPECT_TRUE(rejected.is()); + EXPECT_EQ(reservation.bytes(), 0); + EXPECT_EQ(reporter.current_bytes(), 40); + + ASSERT_TRUE(reservation.set_bytes(60).ok()); + EXPECT_EQ(reporter.current_bytes(), 100); + reservation.reset(); + reporter.report(-40); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporterTest, SpillThresholdReservationsRemainAccountedAboveCap) { + MemoryReporter reporter(nullptr, 100, MemoryReporter::CapPolicy::kSpillThreshold); + reporter.report(120); + EXPECT_TRUE(reporter.over_cap()); + + auto reservation = reporter.make_reservation(); + ASSERT_TRUE(reservation.set_bytes(60).ok()); + EXPECT_EQ(reservation.bytes(), 60); + EXPECT_EQ(reporter.current_bytes(), 180); + EXPECT_TRUE(reporter.over_cap()); + + reservation.reset(); + reporter.report(-120); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporterTest, ReservationMoveTransfersChargeAndUnlimitedCapStillCounts) { + int64_t mirrored_bytes = 0; + MemoryReporter reporter([&](int64_t delta) { mirrored_bytes += delta; }); + + auto source = reporter.make_reservation(); + ASSERT_TRUE(source.set_bytes(512).ok()); + auto destination = std::move(source); + EXPECT_EQ(source.bytes(), 0); + EXPECT_EQ(destination.bytes(), 512); + EXPECT_EQ(reporter.current_bytes(), 512); + + ASSERT_TRUE(destination.set_bytes(1024).ok()); + EXPECT_EQ(reporter.current_bytes(), 1024); + ASSERT_TRUE(destination.set_bytes(17).ok()); + EXPECT_EQ(reporter.current_bytes(), 17); + EXPECT_EQ(mirrored_bytes, 17); + + auto replacement = reporter.make_reservation(); + ASSERT_TRUE(replacement.set_bytes(23).ok()); + destination = std::move(replacement); + EXPECT_EQ(reporter.current_bytes(), 23); + EXPECT_EQ(destination.bytes(), 23); + destination.reset(); + EXPECT_EQ(reporter.current_bytes(), 0); + EXPECT_EQ(mirrored_bytes, 0); +} + +TEST(SniiMemoryReporterTest, ConcurrentReservationsCannotOversubscribeCap) { + std::atomic mirrored_bytes {0}; + MemoryReporter reporter( + [&mirrored_bytes](int64_t delta) { + mirrored_bytes.fetch_add(delta, std::memory_order_relaxed); + }, + 100); + std::array succeeded {}; + std::array rejected_by_limit {}; + std::atomic ready {0}; + std::atomic attempted {0}; + std::atomic start {false}; + std::array threads; + + for (size_t i = 0; i < threads.size(); ++i) { + threads[i] = std::thread([&, i] { + auto reservation = reporter.make_reservation(); + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + const auto status = reservation.set_bytes(60); + succeeded[i] = status.ok(); + rejected_by_limit[i] = status.is(); + attempted.fetch_add(1, std::memory_order_release); + while (attempted.load(std::memory_order_acquire) != threads.size()) { + std::this_thread::yield(); + } + }); + } + + while (ready.load(std::memory_order_acquire) != threads.size()) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (std::thread& thread : threads) { + thread.join(); + } + + EXPECT_EQ(std::count(succeeded.begin(), succeeded.end(), true), 1); + EXPECT_EQ(std::count(rejected_by_limit.begin(), rejected_by_limit.end(), true), 1); + EXPECT_EQ(reporter.current_bytes(), 0); + EXPECT_EQ(mirrored_bytes.load(std::memory_order_relaxed), 0); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_posting_cursor_test.cpp b/be/test/storage/index/snii/compaction/snii_posting_cursor_test.cpp new file mode 100644 index 00000000000000..5db931c18976da --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_posting_cursor_test.cpp @@ -0,0 +1,1036 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/compaction/posting_cursor.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::snii_test; // NOLINT +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using compaction::PostingStream; +using compaction::RemappedPostingChunk; +using compaction::SharedAlignedRegionCache; +using compaction::SniiPostingCursor; +using compaction::SniiPostingReadContext; +using compaction::ValidatedRowIdConversion; +using writer::MemoryReporter; + +static_assert(!std::is_constructible_v< + SniiPostingCursor, SniiPostingReadContext*, format::DictEntry, uint64_t, uint64_t, + uint32_t, std::span>, std::span>); + +constexpr uint64_t kIndexId = 9; +constexpr std::string_view kIndexSuffix = "body"; +constexpr uint32_t kDocCount = 640; +constexpr uint32_t kFreqDroppedDocCount = 65536; +constexpr auto kDeleted = std::pair {std::numeric_limits::max(), + std::numeric_limits::max()}; +constexpr std::array kNoDestinationRows {}; +constexpr std::array kDestinationRows1 {1}; +constexpr std::array kDestinationRows8 {8}; +constexpr std::array kDestinationRows10 {10}; +constexpr std::array kTwoDestinationRows10 {10, 10}; + +struct SourceFixture { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; +}; + +struct TermRef { + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; +}; + +Status build_source(std::vector terms, uint32_t doc_count, + SourceFixture* fixture, bool write_freq = true) { + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = kIndexSuffix; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = doc_count; + input.write_freq = write_freq; + input.terms = std::move(terms); + + writer::SniiCompoundWriter compound(&fixture->file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(&fixture->file, &fixture->segment)); + return fixture->segment.open_index(kIndexId, kIndexSuffix, &fixture->index); +} + +Status build_hybrid_source(std::vector terms, uint32_t doc_count, + SourceFixture* fixture) { + namespace inverted_index = doris::segment_v2::inverted_index; + inverted_index::CommonGramsQueryIdentity identity {.common_grams_dictionary_identity = "dict-a", + .base_analyzer_fingerprint = "base-a", + .common_grams_fingerprint = "grams-a"}; + auto metadata = inverted_index::make_common_grams_segment_metadata(identity); + metadata.common_grams_coverage = inverted_index::CommonGramsCoverage::kMixed; + metadata.scoring_doc_count = doc_count; + metadata.scoring_token_count = 1; + + terms.push_back(make_term("plain", {{.docid = 0, .positions = {0}}})); + std::ranges::sort(terms, [](const auto& lhs, const auto& rhs) { return lhs.term < rhs.term; }); + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = kIndexSuffix; + input.config = format::IndexConfig::kDocsPositionsScoring; + input.doc_count = doc_count; + input.write_freq = true; + input.encoded_norms.assign(doc_count, 1); + input.common_grams_metadata = std::move(metadata); + input.common_grams_posting_policy = format::CommonGramsPostingPolicy::kHybridV1; + input.terms = std::move(terms); + + writer::SniiCompoundWriter compound(&fixture->file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(&fixture->file, &fixture->segment)); + return fixture->segment.open_index(kIndexId, kIndexSuffix, &fixture->index); +} + +writer::TermPostings make_docs_only_gram(std::string left, std::string right, + std::vector docids) { + writer::TermPostings term; + term.term = doris::segment_v2::inverted_index::encode_common_gram(left, right).value(); + term.docids = std::move(docids); + term.retain_positions = false; + return term; +} + +std::vector posting_shapes() { + std::vector terms; + terms.push_back(make_term("inline", + {{.docid = 1, .positions = {2, 7}}, {.docid = 5, .positions = {3}}})); + + std::vector slim; + slim.reserve(500); + for (uint32_t docid = 0; docid < 500; ++docid) { + uint32_t mixed = docid * 2654435761U; + mixed ^= mixed >> 16; + const uint32_t frequency = mixed % 97 + 1; + std::vector positions(frequency); + for (uint32_t i = 0; i < frequency; ++i) { + positions[i] = i * 3; + } + slim.push_back({.docid = docid, .positions = std::move(positions)}); + } + terms.push_back(make_term("slim", std::move(slim))); + + std::vector windowed; + windowed.reserve(600); + for (uint32_t docid = 0; docid < 600; ++docid) { + windowed.push_back({.docid = docid, .positions = {1, 4, 9}}); + } + terms.push_back(make_term("windowed", std::move(windowed))); + return terms; +} + +std::vector freq_dropped_posting_shapes() { + std::vector terms; + terms.push_back(make_term("inline", + {{.docid = 1, .positions = {2, 7}}, {.docid = 5, .positions = {3}}})); + + std::vector slim; + slim.reserve(500); + uint32_t docid = 1; + slim.push_back({.docid = docid, .positions = {2, 7}}); + for (uint32_t i = 1; i < 500; ++i) { + docid += 1 + (i * 73 + i * i * 19) % 230; + slim.push_back({.docid = docid, .positions = {2, 7}}); + } + terms.push_back(make_term("slim", std::move(slim))); + + std::vector windowed; + windowed.reserve(600); + for (uint32_t docid = 0; docid < 600; ++docid) { + windowed.push_back({.docid = docid, .positions = {1, 4, 9}}); + } + terms.push_back(make_term("windowed", std::move(windowed))); + return terms; +} + +std::vector repeated_windowed_terms() { + std::vector terms; + for (std::string term : {"first", "second"}) { + std::vector docs; + docs.reserve(600); + for (uint32_t docid = 0; docid < 600; ++docid) { + docs.push_back({.docid = docid, .positions = {1, 4, 9}}); + } + terms.push_back(make_term(std::move(term), std::move(docs))); + } + return terms; +} + +TermRef lookup_term(const reader::LogicalIndexReader& index, std::string_view term) { + bool found = false; + TermRef ref; + assert_ok(index.lookup(term, &found, &ref.entry, &ref.frq_base, &ref.prx_base)); + EXPECT_TRUE(found) << term; + return ref; +} + +std::vector> deleted_map(uint32_t doc_count = kDocCount) { + return std::vector>(doc_count, kDeleted); +} + +struct PostingCopy { + uint32_t segment = 0; + uint32_t docid = 0; + uint32_t freq = 0; + std::vector positions; + + bool operator==(const PostingCopy&) const = default; +}; + +Status drain(SniiPostingCursor* cursor, std::vector* output) { + RemappedPostingChunk chunk; + bool has_chunk = false; + RETURN_IF_ERROR(cursor->next_chunk(&chunk, &has_chunk)); + while (has_chunk) { + for (size_t ordinal = 0; ordinal < chunk.destination_docids.size(); ++ordinal) { + const uint32_t segment = chunk.destination_segment; + const uint32_t docid = chunk.destination_docids[ordinal]; + uint32_t frequency = 0; + std::vector positions; + if (!chunk.freqs.empty()) { + frequency = chunk.freqs[ordinal]; + const uint32_t position_base = chunk.position_offsets.front(); + const uint32_t begin = chunk.position_offsets[ordinal] - position_base; + const uint32_t end = chunk.position_offsets[ordinal + 1] - position_base; + positions.assign(chunk.positions_flat.begin() + begin, + chunk.positions_flat.begin() + end); + } + output->push_back({segment, docid, frequency, std::move(positions)}); + } + RETURN_IF_ERROR(cursor->next_chunk(&chunk, &has_chunk)); + } + return Status::OK(); +} + +std::unique_ptr make_read_context( + const reader::LogicalIndexReader* index, size_t total_read_ahead_budget_bytes = 128) { + auto context = std::make_unique(index, total_read_ahead_budget_bytes); + assert_ok(context->init()); + return context; +} + +struct CursorMapping { + compaction::RowIdConversionMap conversion; + std::unique_ptr validated; +}; + +const ValidatedRowIdConversion* make_validated_cursor_mapping( + const std::vector>& trans) { + static std::vector> arena; + auto fixture = std::make_unique(); + fixture->conversion.resize(4); + fixture->conversion[3] = trans; + + uint32_t max_segment = 0; + bool has_live_row = false; + for (const auto& [segment, docid] : trans) { + if (segment != std::numeric_limits::max()) { + max_segment = std::max(max_segment, segment); + has_live_row = true; + } + } + std::vector destination_rows(has_live_row ? max_segment + 1 : 0, 0); + for (auto& [segment, docid] : fixture->conversion[3]) { + if (segment != std::numeric_limits::max()) { + docid = destination_rows[segment]++; + } + } + const std::array source_rows {0, 0, 0, static_cast(trans.size())}; + assert_ok(ValidatedRowIdConversion::create(&fixture->conversion, source_rows, destination_rows, + &fixture->validated)); + const ValidatedRowIdConversion* validated = fixture->validated.get(); + arena.push_back(std::move(fixture)); + return validated; +} + +SniiPostingCursor make_cursor(SniiPostingReadContext* read_context, TermRef ref, + const std::vector>& trans, + std::span) { + return SniiPostingCursor(read_context, std::move(ref.entry), ref.frq_base, ref.prx_base, + /*source_ordinal=*/3, make_validated_cursor_mapping(trans)); +} + +TEST(SniiPostingCursorTest, ScratchGrowthChargesOldAndNewAllocationsBeforeReserve) { + MemoryReporter reporter(nullptr, /*cap_bytes=*/35); + { + MemoryFile file; + std::vector bytes(32, 7); + assert_ok(file.append(Slice(bytes))); + + SharedAlignedRegionCache cache(&file, /*region_offset=*/0, bytes.size(), + /*total_budget_bytes=*/16, &reporter); + assert_ok(cache.init()); + ASSERT_EQ(reporter.current_bytes(), 16); + + std::vector scratch; + scratch.reserve(8); + auto scratch_reservation = reporter.make_reservation(); + assert_ok(scratch_reservation.set_bytes(scratch.capacity())); + ASSERT_EQ(reporter.current_bytes(), 24); + + Slice resolved; + const Status status = cache.resolve(PostingStream::kDocs, /*abs_off=*/4, /*len=*/12, + &scratch, &resolved, &scratch_reservation); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(scratch.capacity(), 8); + EXPECT_EQ(scratch_reservation.bytes(), 8); + EXPECT_EQ(reporter.current_bytes(), 24); + EXPECT_TRUE(resolved.empty()); + EXPECT_TRUE(file.reads().empty()); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiPostingCursorTest, ChargesEveryRetainedDecoderWorkspaceBuffer) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "slim"); + auto trans = deleted_map(); + trans[0] = {0, 0}; + + MemoryReporter reporter(nullptr, /*cap_bytes=*/4U << 20); + { + SniiPostingReadContext read_context(&source.index, + /*total_read_ahead_budget_bytes=*/1U << 20, &reporter); + assert_ok(read_context.init()); + SniiPostingCursor cursor = make_cursor(&read_context, ref, trans, kDestinationRows1); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + EXPECT_EQ(static_cast(reporter.current_bytes()), + read_context.resident_read_ahead_capacity_bytes() + + read_context.decoder_workspace_capacity_bytes()); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiPostingCursorTest, PositionWorkspaceLimitPoisonsReadContext) { + std::vector positions(100000); + std::iota(positions.begin(), positions.end(), 0); + SourceFixture source; + assert_ok(build_source({make_term("large", {{.docid = 0, .positions = std::move(positions)}})}, + /*doc_count=*/1, &source)); + const TermRef ref = lookup_term(source.index, "large"); + std::vector> trans = {{0, 0}}; + + MemoryReporter reporter(nullptr, /*cap_bytes=*/192U << 10); + SniiPostingReadContext read_context(&source.index, /*total_read_ahead_budget_bytes=*/128, + &reporter); + assert_ok(read_context.init()); + SniiPostingCursor cursor = make_cursor(&read_context, ref, trans, kDestinationRows1); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + const Status first = cursor.next_chunk(&chunk, &has_chunk); + EXPECT_TRUE(first.is()) << first; + EXPECT_FALSE(has_chunk); + EXPECT_LE(reporter.current_bytes(), static_cast(reporter.cap_bytes())); + + SniiPostingCursor retry = make_cursor(&read_context, ref, trans, kDestinationRows1); + EXPECT_EQ(retry.init(), first); +} + +TEST(SniiPostingCursorTest, DecodesInlineAndSkipsDeletedDocs) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "inline"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kInline); + + auto trans = deleted_map(); + trans[1] = {0, 2}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kDestinationRows10); + assert_ok(cursor.init()); + + std::vector got; + assert_ok(drain(&cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 2, {2, 7}}})); +} + +TEST(SniiPostingCursorTest, DecodesZstdDdWithoutMemoryReporter) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + TermRef ref = lookup_term(source.index, "inline"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kInline); + + const auto freq_begin = ref.entry.frq_bytes.begin() + ref.entry.inline_dd_disk_len; + std::vector freq_bytes(freq_begin, ref.entry.frq_bytes.end()); + ByteSink dd_sink; + format::FrqRegionMeta dd_meta; + const std::array docids {1, 5}; + assert_ok( + format::build_dd_region(docids, /*win_base=*/0, /*zstd_level=*/3, &dd_sink, &dd_meta)); + ASSERT_TRUE(dd_meta.zstd); + + ref.entry.frq_bytes = dd_sink.take(); + ref.entry.frq_bytes.insert(ref.entry.frq_bytes.end(), freq_bytes.begin(), freq_bytes.end()); + ref.entry.dd_meta = dd_meta; + ref.entry.inline_dd_disk_len = dd_meta.disk_len; + ref.entry.frq_len = ref.entry.frq_bytes.size(); + + auto trans = deleted_map(); + trans[1] = {0, 2}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = + make_cursor(read_context.get(), std::move(ref), trans, kDestinationRows10); + assert_ok(cursor.init()); + + std::vector got; + assert_ok(drain(&cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 2, {2, 7}}})); +} + +TEST(SniiPostingCursorTest, DecodesSlimPodRefAcrossDestinations) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "slim"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kPodRef); + ASSERT_EQ(ref.entry.enc, format::DictEntryEnc::kSlim); + + auto trans = deleted_map(); + trans[0] = {0, 3}; + trans[150] = {1, 0}; + trans[499] = {1, 4}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kTwoDestinationRows10); + assert_ok(cursor.init()); + + std::vector got; + assert_ok(drain(&cursor, &got)); + auto expected_positions = [](uint32_t docid) { + uint32_t mixed = docid * 2654435761U; + mixed ^= mixed >> 16; + std::vector positions(mixed % 97 + 1); + for (uint32_t i = 0; i < positions.size(); ++i) { + positions[i] = i * 3; + } + return positions; + }; + const std::vector positions_0 = expected_positions(0); + const std::vector positions_150 = expected_positions(150); + const std::vector positions_499 = expected_positions(499); + EXPECT_EQ(got, (std::vector { + {0, 0, static_cast(positions_0.size()), positions_0}, + {1, 0, static_cast(positions_150.size()), positions_150}, + {1, 1, static_cast(positions_499.size()), positions_499}})); +} + +TEST(SniiPostingCursorTest, DecodesWindowedPodRefSequentially) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "windowed"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kPodRef); + ASSERT_EQ(ref.entry.enc, format::DictEntryEnc::kWindowed); + + auto trans = deleted_map(); + trans[0] = {0, 1}; + trans[511] = {0, 8}; + trans[512] = {1, 1}; + trans[599] = {1, 9}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kTwoDestinationRows10); + assert_ok(cursor.init()); + + std::vector got; + assert_ok(drain(&cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 3, {1, 4, 9}}, + {0, 1, 3, {1, 4, 9}}, + {1, 0, 3, {1, 4, 9}}, + {1, 1, 3, {1, 4, 9}}})); +} + +TEST(SniiPostingCursorTest, EmitsMappedChunksAtDecoderWindowBoundaries) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "windowed"); + auto trans = deleted_map(); + trans[0] = {0, 0}; + trans[511] = {0, 1}; + trans[512] = {1, 0}; + trans[599] = {1, 1}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kTwoDestinationRows10); + assert_ok(cursor.init()); + + std::vector chunk_sizes; + std::vector> keys; + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + while (has_chunk) { + chunk_sizes.push_back(chunk.destination_docids.size()); + for (uint32_t docid : chunk.destination_docids) { + keys.emplace_back(chunk.destination_segment, docid); + } + EXPECT_EQ(chunk.freqs.size(), chunk.destination_docids.size()); + EXPECT_EQ(chunk.position_offsets.size(), chunk.destination_docids.size() + 1); + EXPECT_EQ(chunk.position_offsets.back() - chunk.position_offsets.front(), + chunk.positions_flat.size()); + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + } + + EXPECT_EQ(chunk_sizes, std::vector({1, 1, 2})); + EXPECT_EQ(keys, (std::vector> {{0, 0}, {0, 1}, {1, 0}, {1, 1}})); +} + +TEST(SniiPostingCursorTest, SplitsOneDecodedChunkIntoDestinationHomogeneousRuns) { + SourceFixture source; + assert_ok(build_source({make_term("split", {{.docid = 0, .positions = {2, 7}}, + {.docid = 1, .positions = {3}}})}, + /*doc_count=*/2, &source)); + const TermRef ref = lookup_term(source.index, "split"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kInline); + std::vector> trans {{0, 0}, {1, 0}}; + std::array destination_rows {1, 1}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, destination_rows); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + EXPECT_EQ(chunk.destination_segment, 0); + EXPECT_EQ( + std::vector(chunk.destination_docids.begin(), chunk.destination_docids.end()), + (std::vector {0})); + EXPECT_EQ(std::vector(chunk.position_offsets.begin(), chunk.position_offsets.end()), + (std::vector {0, 2})); + EXPECT_EQ(std::vector(chunk.positions_flat.begin(), chunk.positions_flat.end()), + (std::vector {2, 7})); + + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + EXPECT_EQ(chunk.destination_segment, 1); + EXPECT_EQ( + std::vector(chunk.destination_docids.begin(), chunk.destination_docids.end()), + (std::vector {0})); + EXPECT_EQ(std::vector(chunk.position_offsets.begin(), chunk.position_offsets.end()), + (std::vector {2, 3})); + EXPECT_EQ(std::vector(chunk.positions_flat.begin(), chunk.positions_flat.end()), + (std::vector {3})); + + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + EXPECT_FALSE(has_chunk); +} + +TEST(SniiPostingCursorTest, ExposesBaseRelativePositionSliceForDirectRun) { + SourceFixture source; + assert_ok(build_source({make_term("split", {{.docid = 0, .positions = {2, 7}}, + {.docid = 1, .positions = {3}}})}, + /*doc_count=*/2, &source)); + const TermRef ref = lookup_term(source.index, "split"); + std::vector> trans {{0, 0}, {1, 0}}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = + make_cursor(read_context.get(), ref, trans, std::array {1, 1}); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + ASSERT_EQ(chunk.position_offsets.size(), 2U); + EXPECT_GT(chunk.position_offsets.front(), 0U); + EXPECT_EQ(chunk.position_offsets.back() - chunk.position_offsets.front(), + chunk.positions_flat.size()); + EXPECT_EQ(std::vector(chunk.positions_flat.begin(), chunk.positions_flat.end()), + (std::vector {3})); +} + +TEST(SniiPostingCursorTest, UsesChunkLevelNoDeletionFastPath) { + SourceFixture source; + assert_ok(build_source({make_term("all-live", {{.docid = 1, .positions = {2, 7}}, + {.docid = 5, .positions = {3}}})}, + /*doc_count=*/6, &source)); + const TermRef ref = lookup_term(source.index, "all-live"); + std::vector> trans(6); + for (uint32_t docid = 0; docid < trans.size(); ++docid) { + trans[docid] = {0, docid}; + } + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kDestinationRows10); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + EXPECT_EQ(chunk.destination_segment, 0); + EXPECT_EQ( + std::vector(chunk.destination_docids.begin(), chunk.destination_docids.end()), + (std::vector {1, 5})); + EXPECT_EQ(chunk.freqs.size(), 2U); + EXPECT_EQ(chunk.position_offsets.size(), 3U); + EXPECT_EQ(chunk.positions_flat.size(), 3U); + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + EXPECT_FALSE(has_chunk); +} + +TEST(SniiPostingCursorTest, DecodesFlatAndWindowedDocsOnlyWithoutInventingPayloads) { + constexpr uint32_t kWideDocs = format::kSlimDfThreshold; + std::vector wide_docids(kWideDocs); + std::iota(wide_docids.begin(), wide_docids.end(), 0); + SourceFixture source; + assert_ok(build_hybrid_source({make_docs_only_gram("a", "of", {0, 4}), + make_docs_only_gram("z", "of", std::move(wide_docids))}, + kWideDocs, &source)); + + const TermRef flat_ref = lookup_term( + source.index, doris::segment_v2::inverted_index::encode_common_gram("a", "of").value()); + const TermRef windowed_ref = lookup_term( + source.index, doris::segment_v2::inverted_index::encode_common_gram("z", "of").value()); + ASSERT_EQ(flat_ref.entry.enc, format::DictEntryEnc::kSlim); + ASSERT_EQ(windowed_ref.entry.enc, format::DictEntryEnc::kWindowed); + + auto trans = deleted_map(kWideDocs); + trans[0] = {0, 1}; + trans[4] = {0, 5}; + { + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = + make_cursor(read_context.get(), flat_ref, trans, kDestinationRows10); + assert_ok(cursor.init()); + EXPECT_FALSE(cursor.has_positions()); + std::vector got; + assert_ok(drain(&cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 0, {}}, {0, 1, 0, {}}})); + } + + trans = deleted_map(kWideDocs); + trans[0] = {0, 0}; + trans[kWideDocs - 1] = {0, 9}; + { + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = + make_cursor(read_context.get(), windowed_ref, trans, kDestinationRows10); + assert_ok(cursor.init()); + EXPECT_FALSE(cursor.has_positions()); + std::vector got; + assert_ok(drain(&cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 0, {}}, {0, 1, 0, {}}})); + } +} + +TEST(SniiPostingCursorTest, RejectsOverflowAndOutOfFilePostingRegions) { + format::RegionRef region; + region.offset = std::numeric_limits::max() - 3; + region.length = 8; + Status status = compaction::validate_posting_region(region, UINT64_MAX); + EXPECT_TRUE(status.is()) << status; + + region.offset = 90; + region.length = 11; + status = compaction::validate_posting_region(region, 100); + EXPECT_TRUE(status.is()) << status; + + region.length = 10; + EXPECT_TRUE(compaction::validate_posting_region(region, 100).ok()); +} + +TEST(SniiPostingCursorTest, RejectsSourceDocOutsideIndexDocCount) { + SourceFixture posting_source; + assert_ok(build_source({make_term("bad", {{.docid = 8, .positions = {1}}})}, + /*doc_count=*/9, &posting_source)); + const TermRef ref = lookup_term(posting_source.index, "bad"); + ASSERT_EQ(ref.entry.kind, format::DictEntryKind::kInline); + + SourceFixture bounded_source; + assert_ok(build_source({}, /*doc_count=*/8, &bounded_source)); + auto trans = deleted_map(8); + auto read_context = make_read_context(&bounded_source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, kDestinationRows8); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + const Status status = cursor.next_chunk(&chunk, &has_chunk); + EXPECT_TRUE(status.is()) << status; + EXPECT_FALSE(has_chunk); +} + +TEST(SniiPostingCursorTest, RejectsInconsistentTermStatisticsAndPositions) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + + TermRef stats_ref = lookup_term(source.index, "inline"); + ++stats_ref.entry.ttf_delta; + auto trans = deleted_map(); + trans[1] = {0, 1}; + trans[5] = {0, 2}; + auto stats_read_context = make_read_context(&source.index); + SniiPostingCursor stats_cursor = + make_cursor(stats_read_context.get(), stats_ref, trans, kDestinationRows10); + assert_ok(stats_cursor.init()); + std::vector ignored; + Status status = drain(&stats_cursor, &ignored); + EXPECT_TRUE(status.is()) << status; + const TermRef valid_stats_ref = lookup_term(source.index, "inline"); + SniiPostingCursor stats_retry = + make_cursor(stats_read_context.get(), valid_stats_ref, trans, kDestinationRows10); + EXPECT_EQ(stats_retry.init(), status); + + TermRef positions_ref = lookup_term(source.index, "inline"); + ASSERT_FALSE(positions_ref.entry.prx_bytes.empty()); + positions_ref.entry.prx_bytes.pop_back(); + positions_ref.entry.prx_len = positions_ref.entry.prx_bytes.size(); + auto positions_read_context = make_read_context(&source.index); + SniiPostingCursor positions_cursor = + make_cursor(positions_read_context.get(), positions_ref, trans, kDestinationRows10); + assert_ok(positions_cursor.init()); + status = positions_cursor.next_chunk(nullptr, nullptr); + EXPECT_TRUE(status.is()) << status; + RemappedPostingChunk chunk; + bool has_chunk = false; + status = positions_cursor.next_chunk(&chunk, &has_chunk); + EXPECT_TRUE(status.is()) << status; + const TermRef retry_ref = lookup_term(source.index, "inline"); + SniiPostingCursor retry = + make_cursor(positions_read_context.get(), retry_ref, trans, kDestinationRows10); + const Status retry_status = retry.init(); + EXPECT_EQ(retry_status, status); +} + +TEST(SniiPostingCursorTest, RejectsWindowMaxFrequencyMismatchWithoutRescanningPositions) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef valid_ref = lookup_term(source.index, "windowed"); + uint64_t dd_block_offset = 0; + uint64_t frq_payload_length = 0; + assert_ok(source.index.resolve_frq_window(valid_ref.entry, valid_ref.frq_base, &dd_block_offset, + &frq_payload_length)); + ASSERT_GE(dd_block_offset, valid_ref.entry.prelude_len); + const uint64_t prelude_offset = dd_block_offset - valid_ref.entry.prelude_len; + const Slice original_prelude(source.file.data().data() + prelude_offset, + valid_ref.entry.prelude_len); + format::FrqPreludeReader prelude; + assert_ok(format::FrqPreludeReader::open(original_prelude, &prelude)); + + format::FrqPreludeColumns columns; + columns.has_freq = prelude.has_freq(); + columns.has_prx = prelude.has_prx(); + columns.group_size = 64; + for (uint32_t window = 0; window < prelude.n_windows(); ++window) { + format::WindowMeta meta; + assert_ok(prelude.window(window, &meta)); + columns.windows.push_back(meta); + } + ASSERT_FALSE(columns.windows.empty()); + ++columns.windows.front().max_freq; + ByteSink corrupt_prelude; + assert_ok(format::build_frq_prelude(columns, &corrupt_prelude)); + ASSERT_EQ(corrupt_prelude.size(), valid_ref.entry.prelude_len); + + std::vector corrupt_file_bytes = source.file.data(); + std::copy(corrupt_prelude.buffer().begin(), corrupt_prelude.buffer().end(), + corrupt_file_bytes.begin() + prelude_offset); + SourceFixture corrupt_source; + assert_ok(corrupt_source.file.append(Slice(corrupt_file_bytes))); + assert_ok(corrupt_source.file.finalize()); + assert_ok(reader::SniiSegmentReader::open(&corrupt_source.file, &corrupt_source.segment)); + assert_ok(corrupt_source.segment.open_index(kIndexId, kIndexSuffix, &corrupt_source.index)); + + const TermRef corrupt_ref = lookup_term(corrupt_source.index, "windowed"); + auto trans = deleted_map(); + trans[0] = {0, 0}; + auto read_context = make_read_context(&corrupt_source.index); + SniiPostingCursor cursor = + make_cursor(read_context.get(), corrupt_ref, trans, kDestinationRows1); + assert_ok(cursor.init()); + RemappedPostingChunk chunk; + bool has_chunk = false; + const Status status = cursor.next_chunk(&chunk, &has_chunk); + EXPECT_TRUE(status.is()) << status; + EXPECT_FALSE(has_chunk); +} + +TEST(SniiPostingCursorTest, ReusesReadAheadAcrossConsecutivePodRefTerms) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef slim_ref = lookup_term(source.index, "slim"); + const TermRef windowed_ref = lookup_term(source.index, "windowed"); + ASSERT_EQ(slim_ref.entry.kind, format::DictEntryKind::kPodRef); + ASSERT_EQ(windowed_ref.entry.kind, format::DictEntryKind::kPodRef); + + auto trans = deleted_map(); + source.file.clear_reads(); + auto read_context = + make_read_context(&source.index, /*total_read_ahead_budget_bytes=*/1U << 20); + SniiPostingCursor slim_cursor = + make_cursor(read_context.get(), slim_ref, trans, kNoDestinationRows); + assert_ok(slim_cursor.init()); + std::vector ignored; + assert_ok(drain(&slim_cursor, &ignored)); + const uint64_t physical_reads_after_slim = read_context->physical_read_ranges(); + const uint64_t docs_hits_after_slim = read_context->docs_buffer_hits(); + const uint64_t prx_hits_after_slim = read_context->prx_buffer_hits(); + ASSERT_GT(physical_reads_after_slim, 0); + + SniiPostingCursor windowed_cursor = + make_cursor(read_context.get(), windowed_ref, trans, kNoDestinationRows); + assert_ok(windowed_cursor.init()); + assert_ok(drain(&windowed_cursor, &ignored)); + + EXPECT_EQ(read_context->physical_read_ranges(), physical_reads_after_slim); + EXPECT_GT(read_context->docs_buffer_hits(), docs_hits_after_slim); + EXPECT_GT(read_context->prx_buffer_hits(), prx_hits_after_slim); + + const format::RegionRef& posting_region = source.index.section_refs().posting_region; + EXPECT_EQ(read_context->physical_read_bytes(), source.file.read_bytes()); + EXPECT_LE(read_context->physical_read_bytes(), posting_region.length); + EXPECT_EQ(read_context->physical_read_ranges(), source.file.reads().size()); + EXPECT_LE(read_context->resident_read_ahead_capacity_bytes(), + read_context->total_read_ahead_budget_bytes()); + for (const MemoryFile::Read& range : source.file.reads()) { + EXPECT_GE(range.offset, posting_region.offset); + EXPECT_LE(range.offset - posting_region.offset, posting_region.length); + EXPECT_LE(range.len, posting_region.length - (range.offset - posting_region.offset)); + } +} + +TEST(SniiPostingCursorTest, ReusesDecoderWorkspaceAcrossTerms) { + SourceFixture source; + assert_ok(build_source(repeated_windowed_terms(), kDocCount, &source)); + const TermRef first_ref = lookup_term(source.index, "first"); + const TermRef second_ref = lookup_term(source.index, "second"); + auto trans = deleted_map(); + auto read_context = make_read_context(&source.index, 1U << 20); + std::vector ignored; + + SniiPostingCursor first = make_cursor(read_context.get(), first_ref, trans, kNoDestinationRows); + assert_ok(first.init()); + assert_ok(drain(&first, &ignored)); + const size_t capacity_after_first = read_context->decoder_workspace_capacity_bytes(); + ASSERT_GT(capacity_after_first, 0); + + SniiPostingCursor second = + make_cursor(read_context.get(), second_ref, trans, kNoDestinationRows); + assert_ok(second.init()); + assert_ok(drain(&second, &ignored)); + EXPECT_EQ(read_context->decoder_workspace_capacity_bytes(), capacity_after_first); +} + +TEST(SniiPostingCursorTest, ReleasesLargeDecoderWorkspaceAtTermBoundary) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef slim_ref = lookup_term(source.index, "slim"); + auto trans = deleted_map(); + trans[0] = {0, 0}; + auto read_context = make_read_context(&source.index, 1U << 20); + SniiPostingCursor cursor = make_cursor(read_context.get(), slim_ref, trans, kDestinationRows1); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + ASSERT_GT(read_context->decoder_workspace_capacity_bytes(), + read_context->retained_decoder_workspace_limit_bytes()); + + std::vector ignored; + assert_ok(drain(&cursor, &ignored)); + EXPECT_LE(read_context->decoder_workspace_capacity_bytes(), + read_context->retained_decoder_workspace_limit_bytes()); +} + +TEST(SniiPostingCursorTest, UsesValidatedCapabilityMapping) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef ref = lookup_term(source.index, "inline"); + auto trans = deleted_map(); + trans[1] = {0, 1}; + std::array destination_rows {1}; + auto read_context = make_read_context(&source.index); + SniiPostingCursor cursor = make_cursor(read_context.get(), ref, trans, destination_rows); + assert_ok(cursor.init()); + + RemappedPostingChunk chunk; + bool has_chunk = false; + assert_ok(cursor.next_chunk(&chunk, &has_chunk)); + ASSERT_TRUE(has_chunk); + ASSERT_EQ(chunk.destination_docids.size(), 1U); + EXPECT_EQ(chunk.destination_segment, 0); + EXPECT_EQ(chunk.destination_docids.front(), 0); +} + +TEST(SniiPostingCursorTest, EnforcesReadContextBudgetAndSingleActiveCursor) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef inline_ref = lookup_term(source.index, "inline"); + const TermRef slim_ref = lookup_term(source.index, "slim"); + auto trans = deleted_map(); + + SniiPostingReadContext zero_budget(&source.index, 0); + Status status = zero_budget.init(); + EXPECT_TRUE(status.is()) << status; + SniiPostingReadContext oversized_budget(&source.index, + SniiPostingReadContext::kMaxReadAheadBudgetBytes + 1); + status = oversized_budget.init(); + EXPECT_TRUE(status.is()) << status; + + SniiPostingReadContext read_context(&source.index, 128); + EXPECT_EQ(read_context.total_read_ahead_budget_bytes(), 128); + SniiPostingCursor before_context = + make_cursor(&read_context, inline_ref, trans, kNoDestinationRows); + status = before_context.init(); + EXPECT_TRUE(status.is()) << status; + assert_ok(read_context.init()); + + { + SniiPostingCursor active = + make_cursor(&read_context, inline_ref, trans, kNoDestinationRows); + assert_ok(active.init()); + SniiPostingCursor concurrent = + make_cursor(&read_context, slim_ref, trans, kNoDestinationRows); + status = concurrent.init(); + EXPECT_TRUE(status.is()) << status; + } + + SniiPostingCursor after_release = + make_cursor(&read_context, slim_ref, trans, kNoDestinationRows); + assert_ok(after_release.init()); + std::vector ignored; + assert_ok(drain(&after_release, &ignored)); +} + +TEST(SniiPostingCursorTest, RejectsBackwardTermRangesAndPoisonsReadContext) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef slim_ref = lookup_term(source.index, "slim"); + const TermRef windowed_ref = lookup_term(source.index, "windowed"); + auto trans = deleted_map(); + auto read_context = make_read_context(&source.index, 128); + + SniiPostingCursor later = + make_cursor(read_context.get(), windowed_ref, trans, kNoDestinationRows); + assert_ok(later.init()); + std::vector ignored; + assert_ok(drain(&later, &ignored)); + + SniiPostingCursor backward = + make_cursor(read_context.get(), slim_ref, trans, kNoDestinationRows); + Status status = backward.init(); + EXPECT_TRUE(status.is()) << status; + SniiPostingCursor retry = make_cursor(read_context.get(), slim_ref, trans, kNoDestinationRows); + const Status retry_status = retry.init(); + EXPECT_EQ(retry_status, status); +} + +TEST(SniiPostingCursorTest, CorruptPreludePoisonsContextAndReleasesLease) { + SourceFixture source; + assert_ok(build_source(posting_shapes(), kDocCount, &source)); + const TermRef valid_ref = lookup_term(source.index, "windowed"); + TermRef corrupt_ref = valid_ref; + ASSERT_GT(corrupt_ref.entry.frq_off_delta, 0); + --corrupt_ref.entry.frq_off_delta; + auto trans = deleted_map(); + auto read_context = make_read_context(&source.index, 128); + + SniiPostingCursor corrupt = + make_cursor(read_context.get(), corrupt_ref, trans, kNoDestinationRows); + const Status first = corrupt.init(); + EXPECT_TRUE(first.is()) << first; + const uint64_t physical_bytes_after_failure = read_context->physical_read_bytes(); + + SniiPostingCursor retry = make_cursor(read_context.get(), valid_ref, trans, kNoDestinationRows); + const Status second = retry.init(); + EXPECT_EQ(second, first); + EXPECT_EQ(read_context->physical_read_bytes(), physical_bytes_after_failure); +} + +TEST(SniiPostingCursorTest, DecodesAllShapesWithoutStoredFrequencies) { + SourceFixture source; + assert_ok(build_source(freq_dropped_posting_shapes(), kFreqDroppedDocCount, &source, + /*write_freq=*/false)); + const TermRef inline_ref = lookup_term(source.index, "inline"); + const TermRef slim_ref = lookup_term(source.index, "slim"); + const TermRef windowed_ref = lookup_term(source.index, "windowed"); + ASSERT_EQ(inline_ref.entry.kind, format::DictEntryKind::kInline); + ASSERT_EQ(slim_ref.entry.kind, format::DictEntryKind::kPodRef); + ASSERT_EQ(slim_ref.entry.enc, format::DictEntryEnc::kSlim); + ASSERT_EQ(windowed_ref.entry.kind, format::DictEntryKind::kPodRef); + ASSERT_EQ(windowed_ref.entry.enc, format::DictEntryEnc::kWindowed); + EXPECT_FALSE(inline_ref.entry.term_stats_present); + EXPECT_FALSE(slim_ref.entry.term_stats_present); + EXPECT_FALSE(windowed_ref.entry.term_stats_present); + + auto trans = deleted_map(kFreqDroppedDocCount); + trans[1] = {0, 0}; + auto read_context = make_read_context(&source.index, 1U << 20); + + std::vector got; + SniiPostingCursor inline_cursor = + make_cursor(read_context.get(), inline_ref, trans, kDestinationRows1); + assert_ok(inline_cursor.init()); + assert_ok(drain(&inline_cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 2, {2, 7}}})); + + got.clear(); + SniiPostingCursor slim_cursor = + make_cursor(read_context.get(), slim_ref, trans, kDestinationRows1); + assert_ok(slim_cursor.init()); + assert_ok(drain(&slim_cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 2, {2, 7}}})); + + got.clear(); + SniiPostingCursor windowed_cursor = + make_cursor(read_context.get(), windowed_ref, trans, kDestinationRows1); + assert_ok(windowed_cursor.init()); + assert_ok(drain(&windowed_cursor, &got)); + EXPECT_EQ(got, (std::vector {{0, 0, 3, {1, 4, 9}}})); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_rowid_conversion_test.cpp b/be/test/storage/index/snii/compaction/snii_rowid_conversion_test.cpp new file mode 100644 index 00000000000000..51f629fd43d0c6 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_rowid_conversion_test.cpp @@ -0,0 +1,155 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/compaction/rowid_conversion.h" + +namespace { + +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using doris::snii::compaction::RowIdConversionMap; +using doris::snii::compaction::ValidatedRowIdConversion; +using doris::snii::compaction::validate_rowid_conversion; + +constexpr uint32_t kDeleted = std::numeric_limits::max(); + +void expect_invalid(const Status& status, std::string_view message_fragment) { + EXPECT_TRUE(status.is()) << status; + EXPECT_NE(status.to_string().find(message_fragment), std::string::npos) << status; +} + +TEST(SniiRowIdConversionTest, AcceptsCompleteKWayInterleavedMapping) { + const RowIdConversionMap conversion = { + {{0, 0}, {kDeleted, kDeleted}, {0, 2}, {1, 1}}, + {{0, 1}, {1, 0}}, + }; + + EXPECT_TRUE(validate_rowid_conversion(conversion, {4, 2}, {3, 2}).ok()); +} + +TEST(SniiRowIdConversionTest, CreatesValidatedTokenWithStableShapeAndMappings) { + const RowIdConversionMap conversion = { + {{0, 0}, {kDeleted, kDeleted}, {0, 2}, {1, 1}}, + {{0, 1}, {1, 0}}, + }; + const std::vector source_rows = {4, 2}; + const std::vector destination_rows = {3, 2}; + + std::unique_ptr validated; + EXPECT_TRUE( + ValidatedRowIdConversion::create(&conversion, source_rows, destination_rows, &validated) + .ok()); + ASSERT_NE(validated, nullptr); + EXPECT_EQ(validated->source_segment_count(), 2U); + EXPECT_EQ(validated->source_segment_doc_counts(), source_rows); + EXPECT_EQ(validated->destination_segment_doc_counts(), destination_rows); + ASSERT_EQ(validated->source_mapping(1).size(), 2U); + EXPECT_EQ(validated->source_mapping(1)[0], (std::pair {0, 1})); + EXPECT_TRUE(validated->source_has_deletions(0)); + EXPECT_FALSE(validated->source_has_deletions(1)); +} + +TEST(SniiRowIdConversionTest, RejectsInvalidMappingAndClearsOutputToken) { + const RowIdConversionMap valid_conversion = {{{0, 0}}}; + const std::vector source_rows = {1}; + const std::vector destination_rows = {1}; + std::unique_ptr validated; + ASSERT_TRUE(ValidatedRowIdConversion::create(&valid_conversion, source_rows, destination_rows, + &validated) + .ok()); + ASSERT_NE(validated, nullptr); + + const RowIdConversionMap incomplete_conversion = {{{0, 0}}}; + const std::vector incomplete_destination_rows = {2}; + const Status status = ValidatedRowIdConversion::create(&incomplete_conversion, source_rows, + incomplete_destination_rows, &validated); + expect_invalid(status, "missing destination ordinal"); + EXPECT_EQ(validated, nullptr); +} + +TEST(SniiRowIdConversionTest, AcceptsEmptyDestinationAndDeletedSources) { + const RowIdConversionMap conversion = { + {}, + {{kDeleted, kDeleted}, {kDeleted, kDeleted}}, + }; + + EXPECT_TRUE(validate_rowid_conversion(conversion, {0, 2}, {}).ok()); +} + +TEST(SniiRowIdConversionTest, RejectsSourceShapeMismatch) { + expect_invalid(validate_rowid_conversion({{}}, {}, {}), "source segment count"); + expect_invalid(validate_rowid_conversion({{{kDeleted, kDeleted}}}, {2}, {}), + "source doc count"); +} + +TEST(SniiRowIdConversionTest, RejectsHalfDeletedEntries) { + expect_invalid(validate_rowid_conversion({{{kDeleted, 0}}}, {1}, {1}), "partially deleted"); + expect_invalid(validate_rowid_conversion({{{0, kDeleted}}}, {1}, {1}), "partially deleted"); +} + +TEST(SniiRowIdConversionTest, RejectsLiveDestinationOutsideBounds) { + expect_invalid(validate_rowid_conversion({{{1, 0}}}, {1}, {1}), "destination segment"); + expect_invalid(validate_rowid_conversion({{{0, 1}}}, {1}, {1}), "destination row"); +} + +TEST(SniiRowIdConversionTest, RejectsNonIncreasingSourceOrdinals) { + expect_invalid(validate_rowid_conversion({{{1, 0}, {0, 1}}}, {2}, {2, 1}), + "not strictly increasing"); + expect_invalid(validate_rowid_conversion({{{0, 0}, {0, 0}}}, {2}, {1}), + "not strictly increasing"); +} + +TEST(SniiRowIdConversionTest, RejectsDuplicateDestinationAcrossSources) { + const RowIdConversionMap conversion = { + {{0, 0}}, + {{0, 0}}, + }; + + expect_invalid(validate_rowid_conversion(conversion, {1, 1}, {1}), + "duplicate destination ordinal"); +} + +TEST(SniiRowIdConversionTest, RejectsIncompleteDestinationCoverage) { + const RowIdConversionMap conversion = { + {{0, 0}, {0, 2}}, + {{kDeleted, kDeleted}}, + }; + + expect_invalid(validate_rowid_conversion(conversion, {2, 1}, {3}), + "missing destination ordinal"); +} + +TEST(SniiRowIdConversionTest, UsesWideGlobalOrdinalsAcrossUint32Boundary) { + const RowIdConversionMap conversion = {{{1, 1}}}; + + const Status status = + validate_rowid_conversion(conversion, {1}, {std::numeric_limits::max(), 2}); + expect_invalid(status, "4294967296"); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_streamed_session_test.cpp b/be/test/storage/index/snii/compaction/snii_streamed_session_test.cpp new file mode 100644 index 00000000000000..59080f79f13b34 --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_streamed_session_test.cpp @@ -0,0 +1,1580 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/norms_pod.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/posting_window_emitter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spillable_byte_buffer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_source.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" +#include "storage/index/snii_query_test_util.h" + +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::snii_test; // NOLINT +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using doris::snii::query::phrase_query; +using doris::snii::query::term_query; +using writer::SniiCompoundWriter; +using writer::SniiIndexInput; +using writer::SniiStreamedIndexSession; +using writer::SpanTermPostingSource; +using writer::SpimiTermBuffer; +using writer::SpillableByteBuffer; +using writer::TermPostings; +using writer::StreamedTermPostings; +using writer::TermPostingBuffer; +using writer::TermPostingSource; +namespace inverted_index = doris::segment_v2::inverted_index; + +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); +static_assert(!std::is_move_assignable_v); + +// FileWriter fault used to prove that a failed append cannot be retried into a +// sealable container. The failing call writes one byte before returning an +// error, modeling append implementations that can have partial side effects. +class PartialFailOnAppendFile final : public io::FileWriter { +public: + explicit PartialFailOnAppendFile(size_t fail_on_append) : fail_on_append_(fail_on_append) {} + + Status append(Slice data) override { + ++append_calls_; + if (append_calls_ == fail_on_append_) { + if (!data.empty()) { + RETURN_IF_ERROR(backing_.append(Slice(data.data(), 1))); + } + return Status::Error( + "injected partial append failure"); + } + return backing_.append(data); + } + + Status finalize() override { + ++finalize_calls_; + return backing_.finalize(); + } + + uint64_t bytes_written() const override { return backing_.bytes_written(); } + + size_t append_calls() const { return append_calls_; } + size_t finalize_calls() const { return finalize_calls_; } + bool finalized() const { return backing_.finalized(); } + +private: + MemoryFile backing_; + size_t fail_on_append_ = 0; + size_t append_calls_ = 0; + size_t finalize_calls_ = 0; +}; + +SniiIndexInput empty_input(uint64_t index_id, std::string suffix, uint32_t doc_count = 8) { + SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = std::move(suffix); + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = doc_count; + input.write_freq = false; + return input; +} + +TermPostings make_sparse_slim_term(std::string term) { + std::vector docs; + docs.reserve(300); + for (uint32_t i = 0; i < 300; ++i) { + docs.push_back({.docid = i * 100003U, .positions = {0, 3, 9}}); + } + return make_term(std::move(term), std::move(docs)); +} + +class VectorTermPostingSource final : public TermPostingSource { +public: + explicit VectorTermPostingSource(const TermPostings* postings) : postings_(postings) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + requests.push_back(target_docs); + out_was_empty = out_was_empty && out->empty(); + const size_t remaining = postings_->docids.size() - doc_offset_; + const size_t count = std::min(target_docs, remaining); + if (count == 0) { + *exhausted = true; + return Status::OK(); + } + + const auto docids = + std::span(postings_->docids).subspan(doc_offset_, count); + if (!postings_->retain_positions) { + const auto freqs = postings_->freqs.empty() + ? std::span {} + : std::span(postings_->freqs) + .subspan(doc_offset_, count); + RETURN_IF_ERROR(out->append(docids, freqs, {})); + } else { + const auto freqs = + std::span(postings_->freqs).subspan(doc_offset_, count); + const size_t position_count = + std::accumulate(freqs.begin(), freqs.end(), static_cast(0)); + const auto positions = std::span(postings_->positions_flat) + .subspan(position_offset_, position_count); + RETURN_IF_ERROR(out->append(docids, freqs, positions)); + position_offset_ += position_count; + } + doc_offset_ += count; + *exhausted = doc_offset_ == postings_->docids.size(); + return Status::OK(); + } + + std::vector requests; + bool out_was_empty = true; + +private: + const TermPostings* postings_; + size_t doc_offset_ = 0; + size_t position_offset_ = 0; +}; + +class InvalidTermPostingSource final : public TermPostingSource { +public: + enum class Mode { + kEmptyBeforeEof, + kShortBeforeEof, + kOverfill, + kUnordered, + kOutOfRange, + kError + }; + + explicit InvalidTermPostingSource(Mode mode) : mode_(mode) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + switch (mode_) { + case Mode::kEmptyBeforeEof: + *exhausted = false; + return Status::OK(); + case Mode::kShortBeforeEof: { + const std::array docids {1}; + RETURN_IF_ERROR(out->append(docids, {}, {})); + *exhausted = false; + return Status::OK(); + } + case Mode::kOverfill: { + std::vector docids(target_docs + 1); + std::iota(docids.begin(), docids.end(), 0); + RETURN_IF_ERROR(out->append(docids, {}, {})); + *exhausted = true; + return Status::OK(); + } + case Mode::kUnordered: { + const std::array docids {2, 1}; + RETURN_IF_ERROR(out->append(docids, {}, {})); + *exhausted = true; + return Status::OK(); + } + case Mode::kOutOfRange: { + const std::array docids {32}; + RETURN_IF_ERROR(out->append(docids, {}, {})); + *exhausted = true; + return Status::OK(); + } + case Mode::kError: + return Status::Error("injected source failure"); + } + __builtin_unreachable(); + } + +private: + Mode mode_; +}; + +class InvalidPositionShapeSource final : public TermPostingSource { +public: + explicit InvalidPositionShapeSource(size_t position_count) : position_count_(position_count) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + if (target_docs == 0 || out == nullptr || exhausted == nullptr || !out->empty()) { + return Status::Error( + "invalid position source: invalid fill arguments"); + } + writer::MutableTermPostingSpan destination; + RETURN_IF_ERROR(out->grow_uninitialized(/*document_count=*/1, /*has_freqs=*/true, + position_count_, &destination)); + destination.docids[0] = 0; + destination.freqs[0] = 2; + std::iota(destination.positions_flat.begin(), destination.positions_flat.end(), 0); + *exhausted = true; + return Status::OK(); + } + +private: + size_t position_count_ = 0; +}; + +TermPostings make_uniform_term(std::string term, uint32_t doc_count) { + TermPostings postings; + postings.term = std::move(term); + postings.docids.resize(doc_count); + std::iota(postings.docids.begin(), postings.docids.end(), 0); + postings.freqs.assign(doc_count, 1); + postings.positions_flat.assign(doc_count, 0); + return postings; +} + +TermPostings make_adaptive_boundary_term(uint32_t tail_docs = 512) { + constexpr uint32_t kPrefixDocs = format::kAdaptiveWindowDfThreshold; + constexpr uint32_t kFarPosition = 1U << 28; + const uint32_t doc_count = kPrefixDocs + tail_docs; + TermPostings postings; + postings.term = "adaptive-boundary"; + postings.docids.resize(doc_count); + std::iota(postings.docids.begin(), postings.docids.end(), 0); + postings.freqs.assign(doc_count, 1); + postings.positions_flat.reserve(kPrefixDocs + 2 * tail_docs); + for (uint32_t doc = 0; doc < doc_count; ++doc) { + postings.positions_flat.push_back(0); + if (doc >= kPrefixDocs) { + postings.freqs[doc] = 2; + postings.positions_flat.push_back(kFarPosition); + } + } + return postings; +} + +class FailAfterFirstFillSource final : public TermPostingSource { +public: + explicit FailAfterFirstFillSource(const TermPostings* postings) : delegate_(postings) {} + + Status fill(uint32_t target_docs, TermPostingBuffer* out, bool* exhausted) override { + if (fill_count_++ != 0) { + return Status::Error( + "injected source failure after the first fill"); + } + return delegate_.fill(target_docs, out, exhausted); + } + +private: + VectorTermPostingSource delegate_; + size_t fill_count_ = 0; +}; + +SniiIndexInput representative_input(bool write_freq = false) { + SniiIndexInput input; + input.index_id = 71; + input.index_suffix = "body"; + input.config = format::IndexConfig::kDocsPositions; + input.write_freq = write_freq; + input.target_dict_block_bytes = 256; + + input.terms.push_back(make_term("aa_inline", {{.docid = 7, .positions = {1}}})); + input.terms.push_back(make_sparse_slim_term("bb_slim")); + + std::vector wide_docs; + wide_docs.reserve(600); + for (uint32_t i = 0; i < 600; ++i) { + wide_docs.push_back({.docid = i * 90001U + 1U, .positions = {2}}); + } + const uint32_t max_docid = wide_docs.back().docid; + input.terms.push_back(make_term("cc_windowed", std::move(wide_docs))); + + // Enough tiny vocabulary to cross multiple small DICT blocks. + for (uint32_t i = 0; i < 24; ++i) { + input.terms.push_back(make_term("zz_filler_" + std::to_string(1000 + i), + {{.docid = i, .positions = {4}}})); + } + std::ranges::sort(input.terms, [](const TermPostings& lhs, const TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + input.doc_count = max_docid + 4; + input.null_docids = {input.doc_count - 2, input.doc_count - 1}; + return input; +} + +Status write_ordinary_index(SniiIndexInput input, MemoryFile* file) { + SniiCompoundWriter compound(file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + return compound.finish(); +} + +Status push_materialized(SniiStreamedIndexSession* session, TermPostings postings) { + SpanTermPostingSource source(postings.docids, postings.freqs, postings.positions_flat); + return session->push_term(StreamedTermPostings {.term = std::move(postings.term), + .retain_positions = postings.retain_positions, + .source = &source}); +} + +Status write_streamed_index(SniiIndexInput input, MemoryFile* file) { + std::vector terms = std::move(input.terms); + input.terms.clear(); + + SniiCompoundWriter compound(file); + SniiStreamedIndexSession* session = nullptr; + RETURN_IF_ERROR(compound.begin_streamed_index(std::move(input), &session)); + for (TermPostings& term : terms) { + RETURN_IF_ERROR(push_materialized(session, std::move(term))); + } + RETURN_IF_ERROR(session->finish()); + return compound.finish(); +} + +Status write_source_index(SniiIndexInput input, const TermPostings* postings, MemoryFile* file) { + input.terms.clear(); + VectorTermPostingSource source(postings); + SniiCompoundWriter compound(file); + SniiStreamedIndexSession* session = nullptr; + RETURN_IF_ERROR(compound.begin_streamed_index(std::move(input), &session)); + RETURN_IF_ERROR(session->push_term(StreamedTermPostings { + .term = postings->term, + .retain_positions = postings->retain_positions, + .source = &source, + })); + RETURN_IF_ERROR(session->finish()); + return compound.finish(); +} + +Status begin_scoring_session_from_local_input(SniiCompoundWriter* compound, + SniiStreamedIndexSession** session) { + SniiIndexInput input; + input.index_id = 91; + input.index_suffix = "owned_scoring_body_with_heap_storage"; + input.config = format::IndexConfig::kDocsPositionsScoring; + input.doc_count = 4; + input.null_docids = {1, 3}; + input.encoded_norms = {7, 11, 13, 17}; + input.common_grams_metadata = make_plain_scoring_metadata(input.doc_count, 4); + return compound->begin_streamed_index(std::move(input), session); +} + +SniiIndexInput common_grams_input(uint64_t index_id, std::string suffix) { + SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = std::move(suffix); + input.config = format::IndexConfig::kDocsPositionsScoring; + input.doc_count = 1; + input.encoded_norms = {9}; + inverted_index::CommonGramsQueryIdentity identity {.common_grams_dictionary_identity = "dict-a", + .base_analyzer_fingerprint = "base-a", + .common_grams_fingerprint = "grams-a"}; + input.common_grams_metadata = inverted_index::make_common_grams_segment_metadata(identity); + input.common_grams_metadata->scoring_doc_count = input.doc_count; + return input; +} + +TEST(SniiStreamedWriterSessionTest, OrdinaryAndStreamedImagesAreByteIdentical) { + for (bool write_freq : {false, true}) { + MemoryFile ordinary; + MemoryFile streamed; + assert_ok(write_ordinary_index(representative_input(write_freq), &ordinary)); + assert_ok(write_streamed_index(representative_input(write_freq), &streamed)); + + EXPECT_EQ(streamed.data(), ordinary.data()); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&streamed, &segment)); + assert_ok(segment.open_index(71, "body", &index)); + EXPECT_GT(index.n_dict_blocks(), 1U); + EXPECT_EQ(index.stats().null_count, 2U); + + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("aa_inline", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.kind, format::DictEntryKind::kInline); + EXPECT_EQ(entry.term_stats_present, write_freq); + + assert_ok(index.lookup("bb_slim", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kSlim); + EXPECT_EQ(entry.term_stats_present, write_freq); + + assert_ok(index.lookup("cc_windowed", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kWindowed); + EXPECT_EQ(entry.term_stats_present, write_freq); + } +} + +TEST(SniiStreamedWriterSessionTest, CompleteCommonGramSkipsTermFreqStatsScan) { + SniiIndexInput input = common_grams_input(113, "gram_freq_stats"); + input.terms.push_back(make_term("plain", {{.docid = 0, .positions = {1}}})); + auto gram = inverted_index::encode_common_gram("of", "the"); + ASSERT_TRUE(gram.has_value()); + input.terms.push_back(make_term(std::move(gram.value()), {{.docid = 0, .positions = {0, 2}}})); + std::ranges::sort(input.terms, [](const TermPostings& lhs, const TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::testing::reset_term_freq_scans(); + MemoryFile file; + std::vector terms = std::move(input.terms); + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + assert_ok(session->set_semantic_token_count(1)); + for (TermPostings& term : terms) { + assert_ok(push_materialized(session, std::move(term))); + } + assert_ok(session->finish()); + assert_ok(compound.finish()); + + // The ordinary term still needs ttf/max_freq, while the statless gram can + // derive ttf from its already-known position count and has no max consumer. + EXPECT_EQ(writer::testing::term_freq_scans(), 1U); +} + +TEST(SniiStreamedWriterSessionTest, CompleteWindowedCommonGramSkipsUnusedWandScans) { + constexpr uint32_t kDocs = format::kSlimDfThreshold; + SniiIndexInput input = common_grams_input(117, "gram_window_norm_stats"); + input.doc_count = kDocs; + input.encoded_norms.resize(kDocs); + std::iota(input.encoded_norms.begin(), input.encoded_norms.end(), uint8_t {1}); + input.common_grams_metadata->scoring_doc_count = kDocs; + + std::vector docs; + docs.reserve(kDocs); + for (uint32_t docid = 0; docid < kDocs; ++docid) { + docs.push_back({.docid = docid, .positions = {docid}}); + } + auto gram = inverted_index::encode_common_gram("of", "the"); + ASSERT_TRUE(gram.has_value()); + TermPostings term = make_term(std::move(gram.value()), docs); + + writer::testing::reset_window_norm_doc_visits(); + writer::testing::reset_window_freq_doc_visits(); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + assert_ok(session->set_semantic_token_count(kDocs)); + assert_ok(push_materialized(session, std::move(term))); + assert_ok(session->finish()); + assert_ok(compound.finish()); + + EXPECT_EQ(writer::testing::window_norm_doc_visits(), 0U); + EXPECT_EQ(writer::testing::window_freq_doc_visits(), 0U); +} + +TEST(SniiStreamedWriterSessionTest, CompleteDocsOnlyGramsOmitPrxInEveryEntryShape) { + constexpr uint32_t kDocs = format::kSlimDfThreshold; + SniiIndexInput input = common_grams_input(118, "gram_docs_only_shapes"); + input.config = format::IndexConfig::kDocsPositions; + input.write_freq = false; + input.encoded_norms.clear(); + input.common_grams_metadata->scoring_coverage = inverted_index::ScoringCoverage::kNone; + input.common_grams_metadata->scoring_stats_version = 0; + input.common_grams_metadata->norm_semantics_version = 0; + input.common_grams_metadata->scoring_doc_count = 0; + + const auto make_docs_only_gram = [](std::string left, std::string right, uint32_t count, + uint32_t stride) { + TermPostings term; + term.term = inverted_index::encode_common_gram(left, right).value(); + term.retain_positions = false; + term.docids.resize(count); + for (uint32_t i = 0; i < count; ++i) { + term.docids[i] = i * stride; + } + term.freqs.assign(count, 1U); + return term; + }; + input.terms.push_back(make_docs_only_gram("a", "of", 1, 1)); + input.terms.push_back(make_docs_only_gram("b", "of", 300, 100003)); + input.terms.push_back(make_docs_only_gram("c", "of", kDocs, 1)); + input.doc_count = input.terms[1].docids.back() + 1; + std::ranges::sort(input.terms, [](const TermPostings& lhs, const TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::testing::reset_term_freq_scans(); + writer::testing::reset_window_norm_doc_visits(); + writer::testing::reset_window_freq_doc_visits(); + MemoryFile file; + SniiCompoundWriter compound(&file); + assert_ok(compound.add_logical_index(input)); + assert_ok(compound.finish()); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(118, "gram_docs_only_shapes", &index)); + for (size_t i = 0; i < input.terms.size(); ++i) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup(input.terms[i].term, &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + if (i == 0) { + EXPECT_EQ(entry.kind, format::DictEntryKind::kInline); + EXPECT_TRUE(entry.prx_bytes.empty()); + } else { + EXPECT_EQ(entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(entry.prx_len, 0U); + } + } + EXPECT_EQ(writer::testing::term_freq_scans(), 0U); + EXPECT_EQ(writer::testing::window_norm_doc_visits(), 0U); + EXPECT_EQ(writer::testing::window_freq_doc_visits(), 0U); +} + +TEST(SniiStreamedWriterSessionTest, MixedCommonGramKeepsTermFreqStatsScan) { + SniiIndexInput input = common_grams_input(114, "mixed_gram_freq_stats"); + input.common_grams_metadata->common_grams_coverage = + inverted_index::CommonGramsCoverage::kMixed; + auto gram = inverted_index::encode_common_gram("of", "the"); + ASSERT_TRUE(gram.has_value()); + TermPostings term = make_term(std::move(gram.value()), {{.docid = 0, .positions = {0, 2}}}); + + writer::testing::reset_term_freq_scans(); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + assert_ok(session->set_semantic_token_count(1)); + assert_ok(push_materialized(session, std::move(term))); + assert_ok(session->finish()); + assert_ok(compound.finish()); + + EXPECT_EQ(writer::testing::term_freq_scans(), 1U); +} + +TEST(SniiStreamedWriterSessionTest, StatlessCommonGramRejectsInvalidPositionPartition) { + SniiIndexInput input = common_grams_input(115, "invalid_gram_partition"); + auto gram = inverted_index::encode_common_gram("of", "the"); + ASSERT_TRUE(gram.has_value()); + TermPostings term = make_term(std::move(gram.value()), {{.docid = 0, .positions = {0, 2}}}); + term.freqs[0] = 1; + + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + assert_ok(session->set_semantic_token_count(1)); + EXPECT_TRUE(push_materialized(session, std::move(term)).is()); +} + +TEST(SniiStreamedWriterSessionTest, StatlessWindowedGramRejectsInvalidPositionPartition) { + constexpr uint32_t kDocs = format::kSlimDfThreshold; + SniiIndexInput input = common_grams_input(116, "invalid_windowed_gram_partition"); + input.doc_count = kDocs; + input.encoded_norms.assign(kDocs, 1); + input.common_grams_metadata->scoring_doc_count = kDocs; + std::vector docs; + docs.reserve(kDocs); + for (uint32_t docid = 0; docid < kDocs; ++docid) { + docs.push_back({.docid = docid, .positions = {0}}); + } + auto gram = inverted_index::encode_common_gram("of", "the"); + ASSERT_TRUE(gram.has_value()); + TermPostings term = make_term(std::move(gram.value()), std::move(docs)); + term.freqs.back() = 2; + + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + assert_ok(session->set_semantic_token_count(kDocs)); + EXPECT_TRUE(push_materialized(session, std::move(term)).is()); +} + +TEST(SniiStreamedWriterSessionTest, NullDocidsStorageIsTransferredIntoStreamedWriter) { + SniiIndexInput input = empty_input(94, "null_handoff", 4); + input.null_docids = {1, 3}; + const uint32_t* const original_storage = input.null_docids.data(); + + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + ASSERT_NE(session, nullptr); + EXPECT_EQ(session->writer_->null_docids_.data(), original_storage); + assert_ok(session->finish()); + assert_ok(compound.finish()); +} + +TEST(SniiStreamedWriterSessionTest, NullBitmapFinalizationHonorsReporterCap) { + SniiIndexInput input = empty_input(95, "null_memory_cap", 2); + input.null_docids = {1}; + const uint64_t retained_null_bytes = input.null_docids.capacity() * sizeof(uint32_t); + const uint64_t bitmap_build_bytes = format::NullBitmapWriter::build_memory_upper_bound( + std::span(input.null_docids)); + writer::MemoryReporter reporter(nullptr, retained_null_bytes + bitmap_build_bytes - 1); + input.mem_reporter = &reporter; + + { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + ASSERT_NE(session, nullptr); + const Status status = session->finish(); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(reporter.current_bytes(), retained_null_bytes); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, PublicOverloadChargesNullDocidCapacity) { + SniiIndexInput input = empty_input(95, "null_public_capacity", 2); + input.null_docids.reserve(8); + input.null_docids.push_back(1); + const uint64_t retained_null_bytes = input.null_docids.capacity() * sizeof(uint32_t); + writer::MemoryReporter reporter(nullptr, retained_null_bytes - 1); + input.mem_reporter = &reporter; + + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + const Status status = compound.begin_streamed_index(std::move(input), &session); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(session, nullptr); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, NullBitmapMemoryIsReleasedAfterCompoundAppend) { + SniiIndexInput input = empty_input(96, "null_memory_release", 2); + input.null_docids = {1}; + writer::MemoryReporter reporter(nullptr, 1U << 20); + input.mem_reporter = &reporter; + + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + ASSERT_NE(session, nullptr); + assert_ok(session->finish()); + EXPECT_GT(reporter.current_bytes(), 0); + assert_ok(compound.finish()); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, EmptyTrackedNullCapacityIsReleasedAtFinalize) { + std::vector null_docids; + null_docids.reserve(4); + writer::MemoryReporter reporter(nullptr, null_docids.capacity() * sizeof(uint32_t)); + auto reservation = reporter.make_reservation(); + assert_ok(reservation.set_bytes(null_docids.capacity() * sizeof(uint32_t))); + writer::TrackedNullDocids tracked(std::move(reservation), std::move(null_docids)); + + SniiIndexInput input = empty_input(97, "empty_null_capacity", 0); + input.mem_reporter = &reporter; + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), std::move(tracked), &session)); + ASSERT_NE(session, nullptr); + assert_ok(session->finish()); + EXPECT_EQ(reporter.current_bytes(), 0); + assert_ok(compound.finish()); +} + +TEST(SniiStreamedWriterSessionTest, TightPrxLimitsRecutSlimTermAtDocumentBoundaries) { + SniiIndexInput input; + input.index_id = 72; + input.index_suffix = "tight_prx"; + input.config = format::IndexConfig::kDocsPositions; + input.write_freq = true; + input.doc_count = 3; + input.prx_window_limits = { + .max_docs = 8, + .max_positions = 5, + .max_uncomp_bytes = 64, + }; + input.terms.push_back(make_term("dense", {{.docid = 0, .positions = {0, 1, 2}}, + {.docid = 1, .positions = {0, 1, 2}}, + {.docid = 2, .positions = {0, 1, 2}}})); + input.terms.push_back(make_term("tail", {{.docid = 0, .positions = {3}}, + {.docid = 1, .positions = {3}}, + {.docid = 2, .positions = {3}}})); + + MemoryFile file; + assert_ok(write_streamed_index(std::move(input), &file)); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(72, "tight_prx", &index)); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("dense", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kWindowed); + + std::vector docs; + assert_ok(phrase_query(index, {"dense", "tail"}, &docs)); + EXPECT_EQ(docs, (std::vector {0, 1, 2})); +} + +TEST(SniiStreamedWriterSessionTest, TightByteLimitKeepsReadableLegacySlimFrame) { + SniiIndexInput input; + input.index_id = 73; + input.index_suffix = "tight_prx_readable"; + input.config = format::IndexConfig::kDocsPositions; + input.write_freq = true; + input.doc_count = 3; + input.prx_window_limits = { + .max_docs = 8, + .max_positions = 10, + .max_uncomp_bytes = 64, + }; + input.terms.push_back(make_term("dense", {{.docid = 0, .positions = {0, 1, 2}}, + {.docid = 1, .positions = {0, 1, 2}}, + {.docid = 2, .positions = {0, 1, 2}}})); + + MemoryFile file; + assert_ok(write_streamed_index(std::move(input), &file)); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(73, "tight_prx_readable", &index)); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("dense", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kSlim); +} + +#ifdef BE_TEST +TEST(SniiStreamedWriterSessionTest, AutoCodecChoosesSmallestReadableRawFrame) { + // All three codecs are reader-safe here. Complete frame sizes are PFOR=707, + // ZSTD=659, and RAW=607 bytes, so RAW must win. + EXPECT_EQ(format::testing::select_auto_prx_codec_for_test( + /*pfor_payload_size=*/700, /*plain_payload_size=*/600, + /*compressed_payload_size=*/650, /*max_uncomp_bytes=*/650), + static_cast(format::PrxCodec::kRaw)); +} +#endif + +TEST(SniiStreamedWriterSessionTest, PrxBuilderRejectsUnreadableFramesBeforeAppending) { + ByteSink sink; + sink.put_u8(0xA5); + const std::vector before = sink.buffer(); + const format::PrxWindowLimits position_limit { + .max_docs = 4, + .max_positions = 3, + .max_uncomp_bytes = 64, + }; + const std::vector positions = {0, 1, 0, 1}; + const std::vector freqs = {2, 2}; + format::PrxWindowBuildOutcome outcome = format::PrxWindowBuildOutcome::kBuilt; + assert_ok(format::try_build_prx_window_flat(positions, freqs, -3, position_limit, &sink, + &outcome)); + EXPECT_EQ(outcome, format::PrxWindowBuildOutcome::kNeedsSplit); + EXPECT_EQ(sink.buffer(), before); + EXPECT_TRUE(format::build_prx_window_flat(positions, freqs, -3, position_limit, &sink) + .is()); + EXPECT_EQ(sink.buffer(), before); + + const std::vector unsplittable_positions = {0, 1, 2, 3}; + const std::vector unsplittable_freqs = {4}; + EXPECT_TRUE(format::try_build_prx_window_flat(unsplittable_positions, unsplittable_freqs, -3, + position_limit, &sink, &outcome) + .is()); + EXPECT_EQ(sink.buffer(), before); + + const format::PrxWindowLimits byte_limit { + .max_docs = 4, + .max_positions = 4, + .max_uncomp_bytes = 4, + }; + const std::vector wide_delta_positions = {0, UINT32_MAX}; + const std::vector two_positions = {2}; + EXPECT_TRUE( + format::build_prx_window_flat(wide_delta_positions, two_positions, 0, byte_limit, &sink) + .is()); + EXPECT_EQ(sink.buffer(), before); +} + +TEST(SniiStreamedWriterSessionTest, PrxBuilderRejectsMultiDocWindowWithUnsplittableDocument) { + ByteSink sink; + sink.put_u8(0xA5); + const std::vector before = sink.buffer(); + const format::PrxWindowLimits byte_limit { + .max_docs = 4, + .max_positions = 4, + .max_uncomp_bytes = 4, + }; + const std::vector positions = {0, UINT32_MAX, 0}; + const std::vector freqs = {2, 1}; + format::PrxWindowBuildOutcome outcome = format::PrxWindowBuildOutcome::kBuilt; + + EXPECT_TRUE(format::try_build_prx_window_flat(positions, freqs, 0, byte_limit, &sink, &outcome) + .is()); + EXPECT_EQ(sink.buffer(), before); + EXPECT_EQ(outcome, format::PrxWindowBuildOutcome::kBuilt); + + EXPECT_TRUE(format::try_build_prx_window_flat(positions, freqs, -3, byte_limit, &sink, &outcome) + .is()); + EXPECT_EQ(sink.buffer(), before); + EXPECT_EQ(outcome, format::PrxWindowBuildOutcome::kBuilt); + + const format::PrxWindowLimits shape_and_byte_limit { + .max_docs = 4, + .max_positions = 2, + .max_uncomp_bytes = 4, + }; + EXPECT_TRUE(format::try_build_prx_window_flat(positions, freqs, 0, shape_and_byte_limit, &sink, + &outcome) + .is()); + EXPECT_EQ(sink.buffer(), before); + EXPECT_EQ(outcome, format::PrxWindowBuildOutcome::kBuilt); +} + +TEST(SniiStreamedWriterSessionTest, ExplicitReaderPrxLimitsPreserveDefaultBytes) { + const std::vector positions = {0, 2, 7, 1, 9}; + const std::vector freqs = {3, 2}; + ByteSink legacy; + ByteSink explicit_limits; + assert_ok(format::build_prx_window_flat(positions, freqs, -3, &legacy)); + assert_ok(format::build_prx_window_flat(positions, freqs, -3, format::kReaderPrxWindowLimits, + &explicit_limits)); + EXPECT_EQ(explicit_limits.buffer(), legacy.buffer()); +} + +TEST(SniiStreamedWriterSessionTest, DefaultSingletonRawFrameMatchesGolden) { + const std::vector positions = {0}; + const std::vector freqs = {1}; + ByteSink frame; + assert_ok(format::build_prx_window_flat(positions, freqs, -3, &frame)); + + const std::vector expected = {0x00, 0x03, 0x01, 0x01, 0x00, 0x05, 0xF5, 0xB3, 0x91}; + EXPECT_EQ(frame.buffer(), expected); +} + +TEST(SniiStreamedWriterSessionTest, ByteLimitRecutReusesSourcePositionsAndFinalPlans) { + constexpr uint32_t kDocs = 512; + constexpr uint32_t kFarPosition = 1U << 28; + + TermPostings dense; + dense.term = "dense"; + dense.docids.resize(kDocs); + std::iota(dense.docids.begin(), dense.docids.end(), 0); + dense.freqs.assign(kDocs, 2); + dense.positions_flat.resize(2 * kDocs); + for (size_t index = 0; index < dense.positions_flat.size(); ++index) { + dense.positions_flat[index] = (index & 1U) == 0 ? 0 : kFarPosition; + } + + std::vector tail_docs; + tail_docs.reserve(kDocs); + for (uint32_t docid = 0; docid < kDocs; ++docid) { + tail_docs.push_back({.docid = docid, .positions = {1}}); + } + + SniiIndexInput input; + input.index_id = 74; + input.index_suffix = "byte_recut"; + input.config = format::IndexConfig::kDocsPositions; + input.write_freq = true; + input.doc_count = kDocs; + input.prx_window_limits = { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 8, + }; + input.terms.push_back(std::move(dense)); + input.terms.push_back(make_term("tail", std::move(tail_docs))); + + MemoryFile file; + assert_ok(write_streamed_index(std::move(input), &file)); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(74, "byte_recut", &index)); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("dense", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kWindowed); + + std::vector expected_docs(kDocs); + std::iota(expected_docs.begin(), expected_docs.end(), 0); + std::vector actual_docs; + assert_ok(phrase_query(index, {"dense", "tail"}, &actual_docs)); + EXPECT_EQ(actual_docs, expected_docs); +} + +TEST(SniiStreamedWriterSessionTest, ReleasesResidentDictAfterStreamingItIntoContainer) { + writer::MemoryReporter reporter; + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiIndexInput input = representative_input(/*write_freq=*/false); + std::vector terms = std::move(input.terms); + input.terms.clear(); + input.mem_reporter = &reporter; + + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + for (TermPostings& term : terms) { + assert_ok(push_materialized(session, std::move(term))); + } + assert_ok(session->finish()); + + // finish() has already copied the complete DICT region into the compound + // output. The only remaining tracked bytes are the BSBF that must be laid + // out after every logical index's posting/DICT regions. + EXPECT_GT(reporter.current_bytes(), 0); + assert_ok(compound.finish()); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, TermHashGrowthIsReservedBeforeAllocation) { + writer::MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/7); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiIndexInput input = empty_input(98, "term_hash_budget", 1); + input.mem_reporter = &reporter; + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + + const Status status = + push_materialized(session, make_term("term", {{.docid = 0, .positions = {0}}})); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, FinalizesAfterPersistentVocabularyCrossesSpillThreshold) { + constexpr uint64_t kReporterCap = 384U << 10; + constexpr uint32_t kTermCount = 3000; + for (bool force_spill : {false, true}) { + writer::MemoryReporter reporter( + /*consume_release=*/nullptr, kReporterCap, + writer::MemoryReporter::CapPolicy::kSpillThreshold); + SpimiTermBuffer terms(/*has_positions=*/true, kReporterCap, &reporter); + terms.set_forced_spill_min_arena_bytes(1); + for (uint32_t i = 0; i < kTermCount; ++i) { + if (force_spill && i == kTermCount / 2) { + terms.request_global_spill_for_test(); + } + terms.add_token("term_" + std::to_string(100000 + i), /*docid=*/0, /*pos=*/i); + } + assert_ok(terms.status()); + ASSERT_GT(reporter.current_bytes(), static_cast(kReporterCap)); + if (force_spill) { + ASSERT_GT(terms.run_count_for_test(), 0); + } else { + ASSERT_EQ(terms.run_count_for_test(), 0); + } + + MemoryFile file; + SniiCompoundWriter compound(&file); + const std::string suffix = + force_spill ? "persistent_vocab_spilled" : "persistent_vocab_in_memory"; + const uint64_t index_id = force_spill ? 100 : 99; + SniiIndexInput input = empty_input(index_id, suffix, /*doc_count=*/1); + input.target_dict_block_bytes = 16U << 10; + input.term_source = &terms; + input.mem_reporter = &reporter; + assert_ok(compound.add_logical_index(input)); + assert_ok(compound.finish()); + EXPECT_EQ(reporter.current_bytes(), 0); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(index_id, suffix, &index)); + EXPECT_EQ(index.stats().term_count, kTermCount); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("term_100000", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + assert_ok(index.lookup("term_102999", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + } +} + +TEST(SniiStreamedWriterSessionTest, TerminalDrainReleasesLookupStateWithoutUnderAccountingIds) { + constexpr uint64_t kReporterCap = 384U << 10; + constexpr uint32_t kTermCount = 3000; + for (bool force_spill : {false, true}) { + writer::MemoryReporter reporter( + /*consume_release=*/nullptr, kReporterCap, + writer::MemoryReporter::CapPolicy::kSpillThreshold); + SpimiTermBuffer terms(/*has_positions=*/true, kReporterCap, &reporter); + terms.set_forced_spill_min_arena_bytes(1); + for (uint32_t i = 0; i < kTermCount; ++i) { + if (force_spill && i == kTermCount / 2) { + terms.request_global_spill_for_test(); + } + terms.add_token("term_" + std::to_string(100000 + i), /*docid=*/0, /*pos=*/i); + } + assert_ok(terms.status()); + if (force_spill) { + ASSERT_GT(terms.run_count_for_test(), 0); + } else { + ASSERT_EQ(terms.run_count_for_test(), 0); + } + const uint64_t before_drain = terms.resident_bytes_for_test(); + ASSERT_EQ(reporter.current_bytes(), static_cast(before_drain)); + ASSERT_GT(before_drain, kReporterCap); + + bool saw_first_term = false; + assert_ok(terms.for_each_term_sorted([&](writer::StreamedTermPostings&& source) { + if (!saw_first_term) { + const uint64_t during_drain = terms.resident_bytes_for_test(); + if (force_spill) { + EXPECT_GT(reporter.current_bytes(), static_cast(during_drain)); + } else { + EXPECT_EQ(reporter.current_bytes(), static_cast(during_drain)); + } + EXPECT_LT(during_drain, before_drain); + if (!force_spill) { + EXPECT_GE(during_drain, static_cast(kTermCount) * sizeof(uint32_t)); + } + saw_first_term = true; + } + return writer::consume_streamed_term(std::move(source)); + })); + EXPECT_TRUE(saw_first_term); + EXPECT_EQ(reporter.current_bytes(), 0); + } +} + +TEST(SniiStreamedWriterSessionTest, MoveAppendAccountsRetainedVectorCapacity) { + constexpr size_t kLogicalBytes = 7; + constexpr size_t kReservedBytes = 4096; + std::vector bytes(kLogicalBytes, 0x5A); + bytes.reserve(kReservedBytes); + const size_t retained_capacity = bytes.capacity(); + ASSERT_GT(retained_capacity, bytes.size()); + const std::vector expected = bytes; + + writer::MemoryReporter reporter; + SpillableByteBuffer buffer(/*cap_bytes=*/1U << 20, "capacity_accounting", &reporter); + assert_ok(buffer.append_move(std::move(bytes))); + + EXPECT_EQ(buffer.size(), kLogicalBytes); + EXPECT_EQ(reporter.current_bytes(), static_cast(retained_capacity)); + + MemoryFile output; + assert_ok(buffer.seal()); + assert_ok(buffer.stream_into_and_release(&output)); + EXPECT_EQ(buffer.size(), kLogicalBytes); + EXPECT_EQ(output.data(), expected); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, ReusesTransferCapacityAcrossAdjacentTerms) { + uint64_t positive_reservations = 0; + writer::MemoryReporter reporter([&](int64_t delta) { + if (delta > 0) { + ++positive_reservations; + } + }); + + { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiIndexInput input = empty_input(109, "reused_transfer_capacity", /*doc_count=*/4); + input.config = format::IndexConfig::kDocsPositions; + input.mem_reporter = &reporter; + + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + for (uint32_t i = 0; i < 12; ++i) { + assert_ok(push_materialized(session, make_term("term_" + std::to_string(100 + i), + {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {3}}}))); + } + const uint64_t reservations_after_warmup = positive_reservations; + assert_ok(push_materialized(session, + make_term("term_112", {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {3}}}))); + EXPECT_EQ(positive_reservations, reservations_after_warmup); + + assert_ok(session->finish()); + assert_ok(compound.finish()); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiStreamedWriterSessionTest, StreamsSpillWithSubMegabyteReporterHeadroom) { + constexpr size_t kSpilledBytes = (1U << 20) + 4096; + constexpr uint64_t kReadHeadroom = 64U << 10; + constexpr uint64_t kReporterCap = 2U << 20; + + std::vector bytes(kSpilledBytes); + for (size_t i = 0; i < bytes.size(); ++i) { + bytes[i] = static_cast(i % 251); + } + + writer::MemoryReporter reporter(/*consume_release=*/nullptr, kReporterCap); + auto retained_reservation = reporter.make_reservation(); + assert_ok(retained_reservation.set_bytes(kReporterCap - kReadHeadroom)); + + SpillableByteBuffer buffer(std::numeric_limits::max(), "small_read_headroom", + &reporter); + assert_ok(buffer.append(Slice(bytes))); + ASSERT_TRUE(buffer.spilled()); + assert_ok(buffer.seal()); + + MemoryFile output; + assert_ok(buffer.stream_into_and_release(&output)); + EXPECT_EQ(output.data(), bytes); + EXPECT_EQ(reporter.current_bytes(), static_cast(kReporterCap - kReadHeadroom)); +} + +TEST(SniiStreamedWriterSessionTest, SessionOwnsMovedInputUntilContainerFinish) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(begin_scoring_session_from_local_input(&compound, &session)); + + // Reuse heap storage after the caller-side SniiIndexInput is gone. The + // session must still own encoded_norms and every other referenced vector. + std::vector heap_churn(64, std::string(4096, 'x')); + ASSERT_EQ(heap_churn.front().size(), 4096U); + + assert_ok(push_materialized(session, make_term("alpha", {{.docid = 0, .positions = {0}}, + {.docid = 2, .positions = {0}}}))); + assert_ok(push_materialized(session, make_term("beta", {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {2}}}))); + assert_ok(session->set_semantic_token_count(4)); + assert_ok(session->finish()); + assert_ok(compound.finish()); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(91, "owned_scoring_body_with_heap_storage", &index)); + EXPECT_EQ(index.stats().doc_count, 4U); + EXPECT_EQ(index.stats().indexed_doc_count, 2U); + EXPECT_EQ(index.stats().null_count, 2U); + + format::NormsPodReader norms; + assert_ok(index.open_norms(&norms)); + ASSERT_EQ(norms.doc_count(), 4U); + EXPECT_EQ(norms.encoded_norm(0), 7U); + EXPECT_EQ(norms.encoded_norm(1), 11U); + EXPECT_EQ(norms.encoded_norm(2), 13U); + EXPECT_EQ(norms.encoded_norm(3), 17U); + + std::vector term_docs; + assert_ok(term_query(index, "alpha", &term_docs)); + EXPECT_EQ(term_docs, (std::vector {0, 2})); + std::vector phrase_docs; + assert_ok(phrase_query(index, {"alpha", "beta"}, &phrase_docs)); + EXPECT_EQ(phrase_docs, (std::vector {0})); +} + +TEST(SniiStreamedWriterSessionTest, ActiveAndFinishedSessionLifecycleIsEnforced) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* first = nullptr; + assert_ok(compound.begin_streamed_index(empty_input(101, "first"), &first)); + + EXPECT_TRUE(compound.add_logical_index(empty_input(102, "blocked_add")) + .is()); + SniiStreamedIndexSession* blocked = first; + EXPECT_TRUE(compound.begin_streamed_index(empty_input(103, "blocked_begin"), &blocked) + .is()); + EXPECT_EQ(blocked, nullptr); + EXPECT_TRUE(compound.finish().is()); + + assert_ok(push_materialized(first, make_term("alpha", {{.docid = 0, .positions = {0}}}))); + assert_ok(first->finish()); + EXPECT_TRUE(first->finished()); + EXPECT_TRUE(push_materialized(first, make_term("beta", {{.docid = 1, .positions = {0}}})) + .is()); + EXPECT_TRUE(first->finish().is()); + + // Starting a second session may grow the owner's session vector. The first + // raw handle remains valid and inert for the compound writer's lifetime. + SniiStreamedIndexSession* second = nullptr; + assert_ok(compound.begin_streamed_index(empty_input(104, "second"), &second)); + EXPECT_TRUE(first->finish().is()); + assert_ok(push_materialized(second, make_term("bravo", {{.docid = 2, .positions = {1}}}))); + assert_ok(second->finish()); + assert_ok(compound.finish()); + EXPECT_TRUE(file.finalized()); + + reader::SniiSegmentReader segment; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + bool exists = false; + assert_ok(segment.index_exists(101, "first", &exists)); + EXPECT_TRUE(exists); + assert_ok(segment.index_exists(104, "second", &exists)); + EXPECT_TRUE(exists); + assert_ok(segment.index_exists(102, "blocked_add", &exists)); + EXPECT_FALSE(exists); +} + +TEST(SniiStreamedWriterSessionTest, SemanticTokenCountIsLateBoundExactlyOnceBeforeFinish) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(common_grams_input(105, "common_grams"), &session)); + assert_ok(push_materialized(session, make_term("alpha", {{.docid = 0, .positions = {0, 1}}}))); + + assert_ok(session->set_semantic_token_count(2)); + EXPECT_TRUE(session->set_semantic_token_count(2).is()); + assert_ok(session->finish()); + EXPECT_TRUE(session->set_semantic_token_count(2).is()); + assert_ok(compound.finish()); + + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(105, "common_grams", &index)); + ASSERT_NE(index.common_grams_metadata(), nullptr); + EXPECT_EQ(index.common_grams_metadata()->scoring_doc_count, 1U); + EXPECT_EQ(index.common_grams_metadata()->scoring_token_count, 2U); +} + +TEST(SniiStreamedWriterSessionTest, CommonGramsFinishRequiresSemanticTokenCount) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(common_grams_input(106, "missing_token_count"), + &session)); + assert_ok(push_materialized(session, make_term("alpha", {{.docid = 0, .positions = {0}}}))); + + EXPECT_TRUE(session->finish().is()); + EXPECT_FALSE(session->finished()); +} + +TEST(SniiStreamedWriterSessionTest, TermOrderRejectionPoisonsCompoundSession) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(empty_input(111, "ordered"), &session)); + + assert_ok(push_materialized(session, make_term("bravo", {{.docid = 0, .positions = {0}}}))); + EXPECT_TRUE(push_materialized(session, make_term("alpha", {{.docid = 1, .positions = {0}}})) + .is()); + EXPECT_TRUE(push_materialized(session, make_term("charlie", {{.docid = 2, .positions = {0}}})) + .is()); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_FALSE(file.finalized()); +} + +TEST(SniiStreamedWriterSessionTest, RejectsNullDocidsOutsideTheDocumentDomain) { + const std::vector> invalid_nulls {{0, 0}, {1}}; + for (const std::vector& nulls : invalid_nulls) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiIndexInput input = empty_input(112, "invalid_nulls", /*doc_count=*/1); + input.null_docids = nulls; + SniiStreamedIndexSession* session = nullptr; + EXPECT_TRUE(compound.begin_streamed_index(std::move(input), &session) + .is()); + EXPECT_EQ(session, nullptr); + EXPECT_FALSE(file.finalized()); + } +} + +TEST(SniiStreamedWriterSessionTest, OutOfRangePostingPoisonsCompoundSession) { + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(empty_input(113, "invalid_docid", /*doc_count=*/1), + &session)); + + EXPECT_TRUE(push_materialized(session, make_term("alpha", {{.docid = 1, .positions = {0}}})) + .is()); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_FALSE(file.finalized()); +} + +TEST(SniiStreamedWriterSessionTest, BootstrapAppendFailurePoisonsCompound) { + PartialFailOnAppendFile file(/*fail_on_append=*/1); + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + EXPECT_FALSE(compound.begin_streamed_index(empty_input(121, "bootstrap"), &session).ok()); + EXPECT_EQ(session, nullptr); + EXPECT_EQ(file.append_calls(), 1U); + + SniiStreamedIndexSession* retry = nullptr; + EXPECT_FALSE(compound.begin_streamed_index(empty_input(122, "retry"), &retry).ok()); + EXPECT_EQ(retry, nullptr); + EXPECT_FALSE(compound.finish().ok()); + EXPECT_EQ(file.append_calls(), 1U); + EXPECT_EQ(file.finalize_calls(), 0U); + EXPECT_FALSE(file.finalized()); +} + +TEST(SniiStreamedWriterSessionTest, PostingAppendFailurePoisonsSession) { + PartialFailOnAppendFile file(/*fail_on_append=*/2); + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + constexpr uint32_t kDocCount = 300U * 100003U + 1U; + assert_ok(compound.begin_streamed_index(empty_input(131, "posting", kDocCount), &session)); + + EXPECT_FALSE(push_materialized(session, make_sparse_slim_term("slim")).ok()); + EXPECT_TRUE(push_materialized(session, make_term("zulu", {{.docid = 0, .positions = {0}}})) + .is()); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_EQ(file.append_calls(), 2U); + EXPECT_EQ(file.finalize_calls(), 0U); + EXPECT_FALSE(file.finalized()); +} + +TEST(SniiStreamedWriterSessionTest, DictAppendFailureLeavesContainerUnsealable) { + PartialFailOnAppendFile file(/*fail_on_append=*/2); + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(empty_input(141, "dict"), &session)); + assert_ok(push_materialized(session, make_term("alpha", {{.docid = 0, .positions = {0}}}))); + + // The inline term writes no posting bytes. Logical finalization succeeds, + // then the first DICT-region append (append #2 after the bootstrap) fails. + EXPECT_FALSE(session->finish().ok()); + EXPECT_TRUE(push_materialized(session, make_term("beta", {{.docid = 1, .positions = {0}}})) + .is()); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_EQ(file.append_calls(), 2U); + EXPECT_EQ(file.finalize_calls(), 0U); + EXPECT_FALSE(file.finalized()); +} + +TEST(SniiStreamedWriterSessionTest, SourcePushMatchesMaterializedShapeMatrix) { + uint64_t index_id = 151; + auto expect_match = [&](std::string suffix, TermPostings postings, format::IndexConfig config, + bool write_freq, format::PrxWindowLimits limits) { + SCOPED_TRACE(suffix + (write_freq ? "-freq" : "-no-freq")); + SniiIndexInput input; + input.index_id = index_id++; + input.index_suffix = std::move(suffix); + input.config = config; + input.doc_count = postings.docids.empty() ? 0 : postings.docids.back() + 1; + input.write_freq = write_freq; + input.prx_window_limits = limits; + + SniiIndexInput expected_input = input; + TermPostings materialized = postings; + if (!materialized.retain_positions && materialized.freqs.empty()) { + materialized.freqs.assign(materialized.docids.size(), 1); + } + expected_input.terms.push_back(std::move(materialized)); + MemoryFile expected; + assert_ok(write_ordinary_index(std::move(expected_input), &expected)); + + MemoryFile actual; + assert_ok(write_source_index(std::move(input), &postings, &actual)); + EXPECT_EQ(actual.data(), expected.data()); + }; + + const format::PrxWindowLimits default_limits = format::kReaderPrxWindowLimits; + const format::PrxWindowLimits recut_limits { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 2048, + }; + for (bool write_freq : {false, true}) { + expect_match("empty-key", make_uniform_term("", 1), format::IndexConfig::kDocsPositions, + write_freq, default_limits); + expect_match("inline", make_uniform_term("inline", 1), format::IndexConfig::kDocsPositions, + write_freq, default_limits); + expect_match("slim", make_sparse_slim_term("slim"), format::IndexConfig::kDocsPositions, + write_freq, default_limits); + expect_match("df-511", make_uniform_term("df-511", 511), + format::IndexConfig::kDocsPositions, write_freq, default_limits); + expect_match("df-512", make_uniform_term("df-512", 512), + format::IndexConfig::kDocsPositions, write_freq, default_limits); + expect_match("df-8191", make_uniform_term("df-8191", 8191), + format::IndexConfig::kDocsPositions, write_freq, default_limits); + expect_match("df-8192", make_uniform_term("df-8192", 8192), + format::IndexConfig::kDocsPositions, write_freq, default_limits); + expect_match("recut-full", make_adaptive_boundary_term(1024), + format::IndexConfig::kDocsPositions, write_freq, recut_limits); + expect_match("recut-tail", make_adaptive_boundary_term(512), + format::IndexConfig::kDocsPositions, write_freq, recut_limits); + } + + TermPostings docs_only = make_uniform_term("docs-only", 512); + docs_only.retain_positions = false; + docs_only.freqs.clear(); + docs_only.positions_flat.clear(); + expect_match("docs-only", std::move(docs_only), format::IndexConfig::kDocsOnly, + /*write_freq=*/false, default_limits); + + TermPostings docs_with_stats = make_uniform_term("docs-with-stats", 512); + docs_with_stats.retain_positions = false; + docs_with_stats.freqs.assign(docs_with_stats.docids.size(), 2); + docs_with_stats.positions_flat.clear(); + expect_match("docs-with-stats", std::move(docs_with_stats), format::IndexConfig::kDocsOnly, + /*write_freq=*/false, default_limits); +} + +TEST(SniiStreamedWriterSessionTest, SourcePushMatchesMaterializedAcrossAdaptiveBoundary) { + TermPostings postings = make_adaptive_boundary_term(); + SniiIndexInput expected_input = empty_input(181, "source-boundary", postings.docids.back() + 1); + expected_input.prx_window_limits = { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 2048, + }; + expected_input.terms.push_back(postings); + MemoryFile expected; + assert_ok(write_ordinary_index(std::move(expected_input), &expected)); + + VectorTermPostingSource source(&postings); + SniiIndexInput actual_input = empty_input(181, "source-boundary", postings.docids.back() + 1); + actual_input.prx_window_limits = { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 2048, + }; + MemoryFile actual; + SniiCompoundWriter compound(&actual); + SniiStreamedIndexSession* session = nullptr; + assert_ok(compound.begin_streamed_index(std::move(actual_input), &session)); + assert_ok(session->push_term(StreamedTermPostings { + .term = postings.term, + .retain_positions = true, + .source = &source, + })); + assert_ok(session->finish()); + assert_ok(compound.finish()); + + EXPECT_EQ(actual.data(), expected.data()); + EXPECT_TRUE(source.out_was_empty); + EXPECT_EQ(source.requests, (std::vector {format::kAdaptiveWindowDfThreshold, + format::kAdaptiveWindowDocs})); +} + +TEST(SniiStreamedWriterSessionTest, InvalidSourceContractPoisonsCompoundSession) { + const std::array modes = { + InvalidTermPostingSource::Mode::kEmptyBeforeEof, + InvalidTermPostingSource::Mode::kShortBeforeEof, + InvalidTermPostingSource::Mode::kOverfill, + InvalidTermPostingSource::Mode::kUnordered, + InvalidTermPostingSource::Mode::kOutOfRange, + InvalidTermPostingSource::Mode::kError, + }; + for (size_t i = 0; i < modes.size(); ++i) { + SCOPED_TRACE(i); + InvalidTermPostingSource source(modes[i]); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + SniiIndexInput input = empty_input(191 + i, "invalid-source-" + std::to_string(i), 32); + input.config = format::IndexConfig::kDocsOnly; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + EXPECT_FALSE(session->push_term(StreamedTermPostings { + .term = "alpha", + .retain_positions = false, + .source = &source, + }) + .ok()); + EXPECT_FALSE( + push_materialized(session, make_term("bravo", {{.docid = 1, .positions = {0}}})) + .ok()); + EXPECT_FALSE(session->finish().ok()); + EXPECT_FALSE(compound.finish().ok()); + EXPECT_FALSE(file.finalized()); + } +} + +TEST(SniiStreamedWriterSessionTest, PositionedSourceShapeIsValidatedWhenIndexDropsPrx) { + for (size_t position_count : {1U, 3U}) { + SCOPED_TRACE(position_count); + InvalidPositionShapeSource source(position_count); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + SniiIndexInput input = empty_input(211 + position_count, "invalid-position-shape", 1); + input.config = format::IndexConfig::kDocsOnly; + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + const Status status = session->push_term(StreamedTermPostings { + .term = "alpha", .retain_positions = true, .source = &source}); + EXPECT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find("source positions count must equal sum(freqs)"), + std::string::npos) + << status.to_string(); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_FALSE(file.finalized()); + } +} + +TEST(SniiStreamedWriterSessionTest, LateSourceFailurePoisonsAfterPostingBytesWereWritten) { + TermPostings postings = make_adaptive_boundary_term(); + FailAfterFirstFillSource source(&postings); + MemoryFile file; + SniiCompoundWriter compound(&file); + SniiStreamedIndexSession* session = nullptr; + SniiIndexInput input = empty_input(201, "late-source-failure", postings.docids.back() + 1); + assert_ok(compound.begin_streamed_index(std::move(input), &session)); + const uint64_t bootstrap_bytes = file.bytes_written(); + + EXPECT_TRUE(session->push_term(StreamedTermPostings { + .term = postings.term, + .retain_positions = true, + .source = &source, + }) + .is()); + EXPECT_GT(file.bytes_written(), bootstrap_bytes); + EXPECT_TRUE(session->finish().is()); + EXPECT_TRUE(compound.finish().is()); + EXPECT_FALSE(file.finalized()); +} + +} // namespace diff --git a/be/test/storage/index/snii/compaction/snii_term_cursor_test.cpp b/be/test/storage/index/snii/compaction/snii_term_cursor_test.cpp new file mode 100644 index 00000000000000..e3105fcb73e15f --- /dev/null +++ b/be/test/storage/index/snii/compaction/snii_term_cursor_test.cpp @@ -0,0 +1,942 @@ +// 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. + +// T2.3 unit tests: SniiSegmentTermCursor (pull-model full-dictionary scan), +// TermMergeFrontier (k-way merge by term) and SequentialRegionReader (chunked +// posting-region read-ahead). The load-bearing assertions, per the design: +// - interleaved dictionaries merge into ONE strictly increasing term order, +// equal terms aggregated across sources in ascending source order; +// - reserved phrase-bigram / sentinel terms (full 0x1F marker) abort +// the scan with a distinct error, while user terms merely starting with a +// raw 0x1F byte pass through (marker classification, not prefix-byte); +// - all three posting encodings (inline / slim pod_ref / windowed pod_ref) +// and the kNoTermStats flag are passed through UNINTERPRETED -- the cursor +// yields exactly the DictEntry lookup() yields; +// - read-ahead honors chunk boundaries, clamps to the region end, and falls +// back to exact reads for oversized/backward windows without disturbing +// the buffered forward stream; +// - empty and single-source inputs degenerate cleanly. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/compaction/indexed_winner_tree.h" +#include "storage/index/snii/compaction/region_reader.h" +#include "storage/index/snii/compaction/term_cursor.h" +#include "storage/index/snii/compaction/term_merge_frontier.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::snii_test; // NOLINT +namespace ErrorCode = doris::ErrorCode; +using doris::Status; +using compaction::PostingStream; +using compaction::SequentialRegionReader; +using compaction::SharedAlignedRegionCache; +using compaction::SniiSegmentTermCursor; +using compaction::TermMergeFrontier; +using compaction::big_endian_term_prefix; +using compaction::IndexedWinnerTree; + +// One built source segment: the compound bytes plus an opened logical index. +struct SourceFixture { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; +}; + +Status build_source(std::vector terms, uint32_t doc_count, SourceFixture* fx, + bool write_freq = true, uint32_t target_dict_block_bytes = 256, + reader::LogicalIndexOpenMode open_mode = reader::LogicalIndexOpenMode::kQuery) { + writer::SniiIndexInput in; + in.index_id = 7; + in.index_suffix = "body"; + in.config = format::IndexConfig::kDocsPositions; + in.doc_count = doc_count; + in.write_freq = write_freq; + // Small blocks force a multi-block dictionary so the scan crosses block + // boundaries (and the frq/prx bases change between blocks). + in.target_dict_block_bytes = target_dict_block_bytes; + std::ranges::sort(terms, [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + in.terms = std::move(terms); + + writer::SniiCompoundWriter cw(&fx->file); + RETURN_IF_ERROR(cw.add_logical_index(in)); + RETURN_IF_ERROR(cw.finish()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(&fx->file, &fx->segment)); + return fx->segment.open_index(7, "body", &fx->index, open_mode); +} + +// Full DictEntry equality: the cursor contract is UNINTERPRETED passthrough, +// so every locator/stat/inline field must match what lookup() returns. +void expect_entry_eq(const format::DictEntry& got, const format::DictEntry& want) { + EXPECT_EQ(got.term, want.term); + EXPECT_EQ(got.kind, want.kind); + EXPECT_EQ(got.enc, want.enc); + EXPECT_EQ(got.has_sb, want.has_sb); + EXPECT_EQ(got.df, want.df); + EXPECT_EQ(got.ttf_delta, want.ttf_delta); + EXPECT_EQ(got.term_stats_present, want.term_stats_present); + EXPECT_EQ(got.max_freq, want.max_freq); + EXPECT_EQ(got.frq_off_delta, want.frq_off_delta); + EXPECT_EQ(got.frq_len, want.frq_len); + EXPECT_EQ(got.prelude_len, want.prelude_len); + EXPECT_EQ(got.frq_docs_len, want.frq_docs_len); + EXPECT_EQ(got.prx_off_delta, want.prx_off_delta); + EXPECT_EQ(got.prx_len, want.prx_len); + EXPECT_EQ(got.inline_dd_disk_len, want.inline_dd_disk_len); + EXPECT_EQ(got.dd_meta.zstd, want.dd_meta.zstd); + EXPECT_EQ(got.dd_meta.uncomp_len, want.dd_meta.uncomp_len); + EXPECT_EQ(got.dd_meta.disk_len, want.dd_meta.disk_len); + EXPECT_EQ(got.dd_meta.crc, want.dd_meta.crc); + EXPECT_EQ(got.freq_meta.zstd, want.freq_meta.zstd); + EXPECT_EQ(got.freq_meta.uncomp_len, want.freq_meta.uncomp_len); + EXPECT_EQ(got.freq_meta.disk_len, want.freq_meta.disk_len); + EXPECT_EQ(got.freq_meta.crc, want.freq_meta.crc); + EXPECT_EQ(got.frq_bytes, want.frq_bytes); + EXPECT_EQ(got.prx_bytes, want.prx_bytes); +} + +// Corpus spanning all three posting encodings: +// "aa_tiny" df=1 -> inline (encoded bytes <= inline threshold) +// "mid_slim" df=500 -> slim pod_ref (df < 512 but bytes > threshold) +// "zz_wide" df=600 -> windowed pod_ref (df >= kSlimDfThreshold) +// plus filler vocabulary so target_dict_block_bytes=256 yields several blocks. +std::vector three_encoding_terms(uint32_t* doc_count) { + *doc_count = 600; + std::vector terms; + terms.push_back(make_term("aa_tiny", {{.docid = 5, .positions = {1}}})); + std::vector mid; + mid.reserve(500); + for (uint32_t d = 0; d < 500; ++d) { + uint32_t mixed = d * 2654435761U; + mixed ^= mixed >> 16; + const uint32_t frequency = mixed % 97 + 1; + std::vector positions(frequency); + for (uint32_t i = 0; i < frequency; ++i) { + positions[i] = i * 3; + } + mid.push_back({.docid = d, .positions = std::move(positions)}); + } + terms.push_back(make_term("mid_slim", std::move(mid))); + terms.push_back(make_term("zz_wide", docs_with_one_position(0, 600, 2))); + for (int i = 0; i < 24; ++i) { + terms.push_back(make_term("filler_" + std::string(1, static_cast('a' + i)), + docs_with_one_position(0, 40 + i, 1))); + } + return terms; +} + +// Drains a cursor, checking each yielded entry against lookup() (locator +// passthrough) and collecting the term sequence. +void drain_and_check(SourceFixture* fx, std::vector* terms_out) { + SniiSegmentTermCursor cursor(&fx->index, /*source_ordinal=*/0); + bool has_term = false; + ASSERT_TRUE(cursor.next(&has_term).ok()); + while (has_term) { + terms_out->push_back(cursor.term()); + + bool found = false; + format::DictEntry want; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(fx->index.lookup(cursor.term(), &found, &want, &frq_base, &prx_base)); + ASSERT_TRUE(found) << cursor.term(); + expect_entry_eq(cursor.entry(), want); + EXPECT_EQ(cursor.frq_base(), frq_base) << cursor.term(); + EXPECT_EQ(cursor.prx_base(), prx_base) << cursor.term(); + + ASSERT_TRUE(cursor.next(&has_term).ok()); + } +} + +// --------------------------------------------------------------------------- +// SniiSegmentTermCursor +// --------------------------------------------------------------------------- + +TEST(SniiTermCursorTest, FullDictionaryOrderedScanWithLocatorPassthrough) { + uint32_t doc_count = 0; + std::vector terms = three_encoding_terms(&doc_count); + std::vector want_terms; + want_terms.reserve(terms.size()); + for (const auto& tp : terms) { + want_terms.push_back(tp.term); + } + std::ranges::sort(want_terms); + + SourceFixture fx; + assert_ok(build_source(std::move(terms), doc_count, &fx)); + // The scan must cross DICT block boundaries to prove per-block base/state + // tracking; 256-byte target blocks guarantee it for this corpus. + ASSERT_GT(fx.index.n_dict_blocks(), 1U); + + std::vector got_terms; + drain_and_check(&fx, &got_terms); + EXPECT_EQ(got_terms, want_terms); + + // Three-encoding coverage: the corpus really exercised inline, slim + // pod_ref and windowed pod_ref (guards against a corpus regression that + // would silently weaken the passthrough assertions above). + auto entry_of = [&](const std::string& term) { + bool found = false; + format::DictEntry e; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + EXPECT_TRUE(fx.index.lookup(term, &found, &e, &frq_base, &prx_base).ok()); + EXPECT_TRUE(found) << term; + return e; + }; + EXPECT_EQ(entry_of("aa_tiny").kind, format::DictEntryKind::kInline); + const format::DictEntry slim = entry_of("mid_slim"); + EXPECT_EQ(slim.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(slim.enc, format::DictEntryEnc::kSlim); + const format::DictEntry wide = entry_of("zz_wide"); + EXPECT_EQ(wide.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(wide.enc, format::DictEntryEnc::kWindowed); +} + +TEST(SniiTermCursorTest, CompactionOpenSkipsResidentQuerySections) { + SourceFixture query; + assert_ok(build_source({make_term("alpha", {{.docid = 0, .positions = {0}}})}, + /*doc_count=*/1, &query)); + const format::SectionRefs refs = query.index.section_refs(); + ASSERT_GT(refs.dict_region.length, 0U); + ASSERT_GT(refs.bsbf.length, 0U); + + query.file.clear_reads(); + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + assert_ok(reader::SniiSegmentReader::open(&query.file, &segment)); + assert_ok(segment.open_index(7, "body", &index, reader::LogicalIndexOpenMode::kCompaction)); + EXPECT_EQ(index.open_mode(), reader::LogicalIndexOpenMode::kCompaction); + for (const MemoryFile::Read& read : query.file.reads()) { + EXPECT_FALSE(read.offset == refs.dict_region.offset && read.len == refs.dict_region.length); + EXPECT_FALSE(read.offset == refs.bsbf.offset && read.len == refs.bsbf.length); + } + + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup("alpha", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); +} + +TEST(SniiTermCursorTest, DictScanMemoryIsReservedAndLimitErrorIsPreserved) { + SourceFixture source; + assert_ok(build_source({make_term("alpha", {{.docid = 0, .positions = {0}}})}, + /*doc_count=*/1, &source, /*write_freq=*/true, + /*target_dict_block_bytes=*/256, + reader::LogicalIndexOpenMode::kCompaction)); + + writer::MemoryReporter reporter; + { + SniiSegmentTermCursor cursor(&source.index, /*source_ordinal=*/0, &reporter); + bool has_term = false; + assert_ok(cursor.next(&has_term)); + ASSERT_TRUE(has_term); + EXPECT_GT(reporter.current_bytes(), 0); + } + EXPECT_EQ(reporter.current_bytes(), 0); + + writer::MemoryReporter limited_reporter(/*consume_release=*/nullptr, /*cap_bytes=*/1); + { + SniiSegmentTermCursor cursor(&source.index, /*source_ordinal=*/0, &limited_reporter); + bool has_term = false; + const Status status = cursor.next(&has_term); + EXPECT_TRUE(status.is()) << status; + EXPECT_FALSE(has_term); + EXPECT_TRUE(cursor.next(&has_term).is()); + } + EXPECT_EQ(limited_reporter.current_bytes(), 0); +} + +TEST(SniiTermCursorTest, EmptyIndexYieldsNothing) { + SourceFixture fx; + assert_ok(build_source({}, /*doc_count=*/16, &fx)); + EXPECT_EQ(fx.index.n_dict_blocks(), 0U); + + SniiSegmentTermCursor cursor(&fx.index, 0); + bool has_term = true; + ASSERT_TRUE(cursor.next(&has_term).ok()); + EXPECT_FALSE(has_term); + // Exhaustion is stable, not an error. + has_term = true; + ASSERT_TRUE(cursor.next(&has_term).ok()); + EXPECT_FALSE(has_term); +} + +TEST(SniiTermCursorTest, SyntheticMarkerTermAborts) { + // A marker term written through the modern writer path (not the frozen + // image) must abort identically -- the gate is on the term key, not on + // legacy meta. + std::vector terms; + terms.push_back(make_term(std::string(format::kPhraseBigramTermMarker) + "ab", + {{.docid = 0, .positions = {0}}})); + terms.push_back(make_term("alpha", {{.docid = 1, .positions = {0}}})); + + SourceFixture fx; + assert_ok(build_source(std::move(terms), 8, &fx)); + + SniiSegmentTermCursor cursor(&fx.index, 0); + bool has_term = false; + const Status st = cursor.next(&has_term); + ASSERT_FALSE(st.ok()); + EXPECT_TRUE(st.is()) << st.to_string(); +} + +TEST(SniiTermCursorTest, LeadingUnitSeparatorUserTermsPass) { + // 0x1F-corner ruling: classification uses the FULL marker. A user term + // that merely begins with a raw 0x1F byte -- including a strict PREFIX of + // the marker -- is a legitimate dictionary term and must scan through. + const std::string partial_marker = + "\x1F" + "SNII_PHRASE"; // prefix of the marker, not the marker + const std::string unit_sep_user = + "\x1F" + "user"; + std::vector terms; + terms.push_back(make_term(partial_marker, {{.docid = 0, .positions = {0}}})); + terms.push_back(make_term(unit_sep_user, {{.docid = 1, .positions = {0}}})); + terms.push_back(make_term("alpha", {{.docid = 2, .positions = {0}}})); + + SourceFixture fx; + assert_ok(build_source(std::move(terms), 8, &fx)); + + std::vector got; + drain_and_check(&fx, &got); + EXPECT_EQ(got, (std::vector {partial_marker, unit_sep_user, "alpha"})); +} + +TEST(SniiTermCursorTest, NoTermStatsFlagPassthrough) { + // A freq-dropped index (write_freq=false on a positions config) writes + // kNoTermStats DICT blocks; the cursor must surface term_stats_present == + // false so the downstream pump recomputes ttf/max_freq from the actual + // freq stream instead of trusting meaningless defaults (invariant 2). + auto build_terms = [] { + std::vector terms; + terms.push_back(make_term("alpha", docs_with_one_position(0, 20, 0))); + terms.push_back(make_term("bravo", {{.docid = 3, .positions = {1, 4}}})); + return terms; + }; + + SourceFixture dropped; + assert_ok(build_source(build_terms(), 32, &dropped, /*write_freq=*/false)); + SniiSegmentTermCursor cursor(&dropped.index, 0); + bool has_term = false; + size_t seen = 0; + ASSERT_TRUE(cursor.next(&has_term).ok()); + while (has_term) { + EXPECT_FALSE(cursor.entry().term_stats_present) << cursor.term(); + ++seen; + ASSERT_TRUE(cursor.next(&has_term).ok()); + } + EXPECT_EQ(seen, 2U); + + SourceFixture kept; + assert_ok(build_source(build_terms(), 32, &kept, /*write_freq=*/true)); + SniiSegmentTermCursor kept_cursor(&kept.index, 0); + ASSERT_TRUE(kept_cursor.next(&has_term).ok()); + ASSERT_TRUE(has_term); + EXPECT_TRUE(kept_cursor.entry().term_stats_present); +} + +// --------------------------------------------------------------------------- +// TermMergeFrontier +// --------------------------------------------------------------------------- + +TEST(SniiIndexedWinnerTreeTest, StartsAsValidEmptyTree) { + auto before = [](size_t lhs, size_t rhs) { return lhs < rhs; }; + IndexedWinnerTree tree(before); + static_assert(noexcept(tree.empty())); + + EXPECT_TRUE(tree.empty()); + tree.build(0, [](size_t) { return false; }); + EXPECT_TRUE(tree.empty()); + + tree.build(3, [](size_t source) { return source != 1; }); + ASSERT_FALSE(tree.empty()); + EXPECT_EQ(tree.winner(), 0); + EXPECT_EQ(tree.runner_up(), 2); + + tree.update(0, false); + EXPECT_EQ(tree.winner(), 2); + EXPECT_EQ(tree.runner_up(), IndexedWinnerTree::kNoSource); + tree.update(1, true); + EXPECT_EQ(tree.winner(), 1); + EXPECT_EQ(tree.runner_up(), 2); + tree.update(1, false); + tree.update(2, false); + EXPECT_TRUE(tree.empty()); +} + +TEST(SniiTermMergeFrontierTest, InterleavedDictionariesMergeInFullOrder) { + // Three sources with interleaved vocabularies and two shared terms; the + // merged stream must be strictly increasing, with same-term sources + // aggregated in ascending source order and per-source df passed through. + SourceFixture a; + SourceFixture b; + SourceFixture c; + assert_ok(build_source({make_term("apple", docs_with_one_position(0, 3, 0)), + make_term("banana", docs_with_one_position(0, 5, 0)), + make_term("cherry", docs_with_one_position(0, 2, 0)), + make_term("shared", docs_with_one_position(0, 7, 0))}, + 16, &a)); + assert_ok(build_source({make_term("banana", docs_with_one_position(0, 4, 0)), + make_term("date", docs_with_one_position(0, 6, 0)), + make_term("shared", docs_with_one_position(0, 9, 0))}, + 16, &b)); + assert_ok(build_source({make_term("apple", docs_with_one_position(0, 8, 0)), + make_term("elderberry", docs_with_one_position(0, 1, 0))}, + 16, &c)); + + SniiSegmentTermCursor ca(&a.index, 0); + SniiSegmentTermCursor cb(&b.index, 1); + SniiSegmentTermCursor cc(&c.index, 2); + TermMergeFrontier frontier; + assert_ok(frontier.init({&ca, &cb, &cc})); + + struct WantGroup { + std::string term; + std::vector> sources; // (ordinal, df) + }; + const std::vector want = { + {"apple", {{0, 3}, {2, 8}}}, {"banana", {{0, 5}, {1, 4}}}, {"cherry", {{0, 2}}}, + {"date", {{1, 6}}}, {"elderberry", {{2, 1}}}, {"shared", {{0, 7}, {1, 9}}}, + }; + + std::string prev_term; + for (const WantGroup& wg : want) { + ASSERT_FALSE(frontier.empty()); + const std::string group_term = frontier.front()->term(); + EXPECT_EQ(group_term, wg.term); + EXPECT_TRUE(prev_term < group_term); // strictly increasing across groups + prev_term = group_term; + size_t source_index = 0; + while (!frontier.empty() && frontier.front()->term() == group_term) { + SniiSegmentTermCursor* cursor = frontier.front(); + ASSERT_LT(source_index, wg.sources.size()) << wg.term; + EXPECT_EQ(cursor->source_ordinal(), wg.sources[source_index].first) << wg.term; + format::DictEntry entry = cursor->take_entry(); + EXPECT_EQ(entry.df, wg.sources[source_index].second) << wg.term; + EXPECT_EQ(entry.term, wg.term); + ++source_index; + assert_ok(frontier.advance_front()); + } + EXPECT_EQ(source_index, wg.sources.size()) << wg.term; + } + EXPECT_TRUE(frontier.empty()); + // Exhaustion is stable. + EXPECT_TRUE(frontier.empty()); +} + +TEST(SniiTermMergeFrontierTest, SingleSourceDegenerates) { + uint32_t doc_count = 0; + std::vector terms = three_encoding_terms(&doc_count); + std::vector want_terms; + for (const auto& tp : terms) { + want_terms.push_back(tp.term); + } + std::ranges::sort(want_terms); + + SourceFixture fx; + assert_ok(build_source(std::move(terms), doc_count, &fx)); + SniiSegmentTermCursor cursor(&fx.index, 5); + TermMergeFrontier frontier; + assert_ok(frontier.init({&cursor})); + + std::vector got; + while (!frontier.empty()) { + SniiSegmentTermCursor* current = frontier.front(); + EXPECT_EQ(current->source_ordinal(), 5U); + format::DictEntry entry = current->take_entry(); + got.push_back(std::move(entry.term)); + assert_ok(frontier.advance_front()); + } + EXPECT_EQ(got, want_terms); +} + +TEST(SniiTermMergeFrontierTest, EmptyInputsDegenerate) { + // No sources at all. + TermMergeFrontier none; + assert_ok(none.init({})); + EXPECT_TRUE(none.empty()); + + // One source whose dictionary is empty: it never enters the frontier. + SourceFixture empty; + assert_ok(build_source({}, 16, &empty)); + SourceFixture full; + assert_ok(build_source({make_term("only", {{.docid = 0, .positions = {0}}})}, 16, &full)); + SniiSegmentTermCursor ce(&empty.index, 0); + SniiSegmentTermCursor cf(&full.index, 1); + TermMergeFrontier frontier; + assert_ok(frontier.init({&ce, &cf})); + ASSERT_FALSE(frontier.empty()); + EXPECT_EQ(frontier.front()->term(), "only"); + EXPECT_EQ(frontier.front()->source_ordinal(), 1U); + static_cast(frontier.front()->take_entry()); + assert_ok(frontier.advance_front()); + EXPECT_TRUE(frontier.empty()); +} + +TEST(SniiTermMergeFrontierTest, CachedPrefixesPreserveRandomizedStringOrderingAndGrouping) { + EXPECT_EQ(big_endian_term_prefix(""), 0U); + EXPECT_EQ(big_endian_term_prefix("a"), 0x6100000000000000ULL); + EXPECT_EQ(big_endian_term_prefix("abcdefgh"), 0x6162636465666768ULL); + EXPECT_EQ(big_endian_term_prefix("abcdefgh-tail"), 0x6162636465666768ULL); + const std::string binary_prefix {static_cast(0xFF), static_cast(0x80), '\0', 'a'}; + EXPECT_EQ(big_endian_term_prefix(binary_prefix), 0xFF80006100000000ULL); + + std::vector vocabulary = { + "", + "a", + std::string("a\0", 2), + std::string("a\0b", 3), + "abcdefgh", + "abcdefghx", + "abcdefghy", + std::string(1, static_cast(0x80)), + std::string {static_cast(0x80), '\0', static_cast(0xFF)}, + std::string(1, static_cast(0xFF)), + std::string {static_cast(0xFF), '\0'}, + std::string(96, 'z'), + "shared-prefix", + }; + uint64_t state = 0x9E3779B97F4A7C15ULL; + for (size_t ordinal = 0; ordinal < 160; ++ordinal) { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + const size_t suffix_size = 1 + static_cast((state >> 32) % 24); + std::string term = ordinal % 3 == 0 ? "shared-prefix-" : "random-"; + for (size_t i = 0; i < suffix_size; ++i) { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + term.push_back(static_cast(state >> 32)); + } + term.append("-").append(std::to_string(ordinal)); + vocabulary.push_back(std::move(term)); + } + + constexpr size_t kSourceCount = 5; + std::array fixtures; + std::array cursor_ptrs {}; + std::vector> cursors; + cursors.reserve(kSourceCount); + std::map> expected; + for (size_t source = 0; source < kSourceCount; ++source) { + std::vector terms; + for (size_t ordinal = 0; ordinal < vocabulary.size(); ++ordinal) { + if ((ordinal * 17 + source * 11) % 4 == 0) { + continue; + } + terms.push_back(make_term(vocabulary[ordinal], {{.docid = 0, .positions = {1}}})); + expected[vocabulary[ordinal]].push_back(static_cast(source)); + } + assert_ok(build_source(std::move(terms), 1, &fixtures[source], /*write_freq=*/true, + /*target_dict_block_bytes=*/128)); + cursors.push_back(std::make_unique(&fixtures[source].index, + static_cast(source))); + cursor_ptrs[source] = cursors.back().get(); + } + + TermMergeFrontier frontier; + assert_ok(frontier.init( + std::vector(cursor_ptrs.begin(), cursor_ptrs.end()))); + for (const auto& [term, source_ordinals] : expected) { + ASSERT_FALSE(frontier.empty()); + const std::string group_term = frontier.front()->term(); + EXPECT_EQ(group_term, term); + size_t source_index = 0; + while (!frontier.empty() && frontier.front()->term() == group_term) { + SniiSegmentTermCursor* current = frontier.front(); + ASSERT_LT(source_index, source_ordinals.size()); + EXPECT_EQ(current->source_ordinal(), source_ordinals[source_index]); + format::DictEntry entry = current->take_entry(); + EXPECT_EQ(entry.term, term); + ++source_index; + assert_ok(frontier.advance_front()); + } + EXPECT_EQ(source_index, source_ordinals.size()); + } + EXPECT_TRUE(frontier.empty()); +} + +TEST(SniiTermMergeFrontierTest, CursorErrorPropagatesFromInit) { + // A source with a reserved marker term fails during priming, and the + // frontier must surface the cursor's distinct error so the caller aborts + // this column's merge. + SourceFixture invalid; + assert_ok(build_source({make_term(std::string(format::kPhraseBigramTermMarker) + "ab", + {{.docid = 0, .positions = {0}}})}, + 8, &invalid)); + SourceFixture normal; + assert_ok(build_source({make_term("alpha", {{.docid = 0, .positions = {0}}})}, 8, &normal)); + + SniiSegmentTermCursor cl(&invalid.index, 0); + SniiSegmentTermCursor cn(&normal.index, 1); + TermMergeFrontier frontier; + const Status st = frontier.init({&cl, &cn}); + ASSERT_FALSE(st.ok()); + EXPECT_TRUE(st.is()) << st.to_string(); + // Priming may already have advanced an earlier source. The frontier is + // terminal after this error: retrying init could otherwise skip a term. + EXPECT_TRUE(frontier.init({&cl, &cn}).is()); + const Status sticky = frontier.advance_front(); + EXPECT_TRUE(sticky.is()) << sticky.to_string(); +} + +TEST(SniiTermMergeFrontierTest, CursorErrorDuringAdvanceIsSticky) { + std::vector terms; + terms.push_back(make_term("", {{.docid = 0, .positions = {0}}})); + terms.push_back(make_term(std::string(format::kPhraseBigramTermMarker) + "ab", + {{.docid = 1, .positions = {0}}})); + + SourceFixture source; + assert_ok(build_source(std::move(terms), 8, &source)); + SniiSegmentTermCursor cursor(&source.index, 0); + TermMergeFrontier frontier; + assert_ok(frontier.init({&cursor})); + ASSERT_FALSE(frontier.empty()); + EXPECT_TRUE(frontier.front()->term().empty()); + static_cast(frontier.front()->take_entry()); + + const Status first = frontier.advance_front(); + ASSERT_FALSE(first.ok()); + EXPECT_TRUE(first.is()) << first.to_string(); + const Status sticky = frontier.advance_front(); + EXPECT_TRUE(sticky.is()) << sticky.to_string(); + EXPECT_EQ(sticky.to_string(), first.to_string()); +} + +TEST(SniiTermMergeFrontierTest, LifecycleGuards) { + TermMergeFrontier frontier; + EXPECT_TRUE(frontier.advance_front().is()); + assert_ok(frontier.init({})); + EXPECT_TRUE(frontier.init({}).is()); + EXPECT_TRUE(frontier.advance_front().is()); + + TermMergeFrontier with_null; + EXPECT_TRUE(with_null.init({nullptr}).is()); +} + +// --------------------------------------------------------------------------- +// SequentialRegionReader +// --------------------------------------------------------------------------- + +// A file of distinct bytes so slices can be verified by value; the region is a +// strict interior window so clamping to the region (not the file) is provable. +void fill_pattern_file(MemoryFile* file, size_t n) { + std::vector bytes(n); + for (size_t i = 0; i < n; ++i) { + bytes[i] = static_cast(i * 31 + 7); + } + ASSERT_TRUE(file->append(Slice(bytes)).ok()); +} + +void expect_slice_matches(const MemoryFile& file, const Slice& got, uint64_t abs_off, + uint64_t len) { + ASSERT_EQ(got.size(), len); + for (uint64_t i = 0; i < len; ++i) { + ASSERT_EQ(got[i], file.data()[abs_off + i]) << "at " << abs_off + i; + } +} + +TEST(SniiTermCursorRegionReaderTest, SequentialChunkingAndTailClamp) { + MemoryFile file; + fill_pattern_file(&file, 200); + // Region [10, 110), chunk 16: interior on both sides. + SequentialRegionReader rr(&file, /*region_offset=*/10, /*region_length=*/100, + /*chunk_bytes=*/16); + std::vector scratch; + Slice out; + + // First window fills one whole chunk at the window start. + assert_ok(rr.resolve(10, 8, &scratch, &out)); + expect_slice_matches(file, out, 10, 8); + ASSERT_EQ(file.reads().size(), 1U); + EXPECT_EQ(file.reads()[0].offset, 10U); + EXPECT_EQ(file.reads()[0].len, 16U); + + // Fully-buffered window: zero-copy, NO new read. + assert_ok(rr.resolve(18, 8, &scratch, &out)); + expect_slice_matches(file, out, 18, 8); + EXPECT_EQ(file.reads().size(), 1U); + EXPECT_EQ(rr.buffer_hits(), 1U); + + // Window straddling the chunk end: forward miss -> refill AT the window. + assert_ok(rr.resolve(24, 8, &scratch, &out)); + expect_slice_matches(file, out, 24, 8); + ASSERT_EQ(file.reads().size(), 2U); + EXPECT_EQ(file.reads()[1].offset, 24U); + EXPECT_EQ(file.reads()[1].len, 16U); + + // Tail window: the refill is clamped to the REGION end (110), not the + // chunk size and not the file end -- read-ahead never leaves the region. + assert_ok(rr.resolve(104, 6, &scratch, &out)); + expect_slice_matches(file, out, 104, 6); + ASSERT_EQ(file.reads().size(), 3U); + EXPECT_EQ(file.reads()[2].offset, 104U); + EXPECT_EQ(file.reads()[2].len, 6U); + + // Zero-length window: no read, empty slice. + assert_ok(rr.resolve(110, 0, &scratch, &out)); + EXPECT_TRUE(out.empty()); + EXPECT_EQ(file.reads().size(), 3U); + EXPECT_EQ(rr.read_calls(), 3U); +} + +TEST(SniiTermCursorRegionReaderTest, OversizedAndBackwardFallThroughToExactReads) { + MemoryFile file; + fill_pattern_file(&file, 200); + SequentialRegionReader rr(&file, 10, 100, /*chunk_bytes=*/16); + std::vector scratch; + Slice out; + + assert_ok(rr.resolve(10, 8, &scratch, &out)); // buffer = [10, 26) + ASSERT_EQ(file.reads().size(), 1U); + + // Oversized window (> chunk): ONE exact read into scratch; the buffered + // chunk survives untouched. + assert_ok(rr.resolve(30, 40, &scratch, &out)); + expect_slice_matches(file, out, 30, 40); + ASSERT_EQ(file.reads().size(), 2U); + EXPECT_EQ(file.reads()[1].offset, 30U); + EXPECT_EQ(file.reads()[1].len, 40U); + + // The old chunk still serves hits (proves the fallback did not evict it). + assert_ok(rr.resolve(12, 4, &scratch, &out)); + expect_slice_matches(file, out, 12, 4); + EXPECT_EQ(file.reads().size(), 2U); + + // Advance the stream, then a BACKWARD window: exact read, buffer kept. + assert_ok(rr.resolve(50, 8, &scratch, &out)); // buffer = [50, 66) + ASSERT_EQ(file.reads().size(), 3U); + assert_ok(rr.resolve(10, 4, &scratch, &out)); + expect_slice_matches(file, out, 10, 4); + ASSERT_EQ(file.reads().size(), 4U); + EXPECT_EQ(file.reads()[3].offset, 10U); + EXPECT_EQ(file.reads()[3].len, 4U); + assert_ok(rr.resolve(52, 4, &scratch, &out)); // still buffered + expect_slice_matches(file, out, 52, 4); + EXPECT_EQ(file.reads().size(), 4U); +} + +TEST(SniiTermCursorRegionReaderTest, OutOfRegionWindowsRejected) { + MemoryFile file; + fill_pattern_file(&file, 200); + SequentialRegionReader rr(&file, 10, 100, 16); + std::vector scratch; + Slice out; + + EXPECT_TRUE(rr.resolve(9, 4, &scratch, &out).is()); + EXPECT_TRUE(rr.resolve(108, 4, &scratch, &out).is()); + EXPECT_TRUE(rr.resolve(111, 0, &scratch, &out).is()); + EXPECT_EQ(file.reads().size(), 0U); +} + +TEST(SniiTermCursorRegionReaderTest, OverflowingRegionRejected) { + MemoryFile file; + fill_pattern_file(&file, 32); + SequentialRegionReader rr(&file, std::numeric_limits::max() - 3, + /*region_length=*/8, /*chunk_bytes=*/16); + std::vector scratch; + Slice out; + + EXPECT_TRUE(rr.resolve(std::numeric_limits::max() - 2, 1, &scratch, &out) + .is()); + EXPECT_TRUE(file.reads().empty()); +} + +TEST(SniiTermCursorRegionReaderTest, SharedCacheCoalescesLogicalStreamsAndPinsSlices) { + MemoryFile file; + fill_pattern_file(&file, 200); + SharedAlignedRegionCache cache(&file, /*region_offset=*/10, /*region_length=*/100, + /*total_budget_bytes=*/32); + assert_ok(cache.init()); + std::vector docs_scratch; + std::vector prx_scratch; + Slice docs; + Slice prx; + + assert_ok(cache.resolve(PostingStream::kDocs, 12, 4, &docs_scratch, &docs)); + expect_slice_matches(file, docs, 12, 4); + ASSERT_EQ(file.reads().size(), 1U); + EXPECT_EQ(file.reads()[0].offset, 10U); + EXPECT_EQ(file.reads()[0].len, 16U); + + assert_ok(cache.resolve(PostingStream::kPrx, 20, 4, &prx_scratch, &prx)); + expect_slice_matches(file, prx, 20, 4); + expect_slice_matches(file, docs, 12, 4); + EXPECT_EQ(file.reads().size(), 1U); + EXPECT_EQ(cache.buffer_hits(PostingStream::kPrx), 1U); + + assert_ok(cache.resolve(PostingStream::kDocs, 30, 4, &docs_scratch, &docs)); + expect_slice_matches(file, docs, 30, 4); + expect_slice_matches(file, prx, 20, 4); + ASSERT_EQ(file.reads().size(), 2U); + EXPECT_EQ(file.reads()[1].offset, 26U); + EXPECT_EQ(file.reads()[1].len, 16U); + + assert_ok(cache.resolve(PostingStream::kPrx, 34, 4, &prx_scratch, &prx)); + expect_slice_matches(file, prx, 34, 4); + expect_slice_matches(file, docs, 30, 4); + EXPECT_EQ(file.reads().size(), 2U); + EXPECT_EQ(cache.physical_read_ranges(), 2U); + EXPECT_EQ(cache.physical_read_bytes(), 32U); + EXPECT_EQ(cache.read_calls(PostingStream::kDocs), 2U); + EXPECT_EQ(cache.read_calls(PostingStream::kPrx), 0U); + EXPECT_LE(cache.resident_capacity_bytes(), 32U); +} + +TEST(SniiTermCursorRegionReaderTest, SharedCacheUsesExactCrossBlockAndClampsTail) { + MemoryFile file; + fill_pattern_file(&file, 200); + SharedAlignedRegionCache cache(&file, /*region_offset=*/10, /*region_length=*/100, + /*total_budget_bytes=*/32); + assert_ok(cache.init()); + std::vector docs_scratch; + std::vector prx_scratch; + Slice docs; + Slice prx; + + assert_ok(cache.resolve(PostingStream::kDocs, 20, 10, &docs_scratch, &docs)); + expect_slice_matches(file, docs, 20, 10); + ASSERT_EQ(file.reads().size(), 1U); + EXPECT_EQ(file.reads()[0].offset, 20U); + EXPECT_EQ(file.reads()[0].len, 10U); + + assert_ok(cache.resolve(PostingStream::kPrx, 106, 4, &prx_scratch, &prx)); + expect_slice_matches(file, prx, 106, 4); + expect_slice_matches(file, docs, 20, 10); + ASSERT_EQ(file.reads().size(), 2U); + EXPECT_EQ(file.reads()[1].offset, 106U); + EXPECT_EQ(file.reads()[1].len, 4U); + EXPECT_EQ(cache.physical_read_bytes(), 14U); + EXPECT_EQ(cache.physical_read_ranges(), 2U); +} + +class FailBootstrapAppendFile final : public io::FileWriter { +public: + Status append(Slice data) override { + ++append_calls_; + if (append_calls_ == 1) { + if (!data.empty()) ++bytes_written_; + return Status::Error("injected bootstrap append failure"); + } + bytes_written_ += data.size(); + return Status::OK(); + } + + Status finalize() override { + finalized_ = true; + return Status::OK(); + } + + uint64_t bytes_written() const override { return bytes_written_; } + size_t append_calls() const { return append_calls_; } + bool finalized() const { return finalized_; } + +private: + uint64_t bytes_written_ = 0; + size_t append_calls_ = 0; + bool finalized_ = false; +}; + +TEST(SniiStreamedCompactionWriterTest, BootstrapFailurePoisonsCompound) { + FailBootstrapAppendFile file; + writer::SniiCompoundWriter compound(&file); + writer::SniiIndexInput input; + input.index_id = 7; + input.index_suffix = "body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 1; + + writer::SniiStreamedIndexSession* session = nullptr; + const Status first = compound.begin_streamed_index(std::move(input), &session); + ASSERT_FALSE(first.ok()); + EXPECT_TRUE(first.is()) << first.to_string(); + EXPECT_EQ(session, nullptr); + EXPECT_EQ(file.append_calls(), 1U); + + writer::SniiIndexInput retry; + retry.index_id = 7; + retry.index_suffix = "body"; + retry.config = format::IndexConfig::kDocsPositions; + retry.doc_count = 1; + const Status second = compound.begin_streamed_index(std::move(retry), &session); + EXPECT_TRUE(second.is()) << second.to_string(); + EXPECT_EQ(file.append_calls(), 1U); + EXPECT_TRUE(compound.finish().is()); + EXPECT_FALSE(file.finalized()); +} + +// Integration seam: resolve REAL posting windows (from cursor entries) through +// the region reader over the source's posting region -- the exact wiring the +// merge pump uses. Every pod_ref window must resolve to in-region bytes. +TEST(SniiTermCursorRegionReaderTest, ResolvesRealPostingWindowsSequentially) { + uint32_t doc_count = 0; + SourceFixture fx; + assert_ok(build_source(three_encoding_terms(&doc_count), doc_count, &fx)); + + const format::RegionRef& region = fx.index.section_refs().posting_region; + ASSERT_GT(region.length, 0U); + SequentialRegionReader rr(&fx.file, region.offset, region.length, + /*chunk_bytes=*/static_cast(region.length)); + + SniiSegmentTermCursor cursor(&fx.index, 0); + bool has_term = false; + std::vector scratch; + size_t pod_windows = 0; + ASSERT_TRUE(cursor.next(&has_term).ok()); + while (has_term) { + if (cursor.entry().kind == format::DictEntryKind::kPodRef) { + uint64_t abs_off = 0; + uint64_t len = 0; + assert_ok( + fx.index.resolve_frq_window(cursor.entry(), cursor.frq_base(), &abs_off, &len)); + Slice out; + assert_ok(rr.resolve(abs_off, len, &scratch, &out)); + EXPECT_EQ(out.size(), len); + ++pod_windows; + } + ASSERT_TRUE(cursor.next(&has_term).ok()); + } + EXPECT_GT(pod_windows, 0U); + EXPECT_GT(rr.buffer_hits(), 0U); // sequential walk actually amortized reads +} + +} // namespace diff --git a/be/test/storage/index/snii/encoding/byte_sink_test.cpp b/be/test/storage/index/snii/encoding/byte_sink_test.cpp new file mode 100644 index 00000000000000..575fea6a017bf8 --- /dev/null +++ b/be/test/storage/index/snii/encoding/byte_sink_test.cpp @@ -0,0 +1,103 @@ +// 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 "storage/index/snii/encoding/byte_sink.h" + +#include + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/varint.h" + +using namespace doris::snii; + +namespace { + +void append_reference_leb128(uint64_t value, std::vector* out) { + do { + uint8_t octet = static_cast(value % 128); + value /= 128; + if (value != 0) { + octet |= 0x80; + } + out->push_back(octet); + } while (value != 0); +} + +} // namespace + +TEST(SniiByteSink, Fixed32LittleEndian) { + ByteSink s; + s.put_fixed32(0x04030201U); + ASSERT_EQ(s.size(), 4U); + const auto& b = s.buffer(); + EXPECT_EQ(b[0], 0x01); + EXPECT_EQ(b[1], 0x02); + EXPECT_EQ(b[3], 0x04); +} + +TEST(SniiByteSink, Fixed64LittleEndian) { + ByteSink s; + s.put_fixed64(0x0807060504030201ULL); + ASSERT_EQ(s.size(), 8U); + EXPECT_EQ(s.buffer()[0], 0x01); + EXPECT_EQ(s.buffer()[7], 0x08); +} + +TEST(SniiByteSink, VarintThenBytes) { + ByteSink s; + s.put_varint32(300); + const uint8_t payload[] = {0xAA, 0xBB}; + s.put_bytes(Slice(payload, 2)); + EXPECT_EQ(s.size(), varint_len(300) + 2); +} + +TEST(SniiByteSink, ScalarVarintsMatchReferenceAtEveryWidth) { + std::vector values32 = {0, 1}; + for (uint32_t width = 1; width < 5; ++width) { + const uint32_t first_of_next_width = 1U << (width * 7); + values32.push_back(first_of_next_width - 1); + values32.push_back(first_of_next_width); + } + values32.push_back(std::numeric_limits::max() - 1); + values32.push_back(std::numeric_limits::max()); + + std::vector values64 = {0, 1}; + for (uint32_t width = 1; width < 10; ++width) { + const uint64_t first_of_next_width = uint64_t {1} << (width * 7); + values64.push_back(first_of_next_width - 1); + values64.push_back(first_of_next_width); + } + values64.push_back(std::numeric_limits::max() - 1); + values64.push_back(std::numeric_limits::max()); + + ByteSink sink; + std::vector expected; + for (uint32_t value : values32) { + append_reference_leb128(value, &expected); + sink.put_varint32(value); + } + for (uint64_t value : values64) { + append_reference_leb128(value, &expected); + sink.put_varint64(value); + } + + EXPECT_EQ(sink.buffer(), expected); +} diff --git a/be/test/storage/index/snii/encoding/byte_source_test.cpp b/be/test/storage/index/snii/encoding/byte_source_test.cpp new file mode 100644 index 00000000000000..339652482ef3cb --- /dev/null +++ b/be/test/storage/index/snii/encoding/byte_source_test.cpp @@ -0,0 +1,262 @@ +// 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 "storage/index/snii/encoding/byte_source.h" + +#include + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" + +using namespace doris::snii; +using doris::Status; + +TEST(SniiByteSource, RoundTripWithSink) { + ByteSink s; + s.put_fixed32(0xDEADBEEF); + s.put_varint64(123456789); + s.put_zigzag(-42); + ByteSource src(s.view()); + uint32_t a; + uint64_t b; + int64_t c; + ASSERT_TRUE(src.get_fixed32(&a).ok()); + EXPECT_EQ(a, 0xDEADBEEFU); + ASSERT_TRUE(src.get_varint64(&b).ok()); + EXPECT_EQ(b, 123456789U); + ASSERT_TRUE(src.get_zigzag(&c).ok()); + EXPECT_EQ(c, -42); + EXPECT_TRUE(src.eof()); +} + +TEST(SniiByteSource, GetBytesAdvances) { + ByteSink s; + const uint8_t p[] = {1, 2, 3, 4, 5}; + s.put_bytes(Slice(p, 5)); + ByteSource src(s.view()); + Slice got; + ASSERT_TRUE(src.get_bytes(3, &got).ok()); + ASSERT_EQ(got.size(), 3U); + EXPECT_EQ(got[0], 1U); + EXPECT_EQ(src.remaining(), 2U); +} + +TEST(SniiByteSource, OverrunFails) { + uint8_t one[1] = {0x01}; + ByteSource src(Slice(one, 1)); + uint32_t a; + EXPECT_FALSE(src.get_fixed32(&a).ok()); +} + +// skip_varints must advance the cursor EXACTLY as decoding the same varints +// would -- across 1..5-byte encodings and back-to-back values -- so a CSR +// selective decode can skip non-candidate docs' position deltas losslessly. +TEST(SniiByteSource, SkipVarintsAdvancesLikeDecode) { + ByteSink s; + const uint64_t vals[] = {0, 1, 127, 128, 300, 16384, 0xFFFFFFFFULL, 42, 1u << 21, 5}; + for (uint64_t v : vals) s.put_varint64(v); + s.put_u8(0xAB); // sentinel after the varints + + // Decode-advance reference: consume all values, record end position. + ByteSource ref(s.view()); + for (size_t i = 0; i < std::size(vals); ++i) { + uint64_t v = 0; + ASSERT_TRUE(ref.get_varint64(&v).ok()); + EXPECT_EQ(v, vals[i]); + } + const size_t decoded_pos = ref.position(); + + // Skip-advance: skipping the same count must land at the identical position. + ByteSource skp(s.view()); + ASSERT_TRUE(skp.skip_varints(std::size(vals)).ok()); + EXPECT_EQ(skp.position(), decoded_pos); + uint8_t sentinel = 0; + ASSERT_TRUE(skp.get_u8(&sentinel).ok()); + EXPECT_EQ(sentinel, 0xAB); + + // Partial skip then decode the rest interleaves correctly. + ByteSource mix(s.view()); + ASSERT_TRUE(mix.skip_varints(3).ok()); // skip 0,1,127 + uint64_t v = 0; + ASSERT_TRUE(mix.get_varint64(&v).ok()); + EXPECT_EQ(v, 128U); // 4th value decoded intact +} + +// A long run of single-byte varints interleaved with multi-byte values. +// skip_varints must land exactly where a full decode would, for every prefix +// count -- covers long runs and mixed widths in one sweep. +TEST(SniiByteSource, SkipVarintsLongRunMatchesDecode) { + ByteSink s; + std::vector vals; + // 40 one-byte values, then a 5-byte value (forces a straddle), then more. + for (int i = 0; i < 40; ++i) { + vals.push_back(static_cast(i % 100)); + } + vals.push_back(0xFFFFFFFFULL); // 5 bytes + for (int i = 0; i < 40; ++i) { + vals.push_back(static_cast((i * 7) % 120)); + } + vals.push_back(300); // 2 bytes + for (int i = 0; i < 20; ++i) { + vals.push_back(1); + } + for (uint64_t v : vals) { + s.put_varint64(v); + } + s.put_u8(0x3C); // sentinel + + // Reference decode end-position after consuming the first k values. + for (size_t k = 0; k <= vals.size(); ++k) { + ByteSource ref(s.view()); + for (size_t i = 0; i < k; ++i) { + uint64_t v = 0; + ASSERT_TRUE(ref.get_varint64(&v).ok()); + } + const size_t want = ref.position(); + ByteSource skp(s.view()); + ASSERT_TRUE(skp.skip_varints(k).ok()) << "k=" << k; + EXPECT_EQ(skp.position(), want) << "k=" << k; + } + // Full skip then read the sentinel. + ByteSource full(s.view()); + ASSERT_TRUE(full.skip_varints(vals.size()).ok()); + uint8_t sentinel = 0; + ASSERT_TRUE(full.get_u8(&sentinel).ok()); + EXPECT_EQ(sentinel, 0x3C); +} + +TEST(SniiByteSource, SkipVarintsTruncationFails) { + // A dangling continuation byte (high bit set, no terminator) must error, not + // run off the buffer. + uint8_t buf[2] = {0x80, 0x80}; + ByteSource src(Slice(buf, 2)); + EXPECT_FALSE(src.skip_varints(1).ok()); + // Zero-count skip is a no-op and always succeeds. + ByteSource empty(Slice(buf, 0)); + EXPECT_TRUE(empty.skip_varints(0).ok()); +} + +// decode_delta_run must produce the SAME running prefix-sum values as the +// per-value get_varint32 loop it replaces in the CSR reader -- across 1..5-byte +// delta encodings -- and advance the cursor identically. +TEST(SniiByteSource, DecodeDeltaRunMatchesManualPrefixSum) { + ByteSink s; + // Deltas chosen to span every varint width and to sum to a >32-bit-looking + // running value only if mis-summed; the running total stays within uint32. + const uint32_t deltas[] = {5, 0, 1, 127, 128, 300, 16384, 1U << 21, 7, 2}; + for (uint32_t d : deltas) { + s.put_varint64(d); + } + s.put_u8(0xCD); // sentinel + + // Manual reference: decode each delta, prefix-sum from 0. + std::vector expect; + uint32_t running = 0; + for (uint32_t d : deltas) { + running += d; + expect.push_back(running); + } + + std::vector got; + ByteSource src(s.view()); + ASSERT_TRUE(src.decode_delta_run(std::size(deltas), &got).ok()); + EXPECT_EQ(got, expect); + // Cursor lands exactly on the sentinel. + uint8_t sentinel = 0; + ASSERT_TRUE(src.get_u8(&sentinel).ok()); + EXPECT_EQ(sentinel, 0xCD); +} + +TEST(SniiByteSource, DecodeDeltaRunAppendsAndZeroCount) { + ByteSink s; + s.put_varint64(10); + s.put_varint64(20); + std::vector out {999}; // pre-existing content must be preserved + ByteSource src(s.view()); + // Zero-count run is a no-op that leaves the cursor and output untouched. + ASSERT_TRUE(src.decode_delta_run(0, &out).ok()); + EXPECT_EQ(out.size(), 1U); + EXPECT_EQ(src.position(), 0U); + // A real run APPENDS running values after the existing entry. + ASSERT_TRUE(src.decode_delta_run(2, &out).ok()); + ASSERT_EQ(out.size(), 3U); + EXPECT_EQ(out[0], 999U); + EXPECT_EQ(out[1], 10U); + EXPECT_EQ(out[2], 30U); +} + +// get_varint32_fast must decode identically to get_varint32 across the 1-byte +// fast path, the multi-byte fallback, and the boundary at 127/128. +TEST(SniiByteSource, GetVarint32FastMatchesSlow) { + ByteSink s; + const uint32_t vals[] = {0, 1, 127, 128, 129, 255, 300, 16384, 0xFFFFFFFFU, 42}; + for (uint32_t v : vals) { + s.put_varint64(v); + } + s.put_u8(0x5A); // sentinel + + ByteSource fast(s.view()); + for (uint32_t expect : vals) { + uint32_t v = 0; + ASSERT_TRUE(fast.get_varint32_fast(&v).ok()); + EXPECT_EQ(v, expect); + } + uint8_t sentinel = 0; + ASSERT_TRUE(fast.get_u8(&sentinel).ok()); + EXPECT_EQ(sentinel, 0x5A); + + // Empty buffer -> fast path bails to the checked slow decoder, which errors. + ByteSource empty(Slice(nullptr, 0)); + uint32_t v = 0; + EXPECT_FALSE(empty.get_varint32_fast(&v).ok()); +} + +TEST(SniiByteSource, DecodeDeltaRunTruncationFails) { + // Dangling continuation byte -> corruption, no buffer overrun. + uint8_t buf[2] = {0x80, 0x80}; + ByteSource src(Slice(buf, 2)); + std::vector out; + EXPECT_FALSE(src.decode_delta_run(1, &out).ok()); + // Asking for more values than present also fails cleanly. + ByteSink s; + s.put_varint64(3); + ByteSource src2(s.view()); + std::vector out2; + EXPECT_FALSE(src2.decode_delta_run(2, &out2).ok()); +} + +TEST(SniiByteSource, DecodeDeltaRunRejectsVarintAndPrefixSumOverflow) { + ByteSink oversized; + oversized.put_varint64(uint64_t {1} << 32); + std::vector out = {17}; + ByteSource oversized_source(oversized.view()); + const Status oversized_status = oversized_source.decode_delta_run(1, &out); + EXPECT_TRUE(oversized_status.is()) + << oversized_status; + EXPECT_EQ(out, std::vector({17})); + + ByteSink overflowing_sum; + overflowing_sum.put_varint32(std::numeric_limits::max()); + overflowing_sum.put_varint32(1); + ByteSource sum_source(overflowing_sum.view()); + const Status sum_status = sum_source.decode_delta_run(2, &out); + EXPECT_TRUE(sum_status.is()) << sum_status; + EXPECT_EQ(out, std::vector({17})); +} diff --git a/be/test/storage/index/snii/encoding/crc32c_test.cpp b/be/test/storage/index/snii/encoding/crc32c_test.cpp new file mode 100644 index 00000000000000..2913d9b1516ce4 --- /dev/null +++ b/be/test/storage/index/snii/encoding/crc32c_test.cpp @@ -0,0 +1,336 @@ +// 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 "storage/index/snii/encoding/crc32c.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +// Integrated crc32c.h pulls in the thirdparty `namespace crc32c`, so a blanket +// `using namespace doris::snii;` makes a bare crc32c() call ambiguous; pull in only the +// Slice type and qualify the doris::snii::crc32c free functions explicitly. +using doris::snii::Slice; + +// leveldb/RocksDB standard CRC32C(Castagnoli) test vectors. +TEST(SniiCrc32c, KnownVectors) { + std::vector zeros(32, 0x00); + EXPECT_EQ(doris::snii::crc32c(Slice(zeros)), 0x8a9136aaU); + std::vector ff(32, 0xff); + EXPECT_EQ(doris::snii::crc32c(Slice(ff)), 0x62a8ab43U); + std::vector ramp(32); + for (int i = 0; i < 32; ++i) { + ramp[i] = static_cast(i); + } + EXPECT_EQ(doris::snii::crc32c(Slice(ramp)), 0x46dd794eU); +} + +TEST(SniiCrc32c, ExtendEqualsContiguous) { + std::vector v {1, 2, 3, 4, 5, 6, 7, 8}; + uint32_t whole = doris::snii::crc32c(Slice(v)); + uint32_t part = doris::snii::crc32c(Slice(v.data(), 4)); + part = doris::snii::crc32c_extend(part, Slice(v.data() + 4, 4)); + EXPECT_EQ(whole, part); +} + +namespace { + +// Independent byte-at-a-time CRC32C reference (bit-reflected Castagnoli). Used to +// pin the optimized slice-by-8 / hardware path to the canonical scalar result. +uint32_t crc32c_ref(const std::vector& data) { + static constexpr uint32_t kPoly = 0x82F63B78U; + uint32_t crc = ~0U; + for (uint8_t b : data) { + crc ^= b; + for (int k = 0; k < 8; ++k) { + crc = (crc & 1) ? (kPoly ^ (crc >> 1)) : (crc >> 1); + } + } + return ~crc; +} + +} // namespace + +// Sweep every length 0..2048: stresses the 8-byte main loop plus the 0..7 residue +// tail (and the hardware u32/u8 tails) against the scalar reference. Catches any +// slice-by-8 table or unaligned-load bug that the fixed-size vectors above miss. +TEST(SniiCrc32c, MatchesScalarReferenceAllLengths) { + std::vector data; + data.reserve(2048); + uint32_t x = 0x12345678U; + for (size_t len = 0; len <= 2048; ++len) { + EXPECT_EQ(doris::snii::crc32c(Slice(data)), crc32c_ref(data)) << "len=" << len; + // Pseudo-random next byte (xorshift) so the stream is non-trivial. + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + data.push_back(static_cast(x)); + } +} + +// Splitting a buffer at an arbitrary boundary and extending must equal the +// one-shot crc -- the property the windowed/region layout relies on. +TEST(SniiCrc32c, ExtendAcrossArbitrarySplits) { + std::vector data(300); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i * 7 + 1); + } + const uint32_t whole = doris::snii::crc32c(Slice(data)); + for (size_t split : {0U, 1U, 7U, 8U, 9U, 16U, 100U, 299U, 300U}) { + uint32_t c = doris::snii::crc32c(Slice(data.data(), split)); + c = doris::snii::crc32c_extend(c, Slice(data.data() + split, data.size() - split)); + EXPECT_EQ(c, whole) << "split=" << split; + } +} + +// ============================================================================= +// T21 -- CRC32C 3-way interleaved hardware reference paths. +// +// Production crc32c()/crc32c_extend() delegate to the bundled Google crc32c +// thirdparty (already hardware-accelerated + interleaved). The doris::snii::detail +// seam below exposes three reference sub-paths -- portable slice-by-8, serial +// SSE4.2 hardware, and the 3-way interleaved SSE4.2 hardware algorithm with a +// GF(2) shift-combine -- so these tests can prove, byte-for-byte, that the +// production path equals the canonical CRC32C across every size/alignment and that +// the hardware path is engaged. detail::* is compiled only under BE_TEST. +// ============================================================================= + +namespace { + +// Alias only the detail seam; the free crc32c()/crc32c_extend() are qualified at +// every call site below because a bare `crc32c` is ambiguous with the thirdparty +// `crc32c` namespace pulled in by (see the file-header note). +namespace detail = doris::snii::detail; + +// Deterministic pseudo-random buffer with `pad` leading bytes, so a Slice can start +// at an arbitrary alignment offset into the heap allocation. The exact bytes are +// irrelevant to the equivalence asserts (all paths hash the same buffer); the +// seeded engine only keeps runs reproducible. +std::vector make_padded_bytes(uint64_t seed, size_t pad, size_t n) { + std::vector v(pad + n); + std::mt19937_64 rng(seed); + for (auto& b : v) { + b = static_cast(rng() & 0xFFU); + } + return v; +} + +// Sizes spanning the 8-byte main loop, every <8 tail, the 4-byte step, and buffers +// straddling the 1024-byte interleave threshold (below/at/above and exact 3x). +constexpr size_t kSweepSizes[] = {0, 1, 7, 8, 9, 15, 16, 255, + 256, 1023, 1024, 1025, 3072, 4096, 65536}; +constexpr size_t kAlignments[] = {0, 1, 2, 3, 4, 7}; + +} // namespace + +// FV-1: hw3 == slice8 == hw_serial == production crc32c across all sizes/alignments. +TEST(SniiCrc32cTest, Hw3MatchesSlice8AcrossSizesAndAlignments) { + for (size_t n : kSweepSizes) { + for (size_t align : kAlignments) { + const std::vector buf = + make_padded_bytes(0x9E3779B97F4A7C15ULL + n * 131 + align, align, n); + const Slice s(buf.data() + align, n); + const uint32_t lib = doris::snii::crc32c(s); + const uint32_t sw = detail::crc32c_slice8_extend(0, s); + const uint32_t hws = detail::crc32c_hw_serial_extend(0, s); + const uint32_t hw3 = detail::crc32c_hw3_extend(0, s); + EXPECT_EQ(hw3, sw) << "n=" << n << " align=" << align; + EXPECT_EQ(hw3, hws) << "n=" << n << " align=" << align; + EXPECT_EQ(hw3, lib) << "n=" << n << " align=" << align; + } + } +} + +// FV-2: degenerate + exact-divisible boundaries. Empty -> canonical crc32c("")==0 +// on every path; n=3072 is exactly 3*1024 (zero interleave tail). +TEST(SniiCrc32cTest, DegenerateAndExactDivisibleBoundaries) { + const Slice empty(nullptr, 0); + EXPECT_EQ(0U, doris::snii::crc32c(empty)); + EXPECT_EQ(0U, detail::crc32c_slice8_extend(0, empty)); + EXPECT_EQ(0U, detail::crc32c_hw_serial_extend(0, empty)); + EXPECT_EQ(0U, detail::crc32c_hw3_extend(0, empty)); + + for (size_t n : {size_t {1}, size_t {7}, size_t {8}, size_t {3072}}) { + const std::vector buf = make_padded_bytes(0xD1CE5EEDULL + n, 0, n); + const Slice s(buf.data(), n); + const uint32_t lib = doris::snii::crc32c(s); + EXPECT_EQ(detail::crc32c_slice8_extend(0, s), lib) << "n=" << n; + EXPECT_EQ(detail::crc32c_hw_serial_extend(0, s), lib) << "n=" << n; + EXPECT_EQ(detail::crc32c_hw3_extend(0, s), lib) << "n=" << n; + } +} + +// FV-3: seeded-extend chaining equals the one-shot over the whole, including a +// length that crosses the 1024 threshold (4097). Covers the seed chain + the hw3 +// GF(2) combine linearity (the property every framed section relies on). +TEST(SniiCrc32cTest, ExtendEqualsConcatenationAcrossThreshold) { + for (size_t total : {size_t {16}, size_t {1024}, size_t {4097}}) { + const std::vector buf = make_padded_bytes(0xABCDEF01ULL + total, 0, total); + const uint8_t* p = buf.data(); + const uint32_t whole = doris::snii::crc32c(Slice(p, total)); + for (size_t k : + {size_t {0}, size_t {1}, size_t {7}, size_t {8}, total / 2, total - 1, total}) { + const Slice a(p, k); + const Slice b(p + k, total - k); + // Public seeded-extend chaining (delegates to the bundled library). + EXPECT_EQ(doris::snii::crc32c_extend(doris::snii::crc32c(a), b), whole) + << "total=" << total << " k=" << k; + // The hw3 reference threads the same seed through lane A and the combine. + EXPECT_EQ(detail::crc32c_hw3_extend(doris::snii::crc32c(a), b), whole) + << "total=" << total << " k=" << k; + } + } +} + +// FV-4: canonical CRC32C vectors from an external authority. A path that altered +// the value -- and hence the on-disk format -- would break these. +TEST(SniiCrc32cTest, KnownVectorsMatchExternalAuthority) { + const char* digits = "123456789"; + const Slice ds(reinterpret_cast(digits), 9); + EXPECT_EQ(0xE3069283U, doris::snii::crc32c(ds)); + EXPECT_EQ(0xE3069283U, detail::crc32c_slice8_extend(0, ds)); + EXPECT_EQ(0xE3069283U, detail::crc32c_hw_serial_extend(0, ds)); + EXPECT_EQ(0xE3069283U, detail::crc32c_hw3_extend(0, ds)); + + EXPECT_EQ(0U, doris::snii::crc32c(Slice(nullptr, 0))); + + // 4096 x 'a' crosses the interleave threshold; cross-check every path against + // the independent byte-at-a-time reference defined above. + const std::vector aaa(4096, static_cast('a')); + const Slice as(aaa.data(), aaa.size()); + const uint32_t ref = crc32c_ref(aaa); + EXPECT_EQ(ref, doris::snii::crc32c(as)); + EXPECT_EQ(ref, detail::crc32c_slice8_extend(0, as)); + EXPECT_EQ(ref, detail::crc32c_hw_serial_extend(0, as)); + EXPECT_EQ(ref, detail::crc32c_hw3_extend(0, as)); +} + +// FV-5: single-bit-flip regression -- the checksum still distinguishes a one-bit +// change (verify capability not degraded) on both the production and hw3 paths. +TEST(SniiCrc32cTest, SingleBitFlipChangesChecksum) { + std::vector buf = make_padded_bytes(0x5A5A5A5AULL, 0, 2048); // crosses threshold + const Slice s(buf.data(), buf.size()); + const uint32_t base_lib = doris::snii::crc32c(s); + const uint32_t base_hw3 = detail::crc32c_hw3_extend(0, s); + for (size_t bit : {size_t {0}, size_t {1}, size_t {8 * 7 + 3}, size_t {8 * 1000 + 5}, + size_t {8 * 2047 + 7}}) { + const size_t byte = bit / 8; + const auto mask = static_cast(1U << (bit % 8)); + buf[byte] ^= mask; // flip + const Slice fs(buf.data(), buf.size()); + EXPECT_NE(base_lib, doris::snii::crc32c(fs)) << "bit=" << bit; + EXPECT_NE(base_hw3, detail::crc32c_hw3_extend(0, fs)) << "bit=" << bit; + buf[byte] ^= mask; // restore + } +} + +// FV-6 ([perf-deterministic]): the interleave optimization is engaged and correct. +// threshold > 0, and when hardware CRC is present the dispatcher-selected hardware +// paths equal the authoritative library value. +TEST(SniiCrc32cPerfTest, OptimizationEngagedThresholdAndHardwarePath) { + EXPECT_GT(detail::crc32c_interleave_threshold(), 0U); + + const std::vector buf = make_padded_bytes(0xC0FFEEULL, 0, 4096); // >= threshold + const Slice s(buf.data(), buf.size()); + const uint32_t lib = doris::snii::crc32c(s); + // Portable path is always available and correct. + EXPECT_EQ(lib, detail::crc32c_slice8_extend(0, s)); + if (detail::crc32c_has_hw()) { + EXPECT_EQ(lib, detail::crc32c_hw_serial_extend(0, s)); + EXPECT_EQ(lib, detail::crc32c_hw3_extend(0, s)); + } +} + +// [perf-report-only]: wall-clock throughput, never a CI gate. The only assertions +// are deterministic byte-equivalence; the GB/s figures are printed for reference. +TEST(SniiCrc32cPerfTest, Hw3ThroughputReportOnly) { + constexpr size_t kBytes = 1U << 16; // 64 KiB + const std::vector buf = make_padded_bytes(0x1234ULL, 0, kBytes); + const Slice s(buf.data(), buf.size()); + const uint32_t expected = detail::crc32c_slice8_extend(0, s); + ASSERT_EQ(expected, detail::crc32c_hw_serial_extend(0, s)); + ASSERT_EQ(expected, detail::crc32c_hw3_extend(0, s)); + + constexpr int kIters = 2000; + auto time_path = [&](auto fn) { + volatile uint32_t sink = 0; + const auto t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < kIters; ++i) { + sink ^= fn(0, s); + } + const auto t1 = std::chrono::steady_clock::now(); + (void)sink; + return std::chrono::duration(t1 - t0).count(); + }; + const double hw3_s = time_path(detail::crc32c_hw3_extend); + const double hws_s = time_path(detail::crc32c_hw_serial_extend); + const double gigabytes = static_cast(kBytes) * kIters / 1e9; + std::cout << "[ REPORT ] crc32c 64KiB hw3=" << (hw3_s > 0 ? gigabytes / hw3_s : 0.0) + << " GB/s hw_serial=" << (hws_s > 0 ? gigabytes / hws_s : 0.0) + << " GB/s (report-only, not a gate)\n"; + SUCCEED(); +} + +// FV-7: reader black-box regression. build_reader stamps every dict-block / +// prx-window / frq-region crc via the production crc32c(); open_index and the +// queries re-verify them, so assert_ok fails on any Status::Corruption from the +// dict_block / prx_pod / frq_pod verify paths. Results must match the known-good +// sets (identical to SniiPhraseQueryTest), confirming the seam addition perturbs +// nothing and stored CRCs still verify byte-identically. +TEST(SniiCrc32cTest, ReaderPhraseAndTermQueriesVerifyCrc) { + using doris::snii::reader::LogicalIndexReader; + using doris::snii::reader::SniiSegmentReader; + using doris::snii::snii_test::assert_ok; + using doris::snii::snii_test::build_reader; + using doris::snii::snii_test::MemoryFile; + + MemoryFile file; + SniiSegmentReader segment_reader; + LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + // Term lookup drives the dict-block crc verify path. + bool found = false; + doris::snii::format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup("failed", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + + // Phrase query drives the framed prx-window crc verify path. + std::vector phrase_docids; + assert_ok(doris::snii::query::phrase_query(index_reader, {"failed", "order"}, &phrase_docids)); + EXPECT_EQ(phrase_docids, (std::vector {5000, 7000, 8000})); + + // Phrase-prefix query drives term expansion + prx crc paths. + std::vector prefix_docids; + assert_ok(doris::snii::query::phrase_prefix_query(index_reader, {"failed", "ord"}, + &prefix_docids, 10)); + EXPECT_EQ(prefix_docids, (std::vector {5000, 6000, 7000, 8000})); +} diff --git a/be/test/storage/index/snii/encoding/pfor_test.cpp b/be/test/storage/index/snii/encoding/pfor_test.cpp new file mode 100644 index 00000000000000..b5685813fa2a63 --- /dev/null +++ b/be/test/storage/index/snii/encoding/pfor_test.cpp @@ -0,0 +1,488 @@ +// 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 "storage/index/snii/encoding/pfor.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" + +using namespace doris::snii; + +static void roundtrip(const std::vector& v) { + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + ByteSource src(sink.view()); + std::vector out(v.size()); + ASSERT_TRUE(pfor_decode(&src, v.size(), out.data()).ok()); + EXPECT_EQ(out, v); +} + +TEST(SniiPfor, Uniform) { + roundtrip(std::vector(200, 5)); +} + +TEST(SniiPfor, Ramp) { + std::vector v; + for (uint32_t i = 0; i < 256; ++i) { + v.push_back(i); + } + roundtrip(v); +} + +TEST(SniiPfor, WithOutliers) { + std::vector v(128, 3); + v[10] = 1000000; + v[77] = 999; + roundtrip(v); +} + +TEST(SniiPfor, AllZero) { + roundtrip(std::vector(64, 0)); +} + +TEST(SniiPfor, MaxValues) { + std::vector v(32, 0xFFFFFFFFU); + v[0] = 0; + roundtrip(v); +} + +// Exhaustive: every bit width 0..32, capped random values, across a spread of n +// (including non-power-of-two and prime lengths) so the bit-unpacker's tail handling +// (partial trailing 64-bit word / partial byte) is exercised for each width. A +// deterministic LCG keeps it reproducible without Math.random. +TEST(SniiPfor, AllWidthsRandomLengths) { + uint64_t rng = 0x9E3779B97F4A7C15ULL; + auto next = [&rng]() { + rng = rng * 6364136223846793005ULL + 1442695040888963407ULL; + return static_cast(rng >> 32); + }; + const size_t lens[] = {1, 2, 3, 7, 8, 9, 31, 32, 33, 63, 64, 65, 127, 200, 257}; + for (int w = 0; w <= 32; ++w) { + const uint32_t cap = (w >= 32) ? 0xFFFFFFFFU : ((1U << w) - 1U); + for (size_t n : lens) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = (w == 0) ? 0U : (next() & cap); + } + roundtrip(v); // byte-identical decode required for every (w, n) + } + } +} + +// A few exceptions sprinkled into otherwise-narrow data, at boundary indices, to +// cover the exception-patch path alongside the fast unpack. +TEST(SniiPfor, ExceptionsAtBoundaries) { + std::vector v(130, 2); + v[0] = 0x00ABCDEF; // first + v[63] = 0x12345678; // around the 64-value word boundary + v[64] = 0x0BADF00D; + v[129] = 0xFFFFFFFF; // last + roundtrip(v); +} + +// --------------------------------------------------------------------------- +// T11: histogram + suffix-sum choose_width and single-pass pfor_encode. +// +// The oracle helpers below are a faithful transcription of the PRE-T11 encoder +// (O(maxw*n) bits_for width selection + std::vector low/exc split). The optimized +// pfor_encode must (a) select the same bit_width and (b) emit byte-identical +// output for arbitrary inputs -- the strongest guard that the on-disk format is +// unchanged. The deterministic op-count seam proves each value's bit-width is +// evaluated exactly once per run (vs the former ~(maxw+2) times). +// --------------------------------------------------------------------------- +namespace { + +uint8_t ref_bits_for(uint32_t v) { + uint8_t b = 0; + while (v) { + ++b; + v >>= 1; + } + return b; +} + +// Original O(maxw*n) width selection -- identical cost formula and ascending-w +// strict-'<' tie-break. +uint8_t ref_choose_width(const std::vector& v) { + const size_t n = v.size(); + uint8_t maxw = 0; + for (size_t i = 0; i < n; ++i) { + // NOLINTNEXTLINE(readability-use-std-min-max): naive reference mirror of the original + if (ref_bits_for(v[i]) > maxw) { + maxw = ref_bits_for(v[i]); + } + } + uint8_t best = maxw; + size_t best_cost = SIZE_MAX; + for (uint8_t w = 0; w <= maxw; ++w) { + size_t exc = 0; + for (size_t i = 0; i < n; ++i) { + if (ref_bits_for(v[i]) > w) { + ++exc; + } + } + const size_t cost = (static_cast(w) * n + 7) / 8 + exc * 6; + if (cost < best_cost) { + best_cost = cost; + best = w; + } + } + return best; +} + +uint32_t ref_low_mask(uint8_t w) { + return (w >= 32) ? 0xFFFFFFFFU : ((1U << w) - 1U); +} + +// Faithful transcription of the pre-T11 pfor_encode byte layout. +std::vector reference_pfor_encode(const std::vector& values) { + const size_t n = values.size(); + const uint8_t w = ref_choose_width(values); + ByteSink sink; + std::vector> exc; // (index, full value) + std::vector low(values.begin(), values.end()); + for (size_t i = 0; i < n; ++i) { + if (ref_bits_for(values[i]) > w) { + exc.emplace_back(static_cast(i), values[i]); + low[i] = 0; + } + } + sink.put_u8(w); + sink.put_varint32(static_cast(exc.size())); + if (w != 0) { + uint64_t acc = 0; + int filled = 0; + for (size_t i = 0; i < n; ++i) { + acc |= static_cast(low[i] & ref_low_mask(w)) << filled; + filled += w; + while (filled >= 8) { + sink.put_u8(static_cast(acc)); + acc >>= 8; + filled -= 8; + } + } + if (filled > 0) { + sink.put_u8(static_cast(acc)); + } + } + uint32_t prev = 0; + for (const auto& e : exc) { + sink.put_varint32(e.first - prev); + sink.put_varint32(e.second); + prev = e.first; + } + return sink.buffer(); +} + +// Deterministic LCG (same constants as AllWidthsRandomLengths above). +struct Lcg { + uint64_t state; + explicit Lcg(uint64_t seed) : state(seed) {} + uint32_t next() { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + return static_cast(state >> 32); + } +}; + +std::vector encode_to_bytes(const std::vector& v) { + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + return sink.buffer(); +} + +// Representative inputs spanning the layout's branches: all-zero, all-one, +// monotonic delta, freq-like (mostly 1 with bursts), bit31-set large values, and +// narrow data with sparse wide exceptions. +std::vector> representative_inputs() { + std::vector> sets; + sets.emplace_back(256, 0U); // all zero -> w == 0 + sets.emplace_back(256, 1U); // all one -> w == 1 + { + std::vector ramp(256); + for (uint32_t i = 0; i < 256; ++i) { + ramp[i] = i; // monotonic delta + } + sets.push_back(std::move(ramp)); + } + { + std::vector freq(256, 1U); // freq-like: mostly 1, a few bursts + freq[3] = 5; + freq[100] = 9; + freq[255] = 2; + sets.push_back(std::move(freq)); + } + { + std::vector big(64, 7U); // some bit31-set values (width 32) + big[0] = 0x80000000U; + big[31] = 0xFFFFFFFFU; + big[63] = 0x80000001U; + sets.push_back(std::move(big)); + } + { + std::vector mixed(200, 4U); // narrow with sparse wide exceptions + mixed[5] = 70000; + mixed[6] = 70000; + mixed[199] = 1234567; + sets.push_back(std::move(mixed)); + } + return sets; +} + +} // namespace + +// FW-01: random data (mostly small, occasional large to force exceptions) round-trips +// exactly and consumes the whole stream. +TEST(SniiPforTest, RoundTripRandom) { + Lcg rng(0xD1B54A32D192ED03ULL); + for (int trial = 0; trial < 50; ++trial) { + std::vector v(256); + for (uint32_t& vi : v) { + const uint32_t r = rng.next(); + vi = (r % 16 == 0) ? r : (r & 0xFFU); // ~1/16 large, rest in [0,255] + } + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + ByteSource src(sink.view()); + std::vector out(v.size()); + ASSERT_TRUE(pfor_decode(&src, v.size(), out.data()).ok()); + EXPECT_EQ(out, v); + EXPECT_TRUE(src.eof()); + } +} + +// FW-02: the chosen bit_width (encoded as byte 0) matches the original linear-scan +// selection for representative + random inputs. +TEST(SniiPforTest, WidthMatchesLinearScan) { + std::vector> sets = representative_inputs(); + Lcg rng(0x0123456789ABCDEFULL); + for (int t = 0; t < 60; ++t) { + const uint32_t cap_bits = rng.next() % 33; // target max width 0..32 + const uint32_t cap = (cap_bits >= 32) ? 0xFFFFFFFFU : ((1U << cap_bits) - 1U); + const size_t n = 1 + (rng.next() % 256); + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = (cap == 0) ? 0U : (rng.next() & cap); + } + sets.push_back(std::move(v)); + } + for (const auto& v : sets) { + const std::vector buf = encode_to_bytes(v); + ASSERT_FALSE(buf.empty()); + EXPECT_EQ(static_cast(buf[0]), static_cast(ref_choose_width(v))); + } +} + +// FW-03: exact, hand-verified golden bytes freeze the on-disk encoding for the +// clean cases (w==0, and clean bit-packing at w==1/4/8). Any drift in width +// choice or bit-packing is caught immediately. +TEST(SniiPforTest, GoldenByteIdentical) { + using B = std::vector; + EXPECT_EQ(encode_to_bytes(std::vector(8, 0U)), (B {0x00, 0x00})); + EXPECT_EQ(encode_to_bytes(std::vector(8, 1U)), (B {0x01, 0x00, 0xFF})); + EXPECT_EQ(encode_to_bytes(std::vector(4, 10U)), (B {0x04, 0x00, 0xAA, 0xAA})); + EXPECT_EQ(encode_to_bytes(std::vector(4, 200U)), + (B {0x08, 0x00, 0xC8, 0xC8, 0xC8, 0xC8})); +} + +// FW-03 (extended): the optimized encoder is byte-identical to the pre-T11 +// reference for representative inputs AND a large random sweep over mixed widths, +// lengths (including n==0 and n>256), and exception densities. +TEST(SniiPforTest, EncodeMatchesLegacyReference) { + for (const auto& v : representative_inputs()) { + EXPECT_EQ(encode_to_bytes(v), reference_pfor_encode(v)); + } + Lcg rng(0x0F0F0F0F0F0F0F0FULL); + for (int t = 0; t < 300; ++t) { + const uint32_t cap_bits = rng.next() % 33; + const uint32_t cap = (cap_bits >= 32) ? 0xFFFFFFFFU : ((1U << cap_bits) - 1U); + const size_t n = rng.next() % 300; // spans n==0 and the n>256 heap path + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + // ~1/32 chance of a full-width value above the cap -> an exception. + uint32_t vi = 0U; + if (rng.next() % 32 == 0) { + vi = rng.next(); + } else if (cap != 0) { + vi = rng.next() & cap; + } + v[i] = vi; + } + EXPECT_EQ(encode_to_bytes(v), reference_pfor_encode(v)) + << "n=" << n << " cap_bits=" << cap_bits; + } +} + +// FW-04: all zeros -> width 0 (the clz(0) guard), no packed bytes, no exceptions. +TEST(SniiPforTest, AllZeros) { + std::vector v(256, 0U); + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + const std::vector& buf = sink.buffer(); + ASSERT_EQ(buf.size(), 2U); // [w=0][n_exc=0], nothing else + EXPECT_EQ(buf[0], 0U); + EXPECT_EQ(buf[1], 0U); + ByteSource src(sink.view()); + std::vector out(v.size(), 0xDEADBEEFU); + ASSERT_TRUE(pfor_decode(&src, v.size(), out.data()).ok()); + EXPECT_EQ(out, v); + EXPECT_TRUE(src.eof()); +} + +// FW-05: single-element runs (small, zero, max). +TEST(SniiPforTest, SingleElement) { + roundtrip(std::vector {42}); + roundtrip(std::vector {0}); + roundtrip(std::vector {0xFFFFFFFFU}); +} + +// FW-06: empty run -> [w=0][n_exc=0]; decode of 0 values is a no-op and consumes +// the header. nullptr values with n==0 must not be dereferenced. +TEST(SniiPforTest, EmptyRun) { + ByteSink sink; + pfor_encode(nullptr, 0, &sink); + ASSERT_EQ(sink.buffer().size(), 2U); + EXPECT_EQ(sink.buffer()[0], 0U); + EXPECT_EQ(sink.buffer()[1], 0U); + ByteSource src(sink.view()); + std::vector out(1, 0xDEADBEEFU); + ASSERT_TRUE(pfor_decode(&src, 0, out.data()).ok()); + EXPECT_EQ(out[0], 0xDEADBEEFU); // nothing written + EXPECT_TRUE(src.eof()); +} + +// FW-07: values with bit31 set (width 32) round-trip, exercising the width-32 +// path (no `value >> w` UB; the exception split uses the cached width comparison). +TEST(SniiPforTest, TopBitSet) { + std::vector v(64, 5U); + v[0] = 0x80000000U; + v[30] = 0xFFFFFFFFU; + v[63] = 0xC0000000U; + roundtrip(v); + roundtrip(std::vector(40, 0x80000000U)); // chosen width == 32 +} + +// FW-08: bimodal data (mostly narrow, a large minority very wide) keeps the chosen +// width small, sending many values to the exception table. Byte-identical to the +// legacy split and round-trips. +TEST(SniiPforTest, ManyExceptions) { + std::vector v(256, 5U); // width 3 + Lcg rng(0x55AA55AA55AA55AAULL); + for (size_t i = 0; i < v.size(); ++i) { + if (i % 4 == 0) { + v[i] = 0x20000000U | (rng.next() & 0x1FFFFFFFU); // width 30 + } + } + const std::vector buf = encode_to_bytes(v); + EXPECT_EQ(buf, reference_pfor_encode(v)); + size_t exceptions = 0; + for (uint32_t x : v) { + if (ref_bits_for(x) > buf[0]) { + ++exceptions; + } + } + EXPECT_GT(exceptions, 50U); // exception-heavy path actually exercised + roundtrip(v); +} + +// FW-09: runs larger than the 256-element stack buffer exercise the heap fallback; +// byte-identical to the reference and round-trip exact. +TEST(SniiPforTest, LargeRunHeapFallback) { + for (size_t n : {257U, 300U, 512U}) { + Lcg rng(0xC0FFEEULL + n); + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = rng.next() & 0x3FFU; + } + v[n / 2] = 0xABCDEF01U; // an exception + EXPECT_EQ(encode_to_bytes(v), reference_pfor_encode(v)) << "n=" << n; + roundtrip(v); + } +} + +// FW-10: a stream whose exception index is out of range must return Corruption, +// not throw. (Decoder path, unchanged by T11, guarded here against regressions.) +TEST(SniiPforTest, CorruptExceptionIndexReturnsCorruption) { + ByteSink sink; + sink.put_u8(0); // w = 0 -> no packed region + sink.put_varint32(1); // n_exc = 1 + sink.put_varint32(10); // index_delta = 10 -> idx = 10 + sink.put_varint32(7); // exception value + ByteSource src(sink.view()); + std::vector out(4, 0U); // n = 4, so idx 10 is out of range + auto st = pfor_decode(&src, 4, out.data()); + EXPECT_FALSE(st.ok()); + EXPECT_TRUE(st.is()) << st.to_string(); +} + +// Perf (deterministic): exactly one bit-width evaluation per value per run -- the +// single histogram pass, with the encoder reusing the cached widths (no 3rd +// bits_for pass, no per-candidate O(maxw*n) rescan). RED on the un-optimized code +// (which evaluated ~(maxw+2)*n times). +TEST(SniiPforPerfTest, WidthEvalsEqualsNPerRun) { + Lcg rng(0x1357924680ULL); + for (size_t n : {0U, 1U, 7U, 8U, 9U, 64U, 255U, 256U, 257U, 300U}) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = rng.next() & 0xFFFFU; + } + doris::snii::testing::reset_pfor_width_evals(); + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + EXPECT_EQ(doris::snii::testing::pfor_width_evals(), static_cast(n)) << "n=" << n; + } +} + +// Perf (deterministic): cumulative evaluations across multiple runs equal the total +// number of values encoded. +TEST(SniiPforPerfTest, WidthEvalsEqualsTotalValues) { + const std::vector runs = {256, 256, 100, 1, 256}; + Lcg rng(0x2468ACE0ULL); + doris::snii::testing::reset_pfor_width_evals(); + size_t total = 0; + for (size_t n : runs) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = rng.next() & 0xFFFFFFU; + } + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + total += n; + } + EXPECT_EQ(doris::snii::testing::pfor_width_evals(), static_cast(total)); +} + +// Perf (deterministic): worst case for the old O(maxw*n) scan is maxw==32. The old +// encoder evaluated each value ~(maxw+2)==34 times; the histogram path evaluates +// each exactly once, so evals stay == n even when maxw is maximal. +TEST(SniiPforPerfTest, ChooseWidthHasNoInnerRescan) { + const size_t n = 256; + std::vector v(n, 1U); + v[0] = 0x80000000U; // forces maxw == 32 + doris::snii::testing::reset_pfor_width_evals(); + ByteSink sink; + pfor_encode(v.data(), v.size(), &sink); + EXPECT_EQ(doris::snii::testing::pfor_width_evals(), static_cast(n)); + EXPECT_LT(doris::snii::testing::pfor_width_evals(), static_cast(2 * n)); +} diff --git a/be/test/storage/index/snii/encoding/section_framer_test.cpp b/be/test/storage/index/snii/encoding/section_framer_test.cpp new file mode 100644 index 00000000000000..302c16d4f01b47 --- /dev/null +++ b/be/test/storage/index/snii/encoding/section_framer_test.cpp @@ -0,0 +1,76 @@ +// 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 "storage/index/snii/encoding/section_framer.h" + +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" + +using namespace doris::snii; + +TEST(SniiSectionFramer, RoundTrip) { + ByteSink sink; + const uint8_t p[] = {9, 8, 7}; + SectionFramer::write(sink, 0x42, Slice(p, 3)); + ByteSource src(sink.view()); + FramedSection sec; + ASSERT_TRUE(SectionFramer::read(src, &sec).ok()); + EXPECT_EQ(sec.type, 0x42); + ASSERT_EQ(sec.payload.size(), 3U); + EXPECT_EQ(sec.payload[0], 9U); + EXPECT_TRUE(src.eof()); +} + +TEST(SniiSectionFramer, DetectsCorruption) { + ByteSink sink; + const uint8_t p[] = {1, 2, 3, 4}; + SectionFramer::write(sink, 1, Slice(p, 4)); + auto bytes = sink.buffer(); + bytes[3] ^= 0xFF; // flip one byte in the payload + Slice corrupted(bytes); + ByteSource src(corrupted); + FramedSection sec; + EXPECT_FALSE(SectionFramer::read(src, &sec).ok()); +} + +TEST(SniiSectionFramer, SkipMultiple) { + ByteSink sink; + const uint8_t a[] = {1}; + const uint8_t b[] = {2, 2}; + SectionFramer::write(sink, 10, Slice(a, 1)); + SectionFramer::write(sink, 11, Slice(b, 2)); + ByteSource src(sink.view()); + FramedSection s1, s2; + ASSERT_TRUE(SectionFramer::read(src, &s1).ok()); + ASSERT_TRUE(SectionFramer::read(src, &s2).ok()); + EXPECT_EQ(s1.type, 10); + EXPECT_EQ(s2.type, 11); + EXPECT_TRUE(src.eof()); +} + +TEST(SniiSectionFramer, EmptyPayload) { + ByteSink sink; + SectionFramer::write(sink, 7, Slice()); + ByteSource src(sink.view()); + FramedSection sec; + ASSERT_TRUE(SectionFramer::read(src, &sec).ok()); + EXPECT_EQ(sec.type, 7); + EXPECT_EQ(sec.payload.size(), 0U); +} diff --git a/be/test/storage/index/snii/encoding/varint_test.cpp b/be/test/storage/index/snii/encoding/varint_test.cpp new file mode 100644 index 00000000000000..2ae735c733c230 --- /dev/null +++ b/be/test/storage/index/snii/encoding/varint_test.cpp @@ -0,0 +1,73 @@ +// 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 "storage/index/snii/encoding/varint.h" + +#include + +#include + +#include "common/status.h" + +using namespace doris::snii; + +TEST(SniiVarint, RoundTrip32) { + for (uint32_t v : {0U, 1U, 127U, 128U, 300U, 16384U, 0xFFFFFFFFU}) { + uint8_t buf[10]; + size_t n = encode_varint32(v, buf); + EXPECT_EQ(n, varint_len(v)); + uint32_t out; + const uint8_t* next; + ASSERT_TRUE(decode_varint32(buf, buf + n, &out, &next).ok()); + EXPECT_EQ(out, v); + EXPECT_EQ(next, buf + n); + } +} + +TEST(SniiVarint, RoundTrip64) { + for (uint64_t v : {0ULL, 1ULL, 127ULL, 128ULL, 1ULL << 35, 0xFFFFFFFFFFFFFFFFULL}) { + uint8_t buf[10]; + size_t n = encode_varint64(v, buf); + uint64_t out; + const uint8_t* next; + ASSERT_TRUE(decode_varint64(buf, buf + n, &out, &next).ok()); + EXPECT_EQ(out, v); + } +} + +TEST(SniiVarint, TruncatedFails) { + uint8_t buf[1] = {0x80}; // continuation bit set but no subsequent byte + uint32_t out; + const uint8_t* next; + EXPECT_FALSE(decode_varint32(buf, buf + 1, &out, &next).ok()); +} + +TEST(SniiVarint, Overflow32Fails) { + // encode a value > 2^32-1, decoding with decode_varint32 should fail + uint8_t buf[10]; + size_t n = encode_varint64((1ULL << 33), buf); + uint32_t out; + const uint8_t* next; + EXPECT_FALSE(decode_varint32(buf, buf + n, &out, &next).ok()); +} + +TEST(SniiVarint, ZigzagRoundTrip) { + const int64_t cases[] = {0, -1, 1, -1000, 1000, INT64_MIN, INT64_MAX}; + for (int64_t v : cases) { + EXPECT_EQ(zigzag_decode(zigzag_encode(v)), v); + } +} diff --git a/be/test/storage/index/snii/encoding/zstd_codec_test.cpp b/be/test/storage/index/snii/encoding/zstd_codec_test.cpp new file mode 100644 index 00000000000000..86b5a54a86a790 --- /dev/null +++ b/be/test/storage/index/snii/encoding/zstd_codec_test.cpp @@ -0,0 +1,51 @@ +// 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 "storage/index/snii/encoding/zstd_codec.h" + +#include + +#include + +#include "common/status.h" + +using namespace doris::snii; + +TEST(SniiZstd, RoundTrip) { + std::vector in; + for (int i = 0; i < 10000; ++i) { + in.push_back(static_cast(i % 7)); + } + std::vector comp, decomp; + ASSERT_TRUE(zstd_compress(Slice(in), 3, &comp).ok()); + EXPECT_LT(comp.size(), in.size()); + ASSERT_TRUE(zstd_decompress(Slice(comp), in.size(), &decomp).ok()); + EXPECT_EQ(decomp, in); +} + +TEST(SniiZstd, WrongLenFails) { + std::vector in(100, 7), comp, decomp; + ASSERT_TRUE(zstd_compress(Slice(in), 3, &comp).ok()); + EXPECT_FALSE(zstd_decompress(Slice(comp), 99, &decomp).ok()); +} + +TEST(SniiZstd, EmptyInput) { + std::vector in, comp, decomp; + ASSERT_TRUE(zstd_compress(Slice(in), 3, &comp).ok()); + ASSERT_TRUE(zstd_decompress(Slice(comp), 0, &decomp).ok()); + EXPECT_TRUE(decomp.empty()); +} diff --git a/be/test/storage/index/snii/format/bootstrap_header_test.cpp b/be/test/storage/index/snii/format/bootstrap_header_test.cpp new file mode 100644 index 00000000000000..75c510c8b77ac1 --- /dev/null +++ b/be/test/storage/index/snii/format/bootstrap_header_test.cpp @@ -0,0 +1,184 @@ +// 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 "storage/index/snii/format/bootstrap_header.h" + +#include + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using doris::Status; + +namespace { + +// Fixed on-disk size of the bootstrap header (all fields + trailing crc32c). +// u32 magic + u16 format_version + u16 min_reader_version + u32 flags +// + u32 header_length + u8 tail_pointer_size + u32 header_checksum +constexpr size_t kExpectedHeaderBytes = 4 + 2 + 2 + 4 + 4 + 1 + 4; + +} // namespace + +TEST(SniiBootstrapHeader, RoundTripDefaults) { + BootstrapHeader in {}; // defaults: magic/version/min_reader_version + in.flags = 0; + in.tail_pointer_size = 40; + + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + EXPECT_EQ(sink.size(), kExpectedHeaderBytes); + + BootstrapHeader out {}; + ASSERT_TRUE(decode_bootstrap_header(sink.view(), &out).ok()); + EXPECT_EQ(out.magic, in.magic); + EXPECT_EQ(out.format_version, in.format_version); + EXPECT_EQ(out.min_reader_version, in.min_reader_version); + EXPECT_EQ(out.flags, in.flags); + EXPECT_EQ(out.header_length, kExpectedHeaderBytes); + EXPECT_EQ(out.tail_pointer_size, in.tail_pointer_size); +} + +TEST(SniiBootstrapHeader, RoundTripNonDefaultFields) { + BootstrapHeader in {}; + in.flags = 0xDEADBEEFU; + in.tail_pointer_size = 255; + + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + BootstrapHeader out {}; + ASSERT_TRUE(decode_bootstrap_header(sink.view(), &out).ok()); + EXPECT_EQ(out.flags, 0xDEADBEEFU); + EXPECT_EQ(out.tail_pointer_size, 255); + EXPECT_EQ(out.magic, kContainerMagic); + EXPECT_EQ(out.format_version, kFormatVersion); +} + +TEST(SniiBootstrapHeader, EncodeFillsHeaderLength) { + BootstrapHeader in {}; + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + // header_length is written into the buffer; first 4 bytes are magic LE. + const auto& bytes = sink.buffer(); + ASSERT_EQ(bytes.size(), kExpectedHeaderBytes); + uint32_t magic_le = static_cast(bytes[0]) | (static_cast(bytes[1]) << 8) | + (static_cast(bytes[2]) << 16) | + (static_cast(bytes[3]) << 24); + EXPECT_EQ(magic_le, kContainerMagic); +} + +TEST(SniiBootstrapHeader, WrongMagicRejected) { + BootstrapHeader in {}; + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + auto bytes = sink.buffer(); + bytes[0] ^= 0xFF; // corrupt the magic; checksum would also fail, but magic + // check fires first + BootstrapHeader out {}; + Status s = decode_bootstrap_header(Slice(bytes), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, ChecksumMismatchRejected) { + BootstrapHeader in {}; + in.flags = 7; + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + auto bytes = sink.buffer(); + // Flip a flags byte (offset 8 = after magic(4)+fmt(2)+min(2)); magic stays + // valid, so only the checksum can detect this. + bytes[8] ^= 0xFF; + BootstrapHeader out {}; + Status s = decode_bootstrap_header(Slice(bytes), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, MinReaderVersionTooNewUnsupported) { + BootstrapHeader in {}; + // A future writer that demands a reader >= kFormatVersion+1: current reader + // must refuse. + in.min_reader_version = static_cast(kFormatVersion + 1); + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + BootstrapHeader out {}; + Status s = decode_bootstrap_header(sink.view(), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, OlderFormatVersionUnsupported) { + BootstrapHeader in {}; + in.format_version = static_cast(kFormatVersion - 1); + in.min_reader_version = static_cast(kFormatVersion - 1); + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + BootstrapHeader out {}; + Status s = decode_bootstrap_header(sink.view(), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, FutureFormatVersionUnsupported) { + BootstrapHeader in {}; + in.format_version = static_cast(kFormatVersion + 1); + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + BootstrapHeader out {}; + Status s = decode_bootstrap_header(sink.view(), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, TruncatedRejected) { + BootstrapHeader in {}; + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + auto bytes = sink.buffer(); + bytes.pop_back(); // drop one checksum byte + BootstrapHeader out {}; + Status s = decode_bootstrap_header(Slice(bytes), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, EmptyInputRejected) { + BootstrapHeader out {}; + std::vector empty; + Status s = decode_bootstrap_header(Slice(empty), &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiBootstrapHeader, TrailingBytesAfterChecksumRejected) { + BootstrapHeader in {}; + ByteSink sink; + ASSERT_TRUE(encode_bootstrap_header(in, &sink).ok()); + + auto bytes = sink.buffer(); + bytes.push_back(0x00); // extra byte beyond the fixed header + BootstrapHeader out {}; + Status s = decode_bootstrap_header(Slice(bytes), &out); + EXPECT_TRUE(s.is()); +} diff --git a/be/test/storage/index/snii/format/bsbf_test.cpp b/be/test/storage/index/snii/format/bsbf_test.cpp new file mode 100644 index 00000000000000..8282c10822036d --- /dev/null +++ b/be/test/storage/index/snii/format/bsbf_test.cpp @@ -0,0 +1,211 @@ +// 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 "storage/index/snii/format/bsbf.h" + +#include +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/io/local_file.h" + +namespace doris::snii::format { +namespace { + +// Independent Parquet-canonical oracle (scalar): XXH64 seed 0 hash (via bsbf_hash), +// fastrange block index, SALT masks. Verifies BsbfBuilder (which may run AVX2) is +// algorithmically correct AND Parquet-canonical. +struct Oracle { + std::vector w; + uint32_t nblocks = 0; + void init(uint32_t ndv, double fpp) { + const uint32_t nb = bsbf_optimal_num_bytes(ndv, fpp); + nblocks = nb / kBsbfBytesPerBlock; + w.assign(nb / 4, 0U); + } + static void masks(uint32_t key, uint32_t m[8]) { + for (int i = 0; i < 8; ++i) { + m[i] = 1U << ((key * kBsbfSalt[i]) >> 27); + } + } + void insert(uint64_t h) { + const auto b = static_cast(((h >> 32) * nblocks) >> 32); + uint32_t m[8]; + masks(static_cast(h), m); + for (int i = 0; i < 8; ++i) { + w[b * 8 + i] |= m[i]; + } + } + bool find(uint64_t h) const { + const auto b = static_cast(((h >> 32) * nblocks) >> 32); + uint32_t m[8]; + masks(static_cast(h), m); + for (int i = 0; i < 8; ++i) { + if ((w[b * 8 + i] & m[i]) != m[i]) { + return false; + } + } + return true; + } +}; + +std::string tmp_path() { + return "/tmp/snii_bsbf_" + std::to_string(::getpid()) + "_" + std::to_string(std::rand()) + + ".bin"; +} + +TEST(SniiBsbf, BuildProbeMatchesParquetOracle) { + const uint32_t n = 100000; + const double fpp = 0.01; + BsbfBuilder b; + ASSERT_TRUE(BsbfBuilder::create(n, fpp, &b).ok()); + Oracle o; + o.init(n, fpp); + ASSERT_EQ(b.num_blocks(), o.nblocks); + for (uint32_t i = 0; i < n; ++i) { + const uint64_t h = bsbf_hash("t" + std::to_string(i)); + b.insert(h); + o.insert(h); + } + for (uint32_t i = 0; i < n; ++i) { + const uint64_t h = bsbf_hash("t" + std::to_string(i)); + EXPECT_TRUE(b.maybe_contains(h)) << i; // no false negatives + EXPECT_EQ(b.maybe_contains(h), o.find(h)); // == Parquet oracle (AVX2==scalar) + } + uint32_t fp = 0; + for (uint32_t i = 0; i < n; ++i) { + const uint64_t h = bsbf_hash("absent_" + std::to_string(i)); + EXPECT_EQ(b.maybe_contains(h), o.find(h)); + if (b.maybe_contains(h)) { + ++fp; + } + } + EXPECT_LT(fp, static_cast(n * 0.05)); // FPR << 5% (target 1%) +} + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiBsbf, SerializeParseOnDemandProbe) { + const uint32_t n = 50000; + BsbfBuilder b; + ASSERT_TRUE(BsbfBuilder::create(n, 0.01, &b).ok()); + std::vector keys; + for (uint32_t i = 0; i < n; ++i) { + const uint64_t h = bsbf_hash("k" + std::to_string(i)); + keys.push_back(h); + b.insert(h); + } + // Serialize at a non-zero section base (prefix) to exercise the constant offset. + ByteSink sink; + const uint64_t prefix = 777; + for (uint64_t i = 0; i < prefix; ++i) { + sink.put_u8(0xAB); + } + ASSERT_TRUE(b.serialize(&sink).ok()); + EXPECT_EQ(sink.size(), prefix + kBsbfHeaderSize + b.num_bytes()); // 28B header, exact + + const std::string path = tmp_path(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + ASSERT_TRUE(w.append(sink.view()).ok()); + ASSERT_TRUE(w.finalize().ok()); + } + io::LocalFileReader r; + ASSERT_TRUE(r.open(path).ok()); + std::vector hdr; + ASSERT_TRUE(r.read_at(prefix, kBsbfHeaderSize, &hdr).ok()); + BsbfHeader h; + ASSERT_TRUE(BsbfHeader::parse(Slice(hdr), prefix, &h).ok()); + EXPECT_EQ(h.num_bytes, b.num_bytes()); + EXPECT_EQ(h.bitset_base, prefix + kBsbfHeaderSize); + + for (uint32_t i = 0; i < 2000; ++i) { + bool mp = false; + ASSERT_TRUE(bsbf_probe(&r, h, keys[i], &mp).ok()); + EXPECT_TRUE(mp) << i; // present -> maybe + EXPECT_EQ(mp, b.maybe_contains(keys[i])); // on-demand == in-memory + } + for (uint32_t i = 0; i < 2000; ++i) { + const uint64_t a = bsbf_hash("absent_" + std::to_string(i)); + bool mp = false; + ASSERT_TRUE(bsbf_probe(&r, h, a, &mp).ok()); + EXPECT_EQ(mp, b.maybe_contains(a)); // incl. any false positive, identical + } + std::remove(path.c_str()); +} + +TEST(SniiBsbf, HeaderValidation) { + BsbfBuilder b; + ASSERT_TRUE(BsbfBuilder::create(1000, 0.01, &b).ok()); + ByteSink sink; + ASSERT_TRUE(b.serialize(&sink).ok()); + const std::vector bytes(sink.buffer()); + BsbfHeader h; + EXPECT_TRUE(BsbfHeader::parse(Slice(bytes.data(), kBsbfHeaderSize), 0, &h).ok()); + + auto fails = [](std::vector c) { + BsbfHeader hh; + return !BsbfHeader::parse(Slice(c.data(), kBsbfHeaderSize), 0, &hh).ok(); + }; + { + auto c = bytes; + c[0] = 'X'; + EXPECT_TRUE(fails(c)); + } // magic + { + auto c = bytes; + c[4] = 9; + EXPECT_TRUE(fails(c)); + } // version + { + auto c = bytes; + c[5] = 1; + EXPECT_TRUE(fails(c)); + } // hash strategy + { + auto c = bytes; + c[8] ^= 0xFF; + EXPECT_TRUE(fails(c)); + } // field -> header crc + EXPECT_FALSE(BsbfHeader::parse(Slice(bytes.data(), 10), 0, &h).ok()); // short +} + +TEST(SniiBsbf, FastrangeNotMaskVariant) { + // The fastrange index must differ from Doris's `&(n-1)` for some hash, proving we + // use the Parquet-canonical multiply-shift (not the mask variant). + const uint32_t nblocks = 4096; // power of 2 so both formulas are defined + bool differ = false; + for (uint64_t s = 1; s < 200 && !differ; ++s) { + const uint64_t h = bsbf_hash("probe" + std::to_string(s)); + const uint32_t fastrange = bsbf_block_index(h, nblocks); + const uint32_t mask = static_cast(h >> 32) & (nblocks - 1); + if (fastrange != mask) { + differ = true; + } + } + EXPECT_TRUE(differ); +} + +} // namespace +} // namespace doris::snii::format diff --git a/be/test/storage/index/snii/format/core_metadata_test.cpp b/be/test/storage/index/snii/format/core_metadata_test.cpp new file mode 100644 index 00000000000000..35f8224847ea7b --- /dev/null +++ b/be/test/storage/index/snii/format/core_metadata_test.cpp @@ -0,0 +1,319 @@ +// 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 "storage/index/snii/format/core_metadata.h" + +#include + +#include +#include +#include +#include +#include + +#include "gen_cpp/snii.pb.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" + +namespace doris::snii::format { +namespace { + +using segment_v2::inverted_index::CommonGramsCoverage; +using segment_v2::inverted_index::CommonGramsSegmentMetadata; +using segment_v2::inverted_index::PlainTermKeyVersion; +using segment_v2::inverted_index::ScoringCoverage; + +CoreMetadata sample_core(IndexConfig index_config = IndexConfig::kDocsOnly) { + CoreMetadata metadata; + metadata.index_config = index_config; + metadata.stats = {.doc_count = 20, + .indexed_doc_count = 19, + .term_count = 18, + .sum_total_term_freq = 123, + .null_count = 1}; + metadata.section_refs = {.dict_region = {.offset = 10, .length = 11}, + .posting_region = {.offset = 21, .length = 22}, + .norms = {.offset = 31, .length = 32}, + .null_bitmap = {.offset = 41, .length = 42}, + .bsbf = {.offset = 51, .length = 52}}; + return metadata; +} + +CommonGramsSegmentMetadata sample_common_grams(CommonGramsCoverage coverage, + ScoringCoverage scoring_coverage) { + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = coverage; + metadata.common_grams_semantics_version = 1; + metadata.common_grams_key_version = 1; + metadata.common_grams_dictionary_identity = std::string("dictionary\0id", 13); + metadata.base_analyzer_fingerprint = std::string("base\0fingerprint", 16); + metadata.common_grams_fingerprint = std::string("grams\0fingerprint", 17); + metadata.scoring_coverage = scoring_coverage; + metadata.scoring_stats_version = 1; + metadata.norm_semantics_version = 1; + metadata.scoring_doc_count = 20; + metadata.scoring_token_count = 123; + return metadata; +} + +std::vector encode(const CoreMetadata& metadata) { + ByteSink sink; + EXPECT_TRUE(encode_core_metadata(metadata, &sink).ok()); + return sink.buffer(); +} + +std::vector payload_of(const std::vector& framed) { + ByteSource source {Slice(framed)}; + FramedSection section; + EXPECT_TRUE(SectionFramer::read(source, §ion).ok()); + EXPECT_TRUE(source.eof()); + return std::vector(section.payload.data(), + section.payload.data() + section.payload.size()); +} + +std::vector frame_payload( + const std::vector& payload, + uint8_t type = static_cast(SectionType::kCoreMetadataPB)) { + ByteSink sink; + SectionFramer::write(sink, type, Slice(payload)); + return sink.buffer(); +} + +std::vector mutate_core_payload( + const CoreMetadata& metadata, + const std::function& mutation) { + const auto payload = payload_of(encode(metadata)); + doris::snii::SniiCoreMetadataPB core; + EXPECT_TRUE(core.ParseFromArray(payload.data(), static_cast(payload.size()))); + mutation(&core); + std::string mutated; + EXPECT_TRUE(core.SerializeToString(&mutated)); + return std::vector(mutated.begin(), mutated.end()); +} + +void expect_core_eq(const CoreMetadata& expected, const CoreMetadata& actual) { + EXPECT_EQ(expected.index_config, actual.index_config); + EXPECT_EQ(expected.stats.doc_count, actual.stats.doc_count); + EXPECT_EQ(expected.stats.indexed_doc_count, actual.stats.indexed_doc_count); + EXPECT_EQ(expected.stats.term_count, actual.stats.term_count); + EXPECT_EQ(expected.stats.sum_total_term_freq, actual.stats.sum_total_term_freq); + EXPECT_EQ(expected.stats.null_count, actual.stats.null_count); + EXPECT_EQ(expected.section_refs.dict_region.offset, actual.section_refs.dict_region.offset); + EXPECT_EQ(expected.section_refs.dict_region.length, actual.section_refs.dict_region.length); + EXPECT_EQ(expected.section_refs.posting_region.offset, + actual.section_refs.posting_region.offset); + EXPECT_EQ(expected.section_refs.posting_region.length, + actual.section_refs.posting_region.length); + EXPECT_EQ(expected.section_refs.norms.offset, actual.section_refs.norms.offset); + EXPECT_EQ(expected.section_refs.norms.length, actual.section_refs.norms.length); + EXPECT_EQ(expected.section_refs.null_bitmap.offset, actual.section_refs.null_bitmap.offset); + EXPECT_EQ(expected.section_refs.null_bitmap.length, actual.section_refs.null_bitmap.length); + EXPECT_EQ(expected.section_refs.bsbf.offset, actual.section_refs.bsbf.offset); + EXPECT_EQ(expected.section_refs.bsbf.length, actual.section_refs.bsbf.length); + EXPECT_EQ(expected.common_grams_metadata, actual.common_grams_metadata); + EXPECT_EQ(expected.common_grams_posting_policy, actual.common_grams_posting_policy); +} + +TEST(SniiCoreMetadata, RoundTripsDocsOnlyWithAllStatsAndRefs) { + const auto expected = sample_core(); + CoreMetadata actual; + ASSERT_TRUE(decode_core_metadata(Slice(encode(expected)), &actual).ok()); + expect_core_eq(expected, actual); +} + +TEST(SniiCoreMetadata, RoundTripsPositions) { + const auto expected = sample_core(IndexConfig::kDocsPositions); + CoreMetadata actual; + ASSERT_TRUE(decode_core_metadata(Slice(encode(expected)), &actual).ok()); + expect_core_eq(expected, actual); +} + +TEST(SniiCoreMetadata, RoundTripsScoringWithBinaryCommonGramsStrings) { + auto expected = sample_core(IndexConfig::kDocsPositionsScoring); + expected.common_grams_metadata = + sample_common_grams(CommonGramsCoverage::kComplete, ScoringCoverage::kComplete); + + CoreMetadata actual; + ASSERT_TRUE(decode_core_metadata(Slice(encode(expected)), &actual).ok()); + expect_core_eq(expected, actual); +} + +TEST(SniiCoreMetadata, RoundTripsHybridCommonGrams) { + auto expected = sample_core(IndexConfig::kDocsPositions); + expected.common_grams_metadata = + sample_common_grams(CommonGramsCoverage::kMixed, ScoringCoverage::kNone); + expected.common_grams_posting_policy = CommonGramsPostingPolicy::kHybridV1; + + CoreMetadata actual; + ASSERT_TRUE(decode_core_metadata(Slice(encode(expected)), &actual).ok()); + expect_core_eq(expected, actual); +} + +TEST(SniiCoreMetadata, AcceptsUnknownOptionalPbField) { + auto payload = payload_of(encode(sample_core())); + ByteSink unknown_field; + unknown_field.put_varint32((100u << 3) | 0u); + unknown_field.put_varint32(7); + payload.insert(payload.end(), unknown_field.buffer().begin(), unknown_field.buffer().end()); + + CoreMetadata actual; + ASSERT_TRUE(decode_core_metadata(Slice(frame_payload(payload)), &actual).ok()); + expect_core_eq(sample_core(), actual); +} + +TEST(SniiCoreMetadata, RejectsMissingRequiredLogicalField) { + ByteSink payload; + payload.put_varint32((1u << 3) | 0u); + payload.put_varint32(static_cast(IndexConfig::kDocsOnly)); + + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload.buffer())), &actual); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiCoreMetadata, RejectsUnsupportedIndexConfig) { + auto payload = payload_of(encode(sample_core())); + ByteSink field; + field.put_varint32((1u << 3) | 0u); + field.put_varint32(9); + payload.insert(payload.end(), field.buffer().begin(), field.buffer().end()); + + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiCoreMetadata, RejectsUnsupportedPostingPolicy) { + auto payload = payload_of(encode(sample_core())); + ByteSink field; + field.put_varint32((5u << 3) | 0u); + field.put_varint32(9); + payload.insert(payload.end(), field.buffer().begin(), field.buffer().end()); + + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiCoreMetadata, UnsupportedCommonGramsEnums) { + auto metadata = sample_core(IndexConfig::kDocsPositionsScoring); + metadata.common_grams_metadata = + sample_common_grams(CommonGramsCoverage::kComplete, ScoringCoverage::kComplete); + + for (const auto& mutation : + std::array, 3> { + [](auto* common_grams) { common_grams->set_plain_term_key_version(3); }, + [](auto* common_grams) { common_grams->set_common_grams_coverage(3); }, + [](auto* common_grams) { common_grams->set_scoring_coverage(2); }}) { + const auto payload = mutate_core_payload( + metadata, [&mutation](auto* core) { mutation(core->mutable_common_grams()); }); + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; + } +} + +TEST(SniiCoreMetadata, RejectsKnownButContradictoryCommonGramsEnumsAsCorruption) { + auto metadata = sample_core(IndexConfig::kDocsPositions); + metadata.common_grams_metadata = + sample_common_grams(CommonGramsCoverage::kComplete, ScoringCoverage::kNone); + + const auto payload = mutate_core_payload(metadata, [](auto* core) { + core->mutable_common_grams()->set_plain_term_key_version(2); + }); + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiCoreMetadata, RejectsMissingEachStatsField) { + const auto metadata = sample_core(); + for (const auto clear : std::array { + &doris::snii::SniiStatsPB::clear_doc_count, + &doris::snii::SniiStatsPB::clear_indexed_doc_count, + &doris::snii::SniiStatsPB::clear_term_count, + &doris::snii::SniiStatsPB::clear_sum_total_term_freq, + &doris::snii::SniiStatsPB::clear_null_count}) { + const auto payload = mutate_core_payload( + metadata, [clear](auto* core) { (core->mutable_stats()->*clear)(); }); + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; + } +} + +TEST(SniiCoreMetadata, RejectsMissingEachRegionRefOffsetOrLength) { + using RegionGetter = doris::snii::SniiRegionRefPB* (doris::snii::SniiSectionRefsPB::*)(); + const auto metadata = sample_core(); + for (const auto region : + std::array {&doris::snii::SniiSectionRefsPB::mutable_dict_region, + &doris::snii::SniiSectionRefsPB::mutable_posting_region, + &doris::snii::SniiSectionRefsPB::mutable_norms, + &doris::snii::SniiSectionRefsPB::mutable_null_bitmap, + &doris::snii::SniiSectionRefsPB::mutable_bsbf}) { + for (const auto clear : std::array { + &doris::snii::SniiRegionRefPB::clear_offset, + &doris::snii::SniiRegionRefPB::clear_length}) { + const auto payload = mutate_core_payload(metadata, [region, clear](auto* core) { + auto* refs = core->mutable_section_refs(); + ((refs->*region)()->*clear)(); + }); + CoreMetadata actual; + const auto status = decode_core_metadata(Slice(frame_payload(payload)), &actual); + EXPECT_TRUE(status.is()) << status; + } + } +} + +TEST(SniiCoreMetadata, RejectsBadFrameTypeCrcAndTruncation) { + const auto framed = encode(sample_core()); + CoreMetadata actual; + + const auto bad_type = frame_payload(payload_of(framed), /*obsolete non-Core frame type=*/1); + EXPECT_TRUE(decode_core_metadata(Slice(bad_type), &actual) + .is()); + + auto bad_crc = framed; + bad_crc.back() ^= 1; + EXPECT_TRUE(decode_core_metadata(Slice(bad_crc), &actual) + .is()); + + std::vector truncated(framed.begin(), framed.end() - 1); + EXPECT_TRUE(decode_core_metadata(Slice(truncated), &actual) + .is()); +} + +TEST(SniiCoreMetadata, ResetsOutputBeforeDecodeFailureAndNullOutputIsInvalid) { + auto populated = sample_core(IndexConfig::kDocsPositions); + populated.common_grams_metadata = + sample_common_grams(CommonGramsCoverage::kMixed, ScoringCoverage::kNone); + populated.common_grams_posting_policy = CommonGramsPostingPolicy::kHybridV1; + CoreMetadata reused; + ASSERT_TRUE(decode_core_metadata(Slice(encode(populated)), &reused).ok()); + + const auto status = decode_core_metadata(Slice(std::vector {1}), &reused); + EXPECT_TRUE(status.is()) << status; + expect_core_eq(CoreMetadata {}, reused); + EXPECT_TRUE(decode_core_metadata(Slice(encode(sample_core())), nullptr) + .is()); +} + +} // namespace +} // namespace doris::snii::format diff --git a/be/test/storage/index/snii/format/dict_block_directory_test.cpp b/be/test/storage/index/snii/format/dict_block_directory_test.cpp new file mode 100644 index 00000000000000..b7d2356939548a --- /dev/null +++ b/be/test/storage/index/snii/format/dict_block_directory_test.cpp @@ -0,0 +1,352 @@ +// 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 "storage/index/snii/format/dict_block_directory.h" + +#include + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; +using namespace doris::snii::format; + +namespace { + +// Serialize a list of block_refs using the builder and return the framed section bytes. +std::vector Build(const std::vector& refs) { + DictBlockDirectoryBuilder builder; + for (const auto& r : refs) { + builder.add(r); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +// Assert that two BlockRef structs are equal field by field. +void ExpectRefEq(const BlockRef& a, const BlockRef& b) { + EXPECT_EQ(a.offset, b.offset); + EXPECT_EQ(a.length, b.length); + EXPECT_EQ(a.n_entries, b.n_entries); + EXPECT_EQ(a.flags, b.flags); + EXPECT_EQ(a.checksum, b.checksum); + EXPECT_EQ(a.uncomp_len, b.uncomp_len); +} + +} // namespace + +// A zstd-compressed block ref carries uncomp_len; a raw ref does not. Both +// round-trip exactly, and the raw ref's directory bytes stay v1-compact (no +// trailing uncomp_len varint when the kZstd flag is clear). +TEST(SniiDictBlockDirectory, ZstdRefCarriesUncompLen) { + std::vector refs = { + {.offset = 0, + .length = 40000, + .n_entries = 250, + .flags = block_ref_flags::kZstd, + .checksum = 0xABCDEF01U, + .uncomp_len = 65536}, // compressed + {.offset = 40000, + .length = 4096, + .n_entries = 64, + .flags = 0, + .checksum = 0x11223344U, + .uncomp_len = 0}, // raw + }; + auto bytes = Build(refs); + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 2U); + BlockRef z {}, r {}; + ASSERT_TRUE(reader.get(0, &z).ok()); + ASSERT_TRUE(reader.get(1, &r).ok()); + ExpectRefEq(z, refs[0]); + ExpectRefEq(r, refs[1]); + EXPECT_EQ(z.uncomp_len, 65536U); + EXPECT_TRUE((z.flags & block_ref_flags::kZstd) != 0); + EXPECT_EQ(r.uncomp_len, 0U); + EXPECT_FALSE((r.flags & block_ref_flags::kZstd) != 0); +} + +TEST(SniiDictBlockDirectory, RoundTripMultipleRefs) { + std::vector refs = { + {.offset = 0, .length = 4096, .n_entries = 120, .flags = 0x01, .checksum = 0xDEADBEEFU}, + {.offset = 4096, + .length = 8192, + .n_entries = 300, + .flags = 0x05, + .checksum = 0x12345678U}, + {.offset = 12288, + .length = 2048, + .n_entries = 64, + .flags = 0x00, + .checksum = 0xCAFEBABEU}, + }; + auto bytes = Build(refs); + + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 3U); + for (uint32_t i = 0; i < refs.size(); ++i) { + BlockRef out {}; + ASSERT_TRUE(reader.get(i, &out).ok()); + ExpectRefEq(out, refs[i]); + } +} + +TEST(SniiDictBlockDirectory, GetOrdinalCorrectMapping) { + std::vector refs; + for (uint32_t i = 0; i < 50; ++i) { + refs.push_back(BlockRef {.offset = static_cast(i) * 1000, + .length = 1000, + .n_entries = i + 1, + .flags = static_cast(i & 0xFF), + .checksum = i * 7U + 3U}); + } + auto bytes = Build(refs); + + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 50U); + // Sample several ordinals and verify the mapping produces no cross-slot errors. + for (uint32_t ord : {0U, 1U, 17U, 49U}) { + BlockRef out {}; + ASSERT_TRUE(reader.get(ord, &out).ok()); + ExpectRefEq(out, refs[ord]); + } +} + +TEST(SniiDictBlockDirectory, OutOfRangeOrdinalRejected) { + std::vector refs = { + {.offset = 0, .length = 100, .n_entries = 1, .flags = 0, .checksum = 0xAAU}}; + auto bytes = Build(refs); + + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 1U); + BlockRef out {}; + // ordinal == n_blocks is out of range + EXPECT_TRUE(reader.get(1, &out).is()); + // far beyond range + EXPECT_TRUE(reader.get(1000, &out).is()); +} + +TEST(SniiDictBlockDirectory, EmptyDirectory) { + std::vector refs; // 0 blocks + auto bytes = Build(refs); + + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + EXPECT_EQ(reader.n_blocks(), 0U); + BlockRef out {}; + EXPECT_TRUE(reader.get(0, &out).is()); +} + +TEST(SniiDictBlockDirectory, LargeOffsetNear2Pow48) { + const uint64_t kBig = (1ULL << 48) - 1; // near 2^48 + std::vector refs = { + {.offset = 1, + .length = kBig - 1, + .n_entries = 0xFFFFFFFFU, + .flags = 0xFF, + .checksum = 0xFFFFFFFFU}, + {.offset = kBig + 12345, .length = 1, .n_entries = 1, .flags = 0x02, .checksum = 0U}, + }; + auto bytes = Build(refs); + + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 2U); + BlockRef out0 {}; + ASSERT_TRUE(reader.get(0, &out0).ok()); + ExpectRefEq(out0, refs[0]); + BlockRef out1 {}; + ASSERT_TRUE(reader.get(1, &out1).ok()); + ExpectRefEq(out1, refs[1]); +} + +TEST(SniiDictBlockDirectory, RejectsInvalidPhysicalBlockRanges) { + const std::vector> invalid_directories = { + {{.offset = 10, .length = 0, .n_entries = 1}}, + {{.offset = std::numeric_limits::max(), .length = 2, .n_entries = 1}}, + {{.offset = 10, .length = 10, .n_entries = 1}, + {.offset = 19, .length = 10, .n_entries = 1}}, + {{.offset = 20, .length = 10, .n_entries = 1}, + {.offset = 10, .length = 10, .n_entries = 1}}, + }; + + for (const auto& refs : invalid_directories) { + const auto bytes = Build(refs); + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader) + .is()); + } +} + +TEST(SniiDictBlockDirectory, FramedAsDictBlockDirectoryType) { + std::vector refs = { + {.offset = 0, .length = 10, .n_entries = 1, .flags = 0, .checksum = 0}}; + auto bytes = Build(refs); + ASSERT_GE(bytes.size(), 1U); + // The first byte is the SectionFramer type and must be kDictBlockDirectory. + EXPECT_EQ(bytes[0], static_cast(SectionType::kDictBlockDirectory)); +} + +TEST(SniiDictBlockDirectory, DetectsCorruption) { + std::vector refs = { + {.offset = 0, .length = 4096, .n_entries = 120, .flags = 0x01, .checksum = 0xDEADBEEFU}, + {.offset = 4096, + .length = 8192, + .n_entries = 300, + .flags = 0x05, + .checksum = 0x12345678U}, + }; + auto bytes = Build(refs); + ASSERT_GE(bytes.size(), 4U); + // Flip one byte in the payload region (skip the type+len prefix); the section CRC must detect the corruption. + bytes[3] ^= 0xFF; + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader) + .is()); +} + +TEST(SniiDictBlockDirectory, DetectsTruncation) { + std::vector refs = {{.offset = 0, + .length = 4096, + .n_entries = 120, + .flags = 0x01, + .checksum = 0xDEADBEEFU}}; + auto bytes = Build(refs); + bytes.pop_back(); // truncate the last byte + DictBlockDirectoryReader reader; + EXPECT_FALSE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); +} + +TEST(SniiDictBlockDirectory, WrongSectionTypeRejected) { + // Write a section with a type other than kDictBlockDirectory via the framer; open must reject it. + ByteSink sink; + const uint8_t p[] = {0, 1, 2}; + SectionFramer::write(sink, static_cast(SectionType::kSampledTermIndex), Slice(p, 3)); + auto bytes = sink.buffer(); + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader) + .is()); +} + +TEST(SniiDictBlockDirectory, TrailingFramedSectionBytesRejected) { + auto bytes = Build( + {BlockRef {.offset = 0, .length = 10, .n_entries = 1, .flags = 0, .checksum = 123}}); + bytes.push_back(0xEE); + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader) + .is()); +} + +TEST(SniiDictBlockDirectory, TrailingBytesRejected) { + // Extra trailing bytes at the end of the payload should be detected (n_blocks does not match actual data). + ByteSink payload; + payload.put_varint32(1); // n_blocks = 1 + payload.put_varint64(0); // offset + payload.put_varint64(10); // length + payload.put_varint32(1); // n_entries + payload.put_u8(0); // flags + payload.put_fixed32(0); // checksum + payload.put_u8(0xEE); // extra trailing byte + ByteSink sink; + SectionFramer::write(sink, static_cast(SectionType::kDictBlockDirectory), + payload.view()); + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(sink.buffer()), &reader) + .is()); +} + +// An attacker-inflated n_blocks count must hit the reserve-bomb guard in +// decode_payload BEFORE the vector reserve, returning Corruption rather than +// attempting a multi-gigabyte allocation. The payload is framed through +// SectionFramer::write so the section crc is VALID over the crafted bytes: +// this proves the inner n_blocks-vs-capacity check fires, not an outer crc flip. +TEST(SniiDictBlockDirectory, InflatedNBlocksHitsReserveGuard) { + ByteSink payload; + payload.put_varint32(0xFFFFFFFFU); // n_blocks = ~4.29 billion (impossible) + // Only a few real bytes follow; capacity is nowhere near 4.29B refs. + payload.put_u8(0x00); + payload.put_u8(0x00); + payload.put_u8(0x00); + payload.put_u8(0x00); + ByteSink sink; + SectionFramer::write(sink, static_cast(SectionType::kDictBlockDirectory), + payload.view()); + DictBlockDirectoryReader reader; + // remaining()/kMinRefBytes == 4/8 == 0, so n_blocks > 0 trips the guard. + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(sink.buffer()), &reader) + .is()); +} + +// A kZstd ref whose trailing uncomp_len varint is missing (truncated) must be +// rejected as Corruption: decode_ref reads offset/length/n_entries/flags/checksum +// fine, sees the kZstd flag, then overruns when it tries to read uncomp_len. +// n_blocks == 1 with an ~8-byte ref body passes the reserve guard (1 > 8/8 is +// false), so the corruption surfaces from the inner uncomp_len read -- not the +// capacity cap and not the section crc (the frame is crc-valid by construction). +TEST(SniiDictBlockDirectory, ZstdRefTruncatedUncompLenRejected) { + ByteSink payload; + payload.put_varint32(1); // n_blocks = 1 + payload.put_varint64(0); // offset (1 byte) + payload.put_varint64(10); // length (1 byte) + payload.put_varint32(1); // n_entries (1 byte) + payload.put_u8(block_ref_flags::kZstd); // flags: kZstd set + payload.put_fixed32(0xDEADBEEFU); // checksum (4 bytes) + // INTENTIONALLY OMIT the trailing varint64 uncomp_len -> decode_ref overruns. + ByteSink sink; + SectionFramer::write(sink, static_cast(SectionType::kDictBlockDirectory), + payload.view()); + DictBlockDirectoryReader reader; + EXPECT_TRUE(DictBlockDirectoryReader::open(Slice(sink.buffer()), &reader) + .is()); +} + +// A kZstd ref round-trips through the builder/reader preserving uncomp_len +// exactly, including a large 64-bit value, proving the trailing varint is both +// written and read back on the kZstd path. +TEST(SniiDictBlockDirectory, ZstdRefRoundTripPreservesUncompLen) { + const uint64_t kUncomp = (1ULL << 40) + 1234567U; // large, multi-byte varint + std::vector refs = { + {.offset = 1024, + .length = 999, + .n_entries = 7, + .flags = block_ref_flags::kZstd, + .checksum = 0x0BADF00DU, + .uncomp_len = kUncomp}, + }; + auto bytes = Build(refs); + DictBlockDirectoryReader reader; + ASSERT_TRUE(DictBlockDirectoryReader::open(Slice(bytes), &reader).ok()); + ASSERT_EQ(reader.n_blocks(), 1U); + BlockRef out {}; + ASSERT_TRUE(reader.get(0, &out).ok()); + ExpectRefEq(out, refs[0]); + EXPECT_EQ(out.uncomp_len, kUncomp); + EXPECT_TRUE((out.flags & block_ref_flags::kZstd) != 0); +} diff --git a/be/test/storage/index/snii/format/dict_block_owned_finish_test.cpp b/be/test/storage/index/snii/format/dict_block_owned_finish_test.cpp new file mode 100644 index 00000000000000..6b0e7666276591 --- /dev/null +++ b/be/test/storage/index/snii/format/dict_block_owned_finish_test.cpp @@ -0,0 +1,98 @@ +// 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 + +#include +#include +#include +#include + +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/dict_block.h" + +namespace doris::snii::format { +namespace { + +DictEntry MakeInlineEntry(std::string term, uint8_t value, size_t payload_size = 1024) { + DictEntry entry; + entry.term = std::move(term); + entry.kind = DictEntryKind::kInline; + entry.enc = DictEntryEnc::kSlim; + entry.df = 32; + entry.frq_bytes.assign(payload_size, value); + entry.inline_dd_disk_len = entry.frq_bytes.size(); + entry.dd_meta.uncomp_len = entry.frq_bytes.size(); + entry.dd_meta.disk_len = entry.frq_bytes.size(); + entry.dd_meta.verify_crc = false; + return entry; +} + +TEST(DictBlockOwnedFinishTest, PreservesBytesWithBoundedCapacity) { + DictBlockBuilder block(IndexTier::kT1, /*has_positions=*/false, /*frq_base=*/17, + /*prx_base=*/0, /*anchor_interval=*/2); + block.add_entry(MakeInlineEntry("\x1fGcommon", 0x11)); + block.add_entry(MakeInlineEntry("alpha", 0x22)); + block.add_entry(MakeInlineEntry("alphabet", 0x33)); + block.add_entry(MakeInlineEntry("beta", 0x44)); + + ByteSink legacy; + block.finish(&legacy); + std::vector owned = block.finish_owned(); + + EXPECT_EQ(owned, legacy.buffer()); + EXPECT_LT(owned.capacity(), legacy.buffer().capacity()); + + std::vector legacy_compressed; + std::vector owned_compressed; + ASSERT_TRUE(zstd_compress(legacy.view(), /*level=*/3, &legacy_compressed).ok()); + ASSERT_TRUE(zstd_compress(Slice(owned), /*level=*/3, &owned_compressed).ok()); + EXPECT_EQ(owned_compressed, legacy_compressed); + + DictBlockReader reader; + ASSERT_TRUE( + DictBlockReader::open(Slice(owned), IndexTier::kT1, /*has_positions=*/false, &reader) + .ok()); + EXPECT_EQ(reader.n_entries(), 4U); + bool found = false; + DictEntry decoded; + ASSERT_TRUE(reader.find_term("alphabet", &found, &decoded).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(decoded.term, "alphabet"); + EXPECT_EQ(decoded.frq_bytes, std::vector(1024, 0x33)); +} + +TEST(DictBlockOwnedFinishTest, DoesNotReservePrefixCompressionUpperBound) { + DictBlockBuilder block(IndexTier::kT1, /*has_positions=*/false, /*frq_base=*/0, + /*prx_base=*/0, /*anchor_interval=*/16); + const std::string common_prefix(2048, 'x'); + for (char suffix = 'a'; suffix <= 'p'; ++suffix) { + block.add_entry(MakeInlineEntry(common_prefix + suffix, 0x11, /*payload_size=*/1)); + } + + ByteSink legacy; + block.finish(&legacy); + std::vector owned = block.finish_owned(); + + EXPECT_EQ(owned, legacy.buffer()); + EXPECT_LT(owned.capacity(), legacy.buffer().capacity()); + EXPECT_LT(owned.capacity(), block.estimated_bytes()); +} + +} // namespace +} // namespace doris::snii::format diff --git a/be/test/storage/index/snii/format/dict_block_test.cpp b/be/test/storage/index/snii/format/dict_block_test.cpp new file mode 100644 index 00000000000000..818460640c3bef --- /dev/null +++ b/be/test/storage/index/snii/format/dict_block_test.cpp @@ -0,0 +1,643 @@ +// 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 "storage/index/snii/format/dict_block.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; // NOLINT +using namespace doris::snii::format; // NOLINT +using doris::Status; + +namespace { + +// Construct a pod_ref slim dict entry. frq/prx off_delta is already relative to the block base. +DictEntry MakePodRef(std::string term, uint32_t df, uint64_t frq_off, uint64_t prx_off = 0) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kPodRef; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.ttf_delta = df * 2; // written only when tier>=T2 + e.max_freq = 9; // written only when tier>=T2 + e.frq_off_delta = frq_off; + e.frq_len = 128; + e.prx_off_delta = prx_off; // written only when positions are enabled + e.prx_len = 64; // written only when positions are enabled + return e; +} + +DictEntry MakeInline(std::string term, uint32_t df) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kInline; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.frq_bytes = {0x10, 0x20, 0x30}; + return e; +} + +void ExpectCommon(const DictEntry& a, const DictEntry& b) { + EXPECT_EQ(a.term, b.term); + EXPECT_EQ(a.kind, b.kind); + EXPECT_EQ(a.enc, b.enc); + EXPECT_EQ(a.df, b.df); +} + +// Serialize a set of entries into one block using the builder; return the byte buffer. +std::vector BuildBlock(const std::vector& entries, IndexTier tier, + bool has_positions, uint64_t frq_base, uint64_t prx_base, + uint32_t anchor_interval) { + DictBlockBuilder builder(tier, has_positions, frq_base, prx_base, anchor_interval); + for (const auto& e : entries) { + builder.add_entry(e); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +} // namespace + +// Empty block: n_entries=0, open should still succeed, find_term on any target returns not-found. +TEST(SniiDictBlock, EmptyBlock) { + std::vector bytes = + BuildBlock({}, IndexTier::kT1, /*has_positions=*/false, 1000, 0, 16); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT1, + /*has_positions=*/false, &reader) + .ok()); + EXPECT_EQ(reader.n_entries(), 0U); + EXPECT_EQ(reader.frq_base(), 1000U); + + bool found = true; + DictEntry out; + ASSERT_TRUE(reader.find_term("anything", &found, &out).ok()); + EXPECT_FALSE(found); +} + +// Single-entry round-trip. +TEST(SniiDictBlock, SingleEntryRoundTrip) { + DictEntry e = MakePodRef("solo", 7, 0); + std::vector bytes = BuildBlock({e}, IndexTier::kT1, false, 4096, 0, 16); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader).ok()); + EXPECT_EQ(reader.n_entries(), 1U); + EXPECT_EQ(reader.frq_base(), 4096U); + + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term("solo", &found, &out).ok()); + ASSERT_TRUE(found); + ExpectCommon(e, out); + EXPECT_EQ(out.frq_off_delta, e.frq_off_delta); + EXPECT_EQ(out.frq_len, e.frq_len); +} + +// Multi-entry round-trip: all terms found, fields preserved. +TEST(SniiDictBlock, MultiEntryRoundTrip) { + std::vector entries = {MakePodRef("alpha", 3, 0), MakePodRef("beta", 5, 100), + MakeInline("gamma", 2), MakePodRef("delta", 9, 300), + MakePodRef("epsilon", 11, 500)}; + // delta < gamma lexicographically, so reorder entries to be sorted. + entries = {MakePodRef("alpha", 3, 0), MakePodRef("beta", 5, 100), MakePodRef("delta", 9, 300), + MakePodRef("epsilon", 11, 500), MakeInline("gamma", 2)}; + std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 8192, 16384, 16); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader).ok()); + EXPECT_EQ(reader.n_entries(), entries.size()); + + for (const auto& e : entries) { + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term(e.term, &found, &out).ok()) << e.term; + ASSERT_TRUE(found) << e.term; + ExpectCommon(e, out); + } +} + +TEST(SniiDictBlock, DecodeAll) { + std::vector entries = {MakePodRef("alpha", 3, 0), MakePodRef("beta", 5, 100), + MakePodRef("delta", 9, 300), MakePodRef("epsilon", 11, 500), + MakeInline("gamma", 2)}; + std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 8192, 16384, 16); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader).ok()); + + // Null output is rejected, not dereferenced. + EXPECT_FALSE(reader.decode_all(nullptr).ok()); + + // decode_all returns every entry, in ascending order, matching n_entries. + std::vector all; + ASSERT_TRUE(reader.decode_all(&all).ok()); + ASSERT_EQ(all.size(), entries.size()); + for (size_t i = 0; i < entries.size(); ++i) { + EXPECT_EQ(all[i].term, entries[i].term) << i; + ExpectCommon(entries[i], all[i]); + } +} + +// Front-coding across anchors: with a small anchor_interval, terms spanning anchor boundaries are decoded correctly. +TEST(SniiDictBlock, PrefixCompressionAcrossAnchors) { + std::vector entries; + // 21 sorted terms sharing a long common prefix, anchor_interval=4 -> multiple anchors. + std::vector terms = {"interest", "interested", "interesting", "interestingly", + "interests", "internal", "internally", "international", + "internet", "internets", "interplay", "interpose", + "interpret", "interpreted", "interval", "intervene", + "interview", "interviewed", "intestine", "intimate", + "intricate"}; + uint64_t off = 0; + for (const auto& t : terms) { + entries.push_back(MakePodRef(t, 4, off)); + off += 50; + } + std::vector bytes = + BuildBlock(entries, IndexTier::kT1, false, 0, 0, /*anchor_interval=*/4); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader).ok()); + EXPECT_EQ(reader.n_entries(), terms.size()); + + for (size_t i = 0; i < terms.size(); ++i) { + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term(terms[i], &found, &out).ok()) << terms[i]; + ASSERT_TRUE(found) << terms[i]; + EXPECT_EQ(out.term, terms[i]); + EXPECT_EQ(out.frq_off_delta, static_cast(i * 50)); + } +} + +// find_term boundaries: less than first term, greater than last term, exactly an anchor term. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiDictBlock, FindTermBoundaries) { + std::vector terms; + for (int i = 0; i < 40; ++i) { + char buf[8]; + snprintf(buf, sizeof(buf), "t%02d", i); + terms.emplace_back(buf); + } + std::vector entries; + for (size_t i = 0; i < terms.size(); ++i) { + entries.push_back(MakePodRef(terms[i], 1, i * 10)); + } + std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 8); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader).ok()); + + // Less than the first term. + { + bool found = true; + DictEntry out; + ASSERT_TRUE(reader.find_term("\x01", &found, &out).ok()); + EXPECT_FALSE(found); + } + // Greater than the last term. + { + bool found = true; + DictEntry out; + ASSERT_TRUE(reader.find_term("zzzz", &found, &out).ok()); + EXPECT_FALSE(found); + } + // Hit terms that are exactly at anchor positions (anchor_interval=8 -> t08, t16, t24 ...). + for (const char* t : {"t00", "t08", "t16", "t24", "t32"}) { + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term(t, &found, &out).ok()) << t; + ASSERT_TRUE(found) << t; + EXPECT_EQ(out.term, t); + } + // Hit terms that are NOT at anchor positions. + for (const char* t : {"t03", "t11", "t27"}) { + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term(t, &found, &out).ok()) << t; + ASSERT_TRUE(found) << t; + EXPECT_EQ(out.term, t); + } + // A key within the term range that does not exist (gap). + { + bool found = true; + DictEntry out; + ASSERT_TRUE(reader.find_term("t05x", &found, &out).ok()); + EXPECT_FALSE(found); + } +} + +TEST(SniiDictBlock, VisitTermRangeCrossesAnchorsAndHonorsBounds) { + std::vector entries; + for (int i = 0; i < 24; ++i) { + char term[8]; + snprintf(term, sizeof(term), "t%02d", i); + entries.push_back(MakePodRef(term, static_cast(i + 1), i * 10)); + } + const std::vector bytes = + BuildBlock(entries, IndexTier::kT1, false, 0, 0, /*anchor_interval=*/4); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader).ok()); + + std::vector visited; + bool exhausted = false; + ASSERT_TRUE(reader.visit_term_range( + "t03", std::string_view("t18"), + [](std::string_view term) { return term.back() % 2 != 0; }, + [&](DictEntry&& entry, bool* stop) { + visited.push_back(std::move(entry.term)); + *stop = visited.size() == 4; + return Status::OK(); + }, + &exhausted) + .ok()); + EXPECT_EQ(visited, (std::vector {"t03", "t05", "t07", "t09"})); + EXPECT_FALSE(exhausted); + + visited.clear(); + ASSERT_TRUE(reader.visit_term_range( + "t17", std::string_view("t20"), {}, + [&](DictEntry&& entry, bool*) { + visited.push_back(std::move(entry.term)); + return Status::OK(); + }, + &exhausted) + .ok()); + EXPECT_EQ(visited, (std::vector {"t17", "t18", "t19"})); + EXPECT_TRUE(exhausted); +} + +// CRC corruption detection: flipping any byte must cause open() to report Corruption. +TEST(SniiDictBlock, CrcCorruptionDetected) { + std::vector entries = {MakePodRef("aaa", 1, 0), MakePodRef("bbb", 2, 50), + MakePodRef("ccc", 3, 100)}; + std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 100, 200, 16); + // Flip one byte in the middle of the entries region. + bytes[bytes.size() / 2] ^= 0xFF; + DictBlockReader reader; + Status s = DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader); + EXPECT_FALSE(s.ok()); + EXPECT_FALSE(s.ok()); +} + +// A truncated block (shorter than the CRC footer) must report Corruption rather than crash with an out-of-range access. +TEST(SniiDictBlock, TruncatedBlockDetected) { + std::vector entries = {MakePodRef("xxx", 1, 0)}; + std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); + bytes.resize(2); // Only a few bytes of the header remain. + DictBlockReader reader; + Status s = DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader); + EXPECT_FALSE(s.ok()); +} + +// Positions mode: prx_base and prx fields are preserved correctly. +TEST(SniiDictBlock, PositionsPrxRoundTrip) { + std::vector entries = {MakePodRef("phrase", 5, 0, 0), + MakePodRef("query", 6, 200, 80)}; + std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 7000, 9000, 16); + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader).ok()); + EXPECT_EQ(reader.prx_base(), 9000U); + + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term("query", &found, &out).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(out.prx_off_delta, 80U); + EXPECT_EQ(out.prx_len, 64U); + EXPECT_EQ(out.ttf_delta, 12U); +} + +// estimated_bytes is monotonically non-decreasing: grows after each add_entry, and the actual byte +// count after finish() does not exceed the estimate (the estimate is an upper bound for block-splitting decisions). +TEST(SniiDictBlock, EstimatedBytesMonotonic) { + DictBlockBuilder builder(IndexTier::kT1, false, 0, 0, 16); + size_t prev = builder.estimated_bytes(); + std::vector terms = {"aa", "bb", "cc", "dd"}; + for (const auto& t : terms) { + builder.add_entry(MakePodRef(t, 1, 0)); + size_t now = builder.estimated_bytes(); + EXPECT_GE(now, prev); + prev = now; + } + ByteSink sink; + builder.finish(&sink); + EXPECT_LE(sink.size(), builder.estimated_bytes()); +} + +TEST(SniiDictBlock, StatlessEstimateOmitsTermStatsUpperBound) { + DictBlockBuilder statful(IndexTier::kT2, /*has_positions=*/true, 0, 0, 16, + /*term_stats=*/true); + DictBlockBuilder statless(IndexTier::kT2, /*has_positions=*/true, 0, 0, 16, + /*term_stats=*/false); + for (const char* term : {"gram-a", "gram-b"}) { + const DictEntry entry = MakePodRef(term, 7, 0, 0); + statful.add_entry(entry); + statless.add_entry(entry); + } + + EXPECT_EQ(statful.estimated_bytes() - statless.estimated_bytes(), 40U); + + ByteSink statful_sink; + ByteSink statless_sink; + statful.finish(&statful_sink); + statless.finish(&statless_sink); + EXPECT_LE(statful_sink.size(), statful.estimated_bytes()); + EXPECT_LE(statless_sink.size(), statless.estimated_bytes()); + EXPECT_LT(statless_sink.size(), statful_sink.size()); +} + +TEST(SniiDictBlock, StatlessEstimateCoversInlineLocatorMetadata) { + DictEntry entry = MakeInline("gram", std::numeric_limits::max()); + entry.frq_bytes.assign(128, 0); + entry.inline_dd_disk_len = entry.frq_bytes.size(); + entry.dd_meta.uncomp_len = std::numeric_limits::max(); + entry.freq_meta.uncomp_len = std::numeric_limits::max(); + entry.prx_bytes = {0}; + + DictBlockBuilder builder(IndexTier::kT2, /*has_positions=*/true, 0, 0, 16, + /*term_stats=*/false); + builder.add_entry(std::move(entry)); + ByteSink sink; + builder.finish(&sink); + EXPECT_LE(sink.size(), builder.estimated_bytes()); +} + +// Security regression: anchor offsets must be strictly increasing with the first +// anchor at the entries start. A tampered anchor table (offsets swapped) whose crc +// is re-stamped must be rejected by open(); otherwise scan_from_anchor would +// underflow seg_end-seg_begin and read out of bounds (found via ASAN/UBSAN review). +TEST(SniiDictBlock, RejectsNonMonotonicAnchorOffsets) { + std::vector entries = {MakePodRef("aaa", 1, 0), MakePodRef("bbb", 2, 10), + MakePodRef("ccc", 3, 20)}; + // anchor_interval = 1 -> every entry is an anchor (3 anchors). + std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 100, 0, 1); + const size_t n = bytes.size(); + ASSERT_GT(n, 20U); + // Tail layout: [...][anchor_offsets: 3 * u32][n_anchors u32][crc32c u32]. + const size_t off1 = n - 4 /*crc*/ - 4 /*n_anchors*/ - 4 /*offset[2]*/ - 4 /*offset[1]*/; + const size_t off2 = off1 + 4; + for (int k = 0; k < 4; ++k) { + uint8_t tmp = bytes[off1 + k]; + bytes[off1 + k] = bytes[off2 + k]; + bytes[off2 + k] = tmp; + } + // Re-stamp crc32c over the covered region [0, n-4) so corruption is structural. + uint32_t crc = doris::snii::crc32c(Slice(bytes.data(), n - 4)); + for (int k = 0; k < 4; ++k) { + bytes[n - 4 + k] = static_cast(crc >> (8 * k)); + } + + DictBlockReader reader; + Status s = DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader); + EXPECT_FALSE(s.ok()); +} + +// =========================================================================== +// T16: DictBlockBuilder move-entry overload + dead prev_term_ removal. +// The move overload must be byte-for-byte equivalent to the const-ref (copy) +// overload, must actually transfer (not copy) the entry, and removing the dead +// prev_term_ member must leave finish() output unchanged. These live in the +// SniiDictBlockTest suite per the T16 plan. +// =========================================================================== +namespace { + +// Inline entry carrying BOTH non-empty frq_bytes and prx_bytes, so the move path +// has real heap-allocated vectors (and a heap term) to transfer -- exercised by +// the byte-equivalence and moved-from tests below. +DictEntry MakeInlineWithPrx(std::string term, uint32_t df) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kInline; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.ttf_delta = df * 3; // written only when tier>=T2 + e.max_freq = 7; // written only when tier>=T2 + e.frq_bytes = {0x01, 0x02, 0x03, 0x04, 0x05}; + e.prx_bytes = {0xAA, 0xBB, 0xCC}; // written only when positions are enabled + return e; +} + +// Build a block by MOVING each entry into the builder. The caller's vector stays +// intact because each entry is copied once and only the copy is moved in (so the +// returned bytes can be compared against the const-ref BuildBlock path). +std::vector BuildBlockMoved(const std::vector& entries, IndexTier tier, + bool has_positions, uint64_t frq_base, uint64_t prx_base, + uint32_t anchor_interval) { + DictBlockBuilder builder(tier, has_positions, frq_base, prx_base, anchor_interval); + for (const auto& e : entries) { + DictEntry tmp = e; // copy, then move the copy into the builder + builder.add_entry(std::move(tmp)); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +// A mixed, sorted entry set that crosses the default anchor_interval (16): 18 +// entries, inline-with-prx at a few indices, pod_ref elsewhere. Zero-padded term +// names keep lexicographic order == construction order. +std::vector MakeMixedCrossAnchorEntries() { + std::vector entries; + for (int i = 0; i < 18; ++i) { + char buf[16]; + snprintf(buf, sizeof(buf), "term%02d", i); + if (i == 3 || i == 8 || i == 14) { + entries.push_back(MakeInlineWithPrx(buf, static_cast(i + 1))); + } else { + entries.push_back(MakePodRef(buf, static_cast(i + 1), + static_cast(i) * 64, + static_cast(i) * 32)); + } + } + return entries; +} + +} // namespace + +// FC-1 / RED-1: the move overload produces byte-identical block output to the +// const-ref (copy) overload, and the size estimate matches -- the latter only +// holds if estimate_entry_bytes runs BEFORE the move (estimating a moved-from +// entry would undercount entries_est_ and diverge here). +TEST(SniiDictBlockTest, MoveAddProducesByteIdenticalOutput) { + const std::vector entries = MakeMixedCrossAnchorEntries(); + + DictBlockBuilder builder_a(IndexTier::kT2, /*has_positions=*/true, 8192, 16384, 16); + DictBlockBuilder builder_b(IndexTier::kT2, /*has_positions=*/true, 8192, 16384, 16); + for (const auto& e : entries) { + builder_a.add_entry(e); // const-ref copy path + DictEntry tmp = e; + builder_b.add_entry(std::move(tmp)); // move path + } + + // estimate-before-move guard: undercounting on the move path would show here. + EXPECT_EQ(builder_a.estimated_bytes(), builder_b.estimated_bytes()); + EXPECT_EQ(builder_a.n_entries(), builder_b.n_entries()); + + ByteSink sink_a; + ByteSink sink_b; + builder_a.finish(&sink_a); + builder_b.finish(&sink_b); + EXPECT_EQ(sink_a.buffer(), sink_b.buffer()); // on-disk bytes identical +} + +// FC-2 / RED-2: add_entry(DictEntry&&) actually moves -- the source entry's term +// and both byte vectors are left empty (moved-from). A heap-sized term (beyond +// SSO) makes "the buffer was stolen" observable regardless of small-string opt. +TEST(SniiDictBlockTest, MoveAddLeavesSourceMovedFrom) { + DictEntry e = MakeInlineWithPrx("a-deliberately-long-term-well-beyond-sso", 5); + ASSERT_FALSE(e.term.empty()); + ASSERT_FALSE(e.frq_bytes.empty()); + ASSERT_FALSE(e.prx_bytes.empty()); + + DictBlockBuilder builder(IndexTier::kT2, /*has_positions=*/true, 0, 0, 16); + builder.add_entry(std::move(e)); + + EXPECT_TRUE( + e.term.empty()); // NOLINT(bugprone-use-after-move): intentionally inspects moved-from state + EXPECT_TRUE(e.frq_bytes.empty()); // NOLINT(bugprone-use-after-move) + EXPECT_TRUE(e.prx_bytes.empty()); // NOLINT(bugprone-use-after-move) + EXPECT_EQ(builder.n_entries(), 1U); +} + +// FC-3: single-entry block built via the move overload round-trips through +// open()/find_term() (open succeeding also verifies the block CRC). +TEST(SniiDictBlockTest, MoveAddSingleEntryRoundTrip) { + DictEntry e = MakePodRef("solo", 7, 0); + const DictEntry expected = e; // pristine copy for comparison + + DictBlockBuilder builder(IndexTier::kT1, /*has_positions=*/false, 4096, 0, 16); + builder.add_entry(std::move(e)); + ByteSink sink; + builder.finish(&sink); + + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(sink.buffer()), IndexTier::kT1, + /*has_positions=*/false, &reader) + .ok()); + EXPECT_EQ(reader.n_entries(), 1U); + EXPECT_EQ(reader.frq_base(), 4096U); + + bool found = false; + DictEntry out; + ASSERT_TRUE(reader.find_term("solo", &found, &out).ok()); + ASSERT_TRUE(found); + ExpectCommon(expected, out); + EXPECT_EQ(out.frq_off_delta, expected.frq_off_delta); + EXPECT_EQ(out.frq_len, expected.frq_len); +} + +// FC-4: empty-block boundary still serializes and opens cleanly (the move +// overload is irrelevant when nothing is added, but the boundary must hold). +TEST(SniiDictBlockTest, MoveAddEmptyBlock) { + DictBlockBuilder builder(IndexTier::kT1, /*has_positions=*/false, 1000, 0, 16); + ByteSink sink; + builder.finish(&sink); + + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(sink.buffer()), IndexTier::kT1, + /*has_positions=*/false, &reader) + .ok()); + EXPECT_EQ(reader.n_entries(), 0U); + + bool found = true; + DictEntry out; + ASSERT_TRUE(reader.find_term("anything", &found, &out).ok()); + EXPECT_FALSE(found); +} + +// FC-5: anchor-boundary readback. anchor_interval default = 16, so 17 move-added +// entries force a second anchor (index 16); decode_all must reset the +// front-coding base at each anchor segment and return every term in order. +TEST(SniiDictBlockTest, MoveAddAnchorBoundaryDecodeAll) { + std::vector terms; + std::vector entries; + for (int i = 0; i < 17; ++i) { + char buf[16]; + snprintf(buf, sizeof(buf), "key%02d", i); + terms.emplace_back(buf); + entries.push_back( + MakePodRef(buf, static_cast(i + 1), static_cast(i) * 50)); + } + + DictBlockBuilder builder(IndexTier::kT1, /*has_positions=*/false, 0, 0, /*anchor_interval=*/16); + for (const auto& e : entries) { + DictEntry tmp = e; + builder.add_entry(std::move(tmp)); + } + ByteSink sink; + builder.finish(&sink); + + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(sink.buffer()), IndexTier::kT1, + /*has_positions=*/false, &reader) + .ok()); + EXPECT_EQ(reader.n_entries(), 17U); + + std::vector all; + ASSERT_TRUE(reader.decode_all(&all).ok()); + ASSERT_EQ(all.size(), entries.size()); + for (size_t i = 0; i < entries.size(); ++i) { + EXPECT_EQ(all[i].term, terms[i]) << i; + ExpectCommon(entries[i], all[i]); + EXPECT_EQ(all[i].frq_off_delta, static_cast(i) * 50) << i; + } +} + +// RED-3: removing the dead prev_term_ member must not change finish() output. +// A multi-anchor, long-shared-prefix dataset stresses front coding across anchor +// segments -- exactly what finish() rebuilds from its local `prev` (never from +// the removed member). The const-ref path is the pristine golden; the move path +// must reproduce it byte-for-byte, and the block must decode back to the inputs. +TEST(SniiDictBlockTest, FinishUnaffectedByDeadPrevTermRemoval) { + const std::vector terms = { + "interest", "interested", "interesting", "interestingly", "interests", + "internal", "internally", "international", "internet", "internets", + "interplay", "interpose", "interpret", "interpreted", "interval", + "intervene", "interview", "interviewed", "intestine", "intimate", + "intricate"}; + std::vector entries; + uint64_t off = 0; + for (const auto& t : terms) { + entries.push_back(MakePodRef(t, 4, off)); + off += 50; + } + + const std::vector golden = BuildBlock(entries, IndexTier::kT1, /*has_positions=*/false, + 0, 0, /*anchor_interval=*/4); + const std::vector moved = + BuildBlockMoved(entries, IndexTier::kT1, /*has_positions=*/false, 0, 0, 4); + EXPECT_EQ(golden, moved); // dead-field removal + move path leave bytes identical + + DictBlockReader reader; + ASSERT_TRUE(DictBlockReader::open(Slice(golden), IndexTier::kT1, false, &reader).ok()); + std::vector all; + ASSERT_TRUE(reader.decode_all(&all).ok()); + ASSERT_EQ(all.size(), terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + EXPECT_EQ(all[i].term, terms[i]) << i; // front-coding base rebuilt without prev_term_ + } +} diff --git a/be/test/storage/index/snii/format/dict_entry_test.cpp b/be/test/storage/index/snii/format/dict_entry_test.cpp new file mode 100644 index 00000000000000..d25b8ed29d42fb --- /dev/null +++ b/be/test/storage/index/snii/format/dict_entry_test.cpp @@ -0,0 +1,377 @@ +// 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 "storage/index/snii/format/dict_entry.h" + +#include + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; // NOLINT +using namespace doris::snii::format; // NOLINT +using doris::Status; + +namespace { + +// Construct a pod_ref slim dict entry (tier determines which optional fields +// are present). +DictEntry MakePodRefSlim(std::string term, uint32_t df) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kPodRef; + e.enc = DictEntryEnc::kSlim; + e.has_sb = false; + e.df = df; + e.ttf_delta = df * 3; // written only when tier>=T2 + e.max_freq = 7; // written only when tier>=T2 + e.frq_off_delta = 4096; + e.frq_len = 333; + e.frq_docs_len = 200; // slim pod_ref: docs-only prefix (<= frq_len) + e.prx_off_delta = 8192; // written only when positions are stored + e.prx_len = 512; // written only when positions are stored + return e; +} + +DictEntry MakePodRefWindowed(std::string term, uint32_t df) { + DictEntry e = MakePodRefSlim(std::move(term), df); + e.enc = DictEntryEnc::kWindowed; + e.has_sb = true; + e.prelude_len = 64; // written only for windowed encoding + e.frq_docs_len = 240; // windowed docs-only prefix: [prelude][dd-block] + return e; +} + +DictEntry MakeInline(std::string term, uint32_t df) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kInline; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.ttf_delta = df * 2; + e.max_freq = 5; + e.frq_bytes = {0x01, 0x02, 0x03, 0x04}; + e.prx_bytes = {0xAA, 0xBB}; // written only when positions are stored + return e; +} + +// round-trip: encode then decode, assert that fields retained by the given tier +// match. +DictEntry RoundTrip(const DictEntry& in, std::string_view prev, IndexTier tier) { + ByteSink sink; + EXPECT_TRUE(encode_dict_entry(in, prev, tier, &sink).ok()); + ByteSource src(sink.view()); + DictEntry out; + Status s = decode_dict_entry(&src, prev, tier, &out); + EXPECT_TRUE(s.ok()) << s.to_string(); + EXPECT_EQ(src.remaining(), 0U) << "decode did not consume the entire entry"; + return out; +} + +void ExpectCommon(const DictEntry& a, const DictEntry& b) { + EXPECT_EQ(a.term, b.term); + EXPECT_EQ(a.kind, b.kind); + EXPECT_EQ(a.enc, b.enc); + EXPECT_EQ(a.has_sb, b.has_sb); + EXPECT_EQ(a.df, b.df); +} + +} // namespace + +TEST(SniiDictEntry, RejectsAliasedOutputAndScratch) { + DictEntry entry = MakePodRefSlim("alias", 1); + ByteSink sink; + Status status = encode_dict_entry(entry, "", IndexTier::kT2, &sink, true, &sink); + EXPECT_TRUE(status.is()); +} + +TEST(SniiDictEntry, PodRefSlimTier1RoundTrip) { + DictEntry in = MakePodRefSlim("apple", 42); + DictEntry out = RoundTrip(in, "", IndexTier::kT1); + ExpectCommon(in, out); + EXPECT_EQ(out.frq_off_delta, in.frq_off_delta); + EXPECT_EQ(out.frq_len, in.frq_len); + EXPECT_EQ(out.frq_docs_len, in.frq_docs_len); // slim docs-only prefix preserved + // tier1: ttf/max_freq/prx are not written; decode restores them to default 0. + EXPECT_EQ(out.ttf_delta, 0U); + EXPECT_EQ(out.max_freq, 0U); + EXPECT_EQ(out.prx_len, 0U); +} + +TEST(SniiDictEntry, PodRefSlimTier2RoundTrip) { + DictEntry in = MakePodRefSlim("banana", 100); + DictEntry out = RoundTrip(in, "", IndexTier::kT2); + ExpectCommon(in, out); + EXPECT_EQ(out.frq_docs_len, in.frq_docs_len); + EXPECT_EQ(out.ttf_delta, in.ttf_delta); + EXPECT_EQ(out.max_freq, in.max_freq); + EXPECT_EQ(out.prx_off_delta, in.prx_off_delta); + EXPECT_EQ(out.prx_len, in.prx_len); +} + +// A slim pod_ref whose frq_docs_len exceeds frq_len is rejected on decode. +TEST(SniiDictEntry, PodRefSlimFrqDocsLenExceedsFrqLenRejected) { + DictEntry in = MakePodRefSlim("kiwi", 50); + in.frq_docs_len = in.frq_len + 1; // impossible prefix + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(in, "", IndexTier::kT2, &sink).ok()); + ByteSource src(sink.view()); + DictEntry out; + Status s = decode_dict_entry(&src, "", IndexTier::kT2, &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiDictEntry, PodRefSlimTier3RoundTrip) { + DictEntry in = MakePodRefSlim("cherry", 500); + DictEntry out = RoundTrip(in, "", IndexTier::kT3); + ExpectCommon(in, out); + EXPECT_EQ(out.ttf_delta, in.ttf_delta); + EXPECT_EQ(out.max_freq, in.max_freq); + EXPECT_EQ(out.prx_off_delta, in.prx_off_delta); + EXPECT_EQ(out.prx_len, in.prx_len); +} + +TEST(SniiDictEntry, PodRefWindowedTier2RoundTrip) { + DictEntry in = MakePodRefWindowed("durian", 2000); + DictEntry out = RoundTrip(in, "", IndexTier::kT2); + ExpectCommon(in, out); + EXPECT_EQ(out.enc, DictEntryEnc::kWindowed); + EXPECT_EQ(out.has_sb, true); + EXPECT_EQ(out.prelude_len, in.prelude_len); + EXPECT_EQ(out.frq_docs_len, in.frq_docs_len); + EXPECT_EQ(out.prx_len, in.prx_len); +} + +TEST(SniiDictEntry, PodRefWindowedTier1RoundTrip) { + DictEntry in = MakePodRefWindowed("elderberry", 1500); + DictEntry out = RoundTrip(in, "", IndexTier::kT1); + ExpectCommon(in, out); + EXPECT_EQ(out.prelude_len, in.prelude_len); + EXPECT_EQ(out.frq_docs_len, in.frq_docs_len); + EXPECT_EQ(out.prx_len, 0U); // tier1 has no prx +} + +TEST(SniiDictEntry, PodRefWindowedDocsPrefixBoundsRejected) { + DictEntry in = MakePodRefWindowed("jackfruit", 1500); + in.frq_docs_len = in.prelude_len - 1; + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(in, "", IndexTier::kT2, &sink).ok()); + ByteSource src(sink.view()); + DictEntry out; + Status s = decode_dict_entry(&src, "", IndexTier::kT2, &out); + EXPECT_TRUE(s.is()); + + in.frq_docs_len = in.frq_len + 1; + ByteSink sink2; + ASSERT_TRUE(encode_dict_entry(in, "", IndexTier::kT2, &sink2).ok()); + ByteSource src2(sink2.view()); + s = decode_dict_entry(&src2, "", IndexTier::kT2, &out); + EXPECT_TRUE(s.is()); +} + +TEST(SniiDictEntry, InlineTier1RoundTrip) { + DictEntry in = MakeInline("fig", 3); + DictEntry out = RoundTrip(in, "", IndexTier::kT1); + ExpectCommon(in, out); + EXPECT_EQ(out.frq_bytes, in.frq_bytes); + EXPECT_TRUE(out.prx_bytes.empty()); // tier1 has no prx +} + +TEST(SniiDictEntry, InlineTier2RoundTrip) { + DictEntry in = MakeInline("grape", 8); + DictEntry out = RoundTrip(in, "", IndexTier::kT2); + ExpectCommon(in, out); + EXPECT_EQ(out.frq_bytes, in.frq_bytes); + EXPECT_EQ(out.prx_bytes, in.prx_bytes); + EXPECT_EQ(out.ttf_delta, in.ttf_delta); + EXPECT_EQ(out.max_freq, in.max_freq); +} + +TEST(SniiDictEntry, InlineTier3RoundTrip) { + DictEntry in = MakeInline("honeydew", 12); + DictEntry out = RoundTrip(in, "", IndexTier::kT3); + ExpectCommon(in, out); + EXPECT_EQ(out.frq_bytes, in.frq_bytes); + EXPECT_EQ(out.prx_bytes, in.prx_bytes); +} + +// Format v2: an INLINE entry omits the redundant per-region crc32c bytes +// (dd_crc + freq_crc) -- its region bytes are covered by the dict block crc32c. +// Decoded inline metas carry verify_crc=false so region decode skips the check. +// A slim POD-ref entry KEEPS its crc (regions live in the separately-fetched +// .frq POD), so it is exactly 8 bytes larger for the same df. +TEST(SniiDictEntry, InlineOmitsRedundantRegionCrc) { + DictEntry inl = MakeInline("melon", 8); + inl.dd_meta.crc = 0xDEADBEEF; // would be written by v1; must be omitted now + inl.freq_meta.crc = 0xCAFEF00D; // ditto + ByteSink inl_sink; + ASSERT_TRUE(encode_dict_entry(inl, "", IndexTier::kT2, &inl_sink).ok()); + + // Same payload but encoded as a slim POD-ref keeps both crcs (dd + freq = + // 8B). + DictEntry ref = inl; + ref.kind = DictEntryKind::kPodRef; + ref.frq_len = inl.frq_bytes.size(); + ref.frq_docs_len = inl.inline_dd_disk_len; + ByteSink ref_sink; + ASSERT_TRUE(encode_dict_entry(ref, "", IndexTier::kT2, &ref_sink).ok()); + // The POD-ref locator differs structurally; assert the inline encoding does + // not carry the 8 crc bytes by decoding and checking verify_crc is OFF. + ByteSource src(inl_sink.view()); + DictEntry out; + ASSERT_TRUE(decode_dict_entry(&src, "", IndexTier::kT2, &out).ok()); + EXPECT_FALSE(out.dd_meta.verify_crc); + EXPECT_FALSE(out.freq_meta.verify_crc); + EXPECT_EQ(out.dd_meta.crc, 0U); // crc not stored -> decoded as 0 + EXPECT_EQ(out.freq_meta.crc, 0U); +} + +// A POD-ref slim entry decodes with verify_crc=true (its regions carry a crc). +TEST(SniiDictEntry, PodRefSlimKeepsRegionCrc) { + DictEntry in = MakePodRefSlim("nectarine", 50); + in.dd_meta.crc = 0x12345678; + in.freq_meta.crc = 0x9ABCDEF0; + DictEntry out = RoundTrip(in, "", IndexTier::kT2); + EXPECT_TRUE(out.dd_meta.verify_crc); + EXPECT_TRUE(out.freq_meta.verify_crc); + EXPECT_EQ(out.dd_meta.crc, 0x12345678U); + EXPECT_EQ(out.freq_meta.crc, 0x9ABCDEF0U); +} + +// Front coding: consecutive encode/decode of sorted terms; suffix stores only +// the differing part. +TEST(SniiDictEntry, PrefixCompressionSharedPrefix) { + // "interest" and "interesting" share the prefix "interest". + DictEntry a = MakePodRefSlim("interest", 10); + DictEntry b = MakePodRefSlim("interesting", 11); + + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(a, "", IndexTier::kT2, &sink).ok()); + size_t a_end = sink.size(); + ASSERT_TRUE(encode_dict_entry(b, a.term, IndexTier::kT2, &sink).ok()); + + ByteSource src(sink.view()); + DictEntry oa; + ASSERT_TRUE(decode_dict_entry(&src, "", IndexTier::kT2, &oa).ok()); + EXPECT_EQ(src.position(), a_end); + EXPECT_EQ(oa.term, "interest"); + + DictEntry ob; + ASSERT_TRUE(decode_dict_entry(&src, oa.term, IndexTier::kT2, &ob).ok()); + EXPECT_EQ(ob.term, "interesting"); + EXPECT_TRUE(src.eof()); +} + +// Encode/decode three sorted terms consecutively to verify the prefix chain. +TEST(SniiDictEntry, PrefixCompressionChain) { + std::vector terms = {"app", "apple", "application"}; + ByteSink sink; + std::string prev; + for (const auto& t : terms) { + DictEntry e = MakePodRefSlim(t, 5); + ASSERT_TRUE(encode_dict_entry(e, prev, IndexTier::kT1, &sink).ok()); + prev = t; + } + ByteSource src(sink.view()); + prev.clear(); + for (const auto& t : terms) { + DictEntry out; + ASSERT_TRUE(decode_dict_entry(&src, prev, IndexTier::kT1, &out).ok()); + EXPECT_EQ(out.term, t); + prev = out.term; + } + EXPECT_TRUE(src.eof()); +} + +// entry_len must allow a reader to skip the entire entry without parsing its +// internal fields. +TEST(SniiDictEntry, EntryLenAllowsSkip) { + DictEntry a = MakePodRefSlim("first", 9); + DictEntry b = MakeInline("second", 4); + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(a, "", IndexTier::kT2, &sink).ok()); + ASSERT_TRUE(encode_dict_entry(b, a.term, IndexTier::kT2, &sink).ok()); + + ByteSource src(sink.view()); + // Skip the first entry using only entry_len. + ASSERT_TRUE(skip_dict_entry(&src).ok()); + DictEntry out; + ASSERT_TRUE(decode_dict_entry(&src, a.term, IndexTier::kT2, &out).ok()); + EXPECT_EQ(out.term, "second"); + EXPECT_TRUE(src.eof()); +} + +// Edge case: empty suffix (term is identical to prev). +TEST(SniiDictEntry, EmptySuffixEqualsPrev) { + DictEntry in = MakePodRefSlim("same", 1); + DictEntry out = RoundTrip(in, "same", IndexTier::kT1); + EXPECT_EQ(out.term, "same"); +} + +// Edge case: df = 0. +TEST(SniiDictEntry, ZeroDf) { + DictEntry in = MakePodRefSlim("zero", 0); + DictEntry out = RoundTrip(in, "", IndexTier::kT2); + EXPECT_EQ(out.df, 0U); +} + +// Edge case: empty term (first entry, prev is empty, suffix is also empty). +TEST(SniiDictEntry, EmptyTerm) { + DictEntry in = MakePodRefSlim("", 1); + DictEntry out = RoundTrip(in, "", IndexTier::kT1); + EXPECT_EQ(out.term, ""); +} + +// Structural integrity: no CRC at the entry level (CRC is at the DICT block +// level), but entry_len and the actual body length must match; tampering with +// entry_len to make it smaller than the real body -> decode must report +// Corruption. +TEST(SniiDictEntry, EntryLenMismatchDetected) { + DictEntry in = MakePodRefSlim("payload", 99); + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(in, "", IndexTier::kT2, &sink).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes[0], 1U); // single-byte entry_len varint + bytes[0] -= 1; // claim the body is 1 byte shorter than it actually is + + Slice corrupt(bytes); + ByteSource src(corrupt); + DictEntry out; + Status s = decode_dict_entry(&src, "", IndexTier::kT2, &out); + EXPECT_FALSE(s.ok()); + EXPECT_TRUE(s.is()); +} + +// prefix_len exceeding the length of prev must be rejected (guard against +// malformed input). +TEST(SniiDictEntry, RejectPrefixLongerThanPrev) { + DictEntry in = MakePodRefSlim("abcdef", 1); + ByteSink sink; + ASSERT_TRUE(encode_dict_entry(in, "abc", IndexTier::kT1, &sink).ok()); + // Decode with a shorter prev: prefix_len(=3) > prev.size()(=2) -> Corruption. + ByteSource src(sink.view()); + DictEntry out; + Status s = decode_dict_entry(&src, "ab", IndexTier::kT1, &out); + EXPECT_FALSE(s.ok()); +} diff --git a/be/test/storage/index/snii/format/format_constants_test.cpp b/be/test/storage/index/snii/format/format_constants_test.cpp new file mode 100644 index 00000000000000..7de6fb1b5397ee --- /dev/null +++ b/be/test/storage/index/snii/format/format_constants_test.cpp @@ -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. + +#include "storage/index/snii/format/format_constants.h" + +#include + +#include "common/status.h" + +using namespace doris::snii::format; + +// Lock down on-disk contract values to prevent accidental changes from breaking +// readability of already-written files. +TEST(SniiFormatConstants, MagicAndVersionStable) { + EXPECT_EQ(kContainerMagic, 0x49494E53U); + EXPECT_EQ(kTailMagic, 0x4C494154U); + EXPECT_EQ(kFormatVersion, 1); + EXPECT_EQ(kMinReaderVersion, 1); +} + +TEST(SniiFormatConstants, ConfigToTierMapping) { + EXPECT_EQ(tier_of(IndexConfig::kDocsOnly), IndexTier::kT1); + EXPECT_EQ(tier_of(IndexConfig::kDocsPositions), IndexTier::kT2); + EXPECT_EQ(tier_of(IndexConfig::kDocsPositionsScoring), IndexTier::kT3); +} + +TEST(SniiFormatConstants, CapabilityPredicates) { + EXPECT_FALSE(has_positions(IndexConfig::kDocsOnly)); + EXPECT_TRUE(has_positions(IndexConfig::kDocsPositions)); + EXPECT_TRUE(has_scoring(IndexConfig::kDocsPositionsScoring)); + EXPECT_FALSE(has_scoring(IndexConfig::kDocsPositions)); +} + +TEST(SniiFormatConstants, DictFlagBitsDistinct) { + EXPECT_EQ(dict_flags::kKind, 0x01); + EXPECT_EQ(dict_flags::kEnc, 0x02); + EXPECT_EQ(dict_flags::kHasSb, 0x04); +} diff --git a/be/test/storage/index/snii/format/frq_pod_test.cpp b/be/test/storage/index/snii/format/frq_pod_test.cpp new file mode 100644 index 00000000000000..65fd842685b8d2 --- /dev/null +++ b/be/test/storage/index/snii/format/frq_pod_test.cpp @@ -0,0 +1,621 @@ +// 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 "storage/index/snii/format/frq_pod.h" + +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/prx_decode_stats.h" + +using doris::snii::ByteSink; +using doris::snii::Slice; +using doris::Status; +using doris::snii::format::build_dd_region; +using doris::snii::format::build_freq_region; +using doris::snii::format::decode_dd_region; +using doris::snii::format::decode_freq_region; +using doris::snii::format::FrqRegionMeta; +using doris::snii::format::PrxCsrAllocationGate; + +namespace { + +using U32Vec = std::vector; + +class RejectingFrqAllocationGate final : public PrxCsrAllocationGate { +public: + Status reserve_csr(std::vector*, size_t, std::vector*, size_t) override { + return Status::OK(); + } + + Status reserve_decompression(size_t bytes, std::vector**) override { + ++calls; + requested_bytes = bytes; + return Status::Error( + "test rejects FRQ decompression allocation"); + } + + size_t calls = 0; + size_t requested_bytes = 0; +}; + +std::pair, FrqRegionMeta> MakeRawDdFromDeltas( + const std::vector& deltas) { + ByteSink sink; + sink.put_varint32(static_cast(deltas.size())); + for (size_t offset = 0; offset < deltas.size(); offset += doris::snii::format::kFrqBaseUnit) { + const size_t count = std::min(deltas.size() - offset, + static_cast(doris::snii::format::kFrqBaseUnit)); + doris::snii::pfor_encode(deltas.data() + offset, count, &sink); + } + FrqRegionMeta meta {.zstd = false, + .uncomp_len = sink.size(), + .disk_len = sink.size(), + .crc = doris::snii::crc32c(sink.view()), + .verify_crc = true}; + return {sink.take(), meta}; +} + +// Round-trips a window's separately-encoded dd + freq regions: build each +// region, then decode the dd region (docs-only) and the freq region, comparing +// originals. +Status RoundTrip(const U32Vec& docs, const U32Vec& freqs, uint64_t win_base, bool has_freq, + int level, U32Vec* out_docs, U32Vec* out_freqs) { + ByteSink dd_sink; + FrqRegionMeta dd_meta; + RETURN_IF_ERROR(build_dd_region(docs, win_base, level, &dd_sink, &dd_meta)); + RETURN_IF_ERROR(decode_dd_region(Slice(dd_sink.buffer()), dd_meta, win_base, out_docs)); + if (!has_freq) { + out_freqs->clear(); + return Status::OK(); + } + ByteSink freq_sink; + FrqRegionMeta freq_meta; + RETURN_IF_ERROR(build_freq_region(freqs, level, &freq_sink, &freq_meta)); + return decode_freq_region(Slice(freq_sink.buffer()), freq_meta, out_docs->size(), out_freqs); +} + +} // namespace + +// Basic round-trip: first window win_base=0, both dd and freq are restored. +TEST(SniiFrqPod, BasicDocFreqRoundTrip) { + U32Vec docs = {0, 3, 5, 10, 11, 50, 200}; + U32Vec freqs = {1, 2, 1, 7, 3, 1, 9}; + U32Vec out_docs, out_freqs; + ASSERT_TRUE(RoundTrip(docs, freqs, /*win_base=*/0, /*has_freq=*/true, -1, &out_docs, &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// Non-first window: win_base != 0, dd[0]=first-win_base, cross-window delta +// rebuild. +TEST(SniiFrqPod, NonFirstWindowDeltaRebuild) { + uint64_t win_base = 1000; + U32Vec docs = {1001, 1005, 1006, 2000, 2001}; + U32Vec freqs = {3, 1, 1, 2, 8}; + U32Vec out_docs, out_freqs; + ASSERT_TRUE( + RoundTrip(docs, freqs, win_base, /*has_freq=*/true, -1, &out_docs, &out_freqs).ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// dd region decodes WITHOUT the freq region present (docs-only path). +TEST(SniiFrqPod, DdRegionDecodesWithoutFreq) { + uint64_t win_base = 500; + U32Vec docs = {600, 700, 701, 900}; + ByteSink dd_sink; + FrqRegionMeta dd_meta; + ASSERT_TRUE(build_dd_region(docs, win_base, -1, &dd_sink, &dd_meta).ok()); + U32Vec out_docs; + ASSERT_TRUE(decode_dd_region(Slice(dd_sink.buffer()), dd_meta, win_base, &out_docs).ok()); + EXPECT_EQ(out_docs, docs); +} + +// Single-doc window. +TEST(SniiFrqPod, SingleDocWindow) { + U32Vec docs = {42}; + U32Vec freqs = {7}; + U32Vec out_docs, out_freqs; + ASSERT_TRUE(RoundTrip(docs, freqs, /*win_base=*/0, /*has_freq=*/true, -1, &out_docs, &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// Large window (2048 docs): auto mode triggers zstd on each region and is +// lossless. +TEST(SniiFrqPod, LargeWindowAutoZstdRoundTrip) { + U32Vec docs, freqs; + docs.reserve(2048); + freqs.reserve(2048); + uint32_t cur = 0; + for (int i = 0; i < 2048; ++i) { + cur += 1 + (i % 4); + docs.push_back(cur); + freqs.push_back(1 + (i % 3)); + } + ByteSink dd_sink; + FrqRegionMeta dd_meta; + ASSERT_TRUE(build_dd_region(docs, 0, -1, &dd_sink, &dd_meta).ok()); + EXPECT_TRUE(dd_meta.zstd); // dd region large enough for zstd + + U32Vec out_docs, out_freqs; + ASSERT_TRUE(RoundTrip(docs, freqs, /*win_base=*/0, /*has_freq=*/true, -1, &out_docs, &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// Small window level<0 auto -> raw (dd region not compressed). +TEST(SniiFrqPod, SmallWindowUsesRaw) { + U32Vec docs = {0, 1, 2}; + ByteSink dd_sink; + FrqRegionMeta dd_meta; + ASSERT_TRUE(build_dd_region(docs, 0, -1, &dd_sink, &dd_meta).ok()); + EXPECT_FALSE(dd_meta.zstd); + EXPECT_EQ(dd_meta.disk_len, dd_meta.uncomp_len); // raw => disk == uncomp +} + +// Explicit level=0 forces raw on both regions; large window still lossless. +TEST(SniiFrqPod, ExplicitRawLargeWindowRoundTrip) { + U32Vec docs, freqs; + uint32_t cur = 0; + for (int i = 0; i < 600; ++i) { + cur += 1 + (i % 7); + docs.push_back(cur); + freqs.push_back(1 + (i % 5)); + } + ByteSink dd_sink, freq_sink; + FrqRegionMeta dd_meta, freq_meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/0, &dd_sink, &dd_meta).ok()); + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &freq_sink, &freq_meta).ok()); + EXPECT_FALSE(dd_meta.zstd); + EXPECT_FALSE(freq_meta.zstd); + + U32Vec out_docs, out_freqs; + ASSERT_TRUE(decode_dd_region(Slice(dd_sink.buffer()), dd_meta, 0, &out_docs).ok()); + ASSERT_TRUE( + decode_freq_region(Slice(freq_sink.buffer()), freq_meta, out_docs.size(), &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// Explicit level>0 forces zstd, still lossless. +TEST(SniiFrqPod, ExplicitZstdRoundTrip) { + U32Vec docs, freqs; + for (uint32_t i = 0; i < 300; ++i) { + docs.push_back(i * 2); + freqs.push_back(2); + } + U32Vec out_docs, out_freqs; + ASSERT_TRUE(RoundTrip(docs, freqs, /*win_base=*/0, /*has_freq=*/true, + /*level=*/5, &out_docs, &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// Non-ascending docids must be rejected (InvalidArgument). +TEST(SniiFrqPod, NonAscendingDocsRejected) { + U32Vec docs = {0, 5, 3}; + ByteSink sink; + FrqRegionMeta meta; + Status s = build_dd_region(docs, 0, -1, &sink, &meta); + EXPECT_TRUE(s.is()); +} + +TEST(SniiFrqPod, DecodeRejectsZeroInteriorDeltaAndDocidOverflow) { + auto [duplicate_bytes, duplicate_meta] = MakeRawDdFromDeltas({0, 0}); + U32Vec docids; + const Status duplicate = + decode_dd_region(Slice(duplicate_bytes), duplicate_meta, /*win_base=*/0, &docids); + EXPECT_TRUE(duplicate.is()) << duplicate; + + auto [overflow_bytes, overflow_meta] = + MakeRawDdFromDeltas({std::numeric_limits::max()}); + const Status overflow = + decode_dd_region(Slice(overflow_bytes), overflow_meta, /*win_base=*/1, &docids); + EXPECT_TRUE(overflow.is()) << overflow; +} + +// first_docid < win_base (dd[0] would underflow) -> InvalidArgument. +TEST(SniiFrqPod, FirstDocBelowWinBaseRejected) { + U32Vec docs = {100, 200}; + ByteSink sink; + FrqRegionMeta meta; + Status s = build_dd_region(docs, 500, -1, &sink, &meta); + EXPECT_TRUE(s.is()); +} + +// CRC corruption inside the dd region is caught by crc_dd on decode. +TEST(SniiFrqPod, DdRegionCrcCorruptionDetected) { + U32Vec docs = {0, 1, 2, 3, 4, 5}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, -1, &sink, &meta).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 1U); + bytes.back() ^= 0xFF; + U32Vec out_docs; + Status s = decode_dd_region(Slice(bytes), meta, 0, &out_docs); + EXPECT_TRUE(s.is()); +} + +// CRC corruption inside the freq region is caught by crc_freq on decode. +TEST(SniiFrqPod, FreqRegionCrcCorruptionDetected) { + U32Vec freqs = {1, 2, 3, 4, 5, 6}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_freq_region(freqs, -1, &sink, &meta).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 1U); + bytes.back() ^= 0xFF; + U32Vec out_freqs; + Status s = decode_freq_region(Slice(bytes), meta, freqs.size(), &out_freqs); + EXPECT_TRUE(s.is()); +} + +// A region slice whose length disagrees with meta.disk_len is rejected +// (anti-DoS). +TEST(SniiFrqPod, RegionSliceLengthMismatchRejected) { + U32Vec docs = {0, 1, 2, 3}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/0, &sink, &meta).ok()); + std::vector bytes = sink.buffer(); + bytes.pop_back(); // slice now shorter than disk_len + U32Vec out_docs; + Status s = decode_dd_region(Slice(bytes), meta, 0, &out_docs); + EXPECT_TRUE(s.is()); +} + +// Empty window (0 docs) round-trip: dd region has n=0, freq region is empty. +TEST(SniiFrqPod, EmptyWindowRoundTrip) { + U32Vec docs, freqs; + U32Vec out_docs, out_freqs; + ASSERT_TRUE(RoundTrip(docs, freqs, /*win_base=*/0, /*has_freq=*/true, -1, &out_docs, &out_freqs) + .ok()); + EXPECT_TRUE(out_docs.empty()); + EXPECT_TRUE(out_freqs.empty()); +} + +// dd region on-disk length is strictly smaller than dd+freq combined (freq +// carries real bytes), proving the docs-only prefix saving at the region level. +TEST(SniiFrqPod, DdRegionSmallerThanCombined) { + U32Vec docs = {0, 3, 5, 10, 11, 50, 200}; + U32Vec freqs = {1, 2, 1, 7, 3, 1, 9}; + ByteSink dd_sink, freq_sink; + FrqRegionMeta dd_meta, freq_meta; + ASSERT_TRUE(build_dd_region(docs, 0, -1, &dd_sink, &dd_meta).ok()); + ASSERT_TRUE(build_freq_region(freqs, -1, &freq_sink, &freq_meta).ok()); + EXPECT_GT(dd_meta.disk_len, 0U); + EXPECT_GT(freq_meta.disk_len, 0U); + EXPECT_LT(dd_meta.disk_len, dd_meta.disk_len + freq_meta.disk_len); +} + +// Null argument guard. +TEST(SniiFrqPod, NullArgsRejected) { + U32Vec docs = {0, 1}; + FrqRegionMeta meta; + EXPECT_TRUE( + build_dd_region(docs, 0, -1, nullptr, &meta).is()); +} + +// N=0 empty-term prelude round-trips: the dd region encodes VInt n=0 (no PFOR +// runs) and the freq region is zero-length; both decode back to empty vectors. +TEST(SniiFrqPod, EmptyTermZeroDocsRoundTrip) { + U32Vec docs; // N = 0 (empty term) + U32Vec freqs; // no freqs + ByteSink dd_sink, freq_sink; + FrqRegionMeta dd_meta, freq_meta; + ASSERT_TRUE(build_dd_region(docs, /*win_base=*/0, /*level=*/0, &dd_sink, &dd_meta).ok()); + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &freq_sink, &freq_meta).ok()); + EXPECT_EQ(freq_meta.uncomp_len, 0U); // empty freq region carries no bytes + + U32Vec out_docs, out_freqs; + ASSERT_TRUE(decode_dd_region(Slice(dd_sink.buffer()), dd_meta, 0, &out_docs).ok()); + EXPECT_TRUE(out_docs.empty()); + ASSERT_TRUE( + decode_freq_region(Slice(freq_sink.buffer()), freq_meta, out_docs.size(), &out_freqs) + .ok()); + EXPECT_TRUE(out_freqs.empty()); +} + +// A truncated freq region (slice shorter than meta.disk_len) is rejected: the +// length-mismatch guard in open_region fires before any payload decode. +TEST(SniiFrqPod, TruncatedFreqRegionRejected) { + U32Vec freqs = {1, 2, 3, 4, 5, 6, 7, 8}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &sink, &meta).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 1U); + bytes.pop_back(); // slice now shorter than meta.disk_len + U32Vec out_freqs; + Status s = decode_freq_region(Slice(bytes), meta, freqs.size(), &out_freqs); + EXPECT_TRUE(s.is()); +} + +// A dd region with extra trailing bytes (crc + length kept VALID for the longer +// slice) must be rejected by the post-decode eof() check, not the crc/length +// guards. We craft a raw region, append a garbage byte, then fix disk_len, +// uncomp_len, and crc so every earlier guard passes and only the inner +// "trailing bytes" assertion can fire. +TEST(SniiFrqPod, DdRegionTrailingBytesRejected) { + U32Vec docs = {0, 1, 2, 3, 4, 5}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/0, &sink, &meta).ok()); // raw + ASSERT_FALSE(meta.zstd); + std::vector bytes = sink.buffer(); + bytes.push_back(0x00); // garbage trailing byte appended to the plaintext + // Keep raw invariant uncomp_len == disk_len and a matching crc so open_region + // accepts the slice; only the payload-level eof() check can reject it. + meta.disk_len = bytes.size(); + meta.uncomp_len = bytes.size(); + meta.crc = doris::snii::crc32c(Slice(bytes)); + U32Vec out_docs; + Status s = decode_dd_region(Slice(bytes), meta, 0, &out_docs); + EXPECT_TRUE(s.is()); +} + +// Compaction pre-reserves the destination vector from dictionary/prelude +// metadata. A corrupt encoded count must be rejected before the decoder can +// resize that vector beyond the charged capacity. +TEST(SniiFrqPod, ExpectedDdCountMismatchDoesNotMutateDestination) { + U32Vec docs = {3, 7}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/0, &sink, &meta).ok()); + + U32Vec out_docs = {99}; + const size_t original_capacity = out_docs.capacity(); + Status s = decode_dd_region(Slice(sink.buffer()), meta, 0, /*expected_doc_count=*/1, &out_docs); + + EXPECT_TRUE(s.is()); + EXPECT_EQ(out_docs, U32Vec({99})); + EXPECT_EQ(out_docs.capacity(), original_capacity); +} + +TEST(SniiFrqPod, AllocationGateRejectsZstdBeforeDdMutation) { + U32Vec docs(2048); + for (uint32_t i = 0; i < docs.size(); ++i) { + docs[i] = i * 3; + } + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/3, &sink, &meta).ok()); + ASSERT_TRUE(meta.zstd); + + RejectingFrqAllocationGate gate; + U32Vec out_docs = {99}; + const size_t original_capacity = out_docs.capacity(); + const Status status = decode_dd_region(Slice(sink.buffer()), meta, 0, + /*expected_doc_count=*/docs.size(), &gate, &out_docs); + + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(gate.calls, 1); + EXPECT_EQ(gate.requested_bytes, meta.uncomp_len); + EXPECT_EQ(out_docs, U32Vec({99})); + EXPECT_EQ(out_docs.capacity(), original_capacity); +} + +// A freq region with extra trailing bytes (crc + length kept VALID) must be +// rejected by the post-decode eof() check. Same crafting strategy as above. +TEST(SniiFrqPod, FreqRegionTrailingBytesRejected) { + U32Vec freqs = {3, 1, 4, 1, 5, 9}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &sink, &meta).ok()); // raw + ASSERT_FALSE(meta.zstd); + std::vector bytes = sink.buffer(); + bytes.push_back(0x00); // garbage trailing byte appended to the plaintext + meta.disk_len = bytes.size(); + meta.uncomp_len = bytes.size(); + meta.crc = doris::snii::crc32c(Slice(bytes)); + U32Vec out_freqs; + Status s = decode_freq_region(Slice(bytes), meta, freqs.size(), &out_freqs); + EXPECT_TRUE(s.is()); +} + +// The uncomp_len sanity cap (kMaxRegionUncompBytes = 256 MiB) guards against a +// corrupted length that would inflate to a giant allocation. We keep the +// on-disk slice + crc valid (so the length/crc guards pass) and only override +// uncomp_len to an absurd value, proving the cap branch is what rejects it. +TEST(SniiFrqPod, UncompLenCapRejected) { + U32Vec docs = {0, 1, 2, 3, 4, 5}; + ByteSink sink; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, 0, /*level=*/0, &sink, &meta).ok()); + std::vector bytes = sink.buffer(); + // kMaxRegionUncompBytes is an internal (anonymous-namespace) constant; use + // its literal value (256 MiB) + 1 to step just past the cap. + meta.uncomp_len = static_cast(256U * 1024 * 1024) + 1; + U32Vec out_docs; + Status s = decode_dd_region(Slice(bytes), meta, 0, &out_docs); + EXPECT_TRUE(s.is()); +} + +// =========================================================================== +// T22 -- emit_region single-copy (raw branch writes the plaintext view straight +// to `out` with no temp `disk` vector). The raw region's on-disk bytes must equal +// the plaintext EXACTLY (regions carry no header / no trailing crc), and the meta +// (disk_len / crc / zstd / uncomp_len) must be unchanged. The zstd branch keeps +// its own buffer and stays byte-identical. All bytes MUST match the pre-refactor +// output. +// =========================================================================== + +namespace { + +// PFOR_runs(values) wire form: fixed runs of kFrqBaseUnit, run count derived from +// n (mirrors the in-source encode_pfor_runs so a test can rebuild region bytes). +void AppendPforRuns(const U32Vec& values, ByteSink* out) { + const size_t n = values.size(); + const size_t unit = doris::snii::format::kFrqBaseUnit; + for (size_t off = 0; off < n; off += unit) { + const size_t run = (n - off < unit) ? (n - off) : unit; + doris::snii::pfor_encode(values.data() + off, run, out); + } +} + +// dd_region plaintext (VInt n ++ PFOR_runs(doc_delta)) -- the exact bytes +// emit_region writes for a raw dd region. +ByteSink DdPlaintext(const U32Vec& docs, uint64_t win_base) { + U32Vec dd(docs.size()); + uint64_t prev = win_base; + for (size_t i = 0; i < docs.size(); ++i) { + dd[i] = static_cast(static_cast(docs[i]) - prev); + prev = docs[i]; + } + ByteSink plain; + plain.put_varint32(static_cast(docs.size())); + AppendPforRuns(dd, &plain); + return plain; +} + +} // namespace + +// [perf-deterministic] FRQ-BYTE-DD: a raw (level=0) dd region's on-disk bytes +// equal the plaintext exactly, and the meta invariants hold: disk_len == +// plain.size(), uncomp_len == plain.size(), crc == crc32c(plain), zstd == false. +TEST(SniiFrqPodTest, DdRegionRawProducesByteIdenticalOutputAndMeta) { + U32Vec docs = {0, 3, 5, 10, 11, 50, 200}; + const uint64_t win_base = 0; + ByteSink out; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, win_base, /*level=*/0, &out, &meta).ok()); + ASSERT_FALSE(meta.zstd); + + const ByteSink plain = DdPlaintext(docs, win_base); + EXPECT_EQ(out.buffer(), plain.buffer()); // raw region == plaintext, byte-for-byte + EXPECT_EQ(meta.disk_len, plain.view().size()); + EXPECT_EQ(meta.uncomp_len, plain.view().size()); + EXPECT_EQ(meta.crc, doris::snii::crc32c(plain.view())); + EXPECT_EQ(meta.crc, doris::snii::crc32c(Slice(out.buffer()))); + + U32Vec got; + ASSERT_TRUE(decode_dd_region(Slice(out.buffer()), meta, win_base, &got).ok()); + EXPECT_EQ(got, docs); +} + +// [perf-deterministic] FRQ-BYTE-FREQ: a raw (level=0) freq region's on-disk bytes +// equal PFOR_runs(freqs) exactly (no count prefix), with matching meta. +TEST(SniiFrqPodTest, FreqRegionRawProducesByteIdenticalOutputAndMeta) { + U32Vec freqs = {1, 2, 1, 7, 3, 1, 9}; + ByteSink out; + FrqRegionMeta meta; + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &out, &meta).ok()); + ASSERT_FALSE(meta.zstd); + + ByteSink plain; // freq_region plaintext == PFOR_runs(freqs) + AppendPforRuns(freqs, &plain); + EXPECT_EQ(out.buffer(), plain.buffer()); + EXPECT_EQ(meta.disk_len, plain.view().size()); + EXPECT_EQ(meta.uncomp_len, plain.view().size()); + EXPECT_EQ(meta.crc, doris::snii::crc32c(plain.view())); + + U32Vec got; + ASSERT_TRUE(decode_freq_region(Slice(out.buffer()), meta, freqs.size(), &got).ok()); + EXPECT_EQ(got, freqs); +} + +// [perf-deterministic] FRQ-ZSTD-PATH: the zstd branch (level>0) is preserved +// byte-identical -- the on-disk bytes equal an independent zstd(plaintext) at the +// same level, with meta.zstd == true, disk_len == comp.size(), uncomp_len == +// plaintext size, and a lossless round-trip. +TEST(SniiFrqPodTest, DdRegionZstdBranchPreservedByteIdentical) { + U32Vec docs; + uint32_t cur = 0; + for (uint32_t i = 0; i < 2048; ++i) { + cur += 1 + (i % 4); + docs.push_back(cur); + } + const uint64_t win_base = 0; + ByteSink out; + FrqRegionMeta meta; + ASSERT_TRUE(build_dd_region(docs, win_base, /*level=*/5, &out, &meta).ok()); + ASSERT_TRUE(meta.zstd); + + const ByteSink plain = DdPlaintext(docs, win_base); + std::vector comp; + ASSERT_TRUE(doris::snii::zstd_compress(plain.view(), /*level=*/5, &comp).ok()); + EXPECT_EQ(out.buffer(), comp); + EXPECT_EQ(meta.disk_len, comp.size()); + EXPECT_EQ(meta.uncomp_len, plain.view().size()); + EXPECT_EQ(meta.crc, doris::snii::crc32c(Slice(comp))); + + U32Vec got; + ASSERT_TRUE(decode_dd_region(Slice(out.buffer()), meta, win_base, &got).ok()); + EXPECT_EQ(got, docs); +} + +// [functional] FRQ-RT: a raw dd + freq round-trip with the meta raw invariant +// (disk_len == uncomp_len) checked on both regions -- the single-copy raw path +// must keep uncomp_len == disk_len so open_region accepts the region. +TEST(SniiFrqPodTest, RawRegionRoundTripKeepsDiskEqualsUncomp) { + U32Vec docs = {100, 103, 104, 200, 201}; + U32Vec freqs = {1, 4, 2, 1, 9}; + const uint64_t win_base = 100; + ByteSink dd_sink, freq_sink; + FrqRegionMeta dd_meta, freq_meta; + ASSERT_TRUE(build_dd_region(docs, win_base, /*level=*/0, &dd_sink, &dd_meta).ok()); + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &freq_sink, &freq_meta).ok()); + EXPECT_FALSE(dd_meta.zstd); + EXPECT_FALSE(freq_meta.zstd); + EXPECT_EQ(dd_meta.disk_len, dd_meta.uncomp_len); + EXPECT_EQ(freq_meta.disk_len, freq_meta.uncomp_len); + + U32Vec out_docs, out_freqs; + ASSERT_TRUE(decode_dd_region(Slice(dd_sink.buffer()), dd_meta, win_base, &out_docs).ok()); + ASSERT_TRUE( + decode_freq_region(Slice(freq_sink.buffer()), freq_meta, out_docs.size(), &out_freqs) + .ok()); + EXPECT_EQ(out_docs, docs); + EXPECT_EQ(out_freqs, freqs); +} + +// [functional/boundary] FRQ-DOCCOUNT0: an empty freq region (0 freqs) yields a +// zero-length region with zero-valued meta and decodes to an empty vector. +TEST(SniiFrqPodTest, FreqRegionZeroDocCountDecodesEmpty) { + U32Vec freqs; // empty term + ByteSink out; + FrqRegionMeta meta; + ASSERT_TRUE(build_freq_region(freqs, /*level=*/0, &out, &meta).ok()); + EXPECT_EQ(meta.uncomp_len, 0U); + EXPECT_EQ(meta.disk_len, 0U); + EXPECT_FALSE(meta.zstd); + EXPECT_TRUE(out.buffer().empty()); + + U32Vec got = {123U}; // sentinel proves decode clears it + ASSERT_TRUE(decode_freq_region(Slice(out.buffer()), meta, /*doc_count=*/0, &got).ok()); + EXPECT_TRUE(got.empty()); +} diff --git a/be/test/storage/index/snii/format/frq_prelude_test.cpp b/be/test/storage/index/snii/format/frq_prelude_test.cpp new file mode 100644 index 00000000000000..c088c0d4c0c7da --- /dev/null +++ b/be/test/storage/index/snii/format/frq_prelude_test.cpp @@ -0,0 +1,428 @@ +// 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 "storage/index/snii/format/frq_prelude.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/crc32c.h" + +using doris::snii::ByteSink; +using doris::snii::Slice; +using doris::snii::format::build_frq_prelude; +using doris::snii::format::FrqPreludeColumns; +using doris::snii::format::FrqPreludeReader; +using doris::snii::format::WindowMeta; +using doris::Status; + +namespace { + +// Builds N windows that tile a docid space in fixed strides, with deterministic +// per-window metadata. doc_count is `stride`; absolute last docid of window w is +// (w+1)*stride - 1 so windows are contiguous and strictly ascending. dd and freq +// region offsets are CONTIGUOUS prefix sums (the prelude validates this). +FrqPreludeColumns MakeColumns(uint32_t n, uint32_t group_size, uint32_t stride, bool has_prx) { + FrqPreludeColumns cols; + cols.has_freq = true; + cols.has_prx = has_prx; + cols.group_size = group_size; + uint64_t dd_running = 0, freq_running = 0, prx_running = 0; + for (uint32_t w = 0; w < n; ++w) { + WindowMeta m; + m.last_docid = (w + 1) * stride - 1; + m.doc_count = stride; + m.dd_zstd = (w % 2) == 0; + m.dd_off = dd_running; + m.dd_disk_len = 8 + w; + // T18: a raw region's uncomp_len == disk_len (the row stores uncomp_len only + // for zstd regions; the reader derives it otherwise). Keep the test columns + // on that invariant so the round-trip's derived uncomp_len matches. + m.dd_uncomp_len = m.dd_zstd ? m.dd_disk_len + 2 : m.dd_disk_len; + m.crc_dd = 0xDD000000U + w; + dd_running += m.dd_disk_len; + m.freq_zstd = (w % 3) == 0; + m.freq_off = freq_running; + m.freq_disk_len = 5 + (w % 4); + m.freq_uncomp_len = m.freq_zstd ? m.freq_disk_len + 1 : m.freq_disk_len; + m.crc_freq = 0xEE000000U + w; + freq_running += m.freq_disk_len; + if (has_prx) { + m.prx_off = prx_running; + m.prx_len = 7 + (w % 5); + prx_running += m.prx_len; + } + m.max_freq = w * 3 + 1; + m.max_norm = static_cast((w * 13 + 3) & 0xFF); + cols.windows.push_back(m); + } + return cols; +} + +ByteSink MakeSingleWindowPrelude(uint64_t last_docid_delta, uint64_t doc_count, uint64_t max_freq) { + ByteSink block; + block.put_varint64(last_docid_delta); + block.put_varint64(doc_count); + block.put_u8(0); // win_mode: raw dd/freq (no zstd bit => no uncomp_len fields) + block.put_varint64(4); // dd_disk_len (dd_uncomp_len omitted: raw region) + block.put_fixed32(0xDDU); // crc_dd + block.put_varint64(0); // freq_disk_len (freq_uncomp_len omitted: raw region) + block.put_fixed32(0xEEU); // crc_freq + block.put_varint64(max_freq); + block.put_u8(0); // max_norm + + ByteSink dir; + dir.put_varint64(last_docid_delta); + dir.put_varint64(0); // first block starts at the window region + dir.put_varint64(block.size()); + + ByteSink covered; + covered.put_u8(doris::snii::format::frq_prelude_flags::kHasFreq); + covered.put_varint64(1); // N + covered.put_varint64(1); // G + covered.put_varint64(1); // n_super + covered.put_varint64(dir.size()); + covered.put_bytes(dir.view()); + + ByteSink frame; + frame.put_bytes(covered.view()); + frame.put_fixed32(doris::snii::crc32c(covered.view())); + frame.put_bytes(block.view()); + return frame; +} + +// Round-trips columns through build + open, asserting every window's absolute +// fields match (and win_base chains correctly). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +void ExpectRoundTrip(const FrqPreludeColumns& cols) { + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + + ASSERT_EQ(reader.n_windows(), cols.windows.size()); + EXPECT_EQ(reader.has_freq(), cols.has_freq); + EXPECT_EQ(reader.has_prx(), cols.has_prx); + + uint64_t expect_win_base = 0; + for (uint32_t w = 0; w < reader.n_windows(); ++w) { + WindowMeta got; + ASSERT_TRUE(reader.window(w, &got).ok()) << "window w=" << w; + const WindowMeta& exp = cols.windows[w]; + EXPECT_EQ(got.last_docid, exp.last_docid) << "ldd w=" << w; + EXPECT_EQ(got.win_base, expect_win_base) << "win_base w=" << w; + EXPECT_EQ(got.doc_count, exp.doc_count) << "doc_count w=" << w; + EXPECT_EQ(got.dd_zstd, exp.dd_zstd) << "dd_zstd w=" << w; + EXPECT_EQ(got.dd_off, exp.dd_off) << "dd_off w=" << w; + EXPECT_EQ(got.dd_disk_len, exp.dd_disk_len) << "dd_disk_len w=" << w; + EXPECT_EQ(got.dd_uncomp_len, exp.dd_uncomp_len) << "dd_uncomp_len w=" << w; + EXPECT_EQ(got.crc_dd, exp.crc_dd) << "crc_dd w=" << w; + EXPECT_EQ(got.freq_zstd, exp.freq_zstd) << "freq_zstd w=" << w; + EXPECT_EQ(got.freq_off, exp.freq_off) << "freq_off w=" << w; + EXPECT_EQ(got.freq_disk_len, exp.freq_disk_len) << "freq_disk_len w=" << w; + EXPECT_EQ(got.freq_uncomp_len, exp.freq_uncomp_len) << "freq_uncomp_len w=" << w; + EXPECT_EQ(got.crc_freq, exp.crc_freq) << "crc_freq w=" << w; + EXPECT_EQ(got.max_freq, exp.max_freq) << "max_freq w=" << w; + EXPECT_EQ(got.max_norm, exp.max_norm) << "max_norm w=" << w; + if (cols.has_prx) { + EXPECT_EQ(got.prx_off, exp.prx_off) << "prx_off w=" << w; + EXPECT_EQ(got.prx_len, exp.prx_len) << "prx_len w=" << w; + } + expect_win_base = exp.last_docid; + } +} + +} // namespace + +// Many windows across multiple super-blocks, no prx. +TEST(SniiFrqPrelude, RoundTripManyWindowsNoPrx) { + ExpectRoundTrip(MakeColumns(/*n=*/20, /*group_size=*/4, /*stride=*/256, false)); +} + +// Many windows across multiple super-blocks, with prx. +TEST(SniiFrqPrelude, RoundTripManyWindowsWithPrx) { + ExpectRoundTrip(MakeColumns(/*n=*/17, /*group_size=*/4, /*stride=*/256, true)); +} + +// Single window: N=1 collapses to a single super-block. +TEST(SniiFrqPrelude, SingleWindow) { + FrqPreludeColumns cols = MakeColumns(/*n=*/1, /*group_size=*/64, /*stride=*/300, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + EXPECT_EQ(reader.n_windows(), 1U); + EXPECT_EQ(reader.n_super_blocks(), 1U); + WindowMeta m; + ASSERT_TRUE(reader.window(0, &m).ok()); + EXPECT_EQ(m.win_base, 0U); + EXPECT_EQ(m.last_docid, 299U); +} + +// N exactly == group_size -> exactly one super-block. +TEST(SniiFrqPrelude, ExactlyOneSuperBlock) { + FrqPreludeColumns cols = MakeColumns(/*n=*/4, /*group_size=*/4, /*stride=*/256, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + EXPECT_EQ(reader.n_super_blocks(), 1U); +} + +// Large N exercises many super-blocks and varint streams. +TEST(SniiFrqPrelude, LargeN) { + ExpectRoundTrip(MakeColumns(/*n=*/1000, /*group_size=*/64, /*stride=*/256, true)); +} + +// locate_window finds the covering window at boundaries, inside, and past-end. +TEST(SniiFrqPrelude, LocateWindowBoundaries) { + // 10 windows, stride 100 -> window w covers docids [w*100, w*100+99]. + FrqPreludeColumns cols = MakeColumns(/*n=*/10, /*group_size=*/3, /*stride=*/100, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + + struct Case { + uint32_t docid; + bool found; + uint32_t w; + }; + const Case cases[] = { + {.docid = 0, .found = true, .w = 0}, // first docid of window 0 + {.docid = 99, .found = true, .w = 0}, // last docid of window 0 + {.docid = 100, .found = true, .w = 1}, // first docid of window 1 (boundary) + {.docid = 250, .found = true, .w = 2}, // inside window 2 + {.docid = 299, .found = true, .w = 2}, // last docid of window 2 + {.docid = 300, .found = true, .w = 3}, // boundary into window 3 + {.docid = 999, .found = true, .w = 9}, // last docid overall (window 9 covers [900,999]) + {.docid = 1000, .found = false, .w = 0}, // past the term's last docid + {.docid = 5000, .found = false, .w = 0}, // far past end + }; + for (const auto& c : cases) { + bool found = false; + uint32_t w = 999; + ASSERT_TRUE(reader.locate_window(c.docid, &found, &w).ok()) << "docid=" << c.docid; + EXPECT_EQ(found, c.found) << "docid=" << c.docid; + if (c.found) { + EXPECT_EQ(w, c.w) << "docid=" << c.docid; + } + } +} + +// locate_window must agree with a linear scan over windows for every docid in a +// gappy docid layout (windows not contiguous: gaps between windows). +TEST(SniiFrqPrelude, LocateWindowAgainstLinearScan) { + FrqPreludeColumns cols; + cols.has_prx = true; + cols.group_size = 4; + // Non-contiguous absolute last docids with gaps: 50, 130, 131, 400, 900. + const uint32_t lasts[] = {50, 130, 131, 400, 900}; + uint64_t dd_running = 0, freq_running = 0, prx_running = 0; + for (uint32_t v : lasts) { + WindowMeta m; + m.last_docid = v; + m.doc_count = 1; + m.dd_off = dd_running; + m.dd_disk_len = 12; + m.dd_uncomp_len = 12; + m.crc_dd = v; + dd_running += 12; + m.freq_off = freq_running; + m.freq_disk_len = 6; + m.freq_uncomp_len = 6; + m.crc_freq = v + 1; + freq_running += 6; + m.prx_off = prx_running; + m.prx_len = 6; + prx_running += 6; + cols.windows.push_back(m); + } + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + + for (uint32_t docid = 0; docid <= 950; ++docid) { + // Linear oracle: first window whose absolute last_docid >= docid. + bool exp_found = false; + uint32_t exp_w = 0; + for (uint32_t w = 0; w < std::size(lasts); ++w) { + if (docid <= lasts[w]) { + exp_found = true; + exp_w = w; + break; + } + } + bool found = false; + uint32_t w = 999; + ASSERT_TRUE(reader.locate_window(docid, &found, &w).ok()); + EXPECT_EQ(found, exp_found) << "docid=" << docid; + if (exp_found) { + EXPECT_EQ(w, exp_w) << "docid=" << docid; + } + } +} + +// Builder rejects a null sink. +TEST(SniiFrqPrelude, BuildNullSinkRejected) { + FrqPreludeColumns cols = MakeColumns(3, 4, 100, false); + EXPECT_TRUE(build_frq_prelude(cols, nullptr).is()); +} + +// Builder rejects group_size == 0. +TEST(SniiFrqPrelude, BuildZeroGroupSizeRejected) { + FrqPreludeColumns cols = MakeColumns(3, 4, 100, false); + cols.group_size = 0; + ByteSink sink; + EXPECT_TRUE(build_frq_prelude(cols, &sink).is()); +} + +// Builder rejects non-monotonic last_docid ordering across windows. +TEST(SniiFrqPrelude, BuildNonMonotonicRejected) { + FrqPreludeColumns cols = MakeColumns(3, 4, 100, false); + cols.windows[2].last_docid = cols.windows[1].last_docid - 1; // goes backwards + ByteSink sink; + EXPECT_TRUE(build_frq_prelude(cols, &sink).is()); +} + +// T18: dd_off is DERIVED (running prefix sum), not serialized. Tampering a +// column's dd_off no longer changes the on-disk bytes, and the reader +// reconstructs the contiguous offset regardless -- so the offset is always +// contiguous by construction (pre-T18 this stored + cross-checked an explicit +// dd_off, and a gap was rejected on open). +TEST(SniiFrqPrelude, DdOffsetIsDerivedNotStored) { + FrqPreludeColumns cols = MakeColumns(/*n=*/5, /*group_size=*/4, /*stride=*/256, false); + const uint64_t expect_dd_off_2 = + cols.windows[0].dd_disk_len + cols.windows[1].dd_disk_len; // running sum + cols.windows[2].dd_off += 100; // in-memory tamper has no on-disk effect now + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + WindowMeta m; + ASSERT_TRUE(reader.window(2, &m).ok()); + EXPECT_EQ(m.dd_off, expect_dd_off_2); // derived, not the tampered column value +} + +TEST(SniiFrqPrelude, RejectsDocCountBeyondWindowWidth) { + FrqPreludeColumns cols = MakeColumns(/*n=*/1, /*group_size=*/1, + /*stride=*/10, false); + cols.windows[0].doc_count = 11; + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +TEST(SniiFrqPrelude, RejectsWindowDocCountThatDoesNotFitU32) { + ByteSink sink = MakeSingleWindowPrelude( + /*last_docid_delta=*/0, static_cast(std::numeric_limits::max()) + 1, + /*max_freq=*/1); + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +TEST(SniiFrqPrelude, RejectsWindowMaxFreqThatDoesNotFitU32) { + ByteSink sink = MakeSingleWindowPrelude( + /*last_docid_delta=*/0, /*doc_count=*/1, + static_cast(std::numeric_limits::max()) + 1); + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +// The reader reports dd_block_len / freq_block_len as the sums of per-window region +// disk lengths (the contiguous block sizes a reader range-fetches). +TEST(SniiFrqPrelude, BlockLengthsAreRegionSums) { + FrqPreludeColumns cols = MakeColumns(/*n=*/6, /*group_size=*/4, /*stride=*/256, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + uint64_t dd_sum = 0, freq_sum = 0; + for (const auto& m : cols.windows) { + dd_sum += m.dd_disk_len; + freq_sum += m.freq_disk_len; + } + EXPECT_EQ(reader.dd_block_len(), dd_sum); + EXPECT_EQ(reader.freq_block_len(), freq_sum); +} + +// CRC corruption inside the header / super_block_dir is detected by open(). +TEST(SniiFrqPrelude, CrcCorruptionDetected) { + FrqPreludeColumns cols = MakeColumns(8, 4, 256, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GE(bytes.size(), 4U); + bytes[1] ^= 0xFF; // flip a flags/header byte + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(Slice(bytes), &reader); + EXPECT_TRUE(s.is()); +} + +// Truncated buffer (chopping into a window block) is rejected. +TEST(SniiFrqPrelude, TruncatedRejected) { + FrqPreludeColumns cols = MakeColumns(8, 4, 256, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + std::vector bytes = sink.buffer(); + bytes.resize(bytes.size() - 5); // drop tail bytes of the last window block + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(Slice(bytes), &reader); + EXPECT_FALSE(s.ok()); +} + +// Out-of-range window() returns InvalidArgument (no OOB read). +TEST(SniiFrqPrelude, WindowOutOfRangeRejected) { + FrqPreludeColumns cols = MakeColumns(3, 4, 100, true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + WindowMeta m; + EXPECT_TRUE(reader.window(99, &m).is()); +} + +// Anti-DoS: a crc-valid header declaring an absurd window count N must be +// rejected before any reserve(N). We craft header + super_block_dir bytes by +// hand with a matching crc, so verify_crc passes but the count is rejected. +TEST(SniiFrqPrelude, RejectsOversizedWindowCount) { + ByteSink covered; + covered.put_u8(doris::snii::format::frq_prelude_flags::kHasFreq); + covered.put_varint64(0xFFFFFFFFULL); // N: absurd window count + covered.put_varint64(64); // G + covered.put_varint64(1); // n_super (bogus but small) + covered.put_varint64(0); // sbdir_len = 0 + ByteSink frame; + frame.put_bytes(covered.view()); + frame.put_fixed32(doris::snii::crc32c(covered.view())); + FrqPreludeReader reader; + Status s = FrqPreludeReader::open(frame.view(), &reader); + EXPECT_TRUE(s.is()); +} diff --git a/be/test/storage/index/snii/format/metadata_blob_test.cpp b/be/test/storage/index/snii/format/metadata_blob_test.cpp new file mode 100644 index 00000000000000..d76be48b89f90c --- /dev/null +++ b/be/test/storage/index/snii/format/metadata_blob_test.cpp @@ -0,0 +1,293 @@ +// 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 "storage/index/snii/format/metadata_blob.h" + +#include + +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii_query_test_util.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using doris::snii::snii_test::ScopedEnv; + +namespace { + +std::vector BuildSampled(uint32_t n) { + SampledTermIndexBuilder builder; + for (uint32_t i = 0; i < n; ++i) { + builder.add_block_first_term("term_" + std::to_string(1000000 + i)); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +std::vector BuildDirectory() { + DictBlockDirectoryBuilder builder; + builder.add(BlockRef {.offset = 100, .length = 20, .n_entries = 1, .checksum = 123}); + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +struct MetadataFrameCase { + std::vector raw; + SectionType raw_type; + SectionType compressed_type; +}; + +std::vector MetadataFrameCases() { + return {{BuildSampled(3), SectionType::kSampledTermIndex, SectionType::kSampledTermIndexZstd}, + {BuildDirectory(), SectionType::kDictBlockDirectory, + SectionType::kDictBlockDirectoryZstd}}; +} + +std::vector FramePayload(SectionType type, Slice payload) { + ByteSink sink; + SectionFramer::write(sink, static_cast(type), payload); + return sink.buffer(); +} + +SectionType StoredType(Slice frame) { + ByteSource source(frame); + FramedSection section; + EXPECT_TRUE(SectionFramer::read(source, §ion).ok()); + return static_cast(section.type); +} + +std::vector CarrierPayload(uint64_t uncomp_len, Slice compressed) { + ByteSink payload; + payload.put_varint64(uncomp_len); + payload.put_bytes(compressed); + return payload.buffer(); +} + +void ExpectBytesEq(Slice actual, Slice expected) { + ASSERT_EQ(actual.size(), expected.size()); + EXPECT_EQ(std::memcmp(actual.data(), expected.data(), actual.size()), 0); +} + +} // namespace + +TEST(SniiMetadataBlob, RejectsRawFrameWithWrongType) { + const auto raw = BuildSampled(3); + ByteSink stored; + + EXPECT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kDictBlockDirectory, + SectionType::kDictBlockDirectoryZstd, &stored) + .is()); +} + +TEST(SniiMetadataBlob, RejectsTrailingBytesOnRawEncodeAndMaterialize) { + for (const auto& test : MetadataFrameCases()) { + auto trailing = test.raw; + trailing.push_back(0xEE); + + ByteSink stored; + const doris::Status encode_status = + encode_metadata_blob(Slice(trailing), test.raw_type, test.compressed_type, &stored); + EXPECT_TRUE(encode_status.is()) << encode_status; + + std::vector scratch; + Slice materialized; + const doris::Status materialize_status = materialize_metadata_blob( + Slice(trailing), test.raw_type, test.compressed_type, &scratch, &materialized); + EXPECT_TRUE(materialize_status.is()) + << materialize_status; + } +} + +TEST(SniiMetadataBlob, RejectsTrailingOuterAndMalformedInnerCompressedFrames) { + const auto tests = MetadataFrameCases(); + for (size_t i = 0; i < tests.size(); ++i) { + const auto& test = tests[i]; + std::vector compressed; + ASSERT_TRUE(zstd_compress(Slice(test.raw), 3, &compressed).ok()); + const auto payload = CarrierPayload(test.raw.size(), Slice(compressed)); + auto outer_trailing = FramePayload(test.compressed_type, Slice(payload)); + outer_trailing.push_back(0xEE); + + std::vector scratch; + Slice materialized; + doris::Status status = + materialize_metadata_blob(Slice(outer_trailing), test.raw_type, + test.compressed_type, &scratch, &materialized); + EXPECT_TRUE(status.is()) << status; + + auto inner_trailing = test.raw; + inner_trailing.push_back(0xEE); + ASSERT_TRUE(zstd_compress(Slice(inner_trailing), 3, &compressed).ok()); + const auto inner_trailing_payload = + CarrierPayload(inner_trailing.size(), Slice(compressed)); + const auto inner_trailing_carrier = + FramePayload(test.compressed_type, Slice(inner_trailing_payload)); + status = materialize_metadata_blob(Slice(inner_trailing_carrier), test.raw_type, + test.compressed_type, &scratch, &materialized); + EXPECT_TRUE(status.is()) << status; + + const auto& wrong_inner = tests[(i + 1) % tests.size()].raw; + ASSERT_TRUE(zstd_compress(Slice(wrong_inner), 3, &compressed).ok()); + const auto wrong_type_payload = CarrierPayload(wrong_inner.size(), Slice(compressed)); + const auto wrong_type_carrier = + FramePayload(test.compressed_type, Slice(wrong_type_payload)); + status = materialize_metadata_blob(Slice(wrong_type_carrier), test.raw_type, + test.compressed_type, &scratch, &materialized); + EXPECT_TRUE(status.is()) << status; + } +} + +TEST(SniiMetadataBlob, KeepsSmallRawFrameByteIdentical) { + const auto raw = BuildSampled(3); + ByteSink stored; + ASSERT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &stored) + .ok()); + EXPECT_EQ(stored.buffer(), raw); + + std::vector scratch; + Slice materialized; + ASSERT_TRUE(materialize_metadata_blob(stored.view(), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .ok()); + EXPECT_TRUE(scratch.empty()); + ExpectBytesEq(materialized, Slice(raw)); +} + +TEST(SniiMetadataBlob, ThresholdGateControlsCompression) { + const auto raw = BuildSampled(300); + ASSERT_LT(raw.size(), kMetaSectionCompressMinBytes); + + ByteSink default_stored; + ASSERT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &default_stored) + .ok()); + EXPECT_EQ(default_stored.buffer(), raw); + + ByteSink forced_stored; + { + ScopedEnv force("SNII_META_COMPRESS_MIN", "1"); + ASSERT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &forced_stored) + .ok()); + } + EXPECT_EQ(StoredType(forced_stored.view()), SectionType::kSampledTermIndexZstd); + EXPECT_LT(forced_stored.size(), raw.size()); + + std::vector scratch; + Slice materialized; + ASSERT_TRUE(materialize_metadata_blob(forced_stored.view(), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .ok()); + ExpectBytesEq(materialized, Slice(raw)); +} + +TEST(SniiMetadataBlob, CompressesLargeFrameAndMaterializesByteExactly) { + const auto raw = BuildSampled(2000); + ASSERT_GT(raw.size(), kMetaSectionCompressMinBytes); + ByteSink stored; + ASSERT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &stored) + .ok()); + EXPECT_EQ(StoredType(stored.view()), SectionType::kSampledTermIndexZstd); + EXPECT_LT(stored.size(), raw.size()); + + std::vector scratch; + Slice materialized; + ASSERT_TRUE(materialize_metadata_blob(stored.view(), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .ok()); + ExpectBytesEq(materialized, Slice(raw)); +} + +TEST(SniiMetadataBlob, KeepsIncompressibleFrameRaw) { + std::vector payload(kMetaSectionCompressMinBytes * 2); + uint32_t state = 0x12345678U; + for (uint8_t& byte : payload) { + state = state * 1664525U + 1013904223U; + byte = static_cast(state >> 24U); + } + const auto raw = FramePayload(SectionType::kSampledTermIndex, Slice(payload)); + ByteSink stored; + ASSERT_TRUE(encode_metadata_blob(Slice(raw), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &stored) + .ok()); + EXPECT_EQ(stored.buffer(), raw); +} + +TEST(SniiMetadataBlob, RejectsCorruptZstdPayload) { + std::vector garbage(64, 0xAB); + const auto payload = CarrierPayload(1024, Slice(garbage)); + const auto stored = FramePayload(SectionType::kSampledTermIndexZstd, Slice(payload)); + + std::vector scratch; + Slice materialized; + EXPECT_TRUE(materialize_metadata_blob(Slice(stored), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .is()); +} + +TEST(SniiMetadataBlob, RejectsTruncatedOrMisdeclaredZstdPayload) { + const auto raw = BuildSampled(500); + std::vector compressed; + ASSERT_TRUE(zstd_compress(Slice(raw), 3, &compressed).ok()); + + for (const auto& [uncomp_len, comp] : + {std::pair {raw.size(), Slice(compressed.data(), compressed.size() / 2)}, + std::pair {raw.size() - 1, Slice(compressed)}}) { + const auto payload = CarrierPayload(uncomp_len, comp); + const auto stored = FramePayload(SectionType::kSampledTermIndexZstd, Slice(payload)); + std::vector scratch; + Slice materialized; + EXPECT_TRUE(materialize_metadata_blob(Slice(stored), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .is()); + } +} + +TEST(SniiMetadataBlob, RejectsOutOfRangeZstdUncompressedLength) { + std::vector compressed(16, 0x01); + for (const uint64_t uncomp_len : {uint64_t {0}, uint64_t {256ULL * 1024 * 1024 + 1}}) { + const auto payload = CarrierPayload(uncomp_len, Slice(compressed)); + const auto stored = FramePayload(SectionType::kSampledTermIndexZstd, Slice(payload)); + std::vector scratch; + Slice materialized; + EXPECT_TRUE(materialize_metadata_blob(Slice(stored), SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &scratch, + &materialized) + .is()); + } +} diff --git a/be/test/storage/index/snii/format/metadata_directory_test.cpp b/be/test/storage/index/snii/format/metadata_directory_test.cpp new file mode 100644 index 00000000000000..360fffc237ab4a --- /dev/null +++ b/be/test/storage/index/snii/format/metadata_directory_test.cpp @@ -0,0 +1,226 @@ +// 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 "storage/index/snii/format/metadata_directory.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "gen_cpp/snii.pb.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" + +namespace doris::snii::format { +namespace { + +LogicalIndexMetadataRef sample_entry(uint64_t index_id, std::string suffix, uint64_t base) { + return {.index_id = index_id, + .index_suffix = std::move(suffix), + .core_metadata = {.offset = base, .length = 11}, + .sampled_term_index = {.offset = base + 11, .length = 12}, + .dict_block_directory = {.offset = base + 23, .length = 13}}; +} + +std::vector sample_entries() { + return {sample_entry(7, "primary", 100), sample_entry(8, "secondary", 200)}; +} + +std::vector encode(const std::vector& entries) { + ByteSink sink; + EXPECT_TRUE(encode_metadata_directory(entries, &sink).ok()); + return sink.buffer(); +} + +void expect_entry_eq(const LogicalIndexMetadataRef& expected, + const LogicalIndexMetadataRef& actual) { + EXPECT_EQ(expected.index_id, actual.index_id); + EXPECT_EQ(expected.index_suffix, actual.index_suffix); + EXPECT_EQ(expected.core_metadata.offset, actual.core_metadata.offset); + EXPECT_EQ(expected.core_metadata.length, actual.core_metadata.length); + EXPECT_EQ(expected.sampled_term_index.offset, actual.sampled_term_index.offset); + EXPECT_EQ(expected.sampled_term_index.length, actual.sampled_term_index.length); + EXPECT_EQ(expected.dict_block_directory.offset, actual.dict_block_directory.offset); + EXPECT_EQ(expected.dict_block_directory.length, actual.dict_block_directory.length); +} + +doris::snii::SniiMetadataDirectoryPB valid_pb() { + doris::snii::SniiMetadataDirectoryPB directory; + const auto entry = sample_entry(7, "primary", 100); + auto* index = directory.add_indexes(); + index->set_index_id(entry.index_id); + index->set_index_suffix(entry.index_suffix); + auto* core = index->mutable_core_metadata(); + core->set_offset(entry.core_metadata.offset); + core->set_length(entry.core_metadata.length); + auto* sti = index->mutable_sampled_term_index(); + sti->set_offset(entry.sampled_term_index.offset); + sti->set_length(entry.sampled_term_index.length); + auto* dbd = index->mutable_dict_block_directory(); + dbd->set_offset(entry.dict_block_directory.offset); + dbd->set_length(entry.dict_block_directory.length); + return directory; +} + +std::vector serialize(const doris::snii::SniiMetadataDirectoryPB& directory) { + std::string bytes; + EXPECT_TRUE(directory.SerializeToString(&bytes)); + return {bytes.begin(), bytes.end()}; +} + +void expect_corruption(const std::vector& bytes) { + MetadataDirectory directory; + const auto status = MetadataDirectory::decode(Slice(bytes), &directory); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiMetadataDirectory, RoundTripsMultipleIndexesAndFindsDeterministically) { + const auto expected = sample_entries(); + MetadataDirectory directory; + ASSERT_TRUE(MetadataDirectory::decode(Slice(encode(expected)), &directory).ok()); + ASSERT_EQ(expected.size(), directory.size()); + ASSERT_EQ(expected.size(), directory.entries().size()); + for (size_t i = 0; i < expected.size(); ++i) { + expect_entry_eq(expected[i], directory.entries()[i]); + const auto* found = directory.find(expected[i].index_id, expected[i].index_suffix); + ASSERT_NE(nullptr, found); + expect_entry_eq(expected[i], *found); + } + EXPECT_EQ(nullptr, directory.find(expected[0].index_id, "missing")); + EXPECT_EQ(nullptr, directory.find(99, expected[0].index_suffix)); +} + +TEST(SniiMetadataDirectory, RoundTripsEmptyDirectory) { + MetadataDirectory directory; + ASSERT_TRUE(MetadataDirectory::decode(Slice(encode({})), &directory).ok()); + EXPECT_EQ(0U, directory.size()); + EXPECT_TRUE(directory.entries().empty()); + EXPECT_EQ(nullptr, directory.find(7, "primary")); +} + +TEST(SniiMetadataDirectory, PreservesBinarySuffix) { + const auto expected = sample_entry(7, std::string("suffix\0bytes", 12), 100); + MetadataDirectory directory; + ASSERT_TRUE(MetadataDirectory::decode(Slice(encode({expected})), &directory).ok()); + const auto* found = directory.find(expected.index_id, expected.index_suffix); + ASSERT_NE(nullptr, found); + expect_entry_eq(expected, *found); +} + +TEST(SniiMetadataDirectory, RejectsDuplicateKey) { + auto directory = valid_pb(); + *directory.add_indexes() = directory.indexes(0); + expect_corruption(serialize(directory)); + + const auto duplicate = sample_entry(7, "primary", 200); + ByteSink sink; + const auto status = + encode_metadata_directory({sample_entry(7, "primary", 100), duplicate}, &sink); + EXPECT_TRUE(status.is()) << status; + EXPECT_TRUE(sink.buffer().empty()); +} + +TEST(SniiMetadataDirectory, RejectsMissingLogicalFields) { + using Mutation = std::function; + const std::array mutations { + [](auto* index) { index->clear_index_id(); }, + [](auto* index) { index->clear_index_suffix(); }, + [](auto* index) { index->clear_core_metadata(); }, + [](auto* index) { index->clear_sampled_term_index(); }, + [](auto* index) { index->clear_dict_block_directory(); }, + }; + for (const auto& mutation : mutations) { + auto directory = valid_pb(); + mutation(directory.mutable_indexes(0)); + expect_corruption(serialize(directory)); + } +} + +TEST(SniiMetadataDirectory, RejectsMissingNestedBlobOffsetOrLength) { + using BlobGetter = doris::snii::SniiBlobRefPB* (doris::snii::SniiLogicalIndexMetadataPB::*)(); + for (const auto blob : std::array { + &doris::snii::SniiLogicalIndexMetadataPB::mutable_core_metadata, + &doris::snii::SniiLogicalIndexMetadataPB::mutable_sampled_term_index, + &doris::snii::SniiLogicalIndexMetadataPB::mutable_dict_block_directory}) { + for (const auto clear : std::array { + &doris::snii::SniiBlobRefPB::clear_offset, + &doris::snii::SniiBlobRefPB::clear_length}) { + auto directory = valid_pb(); + ((directory.mutable_indexes(0)->*blob)()->*clear)(); + expect_corruption(serialize(directory)); + } + } +} + +TEST(SniiMetadataDirectory, RejectsZeroLengthMandatoryBlob) { + using BlobGetter = doris::snii::SniiBlobRefPB* (doris::snii::SniiLogicalIndexMetadataPB::*)(); + for (const auto blob : std::array { + &doris::snii::SniiLogicalIndexMetadataPB::mutable_core_metadata, + &doris::snii::SniiLogicalIndexMetadataPB::mutable_sampled_term_index, + &doris::snii::SniiLogicalIndexMetadataPB::mutable_dict_block_directory}) { + auto directory = valid_pb(); + (directory.mutable_indexes(0)->*blob)()->set_length(0); + expect_corruption(serialize(directory)); + } + + auto invalid = sample_entry(7, "primary", 100); + invalid.core_metadata.length = 0; + ByteSink sink; + const auto status = encode_metadata_directory({invalid}, &sink); + EXPECT_TRUE(status.is()) << status; + EXPECT_TRUE(sink.buffer().empty()); +} + +TEST(SniiMetadataDirectory, RejectsTruncatedProtobuf) { + expect_corruption({0x12, 0x01}); +} + +TEST(SniiMetadataDirectory, AcceptsUnknownOptionalField) { + auto bytes = serialize(valid_pb()); + ByteSink unknown_field; + unknown_field.put_varint32((100U << 3) | 0U); + unknown_field.put_varint32(7); + bytes.insert(bytes.end(), unknown_field.buffer().begin(), unknown_field.buffer().end()); + + MetadataDirectory directory; + ASSERT_TRUE(MetadataDirectory::decode(Slice(bytes), &directory).ok()); + ASSERT_EQ(1U, directory.size()); + expect_entry_eq(sample_entry(7, "primary", 100), directory.entries()[0]); +} + +TEST(SniiMetadataDirectory, RejectsUnknownRequiredFeature) { + auto directory = valid_pb(); + directory.add_required_features(1); + MetadataDirectory decoded; + const auto status = MetadataDirectory::decode(Slice(serialize(directory)), &decoded); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiMetadataDirectory, RejectsNullOutputPointers) { + EXPECT_TRUE(MetadataDirectory::decode(Slice(encode(sample_entries())), nullptr) + .is()); + EXPECT_TRUE( + encode_metadata_directory(sample_entries(), nullptr).is()); +} + +} // namespace +} // namespace doris::snii::format diff --git a/be/test/storage/index/snii/format/norms_pod_test.cpp b/be/test/storage/index/snii/format/norms_pod_test.cpp new file mode 100644 index 00000000000000..4e4dc75304d6a2 --- /dev/null +++ b/be/test/storage/index/snii/format/norms_pod_test.cpp @@ -0,0 +1,218 @@ +// 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 "storage/index/snii/format/norms_pod.h" + +#include + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; +using doris::Status; // RETURN_IF_ERROR expands to bare Status +using doris::snii::format::NormsPodReader; +using doris::snii::format::NormsPodWriter; + +namespace { + +// Use writer to encode a sequence of encoded_norms into a framed payload and return the buffer. +std::vector BuildPod(const std::vector& norms) { + NormsPodWriter writer; + for (uint8_t n : norms) { + writer.add(n); + } + ByteSink sink; + writer.finish(&sink); + return sink.buffer(); +} + +} // namespace + +// After writing N norms, read them back per-doc and verify they match. +TEST(SniiNormsPod, RoundTripValues) { + std::vector norms = {0, 1, 7, 42, 128, 200, 255}; + NormsPodWriter writer; + for (uint8_t n : norms) { + writer.add(n); + } + EXPECT_EQ(writer.count(), norms.size()); + + ByteSink sink; + writer.finish(&sink); + + NormsPodReader reader; + ASSERT_TRUE(NormsPodReader::open(sink.view(), &reader).ok()); + ASSERT_EQ(reader.doc_count(), norms.size()); + for (uint32_t docid = 0; docid < norms.size(); ++docid) { + EXPECT_EQ(reader.encoded_norm(docid), norms[docid]) << "docid=" << docid; + } +} + +// Large-scale round-trip covering the multi-byte varint doc_count path. +TEST(SniiNormsPod, RoundTripLarge) { + std::vector norms; + norms.reserve(5000); + for (uint32_t i = 0; i < 5000; ++i) { + norms.push_back(static_cast((i * 31 + 7) & 0xFF)); + } + auto buf = BuildPod(norms); + + NormsPodReader reader; + ASSERT_TRUE(NormsPodReader::open(Slice(buf), &reader).ok()); + ASSERT_EQ(reader.doc_count(), 5000U); + for (uint32_t docid = 0; docid < 5000; ++docid) { + EXPECT_EQ(reader.encoded_norm(docid), norms[docid]) << "docid=" << docid; + } +} + +// Empty POD: count = 0 is valid and open should succeed. +TEST(SniiNormsPod, EmptyPod) { + NormsPodWriter writer; + EXPECT_EQ(writer.count(), 0U); + + ByteSink sink; + writer.finish(&sink); + + NormsPodReader reader; + ASSERT_TRUE(NormsPodReader::open(sink.view(), &reader).ok()); + EXPECT_EQ(reader.doc_count(), 0U); +} + +// CRC corruption is detectable: flipping a byte in the payload causes open to fail. +// The integrated reader frames via SectionFramer, which reports a CRC mismatch as +// INVERTED_INDEX_FILE_CORRUPTED (the standalone test's generic Corruption code was +// replaced during integration with this inverted-index-specific code). +TEST(SniiNormsPod, DetectsCorruption) { + std::vector norms = {10, 20, 30, 40, 50}; + auto buf = BuildPod(norms); + // Flip a byte near the end so the framer CRC no longer matches the payload. + buf[buf.size() - 3] ^= 0xFF; + + NormsPodReader reader; + Status s = NormsPodReader::open(Slice(buf), &reader); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// Truncated input should return an error rather than crash. +TEST(SniiNormsPod, DetectsTruncation) { + std::vector norms = {1, 2, 3, 4, 5, 6, 7, 8}; + auto buf = BuildPod(norms); + buf.resize(buf.size() - 4); // Chop off the trailing CRC region. + + NormsPodReader reader; + Status s = NormsPodReader::open(Slice(buf), &reader); + EXPECT_FALSE(s.ok()); +} + +// A mismatch between the declared doc_count and the actual payload byte count should be detected. +// The integrated reader reports this as INVERTED_INDEX_FILE_CORRUPTED. +TEST(SniiNormsPod, DetectsLengthMismatch) { + // Manually construct: framer payload = [varint doc_count=4][only 2 norm bytes]. + ByteSink payload; + payload.put_varint64(4); + payload.put_u8(11); + payload.put_u8(22); + + ByteSink sink; + // Reuse the framer to ensure a self-consistent CRC, specifically to trigger the length-mismatch branch. + doris::snii::SectionFramer::write( + sink, static_cast(doris::snii::format::SectionType::kNormsPod), + payload.view()); + + NormsPodReader reader; + Status s = NormsPodReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +TEST(SniiNormsPod, RejectsWrongSectionType) { + ByteSink payload; + payload.put_varint64(1); + payload.put_u8(7); + ByteSink sink; + SectionFramer::write(sink, static_cast(format::SectionType::kSampledTermIndex), + payload.view()); + + NormsPodReader reader; + const Status status = NormsPodReader::open(sink.view(), &reader); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +TEST(SniiNormsPod, RejectsStatsBlockFrame) { + ByteSink payload; + payload.put_varint64(4); + ByteSink sink; + SectionFramer::write(sink, /*obsolete stats frame type=*/1, payload.view()); + + NormsPodReader reader; + const Status status = NormsPodReader::open(sink.view(), &reader); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +TEST(SniiNormsPod, RejectsNonCanonicalDocCountVarint) { + ByteSink payload; + for (size_t i = 0; i < 9; ++i) { + payload.put_u8(0x80); + } + payload.put_u8(0x02); + ByteSink sink; + SectionFramer::write(sink, static_cast(format::SectionType::kNormsPod), + payload.view()); + + NormsPodReader reader; + const Status status = NormsPodReader::open(sink.view(), &reader); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +TEST(SniiNormsPod, RejectsTrailingFramedBytes) { + std::vector bytes = BuildPod({7}); + bytes.push_back(0); + + NormsPodReader reader; + const Status status = NormsPodReader::open(Slice(bytes), &reader); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +#ifndef NDEBUG +// In a debug build, an out-of-range docid triggers an assertion (death test). +TEST(SniiNormsPodDeathTest, OutOfRangeDocidAsserts) { + std::vector norms = {3, 6, 9}; + auto buf = BuildPod(norms); + NormsPodReader reader; + ASSERT_TRUE(NormsPodReader::open(Slice(buf), &reader).ok()); + EXPECT_DEATH({ (void)reader.encoded_norm(3); }, ""); +} +#endif + +// Checked access: a valid docid returns the value, an out-of-range docid returns +// InvalidArgument (also effective in Release builds). +TEST(SniiNormsPod, TryEncodedNormChecksBounds) { + std::vector norms = {3, 6, 9}; + auto buf = BuildPod(norms); + NormsPodReader reader; + ASSERT_TRUE(NormsPodReader::open(Slice(buf), &reader).ok()); + uint8_t v = 0; + ASSERT_TRUE(reader.try_encoded_norm(1, &v).ok()); + EXPECT_EQ(v, 6U); + Status s = reader.try_encoded_norm(3, &v); + EXPECT_TRUE(s.is()) << s.to_string(); +} diff --git a/be/test/storage/index/snii/format/null_bitmap_test.cpp b/be/test/storage/index/snii/format/null_bitmap_test.cpp new file mode 100644 index 00000000000000..c3be8c51136e96 --- /dev/null +++ b/be/test/storage/index/snii/format/null_bitmap_test.cpp @@ -0,0 +1,286 @@ +// 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 "storage/index/snii/format/null_bitmap.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +// clang-format off +// CRoaring's public header must precede internal container headers. +#include "roaring/roaring.hh" +#include "roaring/containers/array.h" +#include "roaring/containers/bitset.h" +#include "roaring/containers/run.h" +// clang-format on +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/section_framer.h" + +using namespace doris::snii; +using doris::Status; +using doris::snii::format::NullBitmapReader; +using doris::snii::format::NullBitmapWriter; +using doris::snii::format::kNullBitmapSectionType; + +namespace { + +// Encode a set of null docids into a framed buffer using the writer. +std::vector BuildBitmap(const std::vector& nulls, uint32_t doc_count) { + NullBitmapWriter writer; + for (uint32_t d : nulls) { + writer.add_null(d); + } + ByteSink sink; + EXPECT_TRUE(writer.finish(doc_count, &sink).ok()); + return sink.buffer(); +} + +} // namespace + +// After adding nulls, is_null(docid) must match the input set for every doc. +TEST(SniiNullBitmap, RoundTripPerDoc) { + std::vector nulls = {0, 3, 7, 11, 100, 4000}; + uint32_t doc_count = 5000; + auto buf = BuildBitmap(nulls, doc_count); + + NullBitmapReader reader; + ASSERT_TRUE(NullBitmapReader::open(Slice(buf), &reader).ok()); + EXPECT_EQ(reader.doc_count(), doc_count); + EXPECT_EQ(reader.null_count(), nulls.size()); + + std::vector expected(doc_count, false); + for (uint32_t d : nulls) { + expected[d] = true; + } + for (uint32_t docid = 0; docid < doc_count; ++docid) { + EXPECT_EQ(reader.is_null(docid), expected[docid]) << "docid=" << docid; + } +} + +// Writer null_count reflects the number of distinct null docids added. +TEST(SniiNullBitmap, WriterNullCount) { + NullBitmapWriter writer; + EXPECT_EQ(writer.null_count(), 0U); + writer.add_null(5); + writer.add_null(9); + writer.add_null(5); // duplicate is idempotent in a set + EXPECT_EQ(writer.null_count(), 2U); +} + +TEST(SniiNullBitmap, DenseBuildPeakCoversArrayToBitsetConversion) { + constexpr uint32_t kDenseContainerCount = 8; + constexpr uint32_t kDenseCardinality = roaring::internal::DEFAULT_MAX_SIZE + 1; + std::vector nulls; + nulls.reserve(kDenseContainerCount * kDenseCardinality); + for (uint32_t key = 0; key < kDenseContainerCount; ++key) { + for (uint32_t low = 0; low < kDenseCardinality; ++low) { + nulls.push_back((key << 16) | low); + } + } + + constexpr uint64_t kObservedConversionCapacity = 4165; + constexpr uint64_t kTopEntryBytes = sizeof(void*) + sizeof(uint16_t) + sizeof(uint8_t); + constexpr uint64_t kBitsetBytes = + roaring::internal::BITSET_CONTAINER_SIZE_IN_WORDS * sizeof(uint64_t); + constexpr uint64_t kObservedPeak = + sizeof(roaring::Roaring) + 3 * kDenseContainerCount * kTopEntryBytes + + kDenseContainerCount * (kObservedConversionCapacity * sizeof(uint16_t) + kBitsetBytes + + sizeof(roaring::internal::array_container_t) + + sizeof(roaring::internal::bitset_container_t)); + EXPECT_GE(NullBitmapWriter::build_memory_upper_bound(std::span(nulls)), + kObservedPeak); +} + +// Empty bitmap: no nulls. open succeeds, null_count == 0, nothing is null. +TEST(SniiNullBitmap, EmptyNoNulls) { + auto buf = BuildBitmap({}, 1000); + + NullBitmapReader reader; + ASSERT_TRUE(NullBitmapReader::open(Slice(buf), &reader).ok()); + EXPECT_EQ(reader.doc_count(), 1000U); + EXPECT_EQ(reader.null_count(), 0U); + EXPECT_FALSE(reader.is_null(0)); + EXPECT_FALSE(reader.is_null(999)); +} + +// All-null bitmap: every doc in [0, doc_count) is null. +TEST(SniiNullBitmap, AllNull) { + uint32_t doc_count = 256; + std::vector nulls; + for (uint32_t d = 0; d < doc_count; ++d) { + nulls.push_back(d); + } + auto buf = BuildBitmap(nulls, doc_count); + + NullBitmapReader reader; + ASSERT_TRUE(NullBitmapReader::open(Slice(buf), &reader).ok()); + EXPECT_EQ(reader.null_count(), doc_count); + for (uint32_t docid = 0; docid < doc_count; ++docid) { + EXPECT_TRUE(reader.is_null(docid)) << "docid=" << docid; + } +} + +// doc_count round-trips even when there are no nulls and a large doc_count. +TEST(SniiNullBitmap, DocCountRoundTrips) { + auto buf = BuildBitmap({42}, 1234567); + + NullBitmapReader reader; + ASSERT_TRUE(NullBitmapReader::open(Slice(buf), &reader).ok()); + EXPECT_EQ(reader.doc_count(), 1234567U); + EXPECT_TRUE(reader.is_null(42)); +} + +// is_null beyond doc_count is false (docid not in the null set). +TEST(SniiNullBitmap, IsNullOutsideRangeIsFalse) { + auto buf = BuildBitmap({1, 2, 3}, 10); + + NullBitmapReader reader; + ASSERT_TRUE(NullBitmapReader::open(Slice(buf), &reader).ok()); + EXPECT_FALSE(reader.is_null(10)); + EXPECT_FALSE(reader.is_null(1000000)); +} + +// CRC corruption is detectable: flipping a payload byte fails open with a +// corruption error (SectionFramer stamps the crc over type+len+payload, so a +// flipped payload byte makes the recomputed crc disagree with the stored one). +TEST(SniiNullBitmap, DetectsCorruption) { + std::vector nulls = {2, 4, 6, 8, 10, 12, 14}; + auto buf = BuildBitmap(nulls, 100); + // Flip a byte inside the roaring payload region (near the end, before the trailing CRC). + buf[buf.size() - 5] ^= 0xFF; + + NullBitmapReader reader; + Status s = NullBitmapReader::open(Slice(buf), &reader); + EXPECT_TRUE(s.is()); +} + +// Truncated input returns an error rather than crashing. +TEST(SniiNullBitmap, DetectsTruncation) { + auto buf = BuildBitmap({1, 2, 3, 4, 5}, 100); + buf.resize(buf.size() - 4); // chop trailing CRC region + + NullBitmapReader reader; + Status s = NullBitmapReader::open(Slice(buf), &reader); + EXPECT_FALSE(s.ok()); +} + +// An oversized declared roaring_size (larger than the remaining payload bytes) is rejected (anti-DoS). +TEST(SniiNullBitmap, RejectsOversizedRoaringSize) { + // Manually construct a self-consistent frame whose declared roaring_size + // exceeds the bytes actually present, to drive the guard branch. + ByteSink payload; + payload.put_varint64(100); // doc_count + payload.put_varint64(0xFFFFFFFFULL); // roaring_size: absurdly large + payload.put_u8(0x00); // only 1 byte of roaring data present + + ByteSink sink; + SectionFramer::write(sink, kNullBitmapSectionType, payload.view()); + + NullBitmapReader reader; + Status s = NullBitmapReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +// doc_count overflowing uint32 is rejected. +TEST(SniiNullBitmap, RejectsDocCountOverflow) { + ByteSink payload; + payload.put_varint64(0x1'0000'0000ULL); // doc_count > uint32 max + payload.put_varint64(0); // roaring_size + + ByteSink sink; + SectionFramer::write(sink, kNullBitmapSectionType, payload.view()); + + NullBitmapReader reader; + Status s = NullBitmapReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +// A CRC-valid frame carrying malformed roaring container bytes must be rejected +// gracefully (corruption error) without throwing or aborting. The roaring bytes are +// not a valid portable serialization; SectionFramer stamps a correct crc so the framer +// check passes and the roaring pre-validation (deserialize_size probe) must catch it. +TEST(SniiNullBitmap, RejectsMalformedRoaringContainer) { + ByteSink payload; + payload.put_varint64(10); // doc_count + payload.put_varint64(4); // roaring_size + const uint8_t garbage[] = {0xFF, 0xFF, 0xFF, 0xFF}; // invalid roaring cookie + payload.put_bytes(Slice(garbage, 4)); + + ByteSink sink; + SectionFramer::write(sink, kNullBitmapSectionType, payload.view()); + + NullBitmapReader reader; + Status s = NullBitmapReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()); +} + +TEST(SniiNullBitmap, RejectsNegativeNoRunContainerCount) { + ByteSink roaring; + roaring.put_fixed32(12346); // portable cookie without run containers + roaring.put_fixed32(std::numeric_limits::max()); + + ByteSink payload; + payload.put_varint64(1); // doc_count + payload.put_varint64(roaring.size()); + payload.put_bytes(roaring.view()); + + ByteSink sink; + SectionFramer::write(sink, kNullBitmapSectionType, payload.view()); + + uint64_t decoded_bytes = 0; + Status status = NullBitmapReader::decoded_memory_bytes(sink.view(), &decoded_bytes); + EXPECT_TRUE(status.is()); + EXPECT_NE(status.to_string().find("portable container count out of range"), std::string::npos); +} + +TEST(SniiNullBitmap, DecodedMemoryAccountsEachPortableContainer) { + constexpr uint32_t kContainerCount = 3; + auto buffer = BuildBitmap({0, 1U << 16, 2U << 16}, (2U << 16) + 1); + + ByteSource frame_source {Slice(buffer)}; + FramedSection frame; + ASSERT_TRUE(SectionFramer::read(frame_source, &frame).ok()); + ByteSource payload(frame.payload); + uint64_t doc_count = 0; + uint64_t roaring_bytes = 0; + ASSERT_TRUE(payload.get_varint64(&doc_count).ok()); + ASSERT_TRUE(payload.get_varint64(&roaring_bytes).ok()); + EXPECT_EQ(doc_count, (2U << 16) + 1); + + constexpr uint64_t kContainerObjectBytes = + std::max({sizeof(roaring::internal::array_container_t), + sizeof(roaring::internal::bitset_container_t), + sizeof(roaring::internal::run_container_t)}); + constexpr uint64_t kContainerMetadataBytes = + sizeof(void*) + sizeof(uint16_t) + sizeof(uint8_t) + kContainerObjectBytes; + constexpr uint64_t kFixedBytes = + sizeof(roaring::Roaring) + sizeof(roaring::api::roaring_bitmap_t); + + uint64_t decoded_bytes = 0; + ASSERT_TRUE(NullBitmapReader::decoded_memory_bytes(Slice(buffer), &decoded_bytes).ok()); + EXPECT_EQ(decoded_bytes, + roaring_bytes + kContainerCount * kContainerMetadataBytes + kFixedBytes); +} diff --git a/be/test/storage/index/snii/format/prx_pod_test.cpp b/be/test/storage/index/snii/format/prx_pod_test.cpp new file mode 100644 index 00000000000000..bde677b82c13d1 --- /dev/null +++ b/be/test/storage/index/snii/format/prx_pod_test.cpp @@ -0,0 +1,1857 @@ +// 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 "storage/index/snii/format/prx_pod.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/encoding/zstd_codec.h" +#include "storage/index/snii/format/format_constants.h" +#include "util/defer_op.h" + +using doris::Status; // RETURN_IF_ERROR expands to bare Status +using doris::snii::ByteSink; +using doris::snii::ByteSource; +using doris::snii::pfor_encode; +using doris::snii::Slice; +using doris::snii::format::build_prx_window; +using doris::snii::format::build_prx_window_flat; +using doris::snii::format::kFrqBaseUnit; +using doris::snii::format::PrxCodec; +using doris::snii::format::PrxCsrAllocationGate; +using doris::snii::format::PrxDecodeContext; +using doris::snii::format::PrxDecodeStats; +using doris::snii::format::PrxDecodedShape; +using doris::snii::format::read_prx_window; +using doris::snii::format::read_prx_window_csr; +using doris::snii::format::read_prx_window_csr_for_selection; +using doris::snii::format::read_prx_window_csr_selective; + +// Integrated SNII reports corruption via doris ErrorCode::INVERTED_INDEX_FILE_CORRUPTED +// (the standalone build used doris::snii::StatusCode::kCorruption) and writer-side precondition +// violations via ErrorCode::INVALID_ARGUMENT; the negative tests below assert those codes. + +namespace { + +using PerDoc = std::vector>; + +PerDoc MakeProfileDocs(uint32_t count); + +class RejectingPrxAllocationGate final : public PrxCsrAllocationGate { +public: + Status reserve_csr(std::vector*, size_t, std::vector*, size_t) override { + ++csr_calls; + return Status::OK(); + } + + Status reserve_decompression(size_t bytes, std::vector**) override { + ++decompression_calls; + requested_decompression_bytes = bytes; + return Status::Error( + "test rejects PRX decompression allocation"); + } + + size_t csr_calls = 0; + size_t decompression_calls = 0; + size_t requested_decompression_bytes = 0; +}; + +template +concept HasEncodedBytes = requires(T value) { value.encoded_bytes; }; +template +concept HasPayloadBytes = requires(T value) { value.payload_bytes; }; +template +concept HasCompressedBytes = requires(T value) { value.compressed_bytes; }; +template +concept HasChild128Touches = requires(T value) { value.child_128_touches; }; +template +concept HasChild256Touches = requires(T value) { value.child_256_touches; }; +template +concept HasCrcValidationNs = requires(T value) { value.crc_validation_ns; }; +template +concept HasDecompressNs = requires(T value) { value.decompress_ns; }; +template +concept HasScratchAllocationEvents = requires(T value) { value.scratch_allocation_events; }; +template +concept HasContainerAllocationEvents = requires(T value) { value.container_allocation_events; }; +template +concept HasTraceRecords = requires(T value) { value.trace_records; }; +template +concept CanReadPrxCsr = requires(Args... args) { read_prx_window_csr(args...); }; +template +concept CanReadPrxCsrSelective = requires(Args... args) { read_prx_window_csr_selective(args...); }; + +static_assert(!HasEncodedBytes); +static_assert(!HasPayloadBytes); +static_assert(!HasCompressedBytes); +static_assert(!HasChild128Touches); +static_assert(!HasChild256Touches); +static_assert(!HasCrcValidationNs); +static_assert(!HasDecompressNs); +static_assert(!HasScratchAllocationEvents); +static_assert(!HasContainerAllocationEvents); +static_assert(!HasTraceRecords); + +TEST(PrxPodTest, ProfiledDecodeInterfacesDoNotExposeDeferredState) { + using FullDecodeFn = Status (*)(ByteSource*, std::vector*, std::vector*, + PrxDecodeContext*); + using SelectiveDecodeFn = + Status (*)(ByteSource*, std::span, std::vector*, + std::vector*, PrxDecodeContext*); + + const auto full_decode = static_cast(&read_prx_window_csr); + const auto selective_decode = static_cast(&read_prx_window_csr_selective); + EXPECT_NE(full_decode, nullptr); + EXPECT_NE(selective_decode, nullptr); + + using Flat = std::vector; + static_assert(!CanReadPrxCsr); + static_assert(!CanReadPrxCsr); + static_assert(!CanReadPrxCsrSelective, Flat*, Flat*, + PrxDecodeContext*, std::nullptr_t>); + static_assert(!CanReadPrxCsrSelective, Flat*, Flat*, + PrxDecodeContext*, std::nullptr_t, std::nullptr_t>); +} + +TEST(PrxPodTest, AllocationGateRejectsZstdBeforeDecompressionAndCsrMutation) { + const PerDoc docs = MakeProfileDocs(32); + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, /*zstd_level_or_negative_for_auto=*/3, &sink).ok()); + + RejectingPrxAllocationGate gate; + PrxDecodeContext context {.allocation_gate = &gate}; + std::vector flat = {7, 11}; + std::vector offsets = {0, 2}; + const std::vector flat_before = flat; + const std::vector offsets_before = offsets; + ByteSource source(sink.view()); + + const Status status = read_prx_window_csr(&source, &flat, &offsets, &context); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(gate.decompression_calls, 1); + EXPECT_EQ(gate.csr_calls, 0); + EXPECT_GT(gate.requested_decompression_bytes, 0); + EXPECT_EQ(flat, flat_before); + EXPECT_EQ(offsets, offsets_before); +} + +// Encodes a uint32 array as the PFOR_runs wire form the prx PFOR payload uses: +// fixed-size runs of kFrqBaseUnit values, run count derived by the decoder from +// the total length (so it is not stored). Mirrors the in-source +// encode_pfor_runs so a test can hand-assemble a CRC-VALID PFOR payload and +// exercise the INNER decode branches (sum check / trailing bytes) instead of +// the outer CRC/codec. +void AppendPforRuns(const std::vector& values, ByteSink* out) { + const size_t n = values.size(); + for (size_t off = 0; off < n; off += kFrqBaseUnit) { + const size_t run = (n - off < kFrqBaseUnit) ? (n - off) : kFrqBaseUnit; + pfor_encode(values.data() + off, run, out); + } +} + +// Assembles a complete CRC-valid PFOR window: codec, uncomp_len, payload, crc. +// The payload header fields (doc_count, total_pos) are passed explicitly so a +// test can declare values that disagree with the encoded runs, and `trailing` +// lets a test pad the declared payload with undecodable bytes. The CRC always +// matches the emitted frame, so reads fail (if at all) on an INNER payload +// check, never CRC. +std::vector MakePforWindow(uint32_t doc_count, uint32_t total_pos, + const std::vector& freqs, + const std::vector& deltas, + const std::vector& trailing = {}) { + ByteSink payload; + payload.put_varint32(doc_count); + payload.put_varint32(total_pos); + AppendPforRuns(freqs, &payload); + AppendPforRuns(deltas, &payload); + payload.put_bytes(Slice(trailing)); + + ByteSink framed; + framed.put_u8(static_cast(PrxCodec::kPfor)); + framed.put_varint32(static_cast(payload.view().size())); + framed.put_bytes(payload.view()); + + ByteSink full; + full.put_bytes(framed.view()); + full.put_fixed32(doris::snii::crc32c(framed.view())); + return full.buffer(); +} + +std::vector MakeRawWindowWithPositionDeltas(const std::vector& deltas) { + ByteSink payload; + payload.put_varint32(1); + payload.put_varint32(static_cast(deltas.size())); + for (uint32_t delta : deltas) { + payload.put_varint32(delta); + } + + ByteSink framed; + framed.put_u8(static_cast(PrxCodec::kRaw)); + framed.put_varint32(static_cast(payload.size())); + framed.put_bytes(payload.view()); + const uint32_t crc = doris::snii::crc32c(framed.view()); + framed.put_fixed32(crc); + return framed.take(); +} + +// Flattens per-doc lists into (flat positions, freqs) the way the accumulator +// stores them, so the flat builder can be checked for byte-identity. +void Flatten(const PerDoc& in, std::vector* flat, std::vector* freqs) { + flat->clear(); + freqs->clear(); + for (const auto& doc : in) { + freqs->push_back(static_cast(doc.size())); + flat->insert(flat->end(), doc.begin(), doc.end()); + } +} + +// Build data above/below the raw threshold consistent with the production path; +// provides a stable, controlled round-trip helper. +Status RoundTrip(const PerDoc& in, int level, PerDoc* out) { + ByteSink sink; + RETURN_IF_ERROR(build_prx_window(in, level, &sink)); + ByteSource src(sink.view()); + RETURN_IF_ERROR(read_prx_window(&src, out)); + return Status::OK(); +} + +PerDoc MakeProfileDocs(uint32_t count) { + PerDoc docs; + docs.reserve(count); + for (uint32_t doc = 0; doc < count; ++doc) { + std::vector positions; + for (uint32_t i = 0; i < doc % 3 + 1; ++i) { + positions.push_back(doc * 7 + i * 2); + } + docs.push_back(std::move(positions)); + } + return docs; +} + +PerDoc MakeBoundaryProfileDocs() { + PerDoc docs(257); + docs[0] = {1}; + docs[127] = {3, 7}; + docs[128] = {5}; + docs[255] = {11, 13}; + docs[256] = {17}; + return docs; +} + +uint64_t SelectedPositionCount(const PerDoc& docs, const std::vector& ordinals) { + return std::accumulate( + ordinals.begin(), ordinals.end(), uint64_t {0}, + [&](uint64_t total, uint32_t ordinal) { return total + docs[ordinal].size(); }); +} + +std::vector MakePforWindow(const PerDoc& docs, const std::vector& trailing = {}) { + std::vector freqs; + std::vector deltas; + freqs.reserve(docs.size()); + for (const auto& positions : docs) { + freqs.push_back(static_cast(positions.size())); + uint32_t previous = 0; + for (size_t i = 0; i < positions.size(); ++i) { + deltas.push_back(i == 0 ? positions[i] : positions[i] - previous); + previous = positions[i]; + } + } + return MakePforWindow(static_cast(docs.size()), static_cast(deltas.size()), + freqs, deltas, trailing); +} + +std::vector BuildProfileFrame(const PerDoc& docs, PrxCodec codec) { + if (codec == PrxCodec::kPfor) { + return MakePforWindow(docs); + } + ByteSink sink; + EXPECT_TRUE(build_prx_window(docs, codec == PrxCodec::kRaw ? 0 : 3, &sink).ok()); + return sink.buffer(); +} + +void ExpectedSelectiveCsr(const PerDoc& docs, const std::vector& ordinals, + std::vector* flat, std::vector* offsets) { + flat->clear(); + offsets->clear(); + offsets->push_back(0); + for (uint32_t ordinal : ordinals) { + flat->insert(flat->end(), docs[ordinal].begin(), docs[ordinal].end()); + offsets->push_back(static_cast(flat->size())); + } +} + +void RewriteTrailingCrc(std::vector* frame) { + DCHECK_GE(frame->size(), sizeof(uint32_t)); + const size_t crc_offset = frame->size() - sizeof(uint32_t); + const uint32_t crc = doris::snii::crc32c(Slice(frame->data(), crc_offset)); + (*frame)[crc_offset] = static_cast(crc); + (*frame)[crc_offset + 1] = static_cast(crc >> 8); + (*frame)[crc_offset + 2] = static_cast(crc >> 16); + (*frame)[crc_offset + 3] = static_cast(crc >> 24); +} + +} // namespace + +TEST(PrxPodTest, StatsAggregateCodecsPlaintextAndSelectionBoundaries) { + const PerDoc docs = MakeBoundaryProfileDocs(); + const std::vector ordinals = {0, 127, 128, 255, 256}; + const uint64_t total_positions = std::accumulate( + docs.begin(), docs.end(), uint64_t {0}, + [](uint64_t total, const auto& positions) { return total + positions.size(); }); + + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + for (int level : {0, 3, -1}) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, level, &sink).ok()); + std::vector pos_flat; + std::vector pos_off; + ByteSource source(sink.view()); + ASSERT_TRUE(read_prx_window_csr_selective(&source, ordinals, &pos_flat, &pos_off, &context) + .ok()); + } + + EXPECT_EQ(stats.raw_frames, 1U); + EXPECT_EQ(stats.zstd_frames, 1U); + EXPECT_EQ(stats.pfor_frames, 1U); + EXPECT_EQ(stats.total_docs, docs.size() * 3); + EXPECT_EQ(stats.selected_docs, ordinals.size() * 3); + EXPECT_EQ(stats.total_positions, total_positions * 3); + EXPECT_EQ(stats.selected_positions, SelectedPositionCount(docs, ordinals) * 3); + EXPECT_GT(stats.plaintext_bytes, 0U); + EXPECT_TRUE(stats.is_valid()); +} + +TEST(PrxPodTest, StatsMergeEmptyAndMultipleFrames) { + PrxDecodeStats total; + const PrxDecodeStats empty; + total.merge(empty); + EXPECT_EQ(total, PrxDecodeStats {}); + + PerDoc docs = {{1, 3}, {4}}; + for (int level : {0, -1}) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, level, &sink).ok()); + PrxDecodeStats frame; + PrxDecodeContext context {.stats = &frame}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&source, &flat, &offsets, &context).ok()); + total.merge(frame); + } + EXPECT_EQ(total.raw_frames, 1U); + EXPECT_EQ(total.pfor_frames, 1U); + EXPECT_EQ(total.total_docs, 4U); + EXPECT_EQ(total.selected_docs, 4U); + EXPECT_EQ(total.total_positions, 6U); + EXPECT_EQ(total.selected_positions, 6U); + EXPECT_TRUE(total.is_valid()); +} + +TEST(PrxPodTest, StatsUseFixedPlaintextMeaningByCodec) { + const PerDoc docs = MakeBoundaryProfileDocs(); + for (int level : {0, 3, -1}) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, level, &sink).ok()); + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&source, &flat, &offsets, &context).ok()); + + ByteSource header(sink.view()); + uint8_t codec = 0; + uint32_t plaintext_length = 0; + ASSERT_TRUE(header.get_u8(&codec).ok()); + ASSERT_TRUE(header.get_varint32(&plaintext_length).ok()); + EXPECT_EQ(stats.plaintext_bytes, plaintext_length); + EXPECT_GT(stats.plaintext_bytes, 0U); + + if (level == 0) { + EXPECT_EQ(stats.raw_frames, 1U); + EXPECT_EQ(codec, static_cast(PrxCodec::kRaw)); + } else if (level == 3) { + EXPECT_EQ(stats.zstd_frames, 1U); + EXPECT_EQ(codec, static_cast(PrxCodec::kZstd)); + } else { + EXPECT_EQ(stats.pfor_frames, 1U); + EXPECT_EQ(codec, static_cast(PrxCodec::kPfor)); + } + } + + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, -1, &sink).ok()); + const std::vector ordinals = {0, 128, 256}; + PrxDecodeStats selective_stats; + PrxDecodeContext selective_context {.stats = &selective_stats}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view()); + ASSERT_TRUE( + read_prx_window_csr_selective(&source, ordinals, &flat, &offsets, &selective_context) + .ok()); + EXPECT_EQ(selective_stats.pfor_frames, 1U); + EXPECT_EQ(selective_stats.selected_docs, ordinals.size()); +} + +TEST(PrxPodTest, StatsRemainTransactionalOnCorruptionAndInvalidOrdinals) { + const PerDoc docs = MakeProfileDocs(32); + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, 3, &sink).ok()); + + PrxDecodeStats stats; + stats.raw_frames = 7; + const PrxDecodeStats before = stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + + for (size_t trim : {size_t {1}, size_t {4}, sink.size() - 1}) { + ByteSource truncated(sink.view().subslice(0, sink.size() - trim)); + EXPECT_FALSE(read_prx_window_csr(&truncated, &flat, &offsets, &context).ok()); + EXPECT_EQ(stats, before); + } + + std::vector crc_corrupt = sink.buffer(); + crc_corrupt.back() ^= 0x80; + ByteSource corrupt_source((Slice(crc_corrupt))); + EXPECT_FALSE(read_prx_window_csr(&corrupt_source, &flat, &offsets, &context).ok()); + EXPECT_EQ(stats, before); + + const std::vector duplicate = {1, 1}; + ByteSource duplicate_source(sink.view()); + EXPECT_TRUE( + read_prx_window_csr_selective(&duplicate_source, duplicate, &flat, &offsets, &context) + .is()); + EXPECT_EQ(stats, before); + + const std::vector out_of_order = {2, 1}; + ByteSource out_of_order_source(sink.view()); + EXPECT_TRUE(read_prx_window_csr_selective(&out_of_order_source, out_of_order, &flat, &offsets, + &context) + .is()); + EXPECT_EQ(stats, before); + + const std::vector out_of_range = {static_cast(docs.size())}; + ByteSource out_of_range_source(sink.view()); + EXPECT_TRUE(read_prx_window_csr_selective(&out_of_range_source, out_of_range, &flat, &offsets, + &context) + .is()); + EXPECT_EQ(stats, before); + + const std::vector invalid_payload = + MakePforWindow(/*doc_count=*/2, /*total_pos=*/1, /*freqs=*/ {1, 1}, /*deltas=*/ {1}); + ByteSource invalid_payload_source((Slice(invalid_payload))); + EXPECT_TRUE(read_prx_window_csr(&invalid_payload_source, &flat, &offsets, &context) + .is()); + EXPECT_EQ(stats, before); +} + +TEST(PrxPodTest, CodecSelectionMatrixMatchesCsrOracleAtLowAndHighTf) { + for (uint32_t positions_per_doc : {1U, 17U}) { + PerDoc docs(256); + for (uint32_t doc = 0; doc < docs.size(); ++doc) { + for (uint32_t pos = 0; pos < positions_per_doc; ++pos) { + docs[doc].push_back(doc * 100 + pos * 3); + } + } + std::vector half; + std::vector all; + for (uint32_t doc = 0; doc < docs.size(); ++doc) { + all.push_back(doc); + if (doc % 2 == 0) { + half.push_back(doc); + } + } + const std::vector> selections = { + {}, {128}, {0, 17, 127, 128, 255}, half, all}; + + for (PrxCodec codec : {PrxCodec::kRaw, PrxCodec::kZstd, PrxCodec::kPfor}) { + const std::vector encoded = BuildProfileFrame(docs, codec); + SCOPED_TRACE("tf=" + std::to_string(positions_per_doc) + + " codec=" + std::to_string(static_cast(codec))); + + std::vector expected_full_flat; + std::vector expected_full_offsets; + ExpectedSelectiveCsr(docs, all, &expected_full_flat, &expected_full_offsets); + std::vector actual_flat; + std::vector actual_offsets; + PrxDecodeStats full_stats; + PrxDecodeContext full_context {.stats = &full_stats}; + ByteSource full_source((Slice(encoded))); + ASSERT_TRUE( + read_prx_window_csr(&full_source, &actual_flat, &actual_offsets, &full_context) + .ok()); + EXPECT_EQ(actual_flat, expected_full_flat); + EXPECT_EQ(actual_offsets, expected_full_offsets); + EXPECT_EQ(full_stats.total_docs, docs.size()); + EXPECT_EQ(full_stats.selected_docs, docs.size()); + EXPECT_EQ(full_stats.total_positions, expected_full_flat.size()); + EXPECT_EQ(full_stats.selected_positions, expected_full_flat.size()); + + for (const auto& ordinals : selections) { + SCOPED_TRACE("selection=" + std::to_string(ordinals.size())); + std::vector expected_flat; + std::vector expected_offsets; + ExpectedSelectiveCsr(docs, ordinals, &expected_flat, &expected_offsets); + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + ByteSource source((Slice(encoded))); + ASSERT_TRUE(read_prx_window_csr_selective(&source, ordinals, &actual_flat, + &actual_offsets, &context) + .ok()); + EXPECT_EQ(actual_flat, expected_flat); + EXPECT_EQ(actual_offsets, expected_offsets); + EXPECT_EQ(stats.total_docs, docs.size()); + EXPECT_EQ(stats.selected_docs, ordinals.size()); + EXPECT_EQ(stats.total_positions, expected_full_flat.size()); + EXPECT_EQ(stats.selected_positions, expected_flat.size()); + EXPECT_TRUE(stats.is_valid()); + } + } + } +} + +TEST(PrxPodTest, ZstdDeclaredLengthMismatchCommitsNoStats) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(PerDoc {{1, 3}, {7}}, 3, &sink).ok()); + std::vector encoded = sink.buffer(); + ASSERT_EQ(encoded[0], static_cast(PrxCodec::kZstd)); + ASSERT_LT(encoded[1], 0x7f); + ++encoded[1]; + RewriteTrailingCrc(&encoded); + + PrxDecodeStats stats; + stats.raw_frames = 11; + const PrxDecodeStats before = stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + ByteSource source((Slice(encoded))); + EXPECT_TRUE(read_prx_window_csr(&source, &flat, &offsets, &context) + .is()); + EXPECT_EQ(stats, before); +} + +TEST(PrxPodTest, ZstdHeaderCompletePayloadTruncationCommitsNoStats) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(MakeProfileDocs(32), 3, &sink).ok()); + ASSERT_EQ(sink.buffer().front(), static_cast(PrxCodec::kZstd)); + ASSERT_GT(sink.size(), 5U); + + PrxDecodeStats stats; + stats.raw_frames = 11; + const PrxDecodeStats before = stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view().subslice(0, sink.size() - 5)); + EXPECT_FALSE(read_prx_window_csr(&source, &flat, &offsets, &context).ok()); + EXPECT_EQ(stats, before); +} + +TEST(PrxPodTest, NullContextDoesNotReadClock) { + const PerDoc docs = MakeProfileDocs(32); + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, 3, &sink).ok()); + std::vector flat; + std::vector offsets; + + doris::snii::format::testing::reset_prx_clock_read_count(); + ByteSource legacy_source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&legacy_source, &flat, &offsets).ok()); + EXPECT_EQ(doris::snii::format::testing::prx_clock_read_count(), 0U); + + ByteSource explicit_null_source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&explicit_null_source, &flat, &offsets, nullptr).ok()); + EXPECT_EQ(doris::snii::format::testing::prx_clock_read_count(), 0U); + + PrxDecodeStats stats; + PrxDecodeContext stats_context {.stats = &stats}; + ByteSource stats_source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&stats_source, &flat, &offsets, &stats_context).ok()); + EXPECT_GT(doris::snii::format::testing::prx_clock_read_count(), 0U); +} + +TEST(PrxPodTest, InclusiveDecodeTimeCoversSuccessfulDecode) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(PerDoc {}, 3, &sink).ok()); + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&source, &flat, &offsets, &context).ok()); + EXPECT_GT(stats.decode_ns, 0U); +} + +TEST(PrxPodTest, ProfiledDecodeReadsClockTwiceForEveryCodec) { + const PerDoc docs = MakeProfileDocs(32); + DEFER(doris::snii::format::testing::reset_prx_clock_read_count()); + for (int level : {0, 3, -1}) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, level, &sink).ok()); + const auto expected_codec = level == 0 ? PrxCodec::kRaw + : level == 3 ? PrxCodec::kZstd + : PrxCodec::kPfor; + ASSERT_EQ(sink.buffer().front(), static_cast(expected_codec)); + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + doris::snii::format::testing::reset_prx_clock_read_count(); + ByteSource source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&source, &flat, &offsets, &context).ok()); + EXPECT_EQ(doris::snii::format::testing::prx_clock_read_count(), 2U) << level; + } +} + +TEST(PrxPodTest, DenseFullDecodeReportsLogicalSelection) { + const PerDoc docs = {{1}, {2, 4}, {3}, {5, 7, 9}, {}, {11}, {13, 15}, {17}}; + const std::vector ordinals = {1, 3, 6}; + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, -1, &sink).ok()); + + PrxDecodeStats stats; + PrxDecodeContext context {.stats = &stats}; + std::vector flat; + std::vector offsets; + ByteSource source(sink.view()); + ASSERT_TRUE( + read_prx_window_csr_for_selection(&source, ordinals, &flat, &offsets, &context).ok()); + + std::vector expected_flat; + std::vector expected_offsets; + std::vector all_ordinals(docs.size()); + std::iota(all_ordinals.begin(), all_ordinals.end(), 0U); + ExpectedSelectiveCsr(docs, all_ordinals, &expected_flat, &expected_offsets); + EXPECT_EQ(flat, expected_flat); + EXPECT_EQ(offsets, expected_offsets); + EXPECT_EQ(stats.total_docs, docs.size()); + EXPECT_EQ(stats.selected_docs, ordinals.size()); + EXPECT_EQ(stats.total_positions, expected_flat.size()); + EXPECT_EQ(stats.selected_positions, SelectedPositionCount(docs, ordinals)); +} + +TEST(PrxPodTest, StatsDistinguishFullFromEmptySelectiveSelection) { + const PerDoc docs = MakeProfileDocs(257); + ByteSink sink; + ASSERT_TRUE(build_prx_window(docs, 0, &sink).ok()); + + PrxDecodeStats full_stats; + PrxDecodeContext full_context {.stats = &full_stats}; + std::vector flat; + std::vector offsets; + ByteSource full_source(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&full_source, &flat, &offsets, &full_context).ok()); + EXPECT_EQ(full_stats.selected_docs, docs.size()); + + PrxDecodeStats empty_stats; + PrxDecodeContext empty_context {.stats = &empty_stats}; + const std::vector no_ordinals; + ByteSource empty_source(sink.view()); + ASSERT_TRUE(read_prx_window_csr_selective(&empty_source, no_ordinals, &flat, &offsets, + &empty_context) + .ok()); + EXPECT_EQ(empty_stats.total_docs, docs.size()); + EXPECT_EQ(empty_stats.selected_docs, 0U); + EXPECT_EQ(empty_stats.selected_positions, 0U); +} + +// Single doc with ascending positions: delta-encoded round-trip must be +// lossless. +TEST(SniiPrxPod, SingleDocRoundTrip) { + PerDoc in = {{3, 7, 7, 10, 100}}; // includes duplicate positions (delta=0) + PerDoc out; + ASSERT_TRUE(RoundTrip(in, -1, &out).ok()); + EXPECT_EQ(out, in); +} + +// The FLAT builder (used by the writer to avoid materializing a +// vector-of-vectors for high-df terms) must produce BYTE-IDENTICAL window bytes +// to the per-doc builder for the same logical positions, at every codec level. +// This is the load-bearing guarantee that the flat refactor keeps .prx +// byte-identical. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPrxPod, FlatBuilderMatchesPerDocBytes) { + const std::vector cases = { + {}, // 0 docs + {{}, {3}, {}, {}, {1, 2}}, // empty docs interleaved + {{0, 5, 12}, {1}, {2, 2, 9, 9, 40}}, // duplicate positions (delta 0) + {{3, 7, 7, 10, 100}}, // single doc + }; + for (const auto& in : cases) { + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + for (int level : {-1, 0, 3}) { + ByteSink per_doc_sink, flat_sink; + ASSERT_TRUE(build_prx_window(in, level, &per_doc_sink).ok()); + ASSERT_TRUE(build_prx_window_flat(flat, freqs, level, &flat_sink).ok()); + const Slice a = per_doc_sink.view(); + const Slice b = flat_sink.view(); + ASSERT_EQ(a.size(), b.size()) << "level=" << level; + EXPECT_EQ(0, std::memcmp(a.data(), b.data(), a.size())) << "level=" << level; + // The flat-built window still decodes back to the original per-doc lists. + PerDoc out; + ByteSource src(flat_sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in) << "level=" << level; + } + } +} + +// A large window built flat (auto codec) must equal the per-doc build. The auto +// path picks the smaller of PFOR vs zstd; for these deltas it compresses (not +// raw), proving the flat path is byte-identical to the per-doc path through the +// same auto codec. +TEST(SniiPrxPod, FlatBuilderMatchesPerDocLargePfor) { + PerDoc in; + for (uint32_t d = 0; d < 300; ++d) { + in.push_back({d, d + 1U, d + 2U}); + } + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink per_doc_sink, flat_sink; + ASSERT_TRUE(build_prx_window(in, -1, &per_doc_sink).ok()); + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &flat_sink).ok()); + const Slice a = per_doc_sink.view(); + const Slice b = flat_sink.view(); + ASSERT_EQ(a.size(), b.size()); + EXPECT_EQ(0, std::memcmp(a.data(), b.data(), a.size())); + EXPECT_NE(a.data()[0], static_cast(doris::snii::format::PrxCodec::kRaw)); + // Round-trips losslessly back to the per-doc lists. + PerDoc out; + ByteSource src(flat_sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in); +} + +TEST(SniiPrxPod, CsrPforRoundTripMatchesPerDoc) { + PerDoc in; + for (uint32_t d = 0; d < 300; ++d) { + if (d % 7 == 0) { + in.emplace_back(); + } else if (d % 5 == 0) { + in.push_back({d, d + 2U, d + 9U}); + } else { + in.push_back({d + 1U}); + } + } + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + + std::vector pos_flat = {999U}; + std::vector pos_off = {777U}; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&src, &pos_flat, &pos_off).ok()); + + ASSERT_EQ(pos_off.size(), in.size() + 1); + EXPECT_EQ(pos_off.front(), 0U); + EXPECT_EQ(pos_off.back(), pos_flat.size()); + PerDoc got; + got.reserve(in.size()); + for (size_t i = 0; i < in.size(); ++i) { + got.emplace_back(pos_flat.begin() + pos_off[i], pos_flat.begin() + pos_off[i + 1]); + } + EXPECT_EQ(got, in); +} + +TEST(SniiPrxPod, CsrDecodeReportsShapeAndRejectsPositionOverflow) { + const PerDoc docs = {{1, 3, 9}, {2}, {4, 8}}; + ByteSink valid; + ASSERT_TRUE(build_prx_window(docs, /*zstd_level_or_negative_for_auto=*/0, &valid).ok()); + PrxDecodedShape shape; + PrxDecodeContext context {.shape = &shape}; + std::vector positions; + std::vector offsets; + ByteSource valid_source(valid.view()); + ASSERT_TRUE(read_prx_window_csr(&valid_source, &positions, &offsets, &context).ok()); + EXPECT_EQ(shape.total_docs, 3U); + EXPECT_EQ(shape.total_positions, 6U); + EXPECT_EQ(shape.max_frequency, 3U); + + const std::vector overflow_bytes = + MakeRawWindowWithPositionDeltas({std::numeric_limits::max(), 1}); + ByteSource overflow_source((Slice(overflow_bytes))); + const Status overflow = read_prx_window_csr(&overflow_source, &positions, &offsets, &context); + EXPECT_TRUE(overflow.is()) << overflow; +} + +TEST(SniiPrxPod, SelectiveCsrPforReturnsOnlyRequestedDocs) { + PerDoc in; + for (uint32_t d = 0; d < 320; ++d) { + if (d % 11 == 0) { + in.emplace_back(); + } else if (d % 7 == 0) { + in.push_back({d, d + 3U, d + 8U}); + } else { + in.push_back({d + 1U}); + } + } + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + + const std::vector ordinals = {0, 1, 7, 77, 255, 319}; + std::vector pos_flat = {999U}; + std::vector pos_off = {777U}; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window_csr_selective(&src, ordinals, &pos_flat, &pos_off).ok()); + + ASSERT_EQ(pos_off.size(), ordinals.size() + 1); + EXPECT_EQ(pos_off.front(), 0U); + EXPECT_EQ(pos_off.back(), pos_flat.size()); + PerDoc got; + got.reserve(ordinals.size()); + for (size_t i = 0; i < ordinals.size(); ++i) { + got.emplace_back(pos_flat.begin() + pos_off[i], pos_flat.begin() + pos_off[i + 1]); + } + PerDoc want; + for (uint32_t ordinal : ordinals) { + want.push_back(in[ordinal]); + } + EXPECT_EQ(got, want); +} + +TEST(SniiPrxPod, SelectiveCsrRejectsInvalidOrdinals) { + PerDoc in = {{1}, {2}, {3}}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + + std::vector pos_flat; + std::vector pos_off; + { + const std::vector unsorted = {1, 1}; + ByteSource src(sink.view()); + EXPECT_TRUE(read_prx_window_csr_selective(&src, unsorted, &pos_flat, &pos_off) + .is()); + } + { + const std::vector out_of_range = {3}; + ByteSource src(sink.view()); + EXPECT_TRUE(read_prx_window_csr_selective(&src, out_of_range, &pos_flat, &pos_off) + .is()); + } +} + +// Multiple docs, positions ascending within each doc, no shared baseline across +// docs. +TEST(SniiPrxPod, MultiDocRoundTrip) { + PerDoc in = {{0, 5, 12}, {1}, {2, 2, 9, 9, 40}, {7, 8, 9, 10}}; + PerDoc out; + ASSERT_TRUE(RoundTrip(in, -1, &out).ok()); + EXPECT_EQ(out, in); +} + +// Empty positions: supports both 0 docs and 0 positions within a doc. +TEST(SniiPrxPod, EmptyPositions) { + PerDoc empty_window; // 0 docs + PerDoc out1; + ASSERT_TRUE(RoundTrip(empty_window, -1, &out1).ok()); + EXPECT_EQ(out1, empty_window); + + PerDoc with_empty_docs = {{}, {3}, {}, {}, {1, 2}}; // contains empty docs + PerDoc out2; + ASSERT_TRUE(RoundTrip(with_empty_docs, -1, &out2).ok()); + EXPECT_EQ(out2, with_empty_docs); +} + +// A small multi-position window stays on PFOR: only the proven +// single-doc/freq=1 shape bypasses PFOR, avoiding a general small-window policy +// that could trade import CPU for bytes without workload evidence. +TEST(SniiPrxPod, SmallWindowUsesPfor) { + PerDoc in = {{1, 2, 3}, {4, 5}}; + ByteSink sink; + ASSERT_TRUE(build_prx_window(in, -1, &sink).ok()); + + ByteSource src(sink.view()); + uint8_t codec = 0xFF; + ASSERT_TRUE(src.get_u8(&codec).ok()); + EXPECT_EQ(codec, static_cast(doris::snii::format::PrxCodec::kPfor)); + + PerDoc out; + ByteSource rt(sink.view()); + ASSERT_TRUE(read_prx_window(&rt, &out).ok()); + EXPECT_EQ(out, in); +} + +// Auto mode must account for complete frame overhead, not assume PFOR wins for +// every sub-512-byte payload. For a singleton posting RAW needs only three +// payload varints, while PFOR adds two run headers; selecting RAW saves four or +// five bytes without changing the wire format or reader path. +TEST(SniiPrxPod, AutoSelectsSmallerRawFrameForSingleton) { + for (uint32_t position : {0U, 127U, 128U, 0xFFFFFFFFU}) { + SCOPED_TRACE("position=" + std::to_string(position)); + PerDoc in = {{position}}; + ByteSink raw_sink; + ASSERT_TRUE(build_prx_window(in, /*level=*/0, &raw_sink).ok()); + const std::vector pfor_frame = MakePforWindow(in); + ASSERT_LT(raw_sink.size(), pfor_frame.size()); + EXPECT_TRUE(pfor_frame.size() - raw_sink.size() == 4U || + pfor_frame.size() - raw_sink.size() == 5U); + + doris::snii::testing::reset_pfor_width_evals(); + doris::snii::format::testing::reset_prx_delta_materialization_count(); + ByteSink auto_sink; + ASSERT_TRUE(build_prx_window(in, /*level=*/-1, &auto_sink).ok()); + + ASSERT_EQ(raw_sink.view().data()[0], static_cast(PrxCodec::kRaw)); + EXPECT_EQ(auto_sink.view().data()[0], static_cast(PrxCodec::kRaw)); + EXPECT_EQ(auto_sink.buffer(), raw_sink.buffer()); + EXPECT_EQ(doris::snii::testing::pfor_width_evals(), 0U); + EXPECT_EQ(doris::snii::format::testing::prx_delta_materialization_count(), 0U); + + PerDoc out; + ByteSource src(auto_sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in); + } +} + +// Large window (all-1 deltas) auto-encodes as PFOR; for tiny constant deltas +// the 1-bit-packed PFOR payload is much smaller than the forced raw varint +// encoding. +TEST(SniiPrxPod, LargeWindowTriggersPforAndIsSmaller) { + PerDoc in; + in.reserve(64); + for (int d = 0; d < 64; ++d) { + std::vector doc; + uint32_t p = 0; + for (int i = 0; i < 256; ++i) { + p += 1; // all-1 deltas: pack to ~1 bit each + doc.push_back(p); + } + in.push_back(std::move(doc)); + } + + ByteSink auto_sink; + ASSERT_TRUE(build_prx_window(in, -1, &auto_sink).ok()); + ByteSource probe(auto_sink.view()); + uint8_t codec = 0xFF; + ASSERT_TRUE(probe.get_u8(&codec).ok()); + // Auto path compresses with the smaller of PFOR vs zstd; assert it is not raw. + EXPECT_NE(codec, static_cast(doris::snii::format::PrxCodec::kRaw)); + + ByteSink raw_sink; + ASSERT_TRUE(build_prx_window(in, /*level=*/0, &raw_sink).ok()); + ByteSource raw_probe(raw_sink.view()); + uint8_t raw_codec = 0xFF; + ASSERT_TRUE(raw_probe.get_u8(&raw_codec).ok()); + EXPECT_EQ(raw_codec, static_cast(doris::snii::format::PrxCodec::kRaw)); + + EXPECT_LT(auto_sink.size(), raw_sink.size()); + + // The PFOR path restores losslessly. + PerDoc out; + ByteSource z(auto_sink.view()); + ASSERT_TRUE(read_prx_window(&z, &out).ok()); + EXPECT_EQ(out, in); +} + +// PFOR round-trips arbitrary multi-doc windows (including empty docs, duplicate +// positions, and large jumps that force PFOR exceptions) losslessly. +TEST(SniiPrxPod, PforRoundTripVariety) { + const std::vector cases = { + {{0, 5, 12}, {1}, {2, 2, 9, 9, 40}, {7, 8, 9, 10}}, + {{}, {3}, {}, {}, {1, 2}}, + {{0, 1000000, 1000001}, {5}}, // large jump => PFOR exception + {{3, 7, 7, 10, 100}}, // duplicate positions (delta 0) + }; + for (const auto& in : cases) { + ByteSink sink; + ASSERT_TRUE(build_prx_window(in, -1, &sink).ok()); + ByteSource probe(sink.view()); + uint8_t codec = 0xFF; + ASSERT_TRUE(probe.get_u8(&codec).ok()); + EXPECT_EQ(codec, static_cast(doris::snii::format::PrxCodec::kPfor)); + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in); + } +} + +// level=0 explicitly forces raw (no compression), useful for testing and the +// oversized single-doc degenerate path. +TEST(SniiPrxPod, ExplicitRawLevelRoundTrip) { + PerDoc in = {{10, 20, 30}, {40}}; + PerDoc out; + ASSERT_TRUE(RoundTrip(in, /*level=*/0, &out).ok()); + EXPECT_EQ(out, in); +} + +// CRC corruption is detectable: flipping any byte in the payload causes read to +// return Corruption. +TEST(SniiPrxPod, CrcCorruptionDetected) { + PerDoc in = {{1, 2, 3, 4, 5}, {6, 7, 8}}; + ByteSink sink; + ASSERT_TRUE(build_prx_window(in, -1, &sink).ok()); + + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 1U); + bytes.back() ^= 0xFF; // corrupt the last byte of the payload + + Slice corrupt(bytes); + ByteSource src(corrupt); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_FALSE(s.ok()); + EXPECT_TRUE(s.is()); +} + +// A corrupted codec byte (invalid codec value) should also be rejected. +TEST(SniiPrxPod, InvalidCodecRejected) { + PerDoc in = {{1, 2}}; + ByteSink sink; + ASSERT_TRUE(build_prx_window(in, -1, &sink).ok()); + + std::vector bytes = sink.buffer(); + bytes[0] = 0x7F; // invalid codec (bits 0-5 exceed known values) + + Slice corrupt(bytes); + ByteSource src(corrupt); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_FALSE(s.ok()); +} + +// Truncated input (insufficient payload) should return an error rather than +// crash. +TEST(SniiPrxPod, TruncatedInputRejected) { + PerDoc in = {{1, 2, 3, 4, 5, 6, 7, 8}}; + ByteSink sink; + ASSERT_TRUE(build_prx_window(in, -1, &sink).ok()); + + std::vector bytes = sink.buffer(); + bytes.resize(bytes.size() / 2); // truncate to half + + Slice truncated(bytes); + ByteSource src(truncated); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_FALSE(s.ok()); +} + +// DoS prevention: an uncomp_len corrupted to a huge value must be rejected +// before allocation/decompression. +TEST(SniiPrxPod, OversizedUncompLenRejected) { + ByteSink sink; + sink.put_u8(static_cast(doris::snii::format::PrxCodec::kZstd)); + sink.put_varint32(300U * 1024 * 1024); // > 256MiB window limit + // No need to construct the subsequent comp_len/payload/crc — the cap check + // triggers immediately after reading uncomp_len. + ByteSource src(sink.view()); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_TRUE(s.is()); +} + +// DoS prevention: a CRC-VALID raw frame whose decoded doc_count is absurd must +// return Corruption, not a giant reserve()/assign() -> std::bad_alloc. Distinct +// from the CRC and uncomp_len tests above, which are caught BEFORE the inner +// doc_count is read. +TEST(SniiPrxPod, OversizedDocCountRejected) { + ByteSink payload; + payload.put_varint32(0x02000000U); // doc_count = 33M, > kMaxWindowDocs (1<<24) + ByteSink framed; + framed.put_u8(static_cast(doris::snii::format::PrxCodec::kRaw)); + framed.put_varint32(static_cast(payload.view().size())); // uncomp_len + framed.put_bytes(payload.view()); + ByteSink full; + full.put_bytes(framed.view()); + full.put_fixed32(doris::snii::crc32c(framed.view())); // valid crc over codec+uncomp_len+payload + ByteSource src(full.view()); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// Sanity: the MakePforWindow helper builds a window the production reader +// accepts, so the negative variants below isolate the ONE field they tamper +// with (not a helper bug). doc_count=2, total_pos=3, freqs sum to total_pos, +// deltas match. +TEST(SniiPrxPod, CraftedPforWindowRoundTrips) { + // doc0 = {5, 6} (deltas 5,1), doc1 = {9} (delta 9); pos_counts {2,1}, + // total 3. + auto bytes = MakePforWindow(/*doc_count=*/2, /*total_pos=*/3, + /*freqs=*/ {2, 1}, /*deltas=*/ {5, 1, 9}); + Slice s(bytes); + ByteSource src(s); + PerDoc out; + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + PerDoc expected = {{5, 6}, {9}}; + EXPECT_EQ(out, expected); +} + +// CRC-VALID PFOR frame whose declared total_pos disagrees with sum(pos_counts): +// the reader must reach the INNER "pos_count sum mismatch" branch (after the +// cap checks and after decoding the freqs runs) and return Corruption -- NOT a +// CRC or codec rejection. freqs sum to 3 but total_pos is declared 4. +TEST(SniiPrxPod, PforPosCountSumMismatchRejected) { + auto bytes = MakePforWindow(/*doc_count=*/2, /*total_pos=*/4, + /*freqs=*/ {2, 1}, /*deltas=*/ {5, 1, 9}); + Slice slice(bytes); + ByteSource src(slice); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// CRC-VALID PFOR frame with extra bytes appended INSIDE the declared payload +// (uncomp_len covers them, so the CRC is over them too). After +// decode_pfor_payload consumes the real doc_count/total_pos/freqs/deltas, the +// source is not at EOF, so the INNER "trailing bytes after pfor payload" branch +// must fire -> Corruption. +TEST(SniiPrxPod, PforTrailingBytesRejected) { + auto bytes = MakePforWindow(/*doc_count=*/2, /*total_pos=*/3, + /*freqs=*/ {2, 1}, /*deltas=*/ {5, 1, 9}, + /*trailing=*/ {0xAA, 0xBB, 0xCC}); + Slice slice(bytes); + ByteSource src(slice); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// FLAT builder precondition: a (positions_flat, freqs) mismatch where +// sum(freqs) overruns positions_flat would index flat[off+i] past the span end. +// The guard must return a CLEAN InvalidArgument (never OOB / UB) at every flat +// codec level. Here freqs sum to 5 but only 3 positions are supplied. +TEST(SniiPrxPod, FlatBuilderShortPositionsRejectedNotOob) { + std::vector positions_flat = {1, 2, 3}; // 3 positions + std::vector freqs = {2, 3}; // claims 5 positions + for (int level : {-1, 0, 3}) { // pfor, raw, zstd flat encoders all guard + ByteSink sink; + Status s = build_prx_window_flat(positions_flat, freqs, level, &sink); + EXPECT_TRUE(s.is()) + << "level=" << level << " msg=" << s.to_string(); + EXPECT_EQ(sink.size(), 0U) << "level=" << level; // nothing emitted on reject + } +} + +// FLAT builder precondition (other direction): sum(freqs) < positions_flat +// leaves trailing positions unused -- also a writer-side mismatch. The guard +// requires EXACT equality, so this is rejected too (a silently-dropped position +// is a bug). +TEST(SniiPrxPod, FlatBuilderExtraPositionsRejected) { + std::vector positions_flat = {1, 2, 3, 4, 5}; // 5 positions + std::vector freqs = {2, 1}; // only addresses 3 + for (int level : {-1, 0, 3}) { + ByteSink sink; + Status s = build_prx_window_flat(positions_flat, freqs, level, &sink); + EXPECT_TRUE(s.is()) + << "level=" << level << " msg=" << s.to_string(); + } +} + +// --------------------------------------------------------------------------- +// T14: auto-mode single-encode -- the auto builders derive per-doc deltas once, +// directly emit a provably smaller singleton RAW frame, and only materialize a +// throwaway raw plaintext payload when it is large enough (>= 512 B) for zstd. +// The tests below cover the deterministic raw-build counter seam plus builder +// byte/codec equivalence. +// --------------------------------------------------------------------------- + +namespace { + +using doris::snii::format::testing::prx_raw_build_count; +using doris::snii::format::testing::reset_prx_raw_build_count; + +// Byte length of put_varint32(v); mirrors the in-source varint32_size so the +// brute-force codec recomputation below can size frames exactly. +size_t VarintSize(uint32_t v) { + size_t n = 1; + while (v >= 128) { + v >>= 7; + ++n; + } + return n; +} + +// Auto constants restated from prx_pod.cpp (anonymous there). The brute-force +// codec check must use the SAME threshold/level write_auto_pfor_or_zstd uses. +constexpr size_t kAutoZstdMinBytes = 512; +constexpr int kDefaultZstdLevel = 3; + +// Per-doc -> flat delta stream (first position of each doc absolute, the rest +// delta-within-doc): the exact sequence both .prx payload encoders serialize. +std::vector FlatDeltas(const std::vector& flat, + const std::vector& freqs) { + std::vector deltas; + deltas.reserve(flat.size()); + size_t off = 0; + for (uint32_t fc : freqs) { + uint32_t prev = 0; + for (uint32_t i = 0; i < fc; ++i) { + const uint32_t pos = flat[off + i]; + deltas.push_back(i == 0 ? pos : pos - prev); + prev = pos; + } + off += fc; + } + return deltas; +} + +// Independently recompute the codec byte the auto builder must emit, replicating +// build_prx_window_auto_from_flat for the threshold cases below: try zstd only +// when the EXACT plaintext payload size reaches the threshold and choose the +// smallest complete frame among already-materialized PFOR/ZSTD/RAW candidates. +uint8_t BruteForceAutoCodec(const std::vector& freqs, + const std::vector& deltas) { + ByteSink pfor_payload; + pfor_payload.put_varint32(static_cast(freqs.size())); + pfor_payload.put_varint32(static_cast(deltas.size())); + AppendPforRuns(freqs, &pfor_payload); + AppendPforRuns(deltas, &pfor_payload); + const size_t pfor_size = pfor_payload.view().size(); + + ByteSink plain; + plain.put_varint32(static_cast(freqs.size())); + size_t off = 0; + for (uint32_t fc : freqs) { + plain.put_varint32(fc); + for (uint32_t i = 0; i < fc; ++i) { + plain.put_varint32(deltas[off + i]); + } + off += fc; + } + const size_t plain_size = plain.view().size(); + + const auto kPfor = static_cast(PrxCodec::kPfor); + const auto kZstd = static_cast(PrxCodec::kZstd); + const auto kRaw = static_cast(PrxCodec::kRaw); + const size_t pfor_frame = + 1 + VarintSize(static_cast(pfor_size)) + pfor_size + sizeof(uint32_t); + size_t selected_frame = pfor_frame; + uint8_t selected_codec = kPfor; + if (plain_size >= kAutoZstdMinBytes) { + std::vector comp; + EXPECT_TRUE(doris::snii::zstd_compress(plain.view(), kDefaultZstdLevel, &comp).ok()); + const size_t zstd_frame = 1 + VarintSize(static_cast(plain_size)) + + VarintSize(static_cast(comp.size())) + comp.size() + + sizeof(uint32_t); + if (zstd_frame < selected_frame) { + selected_frame = zstd_frame; + selected_codec = kZstd; + } + } + if (plain_size >= kAutoZstdMinBytes) { + const size_t raw_frame = + 1 + VarintSize(static_cast(plain_size)) + plain_size + sizeof(uint32_t); + if (raw_frame < selected_frame) { + selected_codec = kRaw; + } + } + return selected_codec; +} + +// One doc whose raw plaintext payload is EXACTLY `target` bytes. Positions +// {0,1,...,fc-1} make every delta a 1-byte varint (0 then 1s), so the plaintext +// is varint32_size(doc_count=1) + varint32_size(fc) + fc. With fc in +// [128, 16383] the fc varint is 2 bytes, so plaintext == 3 + fc => fc = target-3. +std::vector SingleDocWithPlaintextSize(size_t target) { + EXPECT_GE(target, 131U); // need fc >= 128 so the count varint is 2 bytes + const auto fc = static_cast(target - 3); + std::vector doc(fc); + for (uint32_t i = 0; i < fc; ++i) { + doc[i] = i; + } + return doc; +} + +} // namespace + +// [perf-deterministic] A sub-threshold window (<512 B plaintext) must NOT +// materialize the throwaway raw plaintext payload: the auto path emits PFOR +// directly. On the pre-optimization code (which always materialized) this counter +// would be 1 -> this is the RED guard for the single-encode change. +TEST(SniiPrxPodCounterTest, SmallWindowSkipsRawPlaintextBuild) { + PerDoc in = {{1}, {3, 5}, {2}}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + + reset_prx_raw_build_count(); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + EXPECT_EQ(prx_raw_build_count(), 0U); + + // codec is PFOR and the window still round-trips. + EXPECT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in); +} + +// [perf-deterministic] A >=512 B window materializes the raw plaintext EXACTLY +// once (needed for the zstd-vs-pfor comparison) -- not zero (would skip a viable +// zstd) and not twice (the old double-encode). +TEST(SniiPrxPodCounterTest, LargeWindowStillBuildsRawPlaintextOnce) { + PerDoc in = {SingleDocWithPlaintextSize(512)}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + + reset_prx_raw_build_count(); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + EXPECT_EQ(prx_raw_build_count(), 1U); +} + +// [perf-deterministic] N consecutive small windows materialize ZERO raw +// plaintext payloads in aggregate (the per-window throwaway ByteSink is the +// allocation this task removes for the common Zipfian case). +TEST(SniiPrxPodCounterTest, ManySmallWindowsBuildZeroRawPlaintext) { + reset_prx_raw_build_count(); + for (uint32_t w = 0; w < 64; ++w) { + PerDoc in = {{w}, {w + 1U, w + 3U}, {w + 2U}}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + } + EXPECT_EQ(prx_raw_build_count(), 0U); +} + +// [perf-deterministic] The per-doc (vector) builder funnels through the SAME auto +// path as the flat builder: small window materializes nothing, large window once. +TEST(SniiPrxPodCounterTest, VectorBuilderSharesSingleEncodePath) { + reset_prx_raw_build_count(); + PerDoc small = {{1, 2, 3}, {4, 5}}; + ByteSink small_sink; + ASSERT_TRUE(build_prx_window(small, -1, &small_sink).ok()); + EXPECT_EQ(prx_raw_build_count(), 0U); + + reset_prx_raw_build_count(); + PerDoc large = {SingleDocWithPlaintextSize(512)}; + ByteSink large_sink; + ASSERT_TRUE(build_prx_window(large, -1, &large_sink).ok()); + EXPECT_EQ(prx_raw_build_count(), 1U); +} + +// [functional/perf] The flat and per-doc auto builders must produce BYTE-IDENTICAL +// windows for the same logical positions, across sizes spanning the 512-byte zstd +// threshold (empty, single, tiny, just-below, exactly-at, well-above). This is the +// load-bearing proof that both entry points apply the same auto-codec policy. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPrxPodTest, FlatAutoProducesByteIdenticalOutputAcrossSizes) { + std::vector cases; + cases.emplace_back(); // empty window + cases.push_back({{7}}); // single doc/pos + cases.push_back({{1}, {3, 5}, {2}}); // tiny (<512 plaintext) + cases.push_back({{}, {3}, {}, {}, {1, 2}}); // empty docs interleaved + cases.push_back({SingleDocWithPlaintextSize(511)}); // just below threshold + cases.push_back({SingleDocWithPlaintextSize(512)}); // exactly at threshold + { + PerDoc big; // well above threshold + for (uint32_t d = 0; d < 280; ++d) { + big.push_back({d, d + 1U, d + 2U}); + } + cases.push_back(std::move(big)); + } + + for (const auto& in : cases) { + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink per_doc_sink, flat_sink; + ASSERT_TRUE(build_prx_window(in, -1, &per_doc_sink).ok()); + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &flat_sink).ok()); + const Slice a = per_doc_sink.view(); + const Slice b = flat_sink.view(); + ASSERT_EQ(a.size(), b.size()) << "doc_count=" << in.size(); + EXPECT_EQ(0, std::memcmp(a.data(), b.data(), a.size())) << "doc_count=" << in.size(); + // Both entry points must make the same codec decision, including the + // singleton RAW fast path. + ASSERT_GT(a.size(), 0U); + // The flat-built window round-trips back to the original per-doc lists. + PerDoc out; + ByteSource src(flat_sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in) << "doc_count=" << in.size(); + } +} + +// [functional/perf] At the EXACT 511- vs 512-byte plaintext boundary the codec +// the builder emits must equal an independent brute-force frame-size comparison. +// This guards the risk that an estimated (rather than exact) plaintext size flips +// the codec choice near the threshold and silently changes the on-disk bytes. +TEST(SniiPrxPodTest, AutoCodecChoiceMatchesBruteForceAtThreshold) { + for (size_t target : {size_t {511}, size_t {512}}) { + PerDoc in = {SingleDocWithPlaintextSize(target)}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + + // Confirm the construction lands on exactly the intended plaintext size. + const std::vector deltas = FlatDeltas(flat, freqs); + size_t plain_size = VarintSize(static_cast(freqs.size())); + for (uint32_t fc : freqs) { + plain_size += VarintSize(fc); + } + for (uint32_t d : deltas) { + plain_size += VarintSize(d); + } + ASSERT_EQ(plain_size, target); + + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + const uint8_t emitted = sink.view().data()[0]; + EXPECT_EQ(emitted, BruteForceAutoCodec(freqs, deltas)) << "target=" << target; + if (target < kAutoZstdMinBytes) { + // Below the threshold zstd is never attempted: must be plain PFOR. + EXPECT_EQ(emitted, static_cast(PrxCodec::kPfor)); + } + // Round-trips regardless of the chosen codec. + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in) << "target=" << target; + } +} + +// Once the >=512-byte path has already materialized RAW for zstd, RAW joins the +// exact complete-frame comparison at no extra encoding cost. It must win only +// when strictly smaller, preserving PFOR on an equal-size frame. +TEST(SniiPrxPodTest, MaterializedRawCandidateMustStrictlyReduceFrame) { + constexpr uint32_t kLimit = 1024; + EXPECT_EQ(doris::snii::format::testing::select_auto_prx_codec_for_test( + /*pfor_payload_size=*/520, /*plain_payload_size=*/512, + /*compressed_payload_size=*/600, kLimit), + static_cast(PrxCodec::kRaw)); + EXPECT_EQ(doris::snii::format::testing::select_auto_prx_codec_for_test( + /*pfor_payload_size=*/512, /*plain_payload_size=*/512, + /*compressed_payload_size=*/600, kLimit), + static_cast(PrxCodec::kPfor)); +} + +// [functional] F-RT-small: a small flat window round-trips through the CSR reader +// (codec=pfor path) back to the original per-doc positions. +TEST(SniiPrxPodTest, SmallFlatWindowRoundTripsViaCsr) { + PerDoc in = {{1}, {3, 5}, {2}}; + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + EXPECT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + + std::vector pos_flat, pos_off; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window_csr(&src, &pos_flat, &pos_off).ok()); + ASSERT_EQ(pos_off.size(), in.size() + 1); + PerDoc out; + for (size_t d = 0; d < in.size(); ++d) { + out.emplace_back(pos_flat.begin() + pos_off[d], pos_flat.begin() + pos_off[d + 1]); + } + EXPECT_EQ(out, in); +} + +// [functional] F-RT-large: a >=512 B flat window (exercising the zstd-vs-pfor +// pick) round-trips losslessly through the per-doc reader. +TEST(SniiPrxPodTest, LargeFlatWindowRoundTrips) { + PerDoc in; + for (uint32_t d = 0; d < 280; ++d) { + in.push_back({d, d + 4U, d + 9U}); + } + std::vector flat, freqs; + Flatten(in, &flat, &freqs); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + EXPECT_NE(sink.view().data()[0], static_cast(PrxCodec::kRaw)); + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_EQ(out, in); +} + +// [functional/boundary] F-EMPTY: a 0-doc window builds OK, materializes no raw +// plaintext, decodes to an empty result. +TEST(SniiPrxPodTest, EmptyFlatWindowRoundTripsWithNoRawBuild) { + std::vector flat, freqs; // 0 docs, 0 positions + reset_prx_raw_build_count(); + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + EXPECT_EQ(prx_raw_build_count(), 0U); + EXPECT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + EXPECT_TRUE(out.empty()); +} + +// [functional/boundary] F-SINGLE: a single doc with a single position round-trips. +TEST(SniiPrxPodTest, SingleDocSinglePositionFlatRoundTrips) { + std::vector flat = {42}; + std::vector freqs = {1}; + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + PerDoc expected = {{42}}; + EXPECT_EQ(out, expected); +} + +// [error] F-ERR-asc: descending positions within a doc are rejected by the +// single delta walk -- the ascending check is not lost when the second encode is +// skipped. Nothing is emitted. +TEST(SniiPrxPodTest, NonAscendingPositionsRejected) { + std::vector flat = {5, 3}; // within one doc, descending + std::vector freqs = {2}; + ByteSink sink; + Status s = build_prx_window_flat(flat, freqs, -1, &sink); + EXPECT_TRUE(s.is()) << s.to_string(); + EXPECT_EQ(sink.size(), 0U); +} + +// [error] F-ERR-part: a (flat, freqs) partition mismatch (sum(freqs) != size) is +// rejected before any indexing -- the partition check is preserved. +TEST(SniiPrxPodTest, PartitionMismatchRejected) { + std::vector flat = {1, 2, 3}; + std::vector freqs = {2, 3}; // sum 5 != 3 + ByteSink sink; + Status s = build_prx_window_flat(flat, freqs, -1, &sink); + EXPECT_TRUE(s.is()) << s.to_string(); + EXPECT_EQ(sink.size(), 0U); +} + +// [error] F-NULL: a null sink is rejected by both auto builders before any work. +TEST(SniiPrxPodTest, NullSinkRejected) { + std::vector flat = {1, 2}; + std::vector freqs = {2}; + Status s = build_prx_window_flat(flat, freqs, -1, nullptr); + EXPECT_TRUE(s.is()) << s.to_string(); + + PerDoc per_doc = {{1, 2}}; + Status s2 = build_prx_window(per_doc, -1, nullptr); + EXPECT_TRUE(s2.is()) << s2.to_string(); +} + +// =========================================================================== +// T22 -- single-copy framing (write_pfor / write_raw / write_zstd_compressed + +// SectionFramer::write) and the bitpack capacity reserve. The framing refactor +// writes [codec/type][len][payload] DIRECTLY into the caller's sink and crc's +// that span in place (no temp ByteSink), so these tests pin the on-disk bytes to +// an INDEPENDENT hand-assembly of the wire frame. A wrong crc range, a view() +// taken at the wrong time, or a wrong start offset shows up as a byte or Status +// mismatch. All bytes MUST stay identical to the pre-refactor output. +// =========================================================================== + +namespace { + +using doris::snii::FramedSection; +using doris::snii::pfor_decode; +using doris::snii::SectionFramer; + +// Appends the 4-byte crc32c(prefix) tail to `prefix`, returning the full on-disk +// frame. Mirrors the reader's crc coverage: the crc is over +// [codec/type][len][payload] ONLY (the bytes already in prefix), never itself. +std::vector WithCrc(const ByteSink& prefix) { + ByteSink full; + full.put_bytes(prefix.view()); + full.put_fixed32(doris::snii::crc32c(prefix.view())); + return full.buffer(); +} + +// The RAW/ZSTD plaintext payload (encode_payload_flat wire form): VInt doc_count, +// then per doc VInt pos_count followed by that doc's position deltas (VInt). +ByteSink RawPlaintext(const std::vector& flat, const std::vector& freqs) { + const std::vector deltas = FlatDeltas(flat, freqs); + ByteSink plain; + plain.put_varint32(static_cast(freqs.size())); + size_t off = 0; + for (uint32_t fc : freqs) { + plain.put_varint32(fc); + for (uint32_t i = 0; i < fc; ++i) { + plain.put_varint32(deltas[off + i]); + } + off += fc; + } + return plain; +} + +// Deterministic pseudo-random uint32 run whose values fit in `w` bits, with a +// couple of exact-`w`-bit values forced so choose_width sees width w. Round-trip +// must be lossless regardless of the width PFOR actually picks. +std::vector MakeWidthRun(uint8_t w, size_t n) { + std::vector run(n, 0U); + uint32_t mask = 0; + if (w >= 32) { + mask = 0xFFFFFFFFU; + } else if (w != 0) { + mask = (1U << w) - 1U; + } + uint32_t state = 0x12345678U ^ (static_cast(w) << 24) ^ static_cast(n); + for (size_t i = 0; i < n; ++i) { + state = state * 1664525U + 1013904223U; // LCG + run[i] = state & mask; + } + if (n > 0 && w > 0) { + run[0] = mask; // force a full-width value + run[n / 2] = mask; // and one mid-run (exercises exceptions if a smaller w is picked) + } + return run; +} + +} // namespace + +// [perf-deterministic] PRX-BYTE-PFOR: an auto (<512 B) window emits PFOR and its +// on-disk bytes equal an independent [kPfor][len][payload][crc] hand-assembly, +// proving write_pfor's single-copy framing is byte-identical. +TEST(SniiPrxPodTest, PforFramingProducesByteIdenticalOutput) { + std::vector flat = {1, 2, 3, 4, 5}; + std::vector freqs = {3, 2}; + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, -1, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kPfor)); + + const std::vector deltas = FlatDeltas(flat, freqs); + ByteSink payload; + payload.put_varint32(static_cast(freqs.size())); + payload.put_varint32(static_cast(deltas.size())); + AppendPforRuns(freqs, &payload); + AppendPforRuns(deltas, &payload); + ByteSink framed; + framed.put_u8(static_cast(PrxCodec::kPfor)); + framed.put_varint32(static_cast(payload.view().size())); + framed.put_bytes(payload.view()); + + EXPECT_EQ(sink.buffer(), WithCrc(framed)); +} + +// [perf-deterministic] PRX-BYTE-RAW: level=0 forces raw; the bytes equal an +// independent [kRaw][len][plaintext][crc] assembly (write_raw single-copy). +TEST(SniiPrxPodTest, RawFramingProducesByteIdenticalOutput) { + std::vector flat = {10, 20, 30, 40}; + std::vector freqs = {3, 1}; + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, /*level=*/0, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kRaw)); + + ByteSink plain = RawPlaintext(flat, freqs); + ByteSink framed; + framed.put_u8(static_cast(PrxCodec::kRaw)); + framed.put_varint32(static_cast(plain.view().size())); + framed.put_bytes(plain.view()); + + EXPECT_EQ(sink.buffer(), WithCrc(framed)); +} + +// [perf-deterministic] PRX-BYTE-ZSTD: level>0 forces zstd; the bytes equal an +// independent [kZstd][uncomp_len][comp_len][zstd(plaintext)][crc] assembly +// (write_zstd_compressed single-copy). zstd is deterministic at a fixed level. +TEST(SniiPrxPodTest, ZstdFramingProducesByteIdenticalOutput) { + std::vector flat, freqs; + for (uint32_t d = 0; d < 200; ++d) { + flat.push_back(d); + freqs.push_back(1); + } + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, /*level=*/3, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kZstd)); + + ByteSink plain = RawPlaintext(flat, freqs); + std::vector comp; + ASSERT_TRUE(doris::snii::zstd_compress(plain.view(), /*level=*/3, &comp).ok()); + ByteSink framed; + framed.put_u8(static_cast(PrxCodec::kZstd)); + framed.put_varint32(static_cast(plain.view().size())); + framed.put_varint32(static_cast(comp.size())); + framed.put_bytes(Slice(comp)); + + EXPECT_EQ(sink.buffer(), WithCrc(framed)); + + // Round-trips losslessly back to the original per-doc positions. + PerDoc out; + ByteSource src(sink.view()); + ASSERT_TRUE(read_prx_window(&src, &out).ok()); + ASSERT_EQ(out.size(), flat.size()); + for (uint32_t d = 0; d < flat.size(); ++d) { + ASSERT_EQ(out[d].size(), 1U); + EXPECT_EQ(out[d][0], flat[d]); + } +} + +// [perf-deterministic] Two windows appended to ONE sink: the second window's crc +// must cover only its own [codec][len][payload] (start offset != 0), and each +// window's bytes must be identical to building it standalone. This is the guard +// that write_raw / write_pfor use `start = sink->size()` (not absolute 0). +TEST(SniiPrxPodTest, MultipleWindowsInOneSinkAreFramedByteIdentical) { + std::vector flat_a = {1, 2, 3}, freqs_a = {2, 1}; // -> raw + std::vector flat_b = {5, 6}, freqs_b = {1, 1}; // -> pfor + ByteSink combined; + ASSERT_TRUE(build_prx_window_flat(flat_a, freqs_a, /*level=*/0, &combined).ok()); + const size_t boundary = combined.size(); + ASSERT_TRUE(build_prx_window_flat(flat_b, freqs_b, -1, &combined).ok()); + + ByteSink only_a, only_b; + ASSERT_TRUE(build_prx_window_flat(flat_a, freqs_a, /*level=*/0, &only_a).ok()); + ASSERT_TRUE(build_prx_window_flat(flat_b, freqs_b, -1, &only_b).ok()); + ASSERT_EQ(boundary, only_a.size()); + ASSERT_EQ(combined.size(), only_a.size() + only_b.size()); + const std::vector& c = combined.buffer(); + EXPECT_EQ(0, std::memcmp(c.data(), only_a.buffer().data(), only_a.size())); + EXPECT_EQ(0, std::memcmp(c.data() + boundary, only_b.buffer().data(), only_b.size())); + + // Both windows read back sequentially (crc verified per window). + ByteSource src(combined.view()); + PerDoc a, b; + ASSERT_TRUE(read_prx_window(&src, &a).ok()); + ASSERT_TRUE(read_prx_window(&src, &b).ok()); + EXPECT_TRUE(src.eof()); + PerDoc want_a = {{1, 2}, {3}}; + PerDoc want_b = {{5}, {6}}; + EXPECT_EQ(a, want_a); + EXPECT_EQ(b, want_b); +} + +// [functional/error] PRX-CRC-CORRUPT: flipping a byte inside the raw payload (the +// crc-covered [codec][len][payload] span) is detected as Corruption -- confirms +// write_raw's crc range is exactly [codec][len][payload]. +TEST(SniiPrxPodTest, RawFrameCrcCorruptionDetected) { + std::vector flat = {10, 20, 30, 40}; + std::vector freqs = {3, 1}; + ByteSink sink; + ASSERT_TRUE(build_prx_window_flat(flat, freqs, /*level=*/0, &sink).ok()); + ASSERT_EQ(sink.view().data()[0], static_cast(PrxCodec::kRaw)); + + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 5U); + bytes[bytes.size() - 5] ^= 0xFF; // last payload byte, just before the 4-byte crc + ByteSource src((Slice(bytes))); + PerDoc out; + Status s = read_prx_window(&src, &out); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// [perf-deterministic] SF-BYTE: SectionFramer::write bytes equal an independent +// [type][varint64 len][payload][crc] hand-assembly (single-copy framing). +TEST(SniiPrxPodTest, SectionFramerProducesByteIdenticalOutput) { + const std::vector payload_bytes = {0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x11, 0x22}; + const uint8_t type = 7; + ByteSink sink; + SectionFramer::write(sink, type, Slice(payload_bytes)); + + ByteSink framed; + framed.put_u8(type); + framed.put_varint64(payload_bytes.size()); + framed.put_bytes(Slice(payload_bytes)); + EXPECT_EQ(sink.buffer(), WithCrc(framed)); +} + +// [perf-deterministic] SectionFramer::write into a NON-empty sink frames ONLY its +// own region: the crc covers [type][len][payload] starting at `start`, not the +// pre-existing prefix. Guards the start-offset in the single-copy rewrite. +TEST(SniiPrxPodTest, SectionFramerFramesOnlyItsOwnRegionInNonEmptySink) { + const std::vector prefix = {0xAA, 0xBB}; + const std::vector payload_bytes = {0x01, 0x02, 0x03, 0x04}; + ByteSink sink; + sink.put_bytes(Slice(prefix)); + SectionFramer::write(sink, /*type=*/5, Slice(payload_bytes)); + + ByteSink framed; + framed.put_u8(5); + framed.put_varint64(payload_bytes.size()); + framed.put_bytes(Slice(payload_bytes)); + const std::vector expected_region = WithCrc(framed); + + const std::vector& got = sink.buffer(); + ASSERT_EQ(got.size(), prefix.size() + expected_region.size()); + EXPECT_EQ(0, std::memcmp(got.data(), prefix.data(), prefix.size())); + EXPECT_EQ(0, std::memcmp(got.data() + prefix.size(), expected_region.data(), + expected_region.size())); +} + +// [functional] SF-RT: two sections written into one sink (second at start != 0) +// both read back with the correct type + payload and the crc verifies. +TEST(SniiPrxPodTest, SectionFramerWriteReadRoundTripAcrossSections) { + const std::vector p0 = {1, 2, 3}; + const std::vector p1 = {9, 8, 7, 6, 5}; + ByteSink sink; + SectionFramer::write(sink, /*type=*/3, Slice(p0)); + SectionFramer::write(sink, /*type=*/9, Slice(p1)); + + ByteSource src(sink.view()); + FramedSection s0, s1; + ASSERT_TRUE(SectionFramer::read(src, &s0).ok()); + ASSERT_TRUE(SectionFramer::read(src, &s1).ok()); + EXPECT_TRUE(src.eof()); + EXPECT_EQ(s0.type, 3); + EXPECT_EQ(s1.type, 9); + ASSERT_EQ(s0.payload.size(), p0.size()); + ASSERT_EQ(s1.payload.size(), p1.size()); + EXPECT_EQ(0, std::memcmp(s0.payload.data(), p0.data(), p0.size())); + EXPECT_EQ(0, std::memcmp(s1.payload.data(), p1.data(), p1.size())); +} + +// [functional/error] A corrupted SectionFramer payload byte is caught by the +// section crc on read (crc range == [type][len][payload]). +TEST(SniiPrxPodTest, SectionFramerCrcCorruptionDetected) { + const std::vector payload_bytes = {0x10, 0x20, 0x30, 0x40, 0x50}; + ByteSink sink; + SectionFramer::write(sink, /*type=*/4, Slice(payload_bytes)); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 5U); + bytes[bytes.size() - 5] ^= 0xFF; // last payload byte, before the crc + ByteSource src((Slice(bytes))); + FramedSection out; + Status s = SectionFramer::read(src, &out); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// [functional/boundary] BITPACK-RT: pfor_encode -> pfor_decode round-trips for +// every value width w in [0,32] at run lengths n = 1, 255, 256 (the kFrqBaseUnit +// boundary). Proves the bitpack reserve (capacity-only) leaves the packed bytes +// bit-exact and the decode path recovers every value. +TEST(SniiPrxPodTest, PforBitpackRoundTripAcrossWidths) { + for (uint8_t w = 0; w <= 32; ++w) { + for (size_t n : {size_t {1}, size_t {255}, size_t {256}}) { + const std::vector run = MakeWidthRun(w, n); + ByteSink sink; + pfor_encode(run.data(), n, &sink); + std::vector decoded(n, 0xFFFFFFFFU); + ByteSource src(sink.view()); + ASSERT_TRUE(pfor_decode(&src, n, decoded.data()).ok()) + << "w=" << static_cast(w) << " n=" << n; + EXPECT_TRUE(src.eof()) << "w=" << static_cast(w) << " n=" << n; + EXPECT_EQ(decoded, run) << "w=" << static_cast(w) << " n=" << n; + } + } +} diff --git a/be/test/storage/index/snii/format/sampled_term_index_test.cpp b/be/test/storage/index/snii/format/sampled_term_index_test.cpp new file mode 100644 index 00000000000000..4536a47ddd3311 --- /dev/null +++ b/be/test/storage/index/snii/format/sampled_term_index_test.cpp @@ -0,0 +1,265 @@ +// 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 "storage/index/snii/format/sampled_term_index.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/format_constants.h" + +using namespace doris::snii; // NOLINT +using namespace doris::snii::format; // NOLINT +using doris::Status; + +namespace { + +// Build a SampledTermIndex byte buffer from an ordered set of first_terms. +std::vector BuildIndex(const std::vector& first_terms) { + SampledTermIndexBuilder builder; + for (const auto& t : first_terms) { + builder.add_block_first_term(t); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +// Convenience wrapper to open a reader. +SampledTermIndexReader OpenOrDie(const std::vector& bytes) { + SampledTermIndexReader reader; + Status s = SampledTermIndexReader::open(Slice(bytes), &reader); + EXPECT_TRUE(s.ok()) << s.to_string(); + return reader; +} + +} // namespace + +// Multiple blocks: locate returns the correct ordinal for each first_term. +TEST(SniiSampledTermIndex, LocateExactFirstTermHitsOrdinal) { + const std::vector terms = {"alpha", "delta", "kappa", "omega"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + ASSERT_EQ(reader.n_blocks(), 4U); + + for (uint32_t i = 0; i < terms.size(); ++i) { + bool maybe_present = false; + uint32_t ord = 0xFFFFFFFFU; + ASSERT_TRUE(reader.locate(terms[i], &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present) << "term=" << terms[i]; + EXPECT_EQ(ord, i) << "term=" << terms[i]; + } +} + +// target falls between two first_terms → returns the ordinal of the lower one. +TEST(SniiSampledTermIndex, LocateBetweenReturnsLowerOrdinal) { + const std::vector terms = {"alpha", "delta", "kappa", "omega"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + + bool maybe_present = false; + uint32_t ord = 0; + // "echo" is between "delta"(1) and "kappa"(2) → should land in block 1. + ASSERT_TRUE(reader.locate("echo", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 1U); + + // "beta" is between "alpha"(0) and "delta"(1) → block 0. + ASSERT_TRUE(reader.locate("beta", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 0U); + + // "zzz" > "omega"(3, the LAST block's first term) routes to the last block: the + // last block can hold terms greater than its first term, so locate must return a + // candidate (find_term then confirms). This is a router, not an exact set. + ASSERT_TRUE(reader.locate("zzz", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 3U); +} + +// target < min_term → maybe_present=false (not-present signal). +TEST(SniiSampledTermIndex, LocateBelowMinIsOutOfRange) { + const std::vector terms = {"banana", "cherry", "mango"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + + bool maybe_present = true; + uint32_t ord = 12345; + ASSERT_TRUE(reader.locate("apple", &maybe_present, &ord).ok()); + EXPECT_FALSE(maybe_present); +} + +// target > the last block's first term routes to the LAST block (that block may +// contain terms greater than its first term; find_term confirms presence). The +// router must NOT reject such a target, or present tail-of-last-block terms would +// be falsely reported absent. +TEST(SniiSampledTermIndex, LocateAboveLastFirstTermRoutesToLastBlock) { + const std::vector terms = {"banana", "cherry", "mango"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + + bool maybe_present = false; + uint32_t ord = 0; + ASSERT_TRUE(reader.locate("zebra", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 2U); +} + +// target == min_term / == max_term boundary cases should both hit and be in-range. +TEST(SniiSampledTermIndex, LocateBoundaryTermsInRange) { + const std::vector terms = {"banana", "cherry", "mango"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + + bool maybe_present = false; + uint32_t ord = 0; + ASSERT_TRUE(reader.locate("banana", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 0U); + + ASSERT_TRUE(reader.locate("mango", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 2U); +} + +// Single block: min==max==the only first_term. +TEST(SniiSampledTermIndex, SingleBlock) { + const std::vector terms = {"solo"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + ASSERT_EQ(reader.n_blocks(), 1U); + + bool maybe_present = false; + uint32_t ord = 99; + ASSERT_TRUE(reader.locate("solo", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 0U); + + // "solz" > "solo" routes to the single (last) block: that block may hold terms + // greater than its first term, so locate returns block 0 (find_term confirms). + ASSERT_TRUE(reader.locate("solz", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 0U); + + // "sola" < "solo" → before the first block, so it cannot exist (out of range). + ASSERT_TRUE(reader.locate("sola", &maybe_present, &ord).ok()); + EXPECT_FALSE(maybe_present); +} + +// Terms sharing a long common prefix (verify front coding round-trip correctness). +TEST(SniiSampledTermIndex, SharedPrefixTermsRoundTrip) { + const std::vector terms = {"international", "internationalize", + "internationalized", "internet", "interoperate"}; + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + ASSERT_EQ(reader.n_blocks(), 5U); + + for (uint32_t i = 0; i < terms.size(); ++i) { + bool maybe_present = false; + uint32_t ord = 0; + ASSERT_TRUE(reader.locate(terms[i], &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present) << "term=" << terms[i]; + EXPECT_EQ(ord, i) << "term=" << terms[i]; + } + + // "internationalizes" is between idx2 and idx3 → lower ordinal 2. + bool maybe_present = false; + uint32_t ord = 0; + ASSERT_TRUE(reader.locate("internationalizes", &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 2U); +} + +// Terms containing null bytes / high-bit bytes, compared by unsigned byte order. +TEST(SniiSampledTermIndex, BinarySafeTerms) { + std::vector terms; + terms.emplace_back("\x01\x00z", 3); + terms.emplace_back("\x80\x00", 2); + terms.emplace_back("\xFF", 1); + auto bytes = BuildIndex(terms); + auto reader = OpenOrDie(bytes); + ASSERT_EQ(reader.n_blocks(), 3U); + + bool maybe_present = false; + uint32_t ord = 0; + ASSERT_TRUE(reader.locate(std::string("\x80\x00", 2), &maybe_present, &ord).ok()); + EXPECT_TRUE(maybe_present); + EXPECT_EQ(ord, 1U); +} + +// The first byte is the SectionFramer type and must be kSampledTermIndex. +TEST(SniiSampledTermIndex, FramedAsSampledTermIndexType) { + auto bytes = BuildIndex({"alpha", "beta"}); + ASSERT_GE(bytes.size(), 1U); + EXPECT_EQ(bytes[0], static_cast(SectionType::kSampledTermIndex)); +} + +// CRC checksum: flipping one byte in the payload → open returns Corruption. +TEST(SniiSampledTermIndex, DetectsCorruption) { + auto bytes = BuildIndex({"alpha", "delta", "kappa"}); + ASSERT_GE(bytes.size(), 4U); + bytes[3] ^= 0xFF; // Flip a byte in the payload region (skip the type+len prefix) + SampledTermIndexReader reader; + Status s = SampledTermIndexReader::open(Slice(bytes), &reader); + EXPECT_FALSE(s.ok()) << s.to_string(); +} + +// Truncation: drop the trailing byte → open fails. +TEST(SniiSampledTermIndex, DetectsTruncation) { + auto bytes = BuildIndex({"alpha", "delta", "kappa"}); + bytes.pop_back(); + SampledTermIndexReader reader; + Status s = SampledTermIndexReader::open(Slice(bytes), &reader); + EXPECT_FALSE(s.ok()); +} + +// A wrong section type should be rejected. +TEST(SniiSampledTermIndex, WrongSectionTypeRejected) { + ByteSink sink; + const uint8_t p[] = {0, 0}; + SectionFramer::write(sink, static_cast(SectionType::kDictBlockDirectory), Slice(p, 2)); + SampledTermIndexReader reader; + Status s = SampledTermIndexReader::open(sink.view(), &reader); + EXPECT_FALSE(s.ok()); +} + +TEST(SniiSampledTermIndex, TrailingFramedSectionBytesRejected) { + auto bytes = BuildIndex({"alpha", "delta"}); + bytes.push_back(0xEE); + SampledTermIndexReader reader; + EXPECT_TRUE(SampledTermIndexReader::open(Slice(bytes), &reader) + .is()); +} + +// Empty builder (n_blocks=0): valid build, locate on any term is out of range. +TEST(SniiSampledTermIndex, EmptyIndexLocateOutOfRange) { + auto bytes = BuildIndex({}); + auto reader = OpenOrDie(bytes); + ASSERT_EQ(reader.n_blocks(), 0U); + + bool maybe_present = true; + uint32_t ord = 7; + ASSERT_TRUE(reader.locate("anything", &maybe_present, &ord).ok()); + EXPECT_FALSE(maybe_present); +} diff --git a/be/test/storage/index/snii/format/tail_pointer_test.cpp b/be/test/storage/index/snii/format/tail_pointer_test.cpp new file mode 100644 index 00000000000000..ee624c246af23d --- /dev/null +++ b/be/test/storage/index/snii/format/tail_pointer_test.cpp @@ -0,0 +1,172 @@ +// 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 "storage/index/snii/format/tail_pointer.h" + +#include + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/format/format_constants.h" + +namespace doris::snii::format { +namespace { + +constexpr size_t kExpectedTailPointerSize = 31; +constexpr size_t kTailChecksumOffset = 27; + +TailPointer typical_tail() { + return TailPointer {.directory_offset = 0x0011223344556677ULL, + .directory_length = 0x1020304050607080ULL, + .directory_crc32c = 0xABCD1234U}; +} + +std::vector encode(const TailPointer& tail) { + ByteSink sink; + EXPECT_TRUE(encode_tail_pointer(tail, &sink).ok()); + return sink.buffer(); +} + +void store_fixed32(uint32_t value, uint8_t* destination) { + ByteSink sink; + sink.put_fixed32(value); + std::copy(sink.buffer().begin(), sink.buffer().end(), destination); +} + +void refresh_tail_crc(std::vector* bytes) { + ASSERT_EQ(kExpectedTailPointerSize, bytes->size()); + store_fixed32(crc32c(Slice(bytes->data(), kTailChecksumOffset)), + bytes->data() + kTailChecksumOffset); +} + +TEST(SniiTailPointer, UsesExactV1ThirtyOneByteFieldLayout) { + const auto bytes = encode(typical_tail()); + ASSERT_EQ(kExpectedTailPointerSize, tail_pointer_size()); + ASSERT_EQ(kExpectedTailPointerSize, bytes.size()); + + ByteSource source {Slice(bytes)}; + uint32_t magic = 0; + uint16_t version = 0; + uint64_t directory_offset = 0; + uint64_t directory_length = 0; + uint32_t directory_crc = 0; + uint8_t encoded_size = 0; + uint32_t tail_crc = 0; + ASSERT_TRUE(source.get_fixed32(&magic).ok()); + ASSERT_TRUE(source.get_fixed16(&version).ok()); + ASSERT_TRUE(source.get_fixed64(&directory_offset).ok()); + ASSERT_TRUE(source.get_fixed64(&directory_length).ok()); + ASSERT_TRUE(source.get_fixed32(&directory_crc).ok()); + ASSERT_TRUE(source.get_u8(&encoded_size).ok()); + ASSERT_TRUE(source.get_fixed32(&tail_crc).ok()); + EXPECT_TRUE(source.eof()); + + EXPECT_EQ(kTailMagic, magic); + EXPECT_EQ(1U, version); + EXPECT_EQ(0x0011223344556677ULL, directory_offset); + EXPECT_EQ(0x1020304050607080ULL, directory_length); + EXPECT_EQ(0xABCD1234U, directory_crc); + EXPECT_EQ(kExpectedTailPointerSize, encoded_size); + EXPECT_EQ(crc32c(Slice(bytes.data(), kTailChecksumOffset)), tail_crc); +} + +TEST(SniiTailPointer, RoundTripsThroughTheExactEncodedBytes) { + const auto expected = encode(typical_tail()); + TailPointer decoded; + ASSERT_TRUE(decode_tail_pointer(Slice(expected), &decoded).ok()); + EXPECT_EQ(expected, encode(decoded)); +} + +TEST(SniiTailPointer, NullEncodeAndDecodeOutputsAreInvalidArgument) { + EXPECT_TRUE(encode_tail_pointer(typical_tail(), nullptr).is()); + const auto bytes = encode(typical_tail()); + EXPECT_TRUE(decode_tail_pointer(Slice(bytes), nullptr).is()); +} + +TEST(SniiTailPointer, RejectsWrongVersionAsUnsupportedWithValidCrc) { + auto bytes = encode(typical_tail()); + ASSERT_EQ(kExpectedTailPointerSize, bytes.size()); + bytes[4] = 2; + bytes[5] = 0; + refresh_tail_crc(&bytes); + + TailPointer decoded; + const Status status = decode_tail_pointer(Slice(bytes), &decoded); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiTailPointer, TailCrcCoversEveryPrecedingFieldByte) { + const auto expected = encode(typical_tail()); + ASSERT_EQ(kExpectedTailPointerSize, expected.size()); + for (size_t offset = 0; offset < kTailChecksumOffset; ++offset) { + auto corrupted = expected; + corrupted[offset] ^= 1; + TailPointer decoded; + const Status status = decode_tail_pointer(Slice(corrupted), &decoded); + EXPECT_TRUE(status.is()) + << "unprotected byte offset=" << offset << " status=" << status; + } +} + +TEST(SniiTailPointer, RejectsMagicSizeAndTailCrcCorruption) { + const auto expected = encode(typical_tail()); + ASSERT_EQ(kExpectedTailPointerSize, expected.size()); + + auto bad_magic = expected; + bad_magic[0] ^= 1; + refresh_tail_crc(&bad_magic); + TailPointer decoded; + EXPECT_TRUE(decode_tail_pointer(Slice(bad_magic), &decoded) + .is()); + + auto bad_size = expected; + bad_size[kTailChecksumOffset - 1] = 30; + refresh_tail_crc(&bad_size); + EXPECT_TRUE(decode_tail_pointer(Slice(bad_size), &decoded) + .is()); + + auto bad_crc = expected; + bad_crc.back() ^= 1; + EXPECT_TRUE(decode_tail_pointer(Slice(bad_crc), &decoded) + .is()); +} + +TEST(SniiTailPointer, RejectsEveryTruncationAndTrailingBytes) { + const auto expected = encode(typical_tail()); + ASSERT_EQ(kExpectedTailPointerSize, expected.size()); + TailPointer decoded; + for (size_t size = 0; size < expected.size(); ++size) { + const Status status = decode_tail_pointer(Slice(expected.data(), size), &decoded); + EXPECT_TRUE(status.is()) + << "accepted truncated size=" << size << " status=" << status; + } + + auto longer = expected; + longer.push_back(0); + EXPECT_TRUE(decode_tail_pointer(Slice(longer), &decoded) + .is()); +} + +} // namespace +} // namespace doris::snii::format diff --git a/be/test/storage/index/snii/io/batch_range_fetcher_test.cpp b/be/test/storage/index/snii/io/batch_range_fetcher_test.cpp new file mode 100644 index 00000000000000..faab93419868e5 --- /dev/null +++ b/be/test/storage/index/snii/io/batch_range_fetcher_test.cpp @@ -0,0 +1,127 @@ +// 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 "storage/index/snii/io/batch_range_fetcher.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" + +using namespace doris::snii; +using doris::Status; +using doris::snii::io::BatchRangeFetcher; +using doris::snii::io::LocalFileReader; +using doris::snii::io::LocalFileWriter; +using doris::snii::io::MeteredFileReader; + +namespace { + +std::string MakeRampFile() { + const std::string path = "/tmp/snii_brf_ramp.bin"; + LocalFileWriter w; + EXPECT_TRUE(w.open(path).ok()); + std::vector data(256); + for (int i = 0; i < 256; ++i) { + data[i] = static_cast(i); + } + EXPECT_TRUE(w.append(Slice(data)).ok()); + EXPECT_TRUE(w.finalize().ok()); + return path; +} + +} // namespace + +// Disjoint ranges: each handle returns its own bytes; the whole fetch is one +// serial round on the metered reader. +TEST(SniiBatchRangeFetcher, DisjointRanges) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + BatchRangeFetcher f(&m); + size_t h0 = f.add(0, 4); + size_t h1 = f.add(100, 4); + size_t h2 = f.add(200, 4); + ASSERT_TRUE(f.fetch().ok()); + + EXPECT_EQ(f.get(h0)[0], 0U); + EXPECT_EQ(f.get(h1)[0], 100U); + EXPECT_EQ(f.get(h2)[3], 203U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); // single batched round +} + +// Overlapping requests coalesce into one physical read; bytes still map back. +TEST(SniiBatchRangeFetcher, OverlappingCoalesced) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + BatchRangeFetcher f(&m); + size_t h0 = f.add(0, 4); // [0,4) + size_t h1 = f.add(2, 4); // [2,6) overlaps + size_t h2 = f.add(5, 3); // [5,8) adjacent/overlaps + ASSERT_TRUE(f.fetch().ok()); + + EXPECT_EQ(f.get(h0)[0], 0U); + EXPECT_EQ(f.get(h1)[0], 2U); + EXPECT_EQ(f.get(h2)[0], 5U); + EXPECT_EQ(f.get(h2)[2], 7U); + // Coalesced into a single physical read -> one read_at_call on the metered reader. + EXPECT_EQ(m.metrics().read_at_calls, 1U); +} + +// clear() lets the fetcher be reused for a new round. +TEST(SniiBatchRangeFetcher, ClearAndReuse) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + BatchRangeFetcher f(&m); + f.add(0, 4); + ASSERT_TRUE(f.fetch().ok()); + f.clear(); + EXPECT_EQ(f.pending(), 0U); + size_t h = f.add(64, 8); + ASSERT_TRUE(f.fetch().ok()); + EXPECT_EQ(f.get(h)[0], 64U); +} + +TEST(SniiBatchRangeFetcher, RejectsOverflowingRangeEnd) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + + BatchRangeFetcher f(&inner); + f.add(std::numeric_limits::max() - 1, 8); + const Status st = f.fetch(); + // Integrated fetch() reports range-end overflow as INVERTED_INDEX_FILE_CORRUPTED. + EXPECT_TRUE(st.is()) << st.to_string(); +} + +TEST(SniiBatchRangeFetcher, RejectsNullReaderAtFetch) { + BatchRangeFetcher f(nullptr); + f.add(0, 1); + const Status st = f.fetch(); + EXPECT_TRUE(st.is()) << st.to_string(); +} diff --git a/be/test/storage/index/snii/io/local_file.cpp b/be/test/storage/index/snii/io/local_file.cpp new file mode 100644 index 00000000000000..ff363f70bf6e12 --- /dev/null +++ b/be/test/storage/index/snii/io/local_file.cpp @@ -0,0 +1,134 @@ +// 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 "storage/index/snii/io/local_file.h" + +#include +#include +#include + +#include +#include + +namespace doris::snii::io { +namespace { + +std::string errno_msg(const char* what) { + return std::string(what) + ": " + std::strerror(errno); +} + +} // namespace + +LocalFileReader::~LocalFileReader() { + if (fd_ >= 0) ::close(fd_); +} + +Status LocalFileReader::open(const std::string& path) { + fd_ = ::open(path.c_str(), O_RDONLY); + if (fd_ < 0) return Status::Error(errno_msg("open")); + struct stat st; + if (::fstat(fd_, &st) != 0) + return Status::Error(errno_msg("fstat")); + size_ = static_cast(st.st_size); + return Status::OK(); +} + +Status LocalFileReader::read_at(uint64_t offset, size_t len, std::vector* out) { + if (fd_ < 0) return Status::Error("read_at on unopened file"); + // Non-wrapping bounds check (offset+len could overflow uint64 on a corrupt arg). + if (offset > size_ || len > size_ - offset) { + return Status::Error( + "read_at past end of file"); + } + out->resize(len); + size_t done = 0; + while (done < len) { + ssize_t n = ::pread(fd_, out->data() + done, len - done, static_cast(offset + done)); + if (n < 0) { + if (errno == EINTR) continue; + return Status::Error(errno_msg("pread")); + } + if (n == 0) + return Status::Error( + "pread returned 0 before len"); + done += static_cast(n); + } + return Status::OK(); +} + +LocalFileWriter::~LocalFileWriter() { + if (fd_ >= 0) ::close(fd_); // best-effort: dtor cannot surface a flush error +} + +Status LocalFileWriter::open(const std::string& path) { + fd_ = ::open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); + if (fd_ < 0) return Status::Error(errno_msg("open")); + buf_.reserve(kBufCapacity); + return Status::OK(); +} + +Status LocalFileWriter::write_all(const uint8_t* data, size_t len) { + size_t done = 0; + while (done < len) { + ssize_t n = ::write(fd_, data + done, len - done); + if (n < 0) { + if (errno == EINTR) continue; + return Status::Error(errno_msg("write")); + } + done += static_cast(n); + } + return Status::OK(); +} + +Status LocalFileWriter::flush_buffer() { + if (buf_.empty()) return Status::OK(); + RETURN_IF_ERROR(write_all(buf_.data(), buf_.size())); + buf_.clear(); + return Status::OK(); +} + +Status LocalFileWriter::append(Slice data) { + if (fd_ < 0) return Status::Error("append on unopened file"); + const size_t len = data.size(); + if (len == 0) return Status::OK(); + // Spans larger than the buffer go straight to the fd (after flushing pending + // bytes) to avoid a pointless copy and an oversized buffer. + if (len >= kBufCapacity) { + RETURN_IF_ERROR(flush_buffer()); + RETURN_IF_ERROR(write_all(data.data(), len)); + bytes_written_ += len; + return Status::OK(); + } + if (buf_.size() + len > kBufCapacity) RETURN_IF_ERROR(flush_buffer()); + buf_.insert(buf_.end(), data.data(), data.data() + len); + bytes_written_ += len; + return Status::OK(); +} + +Status LocalFileWriter::finalize() { + if (fd_ < 0) return Status::Error("finalize on unopened file"); + RETURN_IF_ERROR(flush_buffer()); + if (::fsync(fd_) != 0) return Status::Error(errno_msg("fsync")); + if (::close(fd_) != 0) { + fd_ = -1; + return Status::Error(errno_msg("close")); + } + fd_ = -1; + return Status::OK(); +} + +} // namespace doris::snii::io diff --git a/be/test/storage/index/snii/io/local_file.h b/be/test/storage/index/snii/io/local_file.h new file mode 100644 index 00000000000000..aaa26a8fdc681a --- /dev/null +++ b/be/test/storage/index/snii/io/local_file.h @@ -0,0 +1,84 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/file_writer.h" + +namespace doris::snii::io { + +// Local-filesystem FileReader. Uses pread for positional, thread-safe reads +// (so concurrent batch fetches do not contend on a shared file offset). +class LocalFileReader : public FileReader { +public: + LocalFileReader() = default; + ~LocalFileReader() override; + + LocalFileReader(const LocalFileReader&) = delete; + LocalFileReader& operator=(const LocalFileReader&) = delete; + + Status open(const std::string& path); + Status read_at(uint64_t offset, size_t len, std::vector* out) override; + uint64_t size() const override { return size_; } + +private: + int fd_ = -1; + uint64_t size_ = 0; +}; + +// Local-filesystem append-only FileWriter. Appends accumulate in a fixed +// userspace buffer and are flushed to the fd in large chunks, collapsing the +// many tiny per-append ::write() syscalls of the build path (e.g. ~53k writes +// averaging ~683 B each) into a handful of big writes. The produced file is +// byte-identical to the unbuffered path; only the syscall count drops. +class LocalFileWriter : public FileWriter { +public: + LocalFileWriter() = default; + ~LocalFileWriter() override; + + LocalFileWriter(const LocalFileWriter&) = delete; + LocalFileWriter& operator=(const LocalFileWriter&) = delete; + + Status open(const std::string& path); + Status append(Slice data) override; + Status finalize() override; + uint64_t bytes_written() const override { return bytes_written_; } + +private: + // Userspace write buffer size. 256 KiB amortizes the write() syscall cost over + // many appends while keeping transient RAM negligible vs the index sections. + static constexpr size_t kBufCapacity = 256u * 1024; + + // Flushes the userspace buffer to the fd with a robust partial-write loop. + Status flush_buffer(); + // Writes a raw byte span straight to the fd (used for spans larger than the + // buffer, bypassing a needless copy). + Status write_all(const uint8_t* data, size_t len); + + int fd_ = -1; + uint64_t bytes_written_ = 0; + std::vector buf_; +}; + +} // namespace doris::snii::io diff --git a/be/test/storage/index/snii/io/local_file_test.cpp b/be/test/storage/index/snii/io/local_file_test.cpp new file mode 100644 index 00000000000000..4b575c3e7739dc --- /dev/null +++ b/be/test/storage/index/snii/io/local_file_test.cpp @@ -0,0 +1,143 @@ +// 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 "storage/index/snii/io/local_file.h" + +#include + +#include +#include +#include + +#include "common/status.h" + +using namespace doris::snii; +using doris::snii::io::LocalFileReader; +using doris::snii::io::LocalFileWriter; + +namespace { + +std::string TempPath(const char* name) { + return std::string("/tmp/snii_io_test_") + name + ".bin"; +} + +} // namespace + +TEST(SniiLocalFile, AppendThenReadBack) { + const std::string path = TempPath("append"); + { + LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + const uint8_t a[] = {1, 2, 3, 4}; + const uint8_t b[] = {5, 6, 7, 8}; + ASSERT_TRUE(w.append(Slice(a, 4)).ok()); + ASSERT_TRUE(w.append(Slice(b, 4)).ok()); + EXPECT_EQ(w.bytes_written(), 8U); + ASSERT_TRUE(w.finalize().ok()); + } + + LocalFileReader r; + ASSERT_TRUE(r.open(path).ok()); + EXPECT_EQ(r.size(), 8U); + std::vector out; + ASSERT_TRUE(r.read_at(2, 4, &out).ok()); + ASSERT_EQ(out.size(), 4U); + EXPECT_EQ(out[0], 3U); + EXPECT_EQ(out[3], 6U); +} + +TEST(SniiLocalFile, ReadPastEndFails) { + const std::string path = TempPath("past_end"); + { + LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + const uint8_t a[] = {1, 2, 3}; + ASSERT_TRUE(w.append(Slice(a, 3)).ok()); + ASSERT_TRUE(w.finalize().ok()); + } + LocalFileReader r; + ASSERT_TRUE(r.open(path).ok()); + std::vector out; + EXPECT_FALSE(r.read_at(2, 10, &out).ok()); +} + +TEST(SniiLocalFile, OpenMissingFails) { + LocalFileReader r; + EXPECT_FALSE(r.open("/tmp/snii_io_test_does_not_exist_zzz.bin").ok()); +} + +// The buffered writer must produce byte-identical output regardless of how the +// stream is chopped: tiny appends that overflow the 256 KiB buffer, a single +// over-buffer append (the direct-to-fd fast path), and mixes of the two. +TEST(SniiLocalFile, BufferedAppendsByteIdentical) { + const std::string path = TempPath("buffered"); + // Build a deterministic reference stream larger than the 256 KiB buffer so the + // buffer flushes several times and at least one big append bypasses it. + std::vector expected; + expected.reserve(1U << 20); + uint32_t x = 0xC0FFEEU; + auto next = [&]() { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + return static_cast(x); + }; + for (size_t i = 0; i < (1U << 20); ++i) { + expected.push_back(next()); + } + + { + LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + size_t off = 0; + // Phase 1: many small appends (700 B) -> repeated buffer overflow/flush. + while (off + 700 <= 600U * 1024) { + ASSERT_TRUE(w.append(Slice(expected.data() + off, 700)).ok()); + off += 700; + } + // Phase 2: one append larger than the buffer (direct-to-fd path). + const size_t big = 300U * 1024; + ASSERT_TRUE(w.append(Slice(expected.data() + off, big)).ok()); + off += big; + // Phase 3: drain the remainder in a final small append. + ASSERT_TRUE(w.append(Slice(expected.data() + off, expected.size() - off)).ok()); + EXPECT_EQ(w.bytes_written(), expected.size()); + ASSERT_TRUE(w.finalize().ok()); + } + + LocalFileReader r; + ASSERT_TRUE(r.open(path).ok()); + ASSERT_EQ(r.size(), expected.size()); + std::vector got; + ASSERT_TRUE(r.read_at(0, expected.size(), &got).ok()); + EXPECT_EQ(got, expected); +} + +// Empty appends are no-ops and an unused open->finalize yields an empty file. +TEST(SniiLocalFile, EmptyAppendsAndFinalize) { + const std::string path = TempPath("empty_buf"); + { + LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + ASSERT_TRUE(w.append(Slice(nullptr, 0)).ok()); + EXPECT_EQ(w.bytes_written(), 0U); + ASSERT_TRUE(w.finalize().ok()); + } + LocalFileReader r; + ASSERT_TRUE(r.open(path).ok()); + EXPECT_EQ(r.size(), 0U); +} diff --git a/be/test/storage/index/snii/io/metered_file_reader.cpp b/be/test/storage/index/snii/io/metered_file_reader.cpp new file mode 100644 index 00000000000000..3a0b07abe67d78 --- /dev/null +++ b/be/test/storage/index/snii/io/metered_file_reader.cpp @@ -0,0 +1,139 @@ +// 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 "storage/index/snii/io/metered_file_reader.h" + +#include + +namespace doris::snii::io { +namespace { + +// Inclusive [first, last] block ids touched by a validated [offset, offset+len). +// Empty len touches no block (callers guard len==0 before calling this). +void block_range(uint64_t offset, size_t len, size_t block_size, uint64_t* first, uint64_t* last) { + *first = offset / block_size; + *last = (offset + len - 1) / block_size; +} + +} // namespace + +MeteredFileReader::MeteredFileReader(FileReader* inner, size_t block_size) + : inner_(inner), block_size_(block_size) {} + +void MeteredFileReader::reset_metrics() { + resident_.clear(); + metrics_ = IoMetrics {}; +} + +Status MeteredFileReader::validate_range(uint64_t offset, size_t len) const { + if (inner_ == nullptr) + return Status::Error("metered: null inner reader"); + if (block_size_ == 0) + return Status::Error("metered: zero block size"); + const uint64_t total = inner_->size(); + if (offset > total || len > total - offset) { + return Status::Error( + "metered: read range past end"); + } + return Status::OK(); +} + +// Accounts the FileCache effect of touching [offset, offset+len): newly missed +// blocks become coalesced remote GETs and remote bytes. Returns true iff any +// block missed. (Single contiguous span -> at most one coalesced run.) +bool MeteredFileReader::account_blocks(uint64_t offset, size_t len) { + if (len == 0) return false; + uint64_t first = 0, last = 0; + block_range(offset, len, block_size_, &first, &last); + + bool any_miss = false; + bool in_run = false; // currently inside a contiguous run of missing blocks + const uint64_t total = inner_->size(); + for (uint64_t b = first; b <= last; ++b) { + if (resident_.count(b)) { + in_run = false; + continue; + } + resident_.insert(b); + any_miss = true; + const uint64_t block_start = b * block_size_; + metrics_.remote_bytes += std::min(block_size_, total - block_start); + if (!in_run) { + ++metrics_.range_gets; // start of a new coalesced GET + in_run = true; + } + } + return any_miss; +} + +Status MeteredFileReader::read_at(uint64_t offset, size_t len, std::vector* out) { + if (out == nullptr) + return Status::Error("metered: null out"); + RETURN_IF_ERROR(validate_range(offset, len)); + ++metrics_.read_at_calls; + metrics_.total_request_bytes += len; + // A single blocking read: any miss forces one serial round (the next offset is + // not known until these bytes return). + if (account_blocks(offset, len)) ++metrics_.serial_rounds; + return inner_->read_at(offset, len, out); +} + +Status MeteredFileReader::read_batch(const std::vector& ranges, + std::vector>* outs) { + if (outs == nullptr) + return Status::Error("metered: null batch out"); + for (const Range& r : ranges) { + RETURN_IF_ERROR(validate_range(r.offset, r.len)); + } + + // Gather the union of touched blocks so coalescing spans the whole batch, and + // the entire batch counts as at most one serial round. + std::vector blocks; + for (const Range& r : ranges) { + metrics_.total_request_bytes += r.len; + if (r.len == 0) continue; + uint64_t first = 0, last = 0; + block_range(r.offset, r.len, block_size_, &first, &last); + for (uint64_t b = first; b <= last; ++b) blocks.push_back(b); + } + metrics_.read_at_calls += ranges.size(); + + std::sort(blocks.begin(), blocks.end()); + blocks.erase(std::unique(blocks.begin(), blocks.end()), blocks.end()); + + bool any_miss = false; + const uint64_t total = inner_->size(); + uint64_t prev_miss = 0; + bool have_prev = false; + for (uint64_t b : blocks) { + if (resident_.count(b)) continue; + resident_.insert(b); + any_miss = true; + metrics_.remote_bytes += std::min(block_size_, total - b * block_size_); + if (!have_prev || b != prev_miss + 1) ++metrics_.range_gets; // new run + prev_miss = b; + have_prev = true; + } + if (any_miss) ++metrics_.serial_rounds; + + // Delegate the actual byte fetch to the inner reader's batch path, so a backend + // that fetches a batch concurrently (e.g. S3FileReader) realizes the planned + // round as parallel GETs (matching the single serial round accounted above). + return inner_->read_batch(ranges, outs); +} + +} // namespace doris::snii::io diff --git a/be/test/storage/index/snii/io/metered_file_reader.h b/be/test/storage/index/snii/io/metered_file_reader.h new file mode 100644 index 00000000000000..c53453316b105b --- /dev/null +++ b/be/test/storage/index/snii/io/metered_file_reader.h @@ -0,0 +1,67 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/io_metrics.h" + +namespace doris::snii::io { + +// A FileReader decorator that models an object-storage FileCache: reads are +// aligned to fixed (default 1MiB) blocks; only not-yet-resident blocks become +// remote range GETs (adjacent misses are coalesced). It is the single shared +// "yardstick" through which both single blocking reads and batched concurrent +// reads are measured. +// +// - read_at(): a single blocking read. Any cache miss => +1 serial round +// (the cursor must wait for bytes before the next offset is known). +// - read_batch(): all ranges submitted concurrently => the whole batch is at +// most one serial round (+1 iff any range misses). +class MeteredFileReader : public FileReader { +public: + explicit MeteredFileReader(FileReader* inner, size_t block_size = (1u << 20)); + + Status read_at(uint64_t offset, size_t len, std::vector* out) override; + Status read_batch(const std::vector& ranges, + std::vector>* outs) override; + uint64_t size() const override { return inner_->size(); } + + const IoMetrics& metrics() const { return metrics_; } + const IoMetrics* io_metrics() const override { return &metrics_; } + // Clears counters AND the resident block set, modelling a cold (cache-empty) query. + void reset_metrics(); + +private: + Status validate_range(uint64_t offset, size_t len) const; + + // Accounts the cache effect of touching [offset, offset+len): records misses, + // coalesced GETs, and remote bytes. Returns true iff at least one block missed. + bool account_blocks(uint64_t offset, size_t len); + + FileReader* inner_; + size_t block_size_; + std::unordered_set resident_; + IoMetrics metrics_; +}; + +} // namespace doris::snii::io diff --git a/be/test/storage/index/snii/io/metered_file_reader_test.cpp b/be/test/storage/index/snii/io/metered_file_reader_test.cpp new file mode 100644 index 00000000000000..a9fc344fb56d04 --- /dev/null +++ b/be/test/storage/index/snii/io/metered_file_reader_test.cpp @@ -0,0 +1,181 @@ +// 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 "storage/index/snii/io/metered_file_reader.h" + +#include + +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/local_file.h" + +using namespace doris::snii; +using doris::Status; +using doris::snii::io::LocalFileReader; +using doris::snii::io::LocalFileWriter; +using doris::snii::io::MeteredFileReader; +using doris::snii::io::Range; + +namespace { + +// Writes 256 bytes (byte[i] = i) to a temp file and returns its path. +std::string MakeRampFile() { + const std::string path = "/tmp/snii_metered_ramp.bin"; + LocalFileWriter w; + EXPECT_TRUE(w.open(path).ok()); + std::vector data(256); + for (int i = 0; i < 256; ++i) { + data[i] = static_cast(i); + } + EXPECT_TRUE(w.append(Slice(data)).ok()); + EXPECT_TRUE(w.finalize().ok()); + return path; +} + +} // namespace + +// Single reads: first read to a block is a cache miss (1 round, 1 GET, 1 block of +// remote bytes); a second read to the same 16-byte block is a hit (no new round). +TEST(SniiMeteredFileReader, SingleReadCacheAccounting) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, /*block_size=*/16); + + std::vector out; + ASSERT_TRUE(m.read_at(0, 4, &out).ok()); + EXPECT_EQ(out[0], 0U); + EXPECT_EQ(out[3], 3U); + EXPECT_EQ(m.metrics().read_at_calls, 1U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); + EXPECT_EQ(m.metrics().range_gets, 1U); + EXPECT_EQ(m.metrics().remote_bytes, 16U); + + // Same block (offset 8..11) -> cache hit. + ASSERT_TRUE(m.read_at(8, 4, &out).ok()); + EXPECT_EQ(m.metrics().read_at_calls, 2U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); // unchanged + EXPECT_EQ(m.metrics().range_gets, 1U); + EXPECT_EQ(m.metrics().remote_bytes, 16U); + + // Different block (offset 20 -> block 1) -> miss. + ASSERT_TRUE(m.read_at(20, 4, &out).ok()); + EXPECT_EQ(out[0], 20U); + EXPECT_EQ(m.metrics().serial_rounds, 2U); + EXPECT_EQ(m.metrics().range_gets, 2U); + EXPECT_EQ(m.metrics().remote_bytes, 32U); +} + +// A read spanning 3 contiguous blocks is one round and one coalesced GET. +TEST(SniiMeteredFileReader, SpanMultipleBlocksCoalesced) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector out; + ASSERT_TRUE(m.read_at(0, 40, &out).ok()); // blocks 0,1,2 + EXPECT_EQ(out.size(), 40U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); + EXPECT_EQ(m.metrics().range_gets, 1U); + EXPECT_EQ(m.metrics().remote_bytes, 48U); // 3 * 16 +} + +// A batch of reads to non-adjacent blocks: one serial round, one GET per run. +TEST(SniiMeteredFileReader, BatchNonAdjacent) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector ranges = {{.offset = 0, .len = 4}, + {.offset = 100, .len = 4}, + {.offset = 200, .len = 4}}; // blocks 0, 6, 12 + std::vector> outs; + ASSERT_TRUE(m.read_batch(ranges, &outs).ok()); + ASSERT_EQ(outs.size(), 3U); + EXPECT_EQ(outs[1][0], 100U); + EXPECT_EQ(m.metrics().read_at_calls, 3U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); // one batch = one round + EXPECT_EQ(m.metrics().range_gets, 3U); // 3 disjoint runs + EXPECT_EQ(m.metrics().remote_bytes, 48U); +} + +// A batch of reads to adjacent blocks coalesces into a single GET. +TEST(SniiMeteredFileReader, BatchAdjacentCoalesced) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector ranges = {{.offset = 0, .len = 4}, + {.offset = 16, .len = 4}, + {.offset = 32, .len = 4}}; // blocks 0,1,2 + std::vector> outs; + ASSERT_TRUE(m.read_batch(ranges, &outs).ok()); + EXPECT_EQ(m.metrics().read_at_calls, 3U); + EXPECT_EQ(m.metrics().serial_rounds, 1U); + EXPECT_EQ(m.metrics().range_gets, 1U); // coalesced + EXPECT_EQ(m.metrics().remote_bytes, 48U); +} + +// reset_metrics clears both counters and the resident cache (cold query). +TEST(SniiMeteredFileReader, ResetClearsCacheAndCounters) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector out; + ASSERT_TRUE(m.read_at(0, 4, &out).ok()); + m.reset_metrics(); + EXPECT_EQ(m.metrics().read_at_calls, 0U); + EXPECT_EQ(m.metrics().serial_rounds, 0U); + // Cache cleared -> same read misses again. + ASSERT_TRUE(m.read_at(0, 4, &out).ok()); + EXPECT_EQ(m.metrics().serial_rounds, 1U); + EXPECT_EQ(m.metrics().remote_bytes, 16U); +} + +TEST(SniiMeteredFileReader, InvalidRangeDoesNotPolluteMetrics) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector out; + const Status st = m.read_at(250, 16, &out); + EXPECT_TRUE(st.is()) << st.to_string(); + EXPECT_EQ(m.metrics().read_at_calls, 0U); + EXPECT_EQ(m.metrics().serial_rounds, 0U); + EXPECT_EQ(m.metrics().range_gets, 0U); + EXPECT_EQ(m.metrics().remote_bytes, 0U); + EXPECT_EQ(m.metrics().total_request_bytes, 0U); +} + +TEST(SniiMeteredFileReader, InvalidBatchRangeDoesNotPolluteMetrics) { + LocalFileReader inner; + ASSERT_TRUE(inner.open(MakeRampFile()).ok()); + MeteredFileReader m(&inner, 16); + + std::vector> outs; + const Status st = + m.read_batch({Range {.offset = 0, .len = 4}, Range {.offset = 250, .len = 16}}, &outs); + EXPECT_TRUE(st.is()) << st.to_string(); + EXPECT_EQ(m.metrics().read_at_calls, 0U); + EXPECT_EQ(m.metrics().serial_rounds, 0U); + EXPECT_EQ(m.metrics().range_gets, 0U); + EXPECT_EQ(m.metrics().remote_bytes, 0U); + EXPECT_EQ(m.metrics().total_request_bytes, 0U); +} diff --git a/be/test/storage/index/snii/query/boolean_query_test.cpp b/be/test/storage/index/snii/query/boolean_query_test.cpp new file mode 100644 index 00000000000000..f89533442b1c10 --- /dev/null +++ b/be/test/storage/index/snii/query/boolean_query_test.cpp @@ -0,0 +1,337 @@ +// 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 "storage/index/snii/query/boolean_query.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_boolean_query_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +struct Corpus { + uint32_t doc_count = 900; + std::vector> docs; + + std::vector term_docs(const std::string& term) const { + std::vector out; + for (uint32_t d = 0; d < docs.size(); ++d) { + if (std::find(docs[d].begin(), docs[d].end(), term) != docs[d].end()) { + out.push_back(d); + } + } + return out; + } + + std::vector any_docs(const std::vector& terms) const { + std::set ids; + for (const auto& term : terms) { + const std::vector docs_for_term = term_docs(term); + ids.insert(docs_for_term.begin(), docs_for_term.end()); + } + return {ids.begin(), ids.end()}; + } +}; + +Corpus BuildCorpus() { + Corpus c; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + if (d < 700) { + toks.emplace_back("aa_hot"); + } + if (d < 700) { + toks.emplace_back("aa_warm"); + } + if (d < 700) { + toks.emplace_back("aa_loud"); + } + if (d % 17 == 0) { + toks.emplace_back("aa_rare"); + } + if (d % 41 == 0) { + toks.emplace_back("aa_sparse"); + } + char filler[16]; + std::snprintf(filler, sizeof(filler), "zz_%03u", d % 200); + toks.emplace_back(filler); + } + return c; +} + +Corpus BuildDenseAndCorpus() { + Corpus c; + c.doc_count = 9000; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + toks.emplace_back("aa_dense"); + if (d % 1024 == 0) { + toks.emplace_back("aa_rare_driver"); + } + char filler[16]; + std::snprintf(filler, sizeof(filler), "zz_%03u", d % 200); + toks.emplace_back(filler); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/false); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + const auto& toks = c.docs[d]; + for (uint32_t pos = 0; pos < toks.size(); ++pos) { + buf.add_token(toks[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = doris::snii::format::IndexConfig::kDocsOnly; + in.doc_count = c.doc_count; + in.terms = buf.finalize_sorted(); + in.target_dict_block_bytes = 256; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +LogicalIndexReader OpenIndex(io::LocalFileReader* file, SniiSegmentReader* segment, + const std::string& path) { + EXPECT_TRUE(file->open(path).ok()); + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +LogicalIndexReader OpenMeteredIndex(io::MeteredFileReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +class RecordingSink final : public query::DocIdSink { +public: + Status append_sorted(std::span docids) override { + ++chunks; + max_chunk = std::max(max_chunk, docids.size()); + out.insert(out.end(), docids.begin(), docids.end()); + return Status::OK(); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + ++chunks; + if (last_exclusive > first) { + max_chunk = std::max(max_chunk, static_cast(last_exclusive - first)); + for (uint64_t docid = first; docid < last_exclusive; ++docid) { + out.push_back(static_cast(docid)); + } + } + return Status::OK(); + } + + std::vector out; + size_t chunks = 0; + size_t max_chunk = 0; +}; + +class FailingSink final : public query::DocIdSink { +public: + Status append_sorted(std::span docids) override { + saw_docids = !docids.empty(); + return Status::IOError("sink append failed"); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + saw_docids = saw_docids || (last_exclusive > first); + return Status::IOError("sink append failed"); + } + + bool saw_docids = false; +}; + +} // namespace + +TEST(SniiBooleanOr, ReturnsSortedDeduplicatedUnion) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"aa_hot", "aa_rare", "absent_token", "aa_hot", + "aa_sparse"}; + std::vector got; + ASSERT_TRUE(query::boolean_or(idx, terms, &got).ok()); + + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(got, corpus.any_docs(terms)); + + std::vector single; + ASSERT_TRUE(query::term_query(idx, "aa_hot", &single).ok()); + std::vector single_or; + ASSERT_TRUE(query::boolean_or(idx, {"aa_hot"}, &single_or).ok()); + EXPECT_EQ(single_or, single); + + std::vector absent; + ASSERT_TRUE(query::boolean_or(idx, {"missing_a", "missing_b"}, &absent).ok()); + EXPECT_TRUE(absent.empty()); + + std::remove(path.c_str()); +} + +TEST(SniiBooleanOr, EmitsHighHitResultsInBulkChunks) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + RecordingSink sink; + ASSERT_TRUE(query::boolean_or(idx, {"aa_hot"}, &sink).ok()); + + const std::vector want = corpus.any_docs({"aa_hot"}); + EXPECT_EQ(sink.out, want); + ASSERT_FALSE(want.empty()); + EXPECT_GT(sink.max_chunk, 1U); + EXPECT_LT(sink.chunks, want.size() / 100) + << "high-hit term results must be handed off in bulk, not per doc"; + + std::remove(path.c_str()); +} + +TEST(SniiBooleanOr, PropagatesSinkAppendFailure) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + FailingSink sink; + const Status st = query::boolean_or(idx, {"aa_hot"}, &sink); + EXPECT_TRUE(st.is()) << st.to_string(); + EXPECT_TRUE(sink.saw_docids); + + std::remove(path.c_str()); +} + +TEST(SniiBooleanOr, BatchesWindowedPostingReadsByStage) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/128); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + const std::vector terms = {"aa_hot", "aa_warm", "aa_loud"}; + for (const std::string& term : terms) { + bool found = false; + doris::snii::format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + ASSERT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found) << term; + } + + const uint64_t serial_before = metered.metrics().serial_rounds; + std::vector got; + ASSERT_TRUE(query::boolean_or(idx, terms, &got).ok()); + const uint64_t serial_delta = metered.metrics().serial_rounds - serial_before; + + EXPECT_EQ(got, corpus.any_docs(terms)); + EXPECT_LE(serial_delta, 2U) << "windowed OR should batch posting reads by stage, not by term"; + + std::remove(path.c_str()); +} + +TEST(SniiBooleanAnd, SkipsDenseFullWindowsForRareDriver) { + const Corpus corpus = BuildDenseAndCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/128); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + std::vector hot_docs; + ASSERT_TRUE(query::term_query(idx, "aa_dense", &hot_docs).ok()); + const uint64_t hot_bytes = metered.metrics().total_request_bytes; + ASSERT_EQ(hot_docs, corpus.term_docs("aa_dense")); + ASSERT_GT(hot_bytes, 0U); + + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::boolean_and(idx, {"aa_dense", "aa_rare_driver"}, &got).ok()); + const uint64_t and_bytes = metered.metrics().total_request_bytes; + + std::vector want; + const std::vector rare = corpus.term_docs("aa_rare_driver"); + std::ranges::set_intersection(hot_docs, rare, std::back_inserter(want)); + EXPECT_EQ(got, want); + EXPECT_LT(and_bytes, hot_bytes / 4) + << "dense full-window AND should keep rare-driver candidates without " + "reading the dense term's dd regions"; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/byte_skip_test.cpp b/be/test/storage/index/snii/query/byte_skip_test.cpp new file mode 100644 index 00000000000000..709340285846a2 --- /dev/null +++ b/be/test/storage/index/snii/query/byte_skip_test.cpp @@ -0,0 +1,473 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/scoring_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +// PHASE B differential + byte-reduction test (design spec sections 1.4 & 2). +// +// Builds an index (~4000 docs, positions + scoring) with a HIGH-DF term spanning +// MANY .frq windows plus low/mid-df terms and a planted 5-term phrase led by the +// high-df term. Over a MeteredFileReader it asserts: +// (a) term_query and phrase_query docids equal a brute-force ORACLE (unchanged +// by the freq-skip optimization); +// (b) a docid-only term_query on the high-df term requests STRICTLY FEWER .frq +// bytes than the pre-PhaseB full-window path (sum of per-window frq_len), +// because the freq region is skipped on the wire; +// (c) scoring_query STILL reads the FULL windows (freq region present -> its +// .frq request bytes match the full-window total, strictly above the +// docid-only path) and returns the correct top-K. +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::snii::query::Bm25Params; +using doris::snii::query::ScoredDoc; +using doris::snii::stats::SniiStatsProvider; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_byte_skip_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +constexpr uint32_t kDocCount = 4000; +constexpr uint32_t kHiGap = 1; // aa_hi in (almost) every doc -> many windows +constexpr uint32_t kHiCount = 3600; // df of aa_hi (>>256 -> many windows) +// The 5-term phrase is planted only in the FIRST kPhraseSpan occurrences of +// aa_hi, so a low-df lead term concentrates candidates into the first windows. +constexpr uint32_t kPhraseSpan = 120; + +// aa_hi appears with VARIED, large frequency so its per-window freq region is +// big in absolute terms -- large enough to span whole cache blocks the docs-only +// path never touches, making the freq-skip visible in remote_bytes (not just in +// the raw wire-byte total). Frequencies cycle through a wide range. +uint32_t HiFreq(uint32_t occ) { + return 40U + (occ * 13U + 7U) % 200U; +} + +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; // docs[d] = ordered tokens + std::vector doc_len; // per-doc token count (for norms) + + bool phrase_in_doc(uint32_t d, const std::vector& phrase) const { + if (phrase.empty()) { + return false; + } + const auto& toks = docs[d]; + if (toks.size() < phrase.size()) { + return false; + } + for (size_t i = 0; i + phrase.size() <= toks.size(); ++i) { + bool match = true; + for (size_t k = 0; k < phrase.size(); ++k) { + if (toks[i + k] != phrase[k]) { + match = false; + break; + } + } + if (match) { + return true; + } + } + return false; + } + + std::vector term_oracle(const std::string& term) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + const auto& toks = docs[d]; + if (std::ranges::find(toks, term) != toks.end()) { + out.push_back(d); + } + } + return out; + } + + std::vector phrase_oracle(const std::vector& phrase) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + if (phrase_in_doc(d, phrase)) { + out.push_back(d); + } + } + return out; + } +}; + +Corpus BuildCorpus() { + Corpus c; + c.doc_count = kDocCount; + c.docs.resize(c.doc_count); + c.doc_len.assign(c.doc_count, 0); + const std::vector vocab = {"alpha", "bravo", "charlie", "delta"}; + + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + const bool is_hi = (d % kHiGap == 0) && (d / kHiGap < kHiCount); + if (is_hi) { + toks.emplace_back("aa_hi"); // high-df, position 0 + } + const uint32_t occ = d / kHiGap; // aa_hi occurrence ordinal + + if (is_hi && occ < kPhraseSpan) { + // 5-term phrase led by the high-df term, concentrated in early windows. + toks.emplace_back("aa_quick"); + toks.emplace_back("aa_brown"); + toks.emplace_back("aa_fox"); + toks.emplace_back("aa_jumps"); + } else if (d % 11 == 0) { + toks.emplace_back("aa_lazy"); + toks.emplace_back("aa_dog"); + } + // Mid-df term: every 5th doc (varied freq via repetition for scoring). + if (d % 5 == 0) { + toks.emplace_back("aa_mid"); + if (d % 15 == 0) { + toks.emplace_back("aa_mid"); + } + } + for (uint32_t k = 0; k < 2; ++k) { + toks.push_back(vocab[(d + k) % vocab.size()]); + } + if (d % 97 == 0) { + toks.emplace_back("aa_rare"); + } + // Filler vocabulary to occupy the lexicographic tail blocks. + char nm[16]; + std::snprintf(nm, sizeof(nm), "zz_%03u", d % 400); + toks.emplace_back(nm); + // Extra aa_hi repetitions AFTER position 0's phrase, giving aa_hi a varied, + // larger freq so its per-window freq region carries real bytes. These trail + // the phrase tokens so the consecutive phrase at position 0 is preserved. + if (is_hi) { + const uint32_t extra = HiFreq(occ) - 1U; // total freq = HiFreq(occ) + for (uint32_t k = 0; k < extra; ++k) { + toks.emplace_back("aa_hi"); + } + } + c.doc_len[d] = toks.size(); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.doc_count; ++d) { + const auto& toks = c.docs[d]; + for (uint32_t pos = 0; pos < toks.size(); ++pos) { + buf.add_token(toks[pos], d, pos); + } + } + std::vector terms = buf.finalize_sorted(); + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositionsScoring; // positions (phrase) + norms (scoring) + in.doc_count = c.doc_count; + in.terms = std::move(terms); + in.target_dict_block_bytes = 256; + + in.encoded_norms.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + in.encoded_norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(in).ok()); + ASSERT_TRUE(cw.finish().ok()); +} + +// Grouped-block byte totals of a windowed term (design 1.6): docs = dd-block +// length (sum of per-window dd_disk_len); full = dd-block + freq-block lengths +// (what scoring reads). Also returns prelude_len and window count. +struct FrqByteTotals { + uint64_t full = 0; // dd-block + freq-block on-disk bytes + uint64_t docs = 0; // dd-block on-disk bytes (the contiguous docs-only run) + uint64_t prelude_len = 0; + uint32_t windows = 0; +}; + +FrqByteTotals MeasureFrqTotals(const LogicalIndexReader& idx, const std::string& term) { + FrqByteTotals t; + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + EXPECT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + EXPECT_TRUE(found); + EXPECT_EQ(entry.enc, DictEntryEnc::kWindowed); + t.prelude_len = entry.prelude_len; + FrqPreludeReader prelude; + EXPECT_TRUE(fetch_windowed_prelude(idx, entry, frq_base, &prelude).ok()); + t.windows = prelude.n_windows(); + t.docs = prelude.dd_block_len(); + t.full = prelude.dd_block_len() + prelude.freq_block_len(); + return t; +} + +// Reference BM25 ranking computed directly from the corpus (df/idf/tf + norms). +std::vector ReferenceRanking(const Corpus& c, const std::vector& terms, + uint32_t k, const Bm25Params& params) { + uint64_t sum_ttf = 0; + for (uint64_t dl : c.doc_len) { + sum_ttf += dl; + } + const double avgdl = static_cast(sum_ttf) / std::max(1, c.doc_count); + + std::unordered_map scores; + for (const auto& term : terms) { + std::map plist; // docid -> freq + for (uint32_t d = 0; d < c.doc_count; ++d) { + uint32_t f = 0; + for (const auto& t : c.docs[d]) { + if (t == term) { + ++f; + } + } + if (f > 0) { + plist[d] = f; + } + } + if (plist.empty()) { + continue; + } + const uint64_t df = plist.size(); + const double idf = + std::log(1.0 + (static_cast(c.doc_count) - df + 0.5) / (df + 0.5)); + for (const auto& [docid, freq] : plist) { + // Match the index's quantization (encode -> decode) so scores compare exact. + const double dl = doris::snii::query::decode_norm( + doris::snii::query::encode_norm(c.doc_len[docid])); + const double denom = freq + params.k1 * (1.0 - params.b + params.b * dl / avgdl); + scores[docid] += idf * (freq * (params.k1 + 1.0)) / denom; + } + } + + std::vector all; + all.reserve(scores.size()); + for (const auto& [docid, s] : scores) { + all.push_back({docid, s}); + } + std::ranges::sort(all, [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) { + return a.score > b.score; + } + return a.docid < b.docid; + }); + if (all.size() > k) { + all.resize(k); + } + return all; +} + +} // namespace + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiByteSkip, DocidAndPhraseSkipFreqWhileScoringKeepsIt) { + Corpus c = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + // Fine-grained cache block (< a window's freq region) so the per-window + // freq-skip shows up in remote_bytes too: each window's freq region spans whole + // cache blocks that the docs-only path never touches. + io::MeteredFileReader metered(&local, /*block_size=*/64); + + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + ASSERT_EQ(idx.stats().doc_count, c.doc_count); + + // The high-df term must span many windows for the freq-skip to be meaningful. + const FrqByteTotals hi = MeasureFrqTotals(idx, "aa_hi"); + ASSERT_GE(hi.windows, 4U) << "high-df term should span multiple windows"; + // PhaseA format invariant: the docs-only prefix is strictly smaller than the + // full window total (the freq region occupies real bytes). + ASSERT_LT(hi.docs, hi.full) << "docs-only prefix not smaller: docs=" << hi.docs + << " full=" << hi.full; + + // ---- (a) term_query / phrase_query docids == ORACLE (unchanged) ----------- + const std::vector terms_to_check = {"aa_hi", "aa_mid", "aa_quick", "aa_lazy", + "aa_rare"}; + for (const auto& term : terms_to_check) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, term, &got).ok()) << term; + EXPECT_TRUE(std::ranges::is_sorted(got)) << term; + EXPECT_EQ(got, c.term_oracle(term)) << "term_query oracle mismatch: " << term; + } + + const std::vector> phrases = { + {"aa_hi", "aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // 5-term, hi lead + {"aa_hi", "aa_quick"}, // hi + mid + {"aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // no high-df term + {"aa_lazy", "aa_dog"}, // scattered phrase + {"aa_brown", "aa_fox"}, // sub-phrase + {"aa_fox", "aa_brown"}, // reversed -> absent + {"aa_hi", "aa_lazy"}, // present terms, absent phrase + }; + for (const auto& p : phrases) { + std::string label; + for (const auto& w : p) { + label += w + " "; + } + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, p, &got).ok()) << label; + EXPECT_TRUE(std::ranges::is_sorted(got)) << label; + EXPECT_EQ(got, c.phrase_oracle(p)) << "phrase oracle mismatch: [" << label << "]"; + } + + // ---- (b) docid-only term_query on aa_hi skips the freq region ------------- + // Format-level proof: the docs-only prefixes total far fewer bytes than the + // full windows (the freq region is real and is what the wire-level skip drops). + const uint64_t full_window_frq = hi.full; // pre-PhaseB .frq byte total + EXPECT_LT(hi.docs * 100, full_window_frq * 85) + << "freq-skip format reduction too small: docs=" << hi.docs + << " full=" << full_window_frq; + // Both A and B below share the SAME lookup + prelude overhead; the ONLY + // difference is whether each window fetches its docs-only prefix or the full + // window. So (B - A) isolates exactly the .frq freq-region bytes the docid-only + // path skips. A and B are measured live over the metered reader. + // + // A = PhaseB docid-only term_query (want_freq=false everywhere). + metered.reset_metrics(); + std::vector hi_docs; + ASSERT_TRUE(query::term_query(idx, "aa_hi", &hi_docs).ok()); + const io::IoMetrics a = metered.metrics(); + ASSERT_EQ(hi_docs, c.term_oracle("aa_hi")); + + // B = the pre-PhaseB full-window baseline: same lookup + prelude, then read + // every window IN FULL (want_freq=true). This is exactly what term_query used + // to fetch before the freq-skip landed. + metered.reset_metrics(); + { + DictEntry entry; + uint64_t fb = 0, pb = 0; + bool found = false; + ASSERT_TRUE(idx.lookup("aa_hi", &found, &entry, &fb, &pb).ok()); + ASSERT_TRUE(found); + DecodedPosting full_posting; + ASSERT_TRUE(read_windowed_posting(idx, entry, fb, pb, /*want_positions=*/false, + /*want_freq=*/true, &full_posting) + .ok()); + EXPECT_EQ(full_posting.docids, c.term_oracle("aa_hi")); + EXPECT_EQ(full_posting.freqs.size(), full_posting.docids.size()) + << "want_freq=true must decode per-doc freqs"; + } + const io::IoMetrics b = metered.metrics(); + + // The docid-only path requests STRICTLY FEWER bytes on the wire (freq skipped). + EXPECT_LT(a.total_request_bytes, b.total_request_bytes) + << "docid-only did not skip freq bytes: a=" << a.total_request_bytes + << " b=" << b.total_request_bytes; + // The skipped amount equals the full-vs-docs .frq delta (sum(frq_len-frq_docs_len)). + const uint64_t saved = b.total_request_bytes - a.total_request_bytes; + EXPECT_EQ(saved, hi.full - hi.docs) << "skipped bytes != freq-region size: saved=" << saved + << " freq_region=" << (hi.full - hi.docs); + // remote_bytes (cache-block granular) is also strictly below the full-window + // baseline: the freq region occupies whole cache blocks the docs path never hits. + EXPECT_LT(a.remote_bytes, b.remote_bytes) + << "docid-only remote_bytes not below full-window: a=" << a.remote_bytes + << " b=" << b.remote_bytes; + + // ---- (c) scoring_query STILL reads full windows (freq present) ------------ + SniiStatsProvider stats; + ASSERT_TRUE(SniiStatsProvider::open(&idx, &stats).ok()); + const Bm25Params params; // defaults + const std::vector score_terms = {"aa_hi", "aa_mid", "aa_rare"}; + const uint32_t kTopK = 10; + + std::vector exhaustive, wand; + ASSERT_TRUE(query::scoring_query_exhaustive(idx, stats, score_terms, kTopK, params, &exhaustive) + .ok()); + ASSERT_TRUE(query::scoring_query_wand(idx, stats, score_terms, kTopK, params, &wand).ok()); + + const std::vector ref = ReferenceRanking(c, score_terms, kTopK, params); + ASSERT_EQ(exhaustive.size(), ref.size()); + for (size_t i = 0; i < ref.size(); ++i) { + EXPECT_EQ(exhaustive[i].docid, ref[i].docid) << "scoring docid mismatch at " << i; + EXPECT_NEAR(exhaustive[i].score, ref[i].score, 1e-9) << "scoring score mismatch at " << i; + } + // WAND must equal the exhaustive ranking exactly (docid + score). + ASSERT_EQ(wand.size(), exhaustive.size()); + for (size_t i = 0; i < wand.size(); ++i) { + EXPECT_EQ(wand[i].docid, exhaustive[i].docid) << "wand vs exhaustive at " << i; + EXPECT_NEAR(wand[i].score, exhaustive[i].score, 1e-9) << "wand vs exhaustive at " << i; + } + + // Scoring reads the FULL .frq windows for the high-df term (freq region present), + // so a single-term scoring query requests STRICTLY MORE bytes than the + // docid-only term_query A (which skipped the freq region). Same lookup + .frq + // section; scoring additionally pulls the freq region (and norms/prelude), so + // its request total exceeds the docid-only path by at least the freq region. + metered.reset_metrics(); + std::vector hi_only; + ASSERT_TRUE( + query::scoring_query_exhaustive(idx, stats, {"aa_hi"}, kTopK, params, &hi_only).ok()); + const io::IoMetrics score_io = metered.metrics(); + const uint64_t score_frq_request = score_io.total_request_bytes; + const uint64_t term_frq_request = a.total_request_bytes; + EXPECT_GE(score_frq_request, term_frq_request + (hi.full - hi.docs)) + << "scoring did not read the full freq region: scoring=" << score_frq_request + << " docid=" << term_frq_request << " freq_region=" << (hi.full - hi.docs); + EXPECT_GT(score_frq_request, term_frq_request) + << "scoring should read MORE .frq bytes than docid-only: scoring=" << score_frq_request + << " docid=" << term_frq_request; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/common_grams_dict_plan_test.cpp b/be/test/storage/index/snii/query/common_grams_dict_plan_test.cpp new file mode 100644 index 00000000000000..3e690c76b8f042 --- /dev/null +++ b/be/test/storage/index/snii/query/common_grams_dict_plan_test.cpp @@ -0,0 +1,489 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii::query { +namespace { + +using segment_v2::inverted_index::CommonGramsCoverage; +using segment_v2::inverted_index::CommonGramsPlanCostModel; +using segment_v2::inverted_index::CommonGramsQueryIdentity; +using segment_v2::inverted_index::CommonGramsSegmentMetadata; +using segment_v2::inverted_index::PlainTermKeyVersion; +using segment_v2::inverted_index::ScoringCoverage; +using segment_v2::inverted_index::COMMON_GRAMS_KEY_VERSION_V1; +using segment_v2::inverted_index::COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; +using segment_v2::inverted_index::COMMON_GRAMS_SCORING_STATS_VERSION_V1; +using segment_v2::inverted_index::COMMON_GRAMS_SEMANTICS_VERSION_V1; +using segment_v2::inverted_index::encode_common_gram; +using segment_v2::inverted_index::encode_plain_term; +using snii_test::MemoryFile; +using snii_test::ScopedEnv; +using snii_test::assert_ok; +using snii_test::make_term; + +constexpr uint64_t kIndexId = 51; +constexpr std::string_view kIndexSuffix = "body"; + +class RecordingReader final : public io::FileReader { +public: + explicit RecordingReader(io::FileReader* inner) : inner_(inner) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + ++default_read_calls_; + return inner_->read_at(offset, len, out); + } + + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + ++default_read_calls_; + return inner_->read_batch(ranges, outs); + } + + uint64_t size() const override { return inner_->size(); } + + void reset_counts() { default_read_calls_ = 0; } + + size_t default_read_calls() const { return default_read_calls_; } + +private: + io::FileReader* inner_; + size_t default_read_calls_ = 0; +}; + +struct Fixture { + Fixture() : recording_reader(&file) {} + + MemoryFile file; + RecordingReader recording_reader; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; +}; + +CommonGramsSegmentMetadata complete_metadata() { + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = CommonGramsCoverage::kComplete; + metadata.common_grams_semantics_version = COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = "builtin:lucene_english_stop:v1"; + metadata.base_analyzer_fingerprint = "common-grams-query-base-v1"; + metadata.common_grams_fingerprint = "common-grams-query-v1"; + return metadata; +} + +CommonGramsQueryIdentity complete_query_identity() { + const CommonGramsSegmentMetadata metadata = complete_metadata(); + return {.common_grams_dictionary_identity = metadata.common_grams_dictionary_identity, + .base_analyzer_fingerprint = metadata.base_analyzer_fingerprint, + .common_grams_fingerprint = metadata.common_grams_fingerprint}; +} + +CommonGramsSegmentMetadata complete_scoring_metadata(uint32_t doc_count, + uint64_t scoring_token_count) { + CommonGramsSegmentMetadata metadata = complete_metadata(); + metadata.scoring_coverage = ScoringCoverage::kComplete; + metadata.scoring_stats_version = COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + metadata.scoring_doc_count = doc_count; + metadata.scoring_token_count = scoring_token_count; + return metadata; +} + +CommonGramsPlanCostModel no_hysteresis_cost_model() { + return {.position_verify_factor = 0, .common_grams_cost_ratio_percent = 100, .generation = 0}; +} + +CommonGramsPlanCostModel verify_cost_model() { + return {.position_verify_factor = 7, .common_grams_cost_ratio_percent = 100, .generation = 0}; +} + +std::string plain_key(std::string_view term) { + auto encoded = encode_plain_term(term, PlainTermKeyVersion::kEscapedV1); + EXPECT_TRUE(encoded.has_value()); + return encoded.has_value() ? std::move(encoded.value()) : std::string(); +} + +std::string gram_key(std::string_view left, std::string_view right) { + auto encoded = encode_common_gram(left, right); + EXPECT_TRUE(encoded.has_value()); + return encoded.has_value() ? std::move(encoded.value()) : std::string(); +} + +std::vector plan_terms() { + return { + make_term(plain_key("the"), {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {0}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {0}}, + {.docid = 6, .positions = {0}}, + {.docid = 7, .positions = {0}}}), + make_term(plain_key("gap"), {{.docid = 2, .positions = {1}}, + {.docid = 3, .positions = {1}}, + {.docid = 4, .positions = {1}}, + {.docid = 5, .positions = {1}}, + {.docid = 6, .positions = {1}}, + {.docid = 7, .positions = {1}}}), + make_term(plain_key("wolf"), {{.docid = 0, .positions = {1}}, + {.docid = 2, .positions = {2}}, + {.docid = 3, .positions = {2}}, + {.docid = 4, .positions = {2}}, + {.docid = 5, .positions = {2}}, + {.docid = 6, .positions = {2}}, + {.docid = 7, .positions = {2}}}), + make_term(plain_key("woman"), {{.docid = 1, .positions = {1}}}), + make_term(gram_key("the", "wolf"), {{.docid = 0, .positions = {0}}}), + make_term(gram_key("the", "woman"), {{.docid = 1, .positions = {0}}}), + make_term(gram_key("the", "gap"), {{.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {0}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {0}}, + {.docid = 6, .positions = {0}}, + {.docid = 7, .positions = {0}}}), + }; +} + +Status build_fixture(Fixture* fixture) { + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = std::string(kIndexSuffix); + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 8; + input.terms = plan_terms(); + input.target_dict_block_bytes = 1; + input.common_grams_metadata = complete_metadata(); + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::SniiCompoundWriter compound_writer(&fixture->file); + RETURN_IF_ERROR(compound_writer.add_logical_index(input)); + RETURN_IF_ERROR(compound_writer.finish()); + RETURN_IF_ERROR( + reader::SniiSegmentReader::open(&fixture->recording_reader, &fixture->segment_reader)); + RETURN_IF_ERROR( + fixture->segment_reader.open_index(kIndexId, kIndexSuffix, &fixture->index_reader)); + fixture->recording_reader.reset_counts(); + return Status::OK(); +} + +constexpr uint32_t kScoringDocCount = 520; + +std::vector scoring_plan_terms() { + auto all_the = snii_test::docs_with_one_position(0, kScoringDocCount, 0); + auto all_wolf = snii_test::docs_with_one_position(0, kScoringDocCount, 1); + auto all_the_wolf = snii_test::docs_with_one_position(0, kScoringDocCount, 0); + return { + make_term(plain_key(std::string("\x1f" + "literal")), + {{.docid = 0, .positions = {2}}}), + make_term(gram_key("the", "wolf"), std::move(all_the_wolf)), + make_term(gram_key("the", "woman"), {{.docid = 0, .positions = {0}}}), + make_term(plain_key("the"), std::move(all_the)), + make_term(plain_key("wolf"), std::move(all_wolf)), + make_term(plain_key("woman"), {{.docid = 0, .positions = {1}}}), + make_term(plain_key("zeta"), {{.docid = 0, .positions = {3}}}), + }; +} + +Status build_scoring_fixture(Fixture* fixture, + CommonGramsCoverage coverage = CommonGramsCoverage::kComplete) { + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = std::string(kIndexSuffix); + input.config = format::IndexConfig::kDocsPositionsScoring; + input.doc_count = kScoringDocCount; + input.encoded_norms.assign(kScoringDocCount, 1); + input.terms = scoring_plan_terms(); + input.target_dict_block_bytes = 1U << 20; + input.common_grams_metadata = complete_scoring_metadata(kScoringDocCount, 1043); + input.common_grams_metadata->common_grams_coverage = coverage; + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::SniiCompoundWriter compound_writer(&fixture->file); + RETURN_IF_ERROR(compound_writer.add_logical_index(input)); + RETURN_IF_ERROR(compound_writer.finish()); + RETURN_IF_ERROR( + reader::SniiSegmentReader::open(&fixture->recording_reader, &fixture->segment_reader)); + return fixture->segment_reader.open_index(kIndexId, kIndexSuffix, &fixture->index_reader); +} + +format::DictEntry find_entry(const reader::LogicalIndexReader& index_reader, + std::string_view term) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + EXPECT_TRUE(index_reader.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + EXPECT_TRUE(found) << term; + return entry; +} + +uint64_t frequency_bytes(const format::DictEntry& entry) { + if (entry.kind == format::DictEntryKind::kInline) { + EXPECT_GE(entry.frq_bytes.size(), entry.inline_dd_disk_len); + return entry.frq_bytes.size() - entry.inline_dd_disk_len; + } + EXPECT_GE(entry.frq_len, entry.frq_docs_len); + return entry.frq_len - entry.frq_docs_len; +} + +uint64_t phrase_visible_bytes(const format::DictEntry& entry) { + if (entry.kind == format::DictEntryKind::kInline) { + return entry.inline_dd_disk_len + entry.prx_bytes.size(); + } + return entry.frq_docs_len + entry.prx_len; +} + +uint64_t docid_visible_bytes(const format::DictEntry& entry) { + return entry.kind == format::DictEntryKind::kInline ? entry.inline_dd_disk_len + : entry.frq_docs_len; +} + +segment_v2::InvertedIndexQueryInfo query_info(segment_v2::TermKeyKind kind, std::string term) { + segment_v2::InvertedIndexQueryInfo info; + segment_v2::TermInfo term_info; + term_info.term = std::move(term); + term_info.position = 1; + term_info.key_kind = kind; + info.term_infos.push_back(std::move(term_info)); + return info; +} + +segment_v2::InvertedIndexQueryInfo plain_query_info(std::string first, std::string second) { + segment_v2::InvertedIndexQueryInfo info; + info.term_infos.emplace_back(std::move(first), 1); + info.term_infos.emplace_back(std::move(second), 2); + return info; +} + +// Repeating the query re-plans from scratch and must land on the same plan and the same docs: +// nothing memoizes the choice, so the second run is only as stable as the inputs make it. +TEST(SniiCommonGramsDictPlan, RepeatedExactPlanningUsesDefaultReads) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + Fixture fixture; + assert_ok(build_fixture(&fixture)); + const CommonGramsQueryIdentity identity = complete_query_identity(); + const auto plain = plain_query_info("the", "wolf"); + const auto gram = query_info(segment_v2::TermKeyKind::kCommonGram, gram_key("the", "wolf")); + + std::vector expected; + assert_ok(phrase_query(fixture.index_reader, {"the", "wolf"}, &expected)); + + for (int run = 0; run < 2; ++run) { + SCOPED_TRACE(run); + fixture.recording_reader.reset_counts(); + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, &selected_plan, no_hysteresis_cost_model())); + + EXPECT_EQ(docs, expected); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_GT(fixture.recording_reader.default_read_calls(), 0U); + } +} + +TEST(SniiCommonGramsDictPlan, RepeatedPrefixPlanningUsesDefaultReads) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + Fixture fixture; + assert_ok(build_fixture(&fixture)); + const CommonGramsQueryIdentity identity = complete_query_identity(); + const auto plain = plain_query_info("the", "wo"); + const auto gram = query_info(segment_v2::TermKeyKind::kCommonGram, gram_key("the", "wo")); + + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"the", "wo"}, &expected, + /*max_expansions=*/50)); + + for (int run = 0; run < 2; ++run) { + SCOPED_TRACE(run); + fixture.recording_reader.reset_counts(); + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, plain, gram, &identity, &docs, nullptr, + /*max_expansions=*/50, &selected_plan, no_hysteresis_cost_model())); + + EXPECT_EQ(docs, expected); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_GT(fixture.recording_reader.default_read_calls(), 0U); + } +} +TEST(SniiCommonGramsDictPlan, IncompatibleExactAndPrefixFallbackUseDefaultReads) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + Fixture fixture; + assert_ok(build_fixture(&fixture)); + CommonGramsQueryIdentity incompatible_identity = complete_query_identity(); + incompatible_identity.common_grams_fingerprint = "incompatible-common-grams-query"; + + const auto exact_plain = plain_query_info("the", "wolf"); + const auto exact_gram = + query_info(segment_v2::TermKeyKind::kCommonGram, gram_key("the", "wolf")); + ExactPhrasePlanKind exact_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector exact_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, exact_plain, exact_gram, + &incompatible_identity, &exact_docs, nullptr, &exact_plan, + no_hysteresis_cost_model())); + + EXPECT_EQ(exact_plan, ExactPhrasePlanKind::kPlain); + EXPECT_EQ(exact_docs, (std::vector {0})); + EXPECT_GT(fixture.recording_reader.default_read_calls(), 0U); + + fixture.recording_reader.reset_counts(); + const auto prefix_plain = plain_query_info("the", "wo"); + const auto prefix_gram = + query_info(segment_v2::TermKeyKind::kCommonGram, gram_key("the", "wo")); + PhrasePrefixPlanKind prefix_plan = PhrasePrefixPlanKind::kCommonGrams; + std::vector prefix_docs; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, prefix_plain, prefix_gram, &incompatible_identity, &prefix_docs, + nullptr, /*max_expansions=*/50, &prefix_plan, no_hysteresis_cost_model())); + + EXPECT_EQ(prefix_plan, PhrasePrefixPlanKind::kPlain); + EXPECT_EQ(prefix_docs, (std::vector {0, 1})); + EXPECT_GT(fixture.recording_reader.default_read_calls(), 0U); +} + +TEST(SniiCommonGramsDictPlan, ScoringIndexOmitsOnlyGramFrequencyAndTermStats) { + Fixture fixture; + assert_ok(build_scoring_fixture(&fixture)); + + const std::string escaped_plain = + plain_key(std::string("\x1f" + "literal")); + const std::string dense_gram = gram_key("the", "wolf"); + const std::string small_gram = gram_key("the", "woman"); + ASSERT_LT(escaped_plain, dense_gram); + ASSERT_LT(dense_gram, plain_key("the")); + + const format::DictEntry escaped_entry = find_entry(fixture.index_reader, escaped_plain); + const format::DictEntry dense_gram_entry = find_entry(fixture.index_reader, dense_gram); + const format::DictEntry small_gram_entry = find_entry(fixture.index_reader, small_gram); + const format::DictEntry the_entry = find_entry(fixture.index_reader, plain_key("the")); + + EXPECT_TRUE(escaped_entry.term_stats_present); + EXPECT_GT(frequency_bytes(escaped_entry), 0U); + EXPECT_FALSE(dense_gram_entry.term_stats_present); + EXPECT_EQ(frequency_bytes(dense_gram_entry), 0U); + EXPECT_FALSE(small_gram_entry.term_stats_present); + EXPECT_EQ(frequency_bytes(small_gram_entry), 0U); + EXPECT_TRUE(the_entry.term_stats_present); + EXPECT_GT(frequency_bytes(the_entry), 0U); + + stats::SniiStatsProvider stats_provider; + assert_ok(stats::SniiStatsProvider::open(&fixture.index_reader, &stats_provider)); + uint64_t total_term_freq = 0; + assert_ok(stats_provider.total_term_freq(plain_key("the"), &total_term_freq)); + EXPECT_EQ(total_term_freq, kScoringDocCount); + EXPECT_FALSE(stats_provider.total_term_freq(dense_gram, &total_term_freq).ok()); + + std::vector plain_docs; + assert_ok( + phrase_query(fixture.index_reader, {plain_key("the"), plain_key("wolf")}, &plain_docs)); + EXPECT_EQ(plain_docs.size(), kScoringDocCount); +} + +TEST(SniiCommonGramsDictPlan, IncompleteCoverageKeepsGramFrequencyAndTermStats) { + for (const CommonGramsCoverage coverage : + {CommonGramsCoverage::kMixed, CommonGramsCoverage::kNone}) { + SCOPED_TRACE(static_cast(coverage)); + Fixture fixture; + assert_ok(build_scoring_fixture(&fixture, coverage)); + + const format::DictEntry gram_entry = + find_entry(fixture.index_reader, gram_key("the", "wolf")); + EXPECT_TRUE(gram_entry.term_stats_present); + EXPECT_GT(frequency_bytes(gram_entry), 0U); + } +} + +TEST(SniiCommonGramsDictPlan, ExactPlanCostMatchesExecutorReads) { + Fixture fixture; + assert_ok(build_scoring_fixture(&fixture)); + const std::string the = plain_key("the"); + const std::string wolf = plain_key("wolf"); + const std::string gram = gram_key("the", "wolf"); + const CommonGramsQueryIdentity identity = complete_query_identity(); + + QueryProfile profile; + std::vector planned_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain_query_info("the", "wolf"), + query_info(segment_v2::TermKeyKind::kCommonGram, gram), + &identity, &planned_docs, &profile, nullptr, + verify_cost_model())); + + EXPECT_EQ(planned_docs.size(), kScoringDocCount); + const auto& stats = profile.phrase_query_stats; + EXPECT_EQ(stats.common_grams_plain_posting_bytes, + phrase_visible_bytes(find_entry(fixture.index_reader, the)) + + phrase_visible_bytes(find_entry(fixture.index_reader, wolf))); + EXPECT_EQ(stats.common_grams_gram_posting_bytes, + docid_visible_bytes(find_entry(fixture.index_reader, gram))); + EXPECT_EQ(stats.common_grams_plain_estimated_cost, + stats.common_grams_plain_posting_bytes + + static_cast(kScoringDocCount) * 7 * 2); + EXPECT_EQ(stats.common_grams_gram_estimated_cost, stats.common_grams_gram_posting_bytes); +} + +TEST(SniiCommonGramsDictPlan, TwoTermPrefixGramPlanCostsOnlyDocidUnion) { + Fixture fixture; + assert_ok(build_fixture(&fixture)); + const CommonGramsQueryIdentity identity = complete_query_identity(); + + QueryProfile profile; + std::vector planned_docs; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, plain_query_info("the", "wo"), + query_info(segment_v2::TermKeyKind::kCommonGram, gram_key("the", "wo")), &identity, + &planned_docs, &profile, /*max_expansions=*/50, nullptr, verify_cost_model())); + + EXPECT_EQ(planned_docs, (std::vector {0, 1})); + const auto& stats = profile.phrase_query_stats; + EXPECT_EQ(stats.common_grams_gram_posting_bytes, + docid_visible_bytes(find_entry(fixture.index_reader, gram_key("the", "wolf"))) + + docid_visible_bytes( + find_entry(fixture.index_reader, gram_key("the", "woman")))); + EXPECT_EQ(stats.common_grams_gram_estimated_cost, stats.common_grams_gram_posting_bytes); +} + +} // namespace +} // namespace doris::snii::query diff --git a/be/test/storage/index/snii/query/common_grams_namespace_query_test.cpp b/be/test/storage/index/snii/query/common_grams_namespace_query_test.cpp new file mode 100644 index 00000000000000..c9f6702b092728 --- /dev/null +++ b/be/test/storage/index/snii/query/common_grams_namespace_query_test.cpp @@ -0,0 +1,1563 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/count_query.h" +#include "storage/index/snii/query/internal/plain_term_routing.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" +#include "util/debug_points.h" +#include "util/defer_op.h" + +namespace doris::snii::query { +namespace { + +using segment_v2::inverted_index::CommonGramsCoverage; +using segment_v2::inverted_index::CommonGramsPlanCostModel; +using segment_v2::inverted_index::CommonGramsQueryIdentity; +using segment_v2::inverted_index::CommonGramsSegmentMetadata; +using segment_v2::inverted_index::PlainTermKeyVersion; +using segment_v2::inverted_index::COMMON_GRAMS_KEY_VERSION_V1; +using segment_v2::inverted_index::COMMON_GRAMS_SEMANTICS_VERSION_V1; +using segment_v2::inverted_index::encode_common_gram; +using segment_v2::inverted_index::encode_plain_term; +using snii_test::MemoryFile; +using snii_test::ScopedEnv; +using snii_test::assert_ok; +using snii_test::make_term; + +constexpr uint64_t kIndexId = 41; +constexpr std::string_view kIndexSuffix = "body"; + +struct Fixture { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; +}; + +CommonGramsSegmentMetadata metadata_for(PlainTermKeyVersion version) { + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = version; + metadata.common_grams_coverage = version == PlainTermKeyVersion::kRawNoInternal + ? CommonGramsCoverage::kNone + : CommonGramsCoverage::kMixed; + return metadata; +} + +CommonGramsSegmentMetadata complete_metadata() { + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = CommonGramsCoverage::kComplete; + metadata.common_grams_semantics_version = COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = "builtin:lucene_english_stop:v1"; + metadata.base_analyzer_fingerprint = "common-grams-query-base-v1"; + metadata.common_grams_fingerprint = "common-grams-query-v1"; + return metadata; +} + +CommonGramsSegmentMetadata hybrid_metadata() { + CommonGramsSegmentMetadata metadata = complete_metadata(); + metadata.common_grams_coverage = CommonGramsCoverage::kMixed; + return metadata; +} + +CommonGramsQueryIdentity complete_query_identity() { + const auto metadata = complete_metadata(); + return {.common_grams_dictionary_identity = metadata.common_grams_dictionary_identity, + .base_analyzer_fingerprint = metadata.base_analyzer_fingerprint, + .common_grams_fingerprint = metadata.common_grams_fingerprint}; +} + +CommonGramsPlanCostModel no_hysteresis_cost_model() { + return {.position_verify_factor = 0, .common_grams_cost_ratio_percent = 100, .generation = 0}; +} + +CommonGramsPlanCostModel verification_dominated_cost_model() { + return {.position_verify_factor = 100000, + .common_grams_cost_ratio_percent = 100, + .generation = 0}; +} + +Status build_fixture( + Fixture* fixture, std::vector terms, + std::optional metadata = std::nullopt, + format::CommonGramsPostingPolicy posting_policy = format::CommonGramsPostingPolicy::kNone) { + std::ranges::sort(terms, [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = std::string(kIndexSuffix); + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 8; + input.terms = std::move(terms); + input.target_dict_block_bytes = 64; + input.common_grams_metadata = std::move(metadata); + input.common_grams_posting_policy = posting_policy; + + writer::SniiCompoundWriter compound(&fixture->file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(&fixture->file, &fixture->segment_reader)); + return fixture->segment_reader.open_index(kIndexId, kIndexSuffix, &fixture->index_reader); +} + +Status build_streamed_fixture(Fixture* fixture, writer::SpimiTermBuffer* terms, uint32_t doc_count, + CommonGramsSegmentMetadata metadata, + format::CommonGramsPostingPolicy posting_policy) { + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = std::string(kIndexSuffix); + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = doc_count; + input.term_source = terms; + input.target_dict_block_bytes = 64; + input.common_grams_metadata = std::move(metadata); + input.common_grams_posting_policy = posting_policy; + + writer::SniiCompoundWriter compound(&fixture->file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(&fixture->file, &fixture->segment_reader)); + return fixture->segment_reader.open_index(kIndexId, kIndexSuffix, &fixture->index_reader); +} + +std::string encoded_plain(std::string_view logical, PlainTermKeyVersion version) { + auto encoded = encode_plain_term(logical, version); + EXPECT_TRUE(encoded.has_value()); + return encoded.has_value() ? std::move(encoded.value()) : std::string(); +} + +std::vector escaped_terms() { + const std::string literal_marker = + std::string(segment_v2::inverted_index::CG_V1_MARKER) + "literal"; + auto gram = encode_common_gram("the", "cat"); + EXPECT_TRUE(gram.has_value()); + return { + make_term(encoded_plain(literal_marker, PlainTermKeyVersion::kEscapedV1), + {{.docid = 0, .positions = {0}}}), + make_term(gram.has_value() ? std::move(gram.value()) : std::string(), + {{.docid = 1, .positions = {0}}}), + make_term(std::string(format::kPhraseBigramTermMarker) + "legacy", + {{.docid = 2, .positions = {0}}}), + make_term("alpha", {{.docid = 3, .positions = {0}}}), + make_term("beta", {{.docid = 3, .positions = {1}}, {.docid = 4, .positions = {0}}}), + }; +} + +enum class TestGramPostingShape : uint8_t { + kPositioned, + kHybrid, + kDocsOnly, +}; + +std::vector build_common_grams_terms( + const std::vector>& docs, bool include_grams = true, + TestGramPostingShape posting_shape = TestGramPostingShape::kPositioned) { + writer::SpimiTermBuffer postings(/*has_positions=*/true); + const auto& common_words = + segment_v2::inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + for (uint32_t docid = 0; docid < docs.size(); ++docid) { + const auto& terms = docs[docid]; + for (uint32_t position = 0; position < terms.size(); ++position) { + postings.add_token(encoded_plain(terms[position], PlainTermKeyVersion::kEscapedV1), + docid, position); + if (!include_grams || position + 1 == terms.size()) { + continue; + } + const bool left_common = common_words.contains(terms[position]); + const bool right_common = common_words.contains(terms[position + 1]); + if (!left_common && !right_common) { + continue; + } + auto gram = encode_common_gram(terms[position], terms[position + 1]); + EXPECT_TRUE(gram.has_value()); + if (gram.has_value()) { + const bool retain_positions = posting_shape == TestGramPostingShape::kPositioned || + (posting_shape == TestGramPostingShape::kHybrid && + left_common && right_common); + postings.add_token(std::move(gram.value()), docid, position, retain_positions); + } + } + } + return postings.finalize_sorted(); +} + +format::DictEntry find_entry(const reader::LogicalIndexReader& index_reader, + std::string_view term) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + EXPECT_TRUE(index_reader.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + EXPECT_TRUE(found) << term; + return entry; +} + +bool entry_has_prx(const format::DictEntry& entry) { + return entry.kind == format::DictEntryKind::kInline ? !entry.prx_bytes.empty() + : entry.prx_len != 0; +} + +void expect_all_synthetic_grams_docs_only(const reader::LogicalIndexReader& index_reader, + const std::vector>& docs) { + const auto& common_words = + segment_v2::inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + std::set grams; + for (const auto& terms : docs) { + for (size_t position = 0; position + 1 < terms.size(); ++position) { + if (!common_words.contains(terms[position]) && + !common_words.contains(terms[position + 1])) { + continue; + } + auto gram = encode_common_gram(terms[position], terms[position + 1]); + ASSERT_TRUE(gram.has_value()); + grams.emplace(std::move(gram.value())); + } + } + ASSERT_FALSE(grams.empty()); + for (const std::string& gram : grams) { + EXPECT_FALSE(entry_has_prx(find_entry(index_reader, gram))) << gram; + } +} + +segment_v2::InvertedIndexQueryInfo query_info( + const std::vector>& terms) { + segment_v2::InvertedIndexQueryInfo info; + info.term_infos.reserve(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + segment_v2::TermInfo term_info; + term_info.term = terms[i].second; + term_info.position = static_cast(i + 1); + term_info.key_kind = terms[i].first; + info.term_infos.push_back(std::move(term_info)); + } + return info; +} + +segment_v2::InvertedIndexQueryInfo plain_query_info(const std::vector& terms) { + std::vector> typed_terms; + typed_terms.reserve(terms.size()); + for (const std::string& term : terms) { + typed_terms.emplace_back(segment_v2::TermKeyKind::kPlain, term); + } + return query_info(typed_terms); +} + +std::string gram_key(std::string_view left, std::string_view right) { + auto gram = encode_common_gram(left, right); + EXPECT_TRUE(gram.has_value()); + return gram.has_value() ? std::move(gram.value()) : std::string(); +} + +TEST(SniiCommonGramsNamespaceQuery, WideSpilledDocsOnlyGramDeltaStreamRoundTrips) { + constexpr uint32_t kDocumentCount = 768; + writer::SpimiTermBuffer source(/*has_positions=*/true); + source.enable_common_gram_pair_keys(); + source.set_forced_spill_min_arena_bytes(0); + source.set_max_run_files(2); + const writer::PlainTermId left = source.intern_plain_term("of"); + const writer::PlainTermId right = source.intern_plain_term("world"); + for (uint32_t docid = 0; docid < kDocumentCount; ++docid) { + source.add_common_gram(left, right, docid, /*pos=*/0, /*retain_positions=*/false); + if (docid == 255 || docid == 383 || docid == 511) { + source.request_global_spill_for_test(); + } + } + ASSERT_TRUE(source.status().ok()) << source.status(); + ASSERT_GT(source.run_count_for_test(), 0U); + ASSERT_LE(source.run_count_for_test(), source.max_run_files()); + + Fixture fixture; + assert_ok(build_streamed_fixture(&fixture, &source, kDocumentCount, hybrid_metadata(), + format::CommonGramsPostingPolicy::kDocsOnlyV1)); + + const std::string gram = gram_key("of", "world"); + const format::DictEntry entry = find_entry(fixture.index_reader, gram); + EXPECT_EQ(entry.enc, format::DictEntryEnc::kWindowed); + EXPECT_FALSE(entry_has_prx(entry)); + EXPECT_EQ(entry.df, kDocumentCount); + + std::vector actual; + assert_ok(term_query(fixture.index_reader, gram, &actual)); + ASSERT_EQ(actual.size(), kDocumentCount); + for (uint32_t docid = 0; docid < kDocumentCount; ++docid) { + EXPECT_EQ(actual[docid], docid); + } +} + +void expect_planned_prefix_matches_plain(const reader::LogicalIndexReader& idx, + const segment_v2::InvertedIndexQueryInfo& plain_query, + const segment_v2::InvertedIndexQueryInfo& gram_query, + const CommonGramsQueryIdentity* identity, + int32_t max_expansions, + PhrasePrefixPlanKind expected_plan) { + std::vector plain_terms; + plain_terms.reserve(plain_query.term_infos.size()); + for (const auto& term_info : plain_query.term_infos) { + plain_terms.push_back(term_info.get_single_term()); + } + SCOPED_TRACE(::testing::Message() << "first term: " << plain_terms.front() + << ", term count: " << plain_terms.size()); + + std::vector expected; + assert_ok(phrase_prefix_query(idx, plain_terms, &expected, max_expansions)); + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector actual; + assert_ok(planned_phrase_prefix_query(idx, plain_query, gram_query, identity, &actual, nullptr, + max_expansions, &selected_plan, + no_hysteresis_cost_model())); + EXPECT_EQ(selected_plan, expected_plan); + EXPECT_EQ(actual, expected); +} + +TEST(SniiCommonGramsNamespaceQuery, EscapedExactCountAnyAndAllUsePhysicalPlainKeys) { + Fixture fixture; + assert_ok(build_fixture(&fixture, escaped_terms(), + metadata_for(PlainTermKeyVersion::kEscapedV1))); + const std::string literal_marker = + std::string(segment_v2::inverted_index::CG_V1_MARKER) + "literal"; + + std::string physical_marker; + bool representable = false; + assert_ok(internal::route_plain_query_term(fixture.index_reader, literal_marker, + &physical_marker, &representable)); + ASSERT_TRUE(representable); + EXPECT_NE(physical_marker, literal_marker); + std::vector docs; + assert_ok(term_query(fixture.index_reader, physical_marker, &docs)); + EXPECT_EQ(docs, (std::vector {0})); + + uint64_t count = 0; + assert_ok(count_only_term_df(fixture.index_reader, physical_marker, &count)); + EXPECT_EQ(count, 1U); + + std::vector physical_terms {literal_marker, "alpha"}; + bool all_representable = false; + assert_ok(internal::route_plain_query_terms(fixture.index_reader, {literal_marker, "alpha"}, + &physical_terms, &all_representable)); + ASSERT_TRUE(all_representable); + assert_ok(boolean_or(fixture.index_reader, physical_terms, &docs)); + EXPECT_EQ(docs, (std::vector {0, 3})); + assert_ok(boolean_and(fixture.index_reader, physical_terms, &docs)); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, GeneratedGramRequiresPlannerBeforePhysicalLookup) { + Fixture fixture; + assert_ok(build_fixture(&fixture, escaped_terms(), + metadata_for(PlainTermKeyVersion::kEscapedV1))); + auto gram = encode_common_gram("the", "cat"); + ASSERT_TRUE(gram.has_value()); + + segment_v2::TermInfo term_info; + term_info.term = std::move(gram.value()); + term_info.key_kind = segment_v2::TermKeyKind::kCommonGram; + std::string physical_term; + bool representable = false; + EXPECT_TRUE(internal::route_query_term(fixture.index_reader, term_info, &physical_term, + &representable) + .is()); + + segment_v2::InvertedIndexQueryInfo query_info; + query_info.term_infos.emplace_back(term_info); + query_info.term_infos.emplace_back("plain", 1); + std::vector physical_terms {term_info.get_single_term(), "plain"}; + bool all_representable = false; + EXPECT_TRUE(internal::route_query_terms(fixture.index_reader, query_info, &physical_terms, + &all_representable) + .is()); +} + +TEST(SniiCommonGramsNamespaceQuery, EscapedExpansionDecodesPlainAndSkipsInternalBudget) { + Fixture fixture; + assert_ok(build_fixture(&fixture, escaped_terms(), + metadata_for(PlainTermKeyVersion::kEscapedV1))); + const std::string literal_marker = + std::string(segment_v2::inverted_index::CG_V1_MARKER) + "literal"; + + std::vector docs; + assert_ok(prefix_query(fixture.index_reader, literal_marker, &docs)); + EXPECT_EQ(docs, (std::vector {0})); + + assert_ok(prefix_query(fixture.index_reader, "", &docs, /*max_expansions=*/2)); + EXPECT_EQ(docs, (std::vector {0, 3})); + assert_ok(wildcard_query(fixture.index_reader, "*", &docs, /*max_expansions=*/2)); + EXPECT_EQ(docs, (std::vector {0, 3})); + assert_ok(regexp_query(fixture.index_reader, ".*", &docs, /*max_expansions=*/2)); + EXPECT_EQ(docs, (std::vector {0, 3})); + + assert_ok(phrase_prefix_query(fixture.index_reader, {"alpha", "b"}, &docs, + /*max_expansions=*/2)); + EXPECT_EQ(docs, (std::vector {3})); +} + +TEST(SniiCommonGramsNamespaceQuery, EmptyLeadingRangeSkipsColdDictWhenNamespaceIsFirst) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + Fixture fixture; + auto gram = encode_common_gram("the", "cat"); + ASSERT_TRUE(gram.has_value()); + assert_ok(build_fixture(&fixture, + {make_term(std::move(gram.value()), {{.docid = 1, .positions = {0}}})}, + metadata_for(PlainTermKeyVersion::kEscapedV1))); + + fixture.file.clear_reads(); + size_t visited = 0; + assert_ok(fixture.index_reader.visit_term_range( + {}, segment_v2::inverted_index::INTERNAL_TERM_NAMESPACE_BEGIN, + [&](reader::LogicalIndexReader::PrefixHit&&, bool*) { + ++visited; + return Status::OK(); + })); + + EXPECT_EQ(visited, 0U); + EXPECT_TRUE(fixture.file.reads().empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, LegacyRawWithoutReservedKeysAllowsLeadingExpansion) { + Fixture fixture; + assert_ok(build_fixture(&fixture, {make_term("alpha", {{.docid = 1, .positions = {0}}}), + make_term("screenshot", {{.docid = 2, .positions = {0}}})})); + + std::vector docs; + assert_ok(wildcard_query(fixture.index_reader, "*shot", &docs)); + EXPECT_EQ(docs, (std::vector {2})); + assert_ok(regexp_query(fixture.index_reader, ".*shot", &docs)); + EXPECT_EQ(docs, (std::vector {2})); + assert_ok(prefix_query(fixture.index_reader, "", &docs)); + EXPECT_EQ(docs, (std::vector {1, 2})); +} + +TEST(SniiCommonGramsNamespaceQuery, LegacyLeadingExpansionReadsEachColdDictBlockOnce) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + Fixture fixture; + assert_ok(build_fixture(&fixture, {make_term("alpha", {{.docid = 1, .positions = {0}}}), + make_term("screenshot", {{.docid = 2, .positions = {0}}})})); + + fixture.file.clear_reads(); + size_t visited = 0; + assert_ok(internal::visit_expanded_plain_terms( + fixture.index_reader, "", [](std::string_view) { return true; }, + [&](reader::LogicalIndexReader::PrefixHit&&, bool*) { + ++visited; + return Status::OK(); + })); + + EXPECT_EQ(visited, 2U); + std::set> unique_reads; + for (const auto& read : fixture.file.reads()) { + unique_reads.emplace(read.offset, read.len); + } + EXPECT_EQ(unique_reads.size(), fixture.file.reads().size()); +} + +TEST(SniiCommonGramsNamespaceQuery, LegacyReservedQueriesBypassTheIndex) { + for (const std::string& hidden : + {std::string(format::kPhraseBigramTermMarker) + "legacy", + std::string(segment_v2::inverted_index::CG_V1_MARKER) + "legacy"}) { + Fixture fixture; + assert_ok(build_fixture(&fixture, {make_term(hidden, {{.docid = 1, .positions = {0}}}), + make_term("alpha", {{.docid = 2, .positions = {0}}})})); + + std::string physical; + bool representable = false; + EXPECT_TRUE(internal::route_plain_query_term(fixture.index_reader, hidden, &physical, + &representable) + .is()); + std::vector docs; + EXPECT_TRUE(prefix_query(fixture.index_reader, "", &docs) + .is()); + EXPECT_TRUE(wildcard_query(fixture.index_reader, "*", &docs) + .is()); + EXPECT_TRUE(regexp_query(fixture.index_reader, ".*", &docs) + .is()); + } +} + +TEST(SniiCommonGramsNamespaceQuery, RawNoInternalTreatsReservedBytesAsPlain) { + Fixture fixture; + const std::string marker_literal = std::string(format::kPhraseBigramTermMarker) + "literal"; + assert_ok(build_fixture(&fixture, + {make_term(marker_literal, {{.docid = 5, .positions = {0}}}), + make_term("alpha", {{.docid = 6, .positions = {0}}})}, + metadata_for(PlainTermKeyVersion::kRawNoInternal))); + + std::string physical; + bool representable = false; + assert_ok(internal::route_plain_query_term(fixture.index_reader, marker_literal, &physical, + &representable)); + ASSERT_TRUE(representable); + EXPECT_EQ(physical, marker_literal); + std::vector docs; + assert_ok(term_query(fixture.index_reader, physical, &docs)); + EXPECT_EQ(docs, (std::vector {5})); + assert_ok(prefix_query(fixture.index_reader, "\x1f", &docs)); + EXPECT_EQ(docs, (std::vector {5})); + assert_ok(wildcard_query(fixture.index_reader, "\x1f*", &docs)); + EXPECT_EQ(docs, (std::vector {5})); + assert_ok(regexp_query(fixture.index_reader, "\x1f.*", &docs)); + EXPECT_EQ(docs, (std::vector {5})); +} + +TEST(SniiCommonGramsNamespaceQuery, EscapedUnrepresentableBoundaryIsAuthoritativeAbsent) { + Fixture fixture; + assert_ok(build_fixture(&fixture, {make_term("alpha", {{.docid = 0, .positions = {0}}})}, + metadata_for(PlainTermKeyVersion::kEscapedV1))); + std::string logical(segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + logical.front() = segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX; + + std::string physical; + bool representable = true; + assert_ok(internal::route_plain_query_term(fixture.index_reader, logical, &physical, + &representable)); + EXPECT_FALSE(representable); + EXPECT_TRUE(physical.empty()); + + std::vector docs; + assert_ok(prefix_query(fixture.index_reader, logical, &docs)); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, EscapedAllUnrepresentableAnyTermsResolveEmpty) { + Fixture fixture; + assert_ok(build_fixture(&fixture, {make_term("alpha", {{.docid = 0, .positions = {0}}})}, + metadata_for(PlainTermKeyVersion::kEscapedV1))); + std::string escape_term(segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + escape_term.front() = segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX; + std::string marker_term(segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + marker_term.front() = '\x1f'; + + std::vector physical_terms {escape_term, marker_term}; + bool all_representable = true; + assert_ok(internal::route_plain_query_terms(fixture.index_reader, {escape_term, marker_term}, + &physical_terms, &all_representable)); + EXPECT_FALSE(all_representable); + EXPECT_TRUE(physical_terms.empty()); + + std::vector docs; + assert_ok(boolean_or(fixture.index_reader, physical_terms, &docs)); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, ExactPlannerReusesMatcherForEveryCommonTermShape) { + const std::vector> corpus = { + {"cat", "dog", "fox"}, {"cat", "dog", "the"}, {"cat", "the", "dog"}, + {"cat", "the", "the"}, {"the", "cat", "dog"}, {"the", "cat", "the"}, + {"the", "the", "cat"}, {"the", "the", "the"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const std::vector gram_plans = { + query_info({{Kind::kPlain, "cat"}, {Kind::kPlain, "dog"}, {Kind::kPlain, "fox"}}), + query_info({{Kind::kPlain, "cat"}, {Kind::kCommonGram, gram_key("dog", "the")}}), + query_info({{Kind::kCommonGram, gram_key("cat", "the")}, + {Kind::kCommonGram, gram_key("the", "dog")}}), + query_info({{Kind::kCommonGram, gram_key("cat", "the")}, + {Kind::kCommonGram, gram_key("the", "the")}}), + query_info({{Kind::kCommonGram, gram_key("the", "cat")}, + {Kind::kPlain, "cat"}, + {Kind::kPlain, "dog"}}), + query_info({{Kind::kCommonGram, gram_key("the", "cat")}, + {Kind::kCommonGram, gram_key("cat", "the")}}), + query_info({{Kind::kCommonGram, gram_key("the", "the")}, + {Kind::kCommonGram, gram_key("the", "cat")}}), + query_info({{Kind::kCommonGram, gram_key("the", "the")}, + {Kind::kCommonGram, gram_key("the", "the")}})}; + + for (size_t i = 0; i < corpus.size(); ++i) { + std::vector plain_docs; + assert_ok(phrase_query(fixture.index_reader, corpus[i], &plain_docs)); + ASSERT_EQ(plain_docs, (std::vector {static_cast(i)})); + + std::vector planned_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain_query_info(corpus[i]), + gram_plans[i], &identity, &planned_docs)); + EXPECT_EQ(planned_docs, plain_docs) << "shape " << i; + } +} + +TEST(SniiCommonGramsNamespaceQuery, TwoTermsUseSingleGramTermQueryWithoutPrxDecode) { + const std::vector> corpus = { + {"the", "cat"}, {"the", "gap", "cat"}, {"cat", "the"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "cat"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "cat")}}); + QueryProfile profile; + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + &profile, &selected_plan)); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0})); + EXPECT_EQ(profile.prx_decode_stats.frame_count(), 0U); +} + +TEST(SniiCommonGramsNamespaceQuery, + DocsOnlyAllCommonExactUsesTermDocsetOrPlainPrxVerificationByLength) { + const std::vector> corpus = { + {"the", "of"}, + {"the", "gap", "of"}, + {"the", "of", "a"}, + {"the", "of", "gap", "of", "a"}, + {"the", "of", "a", "in", "the"}, + {"the", "of", "gap", "of", "a", "gap", "a", "in", "gap", "in", "the"}, + }; + Fixture fixture; + assert_ok(build_fixture(&fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, + TestGramPostingShape::kDocsOnly), + hybrid_metadata(), format::CommonGramsPostingPolicy::kDocsOnlyV1)); + expect_all_synthetic_grams_docs_only(fixture.index_reader, corpus); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + for (const std::vector& terms : + {std::vector {"the", "of"}, std::vector {"the", "of", "a"}, + std::vector {"the", "of", "a", "in", "the"}}) { + std::vector gram_keys; + std::vector> gram_terms; + for (size_t i = 0; i + 1 < terms.size(); ++i) { + gram_keys.push_back(gram_key(terms[i], terms[i + 1])); + gram_terms.emplace_back(Kind::kCommonGram, gram_keys.back()); + } + + std::vector expected; + assert_ok(phrase_query(fixture.index_reader, terms, &expected)); + std::vector gram_candidates; + assert_ok(boolean_and(fixture.index_reader, gram_keys, &gram_candidates)); + if (terms.size() == 2) { + ASSERT_EQ(gram_candidates, expected); + } else { + ASSERT_GT(gram_candidates.size(), expected.size()); + } + + QueryProfile profile; + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector actual; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain_query_info(terms), + query_info(gram_terms), &identity, &actual, &profile, + &selected_plan, no_hysteresis_cost_model(), + CommonGramsPlanDebugOverride::kForceCommonGrams)); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(actual, expected) << "term count " << terms.size(); + if (terms.size() == 2) { + EXPECT_EQ(profile.prx_decode_stats.frame_count(), 0U); + } + } +} + +TEST(SniiCommonGramsNamespaceQuery, + DocsOnlyAllCommonPhrasePrefixUsesGramCandidatesThenPlainPrxVerification) { + const std::vector> corpus = { + {"the", "of", "a", "in"}, {"the", "of", "a", "into"}, + {"the", "of", "a", "it"}, {"the", "of", "gap", "of", "a", "gap", "a", "in"}, + {"the", "of", "a", "gap", "in"}, + }; + Fixture fixture; + assert_ok(build_fixture(&fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, + TestGramPostingShape::kDocsOnly), + hybrid_metadata(), format::CommonGramsPostingPolicy::kDocsOnlyV1)); + expect_all_synthetic_grams_docs_only(fixture.index_reader, corpus); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "of", "a", "i"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "of")}, + {Kind::kCommonGram, gram_key("of", "a")}, + {Kind::kCommonGram, gram_key("a", "i")}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"the", "of", "a", "i"}, &expected, + /*max_expansions=*/50)); + ASSERT_EQ(expected, (std::vector {0, 1, 2})); + + std::vector in_gram_candidates; + assert_ok(boolean_and(fixture.index_reader, + {gram_key("the", "of"), gram_key("of", "a"), gram_key("a", "in")}, + &in_gram_candidates)); + ASSERT_TRUE(std::ranges::find(in_gram_candidates, 3) != in_gram_candidates.end()); + + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector actual; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, plain, gram, &identity, &actual, nullptr, + /*max_expansions=*/50, &selected_plan, no_hysteresis_cost_model(), + CommonGramsPlanDebugOverride::kForceCommonGrams)); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(actual, expected); +} + +TEST(SniiCommonGramsNamespaceQuery, RepeatedGramUsesExistingRepeatedTermPhrasePath) { + const std::vector> corpus = { + {"the", "the", "the"}, {"the", "gap", "the"}, {"the", "the", "cat"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "the", "the"}); + const std::string the_the = gram_key("the", "the"); + const auto gram = query_info({{Kind::kCommonGram, the_the}, {Kind::kCommonGram, the_the}}); + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, &selected_plan, no_hysteresis_cost_model())); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0})); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridExactUsesDocsOnlyPrefilterAndMinimalPositionedCoverFreshAndCached) { + const std::vector> corpus = { + {"the", "of", "wolf"}, + {"the", "of", "gap", "of", "wolf"}, + {"the", "gap", "of", "wolf"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string the_of = gram_key("the", "of"); + const std::string of_wolf = gram_key("of", "wolf"); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, the_of))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, of_wolf))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "of", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, the_of}, {Kind::kCommonGram, of_wolf}}); + std::vector expected; + assert_ok(phrase_query(fixture.index_reader, {"the", "of", "wolf"}, &expected)); + ASSERT_EQ(expected, (std::vector {0})); + + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector cost_selected_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, + &cost_selected_docs, nullptr, &selected_plan, + no_hysteresis_cost_model())); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(cost_selected_docs, expected); + + selected_plan = ExactPhrasePlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &fresh_docs, + nullptr, &selected_plan, no_hysteresis_cost_model(), + CommonGramsPlanDebugOverride::kForceCommonGrams)); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); +} + +TEST(SniiCommonGramsNamespaceQuery, HybridExactPrunesDominatedPositionedMiddleEdgeFreshAndCached) { + const std::vector> corpus = { + {"the", "of", "a", "in", "wolf"}, + {"the", "of", "a", "gap", "a", "in", "wolf"}, + {"the", "gap", "of", "a", "in", "wolf"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string the_of = gram_key("the", "of"); + const std::string of_a = gram_key("of", "a"); + const std::string a_in = gram_key("a", "in"); + const std::string in_wolf = gram_key("in", "wolf"); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, the_of))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, of_a))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, a_in))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, in_wolf))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "of", "a", "in", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, the_of}, + {Kind::kCommonGram, of_a}, + {Kind::kCommonGram, a_in}, + {Kind::kCommonGram, in_wolf}}); + std::vector expected; + assert_ok(phrase_query(fixture.index_reader, {"the", "of", "a", "in", "wolf"}, &expected)); + ASSERT_EQ(expected, (std::vector {0})); + + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &fresh_docs, + nullptr, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridExactEmptyCandidateIntersectionIsAuthoritativeFreshAndCached) { + const std::vector> corpus = { + {"wolf", "the", "gap", "fox", "gap", "cat"}, + {"gap", "the", "fox", "gap", "cat", "gap", "wolf"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + {"wolf", "gap", "the", "gap", "fox", "cat"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string wolf_the = gram_key("wolf", "the"); + const std::string the_fox = gram_key("the", "fox"); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, wolf_the))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, the_fox))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"wolf", "the", "fox", "cat"}); + const auto gram = query_info({{Kind::kCommonGram, wolf_the}, + {Kind::kCommonGram, the_fox}, + {Kind::kPlain, "fox"}, + {Kind::kPlain, "cat"}}); + QueryProfile fresh_profile; + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + &fresh_profile, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); + EXPECT_EQ(fresh_profile.phrase_query_stats.common_grams_authoritative_empty, 1U); +} + +TEST(SniiCommonGramsNamespaceQuery, ExactPlannerMatchesPlainAtRequiredPhraseLengths) { + const std::vector lengths = {1, 2, 3, 6, 10}; + std::vector> corpus; + corpus.reserve(lengths.size()); + for (size_t length : lengths) { + corpus.emplace_back(length, "the"); + } + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const std::string the_the = gram_key("the", "the"); + for (size_t length : lengths) { + const std::vector plain_terms(length, "the"); + std::vector> gram_terms; + if (length == 1) { + gram_terms.emplace_back(Kind::kPlain, "the"); + } else { + gram_terms.assign(length - 1, {Kind::kCommonGram, the_the}); + } + + std::vector plain_docs; + assert_ok(phrase_query(fixture.index_reader, plain_terms, &plain_docs)); + std::vector planned_docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain_query_info(plain_terms), + query_info(gram_terms), &identity, &planned_docs)); + EXPECT_EQ(planned_docs, plain_docs) << "length " << length; + } +} + +TEST(SniiCommonGramsNamespaceQuery, CompleteCoverageGramMissIsAuthoritativeEmpty) { + const std::vector> corpus = {{"the", "wolf"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus, /*include_grams=*/false), + complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "wolf")}}); + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, &selected_plan, /*cost_model=*/ {})); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, CompleteCoverageChecksGramMissBeforeCostGate) { + using Kind = segment_v2::TermKeyKind; + const std::string the_of = gram_key("the", "of"); + Fixture fixture; + assert_ok(build_fixture(&fixture, + {make_term("the", {{.docid = 0, .positions = {0}}}), + make_term("of", {{.docid = 0, .positions = {1}}}), + make_term("wolf", {{.docid = 0, .positions = {2}}}), + make_term(the_of, {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {0}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {0}}, + {.docid = 6, .positions = {0}}, + {.docid = 7, .positions = {0}}})}, + complete_metadata())); + const auto plain = plain_query_info({"the", "of", "wolf"}); + const auto gram = + query_info({{Kind::kCommonGram, the_of}, {Kind::kCommonGram, gram_key("of", "wolf")}}); + const auto identity = complete_query_identity(); + + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kPlain; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, &selected_plan)); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, CompleteCoveragePlainMissIsAuthoritativeEmpty) { + using Kind = segment_v2::TermKeyKind; + const std::string the_wolf = gram_key("the", "wolf"); + Fixture fixture; + assert_ok(build_fixture(&fixture, + {make_term("the", {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {0}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {0}}, + {.docid = 6, .positions = {0}}, + {.docid = 7, .positions = {0}}}), + make_term(the_wolf, {{.docid = 0, .positions = {0}}})}, + complete_metadata())); + const auto plain = plain_query_info({"the", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, the_wolf}}); + const auto identity = complete_query_identity(); + + QueryProfile fresh_profile; + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + &fresh_profile, &selected_plan, /*cost_model=*/ {})); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kPlain); + EXPECT_TRUE(docs.empty()); + EXPECT_EQ(fresh_profile.phrase_query_stats.common_grams_plain_plans, 1U); + EXPECT_EQ(fresh_profile.phrase_query_stats.common_grams_authoritative_empty, 1U); + EXPECT_GT(fresh_profile.phrase_query_stats.common_grams_planning_ns, 0U); +} + +TEST(SniiCommonGramsNamespaceQuery, OldAndIncompleteCoverageForcePlainPlan) { + const std::vector> corpus = {{"the", "wolf"}}; + const auto terms = build_common_grams_terms(corpus, /*include_grams=*/false); + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "wolf")}}); + const auto identity = complete_query_identity(); + + for (const auto& metadata : {std::optional {}, + std::optional { + metadata_for(PlainTermKeyVersion::kEscapedV1)}}) { + Fixture fixture; + assert_ok(build_fixture(&fixture, terms, metadata)); + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, &selected_plan)); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kPlain); + EXPECT_EQ(docs, (std::vector {0})); + } +} + +TEST(SniiCommonGramsNamespaceQuery, MismatchedOrMissingQueryIdentityForcesPlainPlan) { + const std::vector> corpus = {{"the", "wolf"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus, /*include_grams=*/false), + complete_metadata())); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "wolf"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "wolf")}}); + auto mismatched = complete_query_identity(); + mismatched.common_grams_fingerprint = "different-common-grams-policy"; + const CommonGramsQueryIdentity* identities[] = {nullptr, &mismatched}; + for (const CommonGramsQueryIdentity* identity : identities) { + ExactPhrasePlanKind selected_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector docs; + assert_ok(planned_exact_phrase_query(fixture.index_reader, plain, gram, identity, &docs, + nullptr, &selected_plan)); + + EXPECT_EQ(selected_plan, ExactPhrasePlanKind::kPlain); + EXPECT_EQ(docs, (std::vector {0})); + } +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerMapsFrozenTailTerms) { + const std::vector> corpus = { + {"the", "wolf"}, {"the", "woman"}, {"the", "gap", "wolf"}, {"alpha", "wolf"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "wo"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "wo")}}); + QueryProfile profile; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &docs, + &profile, /*max_expansions=*/50, &selected_plan)); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0, 1})); + EXPECT_GT(profile.phrase_query_stats.common_grams_planning_ns, 0U); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridPrefixUnionsDocsOnlyMappedTailsBeforePlainVerificationFreshAndCached) { + const std::vector> corpus = { + {"the", "of", "wolf"}, + {"the", "of", "woman"}, + {"the", "of", "gap", "of", "wolf"}, + {"the", "of", "gap", "of", "woman"}, + {"the", "gap", "of", "wolf"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string the_of = gram_key("the", "of"); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, the_of))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, gram_key("of", "wolf")))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, gram_key("of", "woman")))); + + struct PrefixCase { + std::string prefix; + std::vector present_gram_tail_ordinals; + std::vector expected_docs; + }; + const std::vector cases = { + {.prefix = "wo", .present_gram_tail_ordinals = {0, 1}, .expected_docs = {0, 1}}, + {.prefix = "wol", .present_gram_tail_ordinals = {0}, .expected_docs = {0}}, + }; + + using Kind = segment_v2::TermKeyKind; + for (const PrefixCase& test_case : cases) { + SCOPED_TRACE(test_case.prefix); + const auto plain = plain_query_info({"the", "of", test_case.prefix}); + const auto gram = query_info({{Kind::kCommonGram, the_of}, + {Kind::kCommonGram, gram_key("of", test_case.prefix)}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"the", "of", test_case.prefix}, + &expected, /*max_expansions=*/50)); + ASSERT_EQ(expected, test_case.expected_docs); + + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kCommonGrams; + std::vector cost_selected_docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, + &cost_selected_docs, nullptr, /*max_expansions=*/50, + &selected_plan, no_hysteresis_cost_model())); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(cost_selected_docs, expected); + + selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, plain, gram, &identity, &fresh_docs, nullptr, + /*max_expansions=*/50, &selected_plan, no_hysteresis_cost_model(), + CommonGramsPlanDebugOverride::kForceCommonGrams)); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); + } +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridPrefixPositionedTailUsesSparseLogicalOffsetFreshAndCached) { + const std::vector> corpus = {{"one", "of", "the", "in", "the"}}; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string one_of = gram_key("one", "of"); + const std::string of_the = gram_key("of", "the"); + const std::string the_in = gram_key("the", "in"); + const std::string in_the = gram_key("in", "the"); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, one_of))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, of_the))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, the_in))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, in_the))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"one", "of", "the", "in", "th"}); + const auto gram = query_info({{Kind::kCommonGram, one_of}, + {Kind::kCommonGram, of_the}, + {Kind::kCommonGram, the_in}, + {Kind::kCommonGram, gram_key("in", "th")}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"one", "of", "the", "in", "th"}, &expected, + /*max_expansions=*/50)); + ASSERT_EQ(expected, (std::vector {0})); + + QueryProfile fresh_profile; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &fresh_docs, + &fresh_profile, /*max_expansions=*/50, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); + EXPECT_EQ(fresh_profile.prx_decode_stats.selected_positions, 3U); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridPrefixMultiplePositionedTailsUseSparseCoverFreshAndCached) { + const std::vector> corpus = { + {"one", "of", "the", "in", "the"}, + {"one", "of", "the", "in", "this"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string one_of = gram_key("one", "of"); + const std::string of_the = gram_key("of", "the"); + const std::string the_in = gram_key("the", "in"); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, one_of))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, of_the))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, the_in))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, gram_key("in", "the")))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, gram_key("in", "this")))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"one", "of", "the", "in", "th"}); + const auto gram = query_info({{Kind::kCommonGram, one_of}, + {Kind::kCommonGram, of_the}, + {Kind::kCommonGram, the_in}, + {Kind::kCommonGram, gram_key("in", "th")}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"one", "of", "the", "in", "th"}, &expected, + /*max_expansions=*/50)); + ASSERT_EQ(expected, (std::vector {0, 1})); + + QueryProfile fresh_profile; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &fresh_docs, + &fresh_profile, /*max_expansions=*/50, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); + EXPECT_EQ(fresh_profile.prx_decode_stats.selected_positions, 6U); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridPrefixSplitsPositionedAndDocsOnlyMappedTailsFreshAndCached) { + const std::vector> corpus = { + {"one", "of", "the", "in", "the"}, + {"one", "of", "the", "in", "thing"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + const std::string one_of = gram_key("one", "of"); + const std::string of_the = gram_key("of", "the"); + const std::string the_in = gram_key("the", "in"); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, one_of))); + EXPECT_TRUE(entry_has_prx(find_entry(fixture.index_reader, gram_key("in", "the")))); + EXPECT_FALSE(entry_has_prx(find_entry(fixture.index_reader, gram_key("in", "thing")))); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"one", "of", "the", "in", "t"}); + const auto gram = query_info({{Kind::kCommonGram, one_of}, + {Kind::kCommonGram, of_the}, + {Kind::kCommonGram, the_in}, + {Kind::kCommonGram, gram_key("in", "t")}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"one", "of", "the", "in", "t"}, &expected, + /*max_expansions=*/50)); + ASSERT_EQ(expected, (std::vector {0, 1})); + + QueryProfile fresh_profile; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector fresh_docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &fresh_docs, + &fresh_profile, /*max_expansions=*/50, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(fresh_docs, expected); + EXPECT_EQ(fresh_profile.prx_decode_stats.selected_positions, 7U); +} + +TEST(SniiCommonGramsNamespaceQuery, + HybridPrefixEmptyCandidatesAreAuthoritativeForFreshAndCachedPlans) { + const std::vector> corpus = { + {"the", "of", "gap"}, + {"x", "of", "wolf"}, + {"x", "of", "woman"}, + {"the", "x", "of", "x", "wolf", "woman"}, + {"the", "x", "of", "x", "wolf", "woman"}, + {"the", "x", "of", "x", "wolf", "woman"}, + {"the", "x", "of", "x", "wolf", "woman"}, + {"the", "x", "of", "x", "wolf", "woman"}, + }; + Fixture fixture; + assert_ok(build_fixture( + &fixture, + build_common_grams_terms(corpus, /*include_grams=*/true, TestGramPostingShape::kHybrid), + hybrid_metadata(), format::CommonGramsPostingPolicy::kHybridV1)); + const auto identity = complete_query_identity(); + + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "of", "wo"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "of")}, + {Kind::kCommonGram, gram_key("of", "wo")}}); + + QueryProfile fresh_profile; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &docs, + &fresh_profile, /*max_expansions=*/50, &selected_plan, + verification_dominated_cost_model())); + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); + EXPECT_EQ(fresh_profile.phrase_query_stats.common_grams_authoritative_empty, 1U); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerDropsOnlyMissingMappedTails) { + const std::vector> corpus = {{"the", "wolf"}, {"the", "woman"}}; + auto terms = build_common_grams_terms(corpus); + const std::string missing_gram = gram_key("the", "woman"); + std::erase_if(terms, + [&](const writer::TermPostings& term) { return term.term == missing_gram; }); + + Fixture fixture; + assert_ok(build_fixture(&fixture, std::move(terms), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain_query_info({"the", "wo"}), + query_info({{Kind::kCommonGram, gram_key("the", "wo")}}), + &identity, &docs, nullptr, + /*max_expansions=*/50, &selected_plan)); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0})); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerTreatsAllMappedTailMissesAsEmpty) { + const std::vector> corpus = {{"the", "wolf"}, {"the", "woman"}}; + auto terms = build_common_grams_terms(corpus); + const std::string wolf_gram = gram_key("the", "wolf"); + const std::string woman_gram = gram_key("the", "woman"); + std::erase_if(terms, [&](const writer::TermPostings& term) { + return term.term == wolf_gram || term.term == woman_gram; + }); + + Fixture fixture; + assert_ok(build_fixture(&fixture, std::move(terms), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain_query_info({"the", "wo"}), + query_info({{Kind::kCommonGram, gram_key("the", "wo")}}), + &identity, &docs, nullptr, + /*max_expansions=*/50, &selected_plan)); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixRepeatedGramBuildsOneTermPlan) { + const std::vector> corpus = { + {"the", "the", "the"}, {"the", "the", "thing"}, {"the", "gap", "the"}}; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + const std::string the_the = gram_key("the", "the"); + + internal::query_test_counters() = {}; + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query( + fixture.index_reader, plain_query_info({"the", "the", "th"}), + query_info({{Kind::kCommonGram, the_the}, {Kind::kCommonGram, gram_key("the", "th")}}), + &identity, &docs, nullptr, /*max_expansions=*/1, &selected_plan, + no_hysteresis_cost_model())); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0})); + EXPECT_EQ(internal::query_test_counters().resolved_term_entry_moves, 1U); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerReusesExistingMatcherShapes) { + const std::vector> corpus = { + {"foo", "of", "theta"}, {"foo", "of", "thing"}, + {"foo", "gap", "of", "theta"}, {"foo", "of", "gap", "theta"}, + {"the", "bar", "baz"}, {"the", "bar", "batch"}, + {"the", "gap", "bar", "baz"}, {"foo", "theater", "the", "x", "bar", "bask"}, + }; + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + + expect_planned_prefix_matches_plain(fixture.index_reader, plain_query_info({"foo", "of", "th"}), + query_info({{Kind::kCommonGram, gram_key("foo", "of")}, + {Kind::kCommonGram, gram_key("of", "th")}}), + &identity, /*max_expansions=*/50, + PhrasePrefixPlanKind::kCommonGrams); + expect_planned_prefix_matches_plain( + fixture.index_reader, plain_query_info({"the", "bar", "baz"}), + query_info({{Kind::kCommonGram, gram_key("the", "bar")}, + {Kind::kPlain, "bar"}, + {Kind::kPlain, "baz"}}), + &identity, /*max_expansions=*/50, PhrasePrefixPlanKind::kCommonGrams); + + // The prefix crosses the common/non-common boundary: "the" is common, + // "theater" is not. The phrase-prefix analyzer therefore emits no gram. + expect_planned_prefix_matches_plain(fixture.index_reader, plain_query_info({"foo", "the"}), + plain_query_info({"foo", "the"}), &identity, + /*max_expansions=*/50, PhrasePrefixPlanKind::kPlain); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerMatchesRequiredPhraseLengthsAndCaps) { + const std::vector lengths = {1, 2, 3, 6, 10}; + std::vector> corpus; + corpus.reserve(lengths.size()); + for (size_t length : lengths) { + corpus.emplace_back(length, "the"); + } + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + const std::string the_the = gram_key("the", "the"); + + for (size_t length : lengths) { + std::vector plain_terms(length, "the"); + plain_terms.back() = "th"; + std::vector> gram_terms; + if (length == 1) { + gram_terms.emplace_back(Kind::kPlain, "th"); + } else { + gram_terms.assign(length - 1, {Kind::kCommonGram, the_the}); + gram_terms.back().second = gram_key("the", "th"); + } + expect_planned_prefix_matches_plain( + fixture.index_reader, plain_query_info(plain_terms), query_info(gram_terms), + &identity, /*max_expansions=*/50, + length == 1 ? PhrasePrefixPlanKind::kPlain : PhrasePrefixPlanKind::kCommonGrams); + } + + const auto plain = plain_query_info({"the", "th"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "th")}}); + for (int32_t cap : {1, 2, 32, 50, 64}) { + expect_planned_prefix_matches_plain(fixture.index_reader, plain, gram, &identity, cap, + PhrasePrefixPlanKind::kCommonGrams); + } +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixGramPlanEmitsMatchesBeforeEmptyFinalTailGroup) { + std::vector> corpus(8); + corpus[0] = {"foo", "of", "aa_000"}; + for (size_t i = 1; i < 64; ++i) { + corpus[1].push_back("of"); + corpus[1].push_back(fmt::format("aa_{:03}", i)); + } + for (size_t docid = 2; docid < corpus.size(); ++docid) { + corpus[docid] = {"foo", "gap", "of", fmt::format("filler_{}", docid)}; + } + + Fixture fixture; + assert_ok(build_fixture(&fixture, build_common_grams_terms(corpus), complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"foo", "of", "aa_"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("foo", "of")}, + {Kind::kCommonGram, gram_key("of", "aa_")}}); + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kPlain; + std::vector docs; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &docs, + nullptr, /*max_expansions=*/64, &selected_plan, + no_hysteresis_cost_model())); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_EQ(docs, (std::vector {0})); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerFallsBackForUnencodableMappedTail) { + const std::string long_tail(segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES - 1, 'x'); + Fixture fixture; + assert_ok(build_fixture(&fixture, + {make_term(encoded_plain("the", PlainTermKeyVersion::kEscapedV1), + {{.docid = 0, .positions = {0}}}), + make_term(encoded_plain(long_tail, PlainTermKeyVersion::kEscapedV1), + {{.docid = 0, .positions = {1}}})}, + complete_metadata())); + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + const auto plain = plain_query_info({"the", "x"}); + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "x")}}); + std::vector expected; + assert_ok(phrase_prefix_query(fixture.index_reader, {"the", "x"}, &expected, + /*max_expansions=*/50)); + PhrasePrefixPlanKind selected_plan = PhrasePrefixPlanKind::kCommonGrams; + std::vector actual; + assert_ok(planned_phrase_prefix_query(fixture.index_reader, plain, gram, &identity, &actual, + nullptr, /*max_expansions=*/50, &selected_plan, + /*cost_model=*/ {})); + + EXPECT_EQ(selected_plan, PhrasePrefixPlanKind::kPlain); + EXPECT_EQ(actual, expected); +} + +TEST(SniiCommonGramsNamespaceQuery, PhrasePrefixPlannerRequiresCompatibleCompleteCoverage) { + const std::vector> corpus = {{"the", "wolf"}}; + const auto terms = build_common_grams_terms(corpus, /*include_grams=*/false); + const auto plain = plain_query_info({"the", "wo"}); + using Kind = segment_v2::TermKeyKind; + const auto gram = query_info({{Kind::kCommonGram, gram_key("the", "wo")}}); + const auto identity = complete_query_identity(); + + for (const auto& metadata : {std::optional {}, + std::optional { + metadata_for(PlainTermKeyVersion::kEscapedV1)}}) { + Fixture fixture; + assert_ok(build_fixture(&fixture, terms, metadata)); + expect_planned_prefix_matches_plain(fixture.index_reader, plain, gram, &identity, + /*max_expansions=*/50, PhrasePrefixPlanKind::kPlain); + } + + Fixture fixture; + assert_ok(build_fixture(&fixture, terms, complete_metadata())); + auto mismatched = identity; + mismatched.common_grams_fingerprint = "different-common-grams-policy"; + expect_planned_prefix_matches_plain(fixture.index_reader, plain, gram, nullptr, + /*max_expansions=*/50, PhrasePrefixPlanKind::kPlain); + expect_planned_prefix_matches_plain(fixture.index_reader, plain, gram, &mismatched, + /*max_expansions=*/50, PhrasePrefixPlanKind::kPlain); +} + +TEST(SniiCommonGramsNamespaceQuery, DebugForceGramDoesNotBypassPlanEligibility) { + const bool original_enable_debug_points = config::enable_debug_points; + config::enable_debug_points = true; + DebugPoints::instance()->add("snii.common_grams.force_gram_plan"); + Defer restore_debug_points([original_enable_debug_points] { + DebugPoints::instance()->remove("snii.common_grams.force_gram_plan"); + config::enable_debug_points = original_enable_debug_points; + }); + + const std::vector> corpus = {{"the", "wolf"}}; + const auto identity = complete_query_identity(); + using Kind = segment_v2::TermKeyKind; + const auto exact_plain = plain_query_info({"the", "wolf"}); + const auto exact_gram = query_info({{Kind::kCommonGram, gram_key("the", "wolf")}}); + + Fixture incomplete_fixture; + assert_ok(build_fixture(&incomplete_fixture, build_common_grams_terms(corpus), + metadata_for(PlainTermKeyVersion::kEscapedV1))); + ExactPhrasePlanKind exact_plan = ExactPhrasePlanKind::kCommonGrams; + std::vector docs; + assert_ok(planned_exact_phrase_query(incomplete_fixture.index_reader, exact_plain, exact_gram, + &identity, &docs, nullptr, &exact_plan)); + EXPECT_EQ(exact_plan, ExactPhrasePlanKind::kPlain); + EXPECT_EQ(docs, (std::vector {0})); + + Fixture missing_gram_fixture; + assert_ok(build_fixture(&missing_gram_fixture, + build_common_grams_terms(corpus, /*include_grams=*/false), + complete_metadata())); + exact_plan = ExactPhrasePlanKind::kPlain; + docs.clear(); + assert_ok(planned_exact_phrase_query(missing_gram_fixture.index_reader, exact_plain, exact_gram, + &identity, &docs, nullptr, &exact_plan, + /*cost_model=*/ {})); + EXPECT_TRUE(docs.empty()); + + PhrasePrefixPlanKind prefix_plan = PhrasePrefixPlanKind::kCommonGrams; + docs.clear(); + assert_ok(planned_phrase_prefix_query( + incomplete_fixture.index_reader, plain_query_info({"the", "wo"}), + query_info({{Kind::kCommonGram, gram_key("the", "wo")}}), &identity, &docs, nullptr, + /*max_expansions=*/50, &prefix_plan)); + EXPECT_EQ(prefix_plan, PhrasePrefixPlanKind::kPlain); + EXPECT_EQ(docs, (std::vector {0})); + + prefix_plan = PhrasePrefixPlanKind::kPlain; + docs.clear(); + assert_ok(planned_phrase_prefix_query( + missing_gram_fixture.index_reader, plain_query_info({"the", "wo"}), + query_info({{Kind::kCommonGram, gram_key("the", "wo")}}), &identity, &docs, nullptr, + /*max_expansions=*/50, &prefix_plan, /*cost_model=*/ {})); + EXPECT_EQ(prefix_plan, PhrasePrefixPlanKind::kCommonGrams); + EXPECT_TRUE(docs.empty()); +} + +} // namespace +} // namespace doris::snii::query diff --git a/be/test/storage/index/snii/query/count_query_test.cpp b/be/test/storage/index/snii/query/count_query_test.cpp new file mode 100644 index 00000000000000..08b1f2a99eb68f --- /dev/null +++ b/be/test/storage/index/snii/query/count_query_test.cpp @@ -0,0 +1,290 @@ +// 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 "storage/index/snii/query/count_query.h" + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "roaring/roaring.hh" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/null_bitmap.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +// G02 single-term count-only fast path: the count answered from dict-entry df +// alone must equal the full posting decode on the shared 9000-doc fixture. The +// count_fastpath_hits seam pins that hits are counted exactly when a dict-only +// answer was produced and never on the ordinary decode path. +using namespace doris::snii; +using namespace doris::snii::snii_test; +using doris::snii::query::count_only_term_df; +using doris::snii::query::fabricate_null_disjoint_count_bitmap; +using doris::snii::query::term_query; +namespace qinternal = doris::snii::query::internal; + +namespace { + +struct Fixture { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; +}; + +void build_fixture(Fixture* f) { + assert_ok(build_reader(&f->file, &f->segment_reader, &f->index_reader)); +} + +uint64_t full_decode_term_count(const reader::LogicalIndexReader& idx, const std::string& term) { + std::vector docids; + assert_ok(term_query(idx, term, &docids)); + return docids.size(); +} + +void reset_query_counters() { + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; +} + +uint64_t count_hits() { + return qinternal::query_test_counters().count_fastpath_hits; +} + +} // namespace + +// Single-term df from the dict equals the full posting decode for every +// posting layout in the fixture (dense/windowed, sparse, tiny, tail) and for +// an absent term (0). The seam counts one hit per dict-only answer and none +// for the decodes. +TEST(SniiCountQuery, TermDfEqualsFullDecode) { + Fixture plain; + build_fixture(&plain); + + const std::vector terms {"failed", "driver", "almost", "sparse_left", + "sparse_right", "needle", "123", "trace", + "repeat", "nosuchterm"}; + + reset_query_counters(); + uint64_t expected_hits = 0; + for (const std::string& term : terms) { + const uint64_t decode_count = full_decode_term_count(plain.index_reader, term); + EXPECT_EQ(count_hits(), expected_hits) << "decode path must not touch the seam"; + + uint64_t df_count = 0; + assert_ok(count_only_term_df(plain.index_reader, term, &df_count)); + EXPECT_EQ(df_count, decode_count) << "term: " << term; + ++expected_hits; + EXPECT_EQ(count_hits(), expected_hits) << "term: " << term; + } + // Spot-pin the fixture dfs so a fixture drift cannot silently weaken the test. + uint64_t df = 0; + assert_ok(count_only_term_df(plain.index_reader, "failed", &df)); + EXPECT_EQ(df, 9000U); + assert_ok(count_only_term_df(plain.index_reader, "nosuchterm", &df)); + EXPECT_EQ(df, 0U); +} + +// Single-term df remains valid on a positionless index. +TEST(SniiCountQuery, DocsOnlyIndexTermStillCounts) { + MemoryFile file; + writer::SniiIndexInput input; + input.index_id = 3; + input.index_suffix = "DocsOnly"; + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = 100; + input.terms = {make_term("alpha", docs_with_one_position(0, 60, 0)), + make_term("bravo", docs_with_one_position(10, 25, 1))}; + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + writer::SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + reader::LogicalIndexReader idx; + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &idx)); + ASSERT_FALSE(idx.has_positions()); + + reset_query_counters(); + uint64_t df = 0; + assert_ok(count_only_term_df(idx, "alpha", &df)); + EXPECT_EQ(df, full_decode_term_count(idx, "alpha")); + EXPECT_EQ(df, 60U); + assert_ok(count_only_term_df(idx, "bravo", &df)); + EXPECT_EQ(df, 15U); + EXPECT_EQ(count_hits(), 2U); +} + +// --- fabricate_null_disjoint_count_bitmap truth table ----------------------- +// The fabricated bitmap must (a) hold exactly `count` ids, (b) be DISJOINT +// from the null bitmap so the caller-side FunctionMatchBase -> +// InvertedIndexResultBitmap::mask_out_null subtraction is a provable no-op, +// and (c) stay inside [0, count + |nulls|), which never exceeds the segment's +// [0, num_rows) row space because df counts only non-null docs. + +TEST(SniiCountQuery, FabricateNoNullsIsDenseRange) { + roaring::Roaring nulls; + roaring::Roaring out; + assert_ok(fabricate_null_disjoint_count_bitmap(100, nulls, &out)); + roaring::Roaring expected; + expected.addRange(0, 100); + EXPECT_EQ(out, expected); +} + +TEST(SniiCountQuery, FabricateNullPrefixShiftsRange) { + roaring::Roaring nulls; + nulls.addRange(0, 100); + roaring::Roaring out; + assert_ok(fabricate_null_disjoint_count_bitmap(50, nulls, &out)); + roaring::Roaring expected; + expected.addRange(100, 150); + EXPECT_EQ(out, expected); + EXPECT_TRUE((out & nulls).isEmpty()); +} + +TEST(SniiCountQuery, FabricateInterleavedNullsStaysDisjointAndExact) { + roaring::Roaring nulls; + for (uint32_t id = 0; id < 200; id += 2) { + nulls.add(id); // 100 even nulls + } + roaring::Roaring out; + assert_ok(fabricate_null_disjoint_count_bitmap(70, nulls, &out)); + EXPECT_EQ(out.cardinality(), 70U); + EXPECT_TRUE((out & nulls).isEmpty()); + // Exactly the first 70 odd (non-null) ids. + roaring::Roaring expected; + for (uint32_t id = 1; id < 140; id += 2) { + expected.add(id); + } + EXPECT_EQ(out, expected); + // The safety property itself: the null-bitmap subtraction the MATCH + // machinery applies to every index result must not change the count. + roaring::Roaring masked = out; + masked -= nulls; + EXPECT_EQ(masked.cardinality(), 70U); +} + +TEST(SniiCountQuery, FabricateNullsBeyondWindowIrrelevant) { + roaring::Roaring nulls; + nulls.add(1000000); + roaring::Roaring out; + assert_ok(fabricate_null_disjoint_count_bitmap(10, nulls, &out)); + roaring::Roaring expected; + expected.addRange(0, 10); + EXPECT_EQ(out, expected); +} + +TEST(SniiCountQuery, FabricateConsumesEntireNonNullWindow) { + // nulls {0..4}, count 5: every non-null id of the window [0, 10) is used. + roaring::Roaring nulls; + nulls.addRange(0, 5); + roaring::Roaring out; + assert_ok(fabricate_null_disjoint_count_bitmap(5, nulls, &out)); + roaring::Roaring expected; + expected.addRange(5, 10); + EXPECT_EQ(out, expected); +} + +TEST(SniiCountQuery, FabricateZeroCountIsEmpty) { + roaring::Roaring nulls; + nulls.addRange(0, 100); + roaring::Roaring out; + out.add(7); // stale content must be overwritten + assert_ok(fabricate_null_disjoint_count_bitmap(0, nulls, &out)); + EXPECT_TRUE(out.isEmpty()); +} + +TEST(SniiCountQuery, FabricateRejectsDocidDomainOverflow) { + // df + null count beyond the uint32 docid domain can only come from a + // corrupt index: it must surface as an error (the reader glue falls + // through to the row-accurate decode), never as a silently wrong bitmap. + roaring::Roaring nulls; + nulls.add(0); + roaring::Roaring out; + const uint64_t count = uint64_t(std::numeric_limits::max()) + 1; // 2^32 + doris::Status st = fabricate_null_disjoint_count_bitmap(count, nulls, &out); + EXPECT_FALSE(st.ok()); +} + +// A segment WITH a null bitmap end to end through the real writer/reader: +// df from the dict is unchanged by nulls (the writer adds NO postings for a +// null doc, so postings -- and df -- can never include null rows), and the +// fabricated bitmap built against the segment's REAL null bitmap keeps +// cardinality df through the caller-side null subtraction. This pins the +// contract that replaced the old "veto whenever a null section exists" guard. +TEST(SniiCountQuery, NullBitmapSegmentDfCountsAndFabricationSurvivesMasking) { + MemoryFile file; + writer::SniiIndexInput input; + input.index_id = 4; + input.index_suffix = "Nullable"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 100; + // Writer invariant mirrored by the fixture: null docids carry no postings + // (scalar add_nulls adds no tokens; a NULL array row is an empty range). + input.null_docids = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 50, 99}; + input.terms = {make_term("alpha", docs_with_one_position(10, 50, 0))}; + writer::SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + reader::LogicalIndexReader idx; + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &idx)); + ASSERT_GT(idx.section_refs().null_bitmap.length, 0U); + + // df is the exact match count despite the null rows. + uint64_t df = 0; + assert_ok(count_only_term_df(idx, "alpha", &df)); + EXPECT_EQ(df, full_decode_term_count(idx, "alpha")); + EXPECT_EQ(df, 40U); + + // Decode the segment's REAL null bitmap the same way the reader glue does. + const auto& ref = idx.section_refs().null_bitmap; + std::vector bytes; + assert_ok(idx.reader()->read_at(ref.offset, ref.length, &bytes)); + format::NullBitmapReader null_reader; + assert_ok(format::NullBitmapReader::open(Slice(bytes), &null_reader)); + roaring::Roaring nulls; + null_reader.copy_to(&nulls); + ASSERT_EQ(nulls.cardinality(), input.null_docids.size()); + + roaring::Roaring fabricated; + assert_ok(fabricate_null_disjoint_count_bitmap(df, nulls, &fabricated)); + EXPECT_EQ(fabricated.cardinality(), df); + EXPECT_TRUE((fabricated & nulls).isEmpty()); + roaring::Roaring masked = fabricated; + masked -= nulls; // FunctionMatchBase::mask_out_null equivalent + EXPECT_EQ(masked.cardinality(), df); + // With the 10 leading nulls the first 40 non-null ids are exactly [10, 50): + // pinned to catch select / removeRange off-by-ones against a real layout. + roaring::Roaring expected; + expected.addRange(10, 50); + EXPECT_EQ(fabricated, expected); +} diff --git a/be/test/storage/index/snii/query/docid_set_ops_test.cpp b/be/test/storage/index/snii/query/docid_set_ops_test.cpp new file mode 100644 index 00000000000000..eda9bf3dc7da65 --- /dev/null +++ b/be/test/storage/index/snii/query/docid_set_ops_test.cpp @@ -0,0 +1,61 @@ +// 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 "storage/index/snii/query/internal/docid_set_ops.h" + +#include + +#include +#include +#include + +#include "common/status.h" + +namespace { + +std::vector Range(uint32_t begin, uint32_t end, uint32_t step = 1) { + std::vector out; + for (uint32_t v = begin; v < end; v += step) { + out.push_back(v); + } + return out; +} + +} // namespace + +TEST(SniiDocIdSetOps, UnionSortedManyDeduplicatesHighOverlapLists) { + std::vector> lists; + lists.reserve(257); + lists.push_back(Range(0, 10000)); + for (uint32_t i = 0; i < 256; ++i) { + lists.push_back(Range(i % 17, 10000, 17)); + } + + const std::vector got = doris::snii::query::internal::union_sorted_many(lists); + + EXPECT_EQ(got, Range(0, 10000)); + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(std::ranges::adjacent_find(got), got.end()); +} + +TEST(SniiDocIdSetOps, UnionSortedManyMergesDisjointLists) { + const std::vector> lists = {{0, 3, 6}, {1, 4, 7}, {}, {2, 5, 8}}; + + const std::vector got = doris::snii::query::internal::union_sorted_many(lists); + + EXPECT_EQ(got, Range(0, 9)); +} diff --git a/be/test/storage/index/snii/query/pattern_query_test.cpp b/be/test/storage/index/snii/query/pattern_query_test.cpp new file mode 100644 index 00000000000000..ccd04047486470 --- /dev/null +++ b/be/test/storage/index/snii/query/pattern_query_test.cpp @@ -0,0 +1,240 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_pattern_query_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +bool WildcardMatch(std::string_view pattern, std::string_view text) { + std::vector prev(text.size() + 1, 0); + std::vector curr(text.size() + 1, 0); + prev[0] = 1; + for (char p : pattern) { + std::ranges::fill(curr, 0); + if (p == '*') { + curr[0] = prev[0]; + for (size_t i = 1; i <= text.size(); ++i) { + curr[i] = prev[i] || curr[i - 1]; + } + } else { + for (size_t i = 1; i <= text.size(); ++i) { + curr[i] = prev[i - 1] && (p == '?' || p == text[i - 1]); + } + } + prev.swap(curr); + } + return prev[text.size()] != 0; +} + +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; + + std::vector matching_docs( + const std::function& matches) const { + std::set ids; + for (uint32_t d = 0; d < docs.size(); ++d) { + for (const std::string& term : docs[d]) { + if (matches(term)) { + ids.insert(d); + } + } + } + return {ids.begin(), ids.end()}; + } +}; + +Corpus BuildMixedCorpus() { + Corpus c; + c.doc_count = 120; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& terms = c.docs[d]; + if (d < 80) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + terms.emplace_back(term); + } + if (d < 80 && d % 2 == 0) { + terms.emplace_back("aa_even"); + } + if (d < 50) { + char term[16]; + std::snprintf(term, sizeof(term), "ab_%03u", d); + terms.emplace_back(term); + } + char filler[16]; + std::snprintf(filler, sizeof(filler), "zz_%03u", d); + terms.emplace_back(filler); + } + return c; +} + +Corpus BuildLowDfPrefixCorpus() { + Corpus c; + c.doc_count = 96; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + c.docs[d].emplace_back(term); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/false); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + const std::vector& terms = c.docs[d]; + for (uint32_t pos = 0; pos < terms.size(); ++pos) { + buf.add_token(terms[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = doris::snii::format::IndexConfig::kDocsOnly; + in.doc_count = c.doc_count; + in.terms = buf.finalize_sorted(); + in.target_dict_block_bytes = 2048; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +LogicalIndexReader OpenIndex(io::LocalFileReader* file, SniiSegmentReader* segment, + const std::string& path) { + EXPECT_TRUE(file->open(path).ok()); + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +LogicalIndexReader OpenMeteredIndex(io::MeteredFileReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +} // namespace + +TEST(SniiPatternQuery, WildcardMatchesOracle) { + const Corpus corpus = BuildMixedCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + for (const char* pattern : {"aa_*", "aa_00?", "aa_even", "zz_0??", "missing*"}) { + std::vector got; + ASSERT_TRUE(query::wildcard_query(idx, pattern, &got).ok()) << pattern; + EXPECT_TRUE(std::ranges::is_sorted(got)) << pattern; + EXPECT_EQ(got, corpus.matching_docs([&](std::string_view term) { + return WildcardMatch(pattern, term); + })) << pattern; + } + + std::remove(path.c_str()); +} + +TEST(SniiPatternQuery, RegexpMatchesOracle) { + const Corpus corpus = BuildMixedCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + for (const char* pattern : + {"aa_00[0-9]", "aa_even", "ab_0[0-9][0-9]", "zz_1..", "no_match.*"}) { + const std::regex re(pattern); + std::vector got; + ASSERT_TRUE(query::regexp_query(idx, pattern, &got).ok()) << pattern; + EXPECT_TRUE(std::ranges::is_sorted(got)) << pattern; + EXPECT_EQ(got, corpus.matching_docs([&](std::string_view term) { + return std::regex_match(term.begin(), term.end(), re); + })) << pattern; + } + + std::vector ignored; + EXPECT_FALSE(query::regexp_query(idx, "[", &ignored).ok()); + + std::remove(path.c_str()); +} + +TEST(SniiPatternQuery, WideWildcardUsesPrefixEnumerationWithoutPerTermLookup) { + const Corpus corpus = BuildLowDfPrefixCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::wildcard_query(idx, "aa_*", &got).ok()); + + EXPECT_EQ(got, corpus.matching_docs( + [](std::string_view term) { return WildcardMatch("aa_*", term); })); + EXPECT_LT(metered.metrics().read_at_calls, corpus.doc_count / 3) + << "wildcard_query must reuse enumerated entries, not lookup every term again"; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/phrase_prefix_query_test.cpp b/be/test/storage/index/snii/query/phrase_prefix_query_test.cpp new file mode 100644 index 00000000000000..f23b5bbdb6183c --- /dev/null +++ b/be/test/storage/index/snii/query/phrase_prefix_query_test.cpp @@ -0,0 +1,959 @@ +// 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 +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::Status; +namespace snii_test = doris::snii::snii_test; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_phrase_prefix_query_" + std::to_string(getpid()) + "_" + + std::to_string(counter++) + ".idx"; +} + +bool HasPrefix(const std::string& term, const std::string& prefix) { + return term.size() >= prefix.size() && term.starts_with(prefix); +} + +struct Corpus { + std::vector> docs; + + std::vector phrase_prefix_docs(const std::vector& terms) const { + std::vector out; + if (terms.empty()) { + return out; + } + for (uint32_t d = 0; d < docs.size(); ++d) { + const std::vector& doc = docs[d]; + bool match = false; + if (terms.size() == 1) { + for (const std::string& term : doc) { + if (HasPrefix(term, terms.front())) { + match = true; + break; + } + } + } else if (doc.size() >= terms.size()) { + for (size_t start = 0; start + terms.size() <= doc.size(); ++start) { + bool exact = true; + for (size_t i = 0; i + 1 < terms.size(); ++i) { + if (doc[start + i] != terms[i]) { + exact = false; + break; + } + } + if (exact && HasPrefix(doc[start + terms.size() - 1], terms.back())) { + match = true; + break; + } + } + } + if (match) { + out.push_back(d); + } + } + return out; + } + + // Truncation-aware oracle: reproduces the byte-exact max_expansions semantics + // the query must honour -- enumerate the REAL tail terms sharing the prefix in + // lexicographic (dict) order, keep only the first `max_expansions`, then match + // exactly as phrase_prefix_docs but restricted to that surviving tail set. + // (These corpora carry no hidden phrase-bigram terms, so the corpus vocabulary + // with the prefix IS the index's real-term enumeration for it.) + std::vector phrase_prefix_docs_capped(const std::vector& terms, + int32_t max_expansions) const { + if (max_expansions <= 0 || terms.size() < 2) { + return phrase_prefix_docs(terms); + } + std::set vocab; + for (const std::vector& doc : docs) { + for (const std::string& t : doc) { + if (HasPrefix(t, terms.back())) { + vocab.insert(t); + } + } + } + std::set allowed; + int32_t taken = 0; + for (const std::string& t : vocab) { // std::set iterates ascending (dict order) + if (taken >= max_expansions) { + break; + } + allowed.insert(t); + ++taken; + } + std::vector out; + for (uint32_t d = 0; d < docs.size(); ++d) { + const std::vector& doc = docs[d]; + if (doc.size() < terms.size()) { + continue; + } + bool match = false; + for (size_t start = 0; start + terms.size() <= doc.size() && !match; ++start) { + bool exact = true; + for (size_t i = 0; i + 1 < terms.size(); ++i) { + if (doc[start + i] != terms[i]) { + exact = false; + break; + } + } + if (exact && allowed.contains(doc[start + terms.size() - 1])) { + match = true; + } + } + if (match) { + out.push_back(d); + } + } + return out; + } +}; + +Corpus BuildPhraseCorpus() { + Corpus c; + c.docs = {{"quick", "brown", "fox"}, {"quick", "blue", "fox"}, {"quick", "bronze", "fox"}, + {"slow", "brown", "fox"}, {"quick", "brownish"}, {"quick", "brown", "fossil"}, + {"quick", "brown", "fog"}, {"quick", "brown"}, {"brown", "fox", "quick"}}; + return c; +} + +Corpus BuildWideTailCorpus() { + Corpus c; + c.docs.resize(96); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + c.docs[d].emplace_back(d == 0 ? "lead" : "other"); + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + c.docs[d].emplace_back(term); + } + return c; +} + +Corpus BuildRepeatedExactCorpus() { + Corpus c; + c.docs = {{"x", "x", "brown"}, + {"x", "y", "brown"}, + {"x", "brown", "x"}, + {"x", "x", "bronze"}, + {"x", "x", "blue"}}; + return c; +} + +Corpus BuildSharedExactWideTailCorpus() { + Corpus c; + c.docs.resize(768); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + c.docs[d].emplace_back("lead"); + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + c.docs[d].emplace_back(term); + } + return c; +} + +Corpus BuildWideNonAdjacentTailCorpus() { + Corpus c; + c.docs.resize(256); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d % 96); + c.docs[d] = {"lead", "gap", term}; + } + return c; +} + +Corpus BuildLeadingPrefilterCorpus(uint32_t tail_docs) { + DCHECK_LE(tail_docs, 1024U); + Corpus c; + c.docs.resize(1024); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + if (d < tail_docs) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d % 64); + c.docs[d] = {"lead", term}; + } else { + c.docs[d] = {"lead", "other"}; + } + } + return c; +} + +Corpus BuildNearEqualLeadingAndTailDfCorpus() { + Corpus c; + c.docs.resize(2048, {"other"}); + for (uint32_t d = 0; d < 256; ++d) { + if (d == 255) { + c.docs[d] = {"lead", "other"}; + continue; + } + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d % 50); + c.docs[d] = {"lead", term}; + } + return c; +} + +Corpus BuildDefaultExpansionPruneCorpus(uint32_t adjacent_docs) { + Corpus c; + c.docs.resize(1024); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + char term[16]; + if (d < 256) { + std::snprintf(term, sizeof(term), "aa_%03u", d % 32); + c.docs[d] = d < adjacent_docs ? std::vector {"lead", term} + : std::vector {"lead", "gap", term}; + } else { + std::snprintf(term, sizeof(term), "aa_%03u", 32 + (d % 18)); + c.docs[d] = {"lead", "gap", term}; + } + } + return c; +} + +Corpus BuildAllMatchedBeforeFinalGroupCorpus() { + Corpus c; + c.docs.resize(1025); + for (uint32_t d = 0; d < 1024; ++d) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d % 32); + c.docs[d] = {"lead", term}; + } + c.docs.back() = {"other", "aa_032"}; + return c; +} + +Corpus BuildAllMatchedAfterPartialGroupCorpus() { + Corpus c; + c.docs.resize(1376); + for (uint32_t d = 0; d < 1024; ++d) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d < 512 ? d % 32 : 32 + (d % 32)); + c.docs[d] = {"lead", term}; + } + for (uint32_t tail = 64; tail < 416; ++tail) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", tail); + c.docs[1024 + tail - 64] = {"other", term}; + } + return c; +} + +Corpus BuildCrossGroupDuplicateTailCorpus() { + Corpus c; + c.docs.resize(3); + for (uint32_t tail = 0; tail < 65; ++tail) { + c.docs[0].emplace_back("lead"); + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", tail); + c.docs[0].emplace_back(term); + } + c.docs[1] = {"lead", "gap", "aa_000"}; + c.docs[2] = {"other", "aa_064"}; + return c; +} + +// The first 32-tail resident group matches doc 0. The final group contains +// exactly 32 real tails, but none of their postings intersects the expected +// leading-term doc set. collect_merged_tail_matches must still flush the match +// accumulated by the earlier group when it processes this empty final group. +Corpus BuildFinalGroupWithoutCandidatesCorpus() { + Corpus c; + c.docs.resize(64); + for (uint32_t tail = 0; tail < c.docs.size(); ++tail) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", tail); + c.docs[tail] = {tail == 0 ? "lead" : "other", term}; + } + return c; +} + +// Large corpus whose leading exact term and every tail expansion have high df, so +// their postings span MULTIPLE windows and the merge advances its cursors across +// window boundaries. Most docs match ("lead" @0 then a "res_" tail @1); a slice +// injects a filler token so the tail is not adjacent (must NOT match), and a slice +// omits the leading term (must NOT match) -- forcing the cross-window sweep to +// reject as well as accept. +Corpus BuildCrossWindowTailCorpus() { + Corpus c; + c.docs.resize(5000); + const char* const tails[] = {"res_a", "res_b", "res_c"}; + for (uint32_t d = 0; d < c.docs.size(); ++d) { + if (d % 50 == 7) { + c.docs[d] = {"lead", "gap", tails[d % 3]}; // tail not adjacent to lead + } else if (d % 50 == 11) { + c.docs[d] = {"nolead", tails[d % 3]}; // leading term absent + } else { + c.docs[d] = {"lead", tails[d % 3]}; + } + } + return c; +} + +// CJK / multi-byte corpus. Leading term and tail prefix are UTF-8; prefix testing +// is byte-wise, matching the index's dict enumeration order. +Corpus BuildCjkTailCorpus() { + Corpus c; + const char* const tails[] = {"\xE7\xBB\x93\xE6\x9E\x9C\xE7\x94\xB2", // 结果甲 + "\xE7\xBB\x93\xE6\x9E\x9C\xE4\xB9\x99", // 结果乙 + "\xE7\xBB\x93\xE6\x9E\x9C\xE4\xB8\x99"}; // 结果丙 + const std::string lead = "\xE8\xBF\x9E\xE6\x8E\xA5"; // 连接 + c.docs.resize(120); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + if (d % 20 == 3) { + c.docs[d] = {lead, "\xE9\x97\xB4\xE9\x9A\x94", + tails[d % 3]}; // 间隔 filler, not adjacent + } else { + c.docs[d] = {lead, tails[d % 3]}; + } + } + return c; +} + +// Two leading exact terms that both occur but are NEVER adjacent, so the leading +// phrase conjunction yields an empty expected-position set -- the multi-tail +// branch's `expected.docs.empty()` early return -- even though the tail prefix +// expands to several real terms. +Corpus BuildEmptyExpectedCorpus() { + Corpus c; + c.docs = {{"alpha", "x", "beta", "res_a"}, + {"alpha", "y", "beta", "res_b"}, + {"beta", "alpha", "res_c"}, + {"alpha", "z", "beta"}}; + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + const std::vector& terms = c.docs[d]; + for (uint32_t pos = 0; pos < terms.size(); ++pos) { + buf.add_token(terms[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = doris::snii::format::IndexConfig::kDocsPositionsScoring; + in.doc_count = static_cast(c.docs.size()); + in.encoded_norms.assign(c.docs.size(), 1); + in.terms = buf.finalize_sorted(); + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + in.target_dict_block_bytes = 2048; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +void WritePostings(SpimiTermBuffer* buffer, uint32_t doc_count, const std::string& path) { + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "body"; + input.config = doris::snii::format::IndexConfig::kDocsPositionsScoring; + input.doc_count = doc_count; + input.encoded_norms.assign(doc_count, 1); + input.terms = buffer->finalize_sorted(); + uint64_t token_count = 0; + for (const auto& term : input.terms) { + token_count += term.positions_flat.size(); + } + input.common_grams_metadata = + snii_test::make_plain_scoring_metadata(input.doc_count, token_count); + input.target_dict_block_bytes = 256; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(input).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +LogicalIndexReader OpenIndex(io::LocalFileReader* file, SniiSegmentReader* segment, + const std::string& path) { + EXPECT_TRUE(file->open(path).ok()); + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +LogicalIndexReader OpenMeteredIndex(io::MeteredFileReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +} // namespace + +TEST(SniiPhrasePrefixQuery, MatchesPositionOracle) { + const Corpus corpus = BuildPhraseCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector> cases = { + {"quick", "bro"}, {"quick", "brown", "fo"}, {"slow", "bro"}, + {"absent", "bro"}, {"quick", "missing"}, {"bro"}}; + for (const std::vector& terms : cases) { + std::vector got; + const Status st = query::phrase_prefix_query(idx, terms, &got); + ASSERT_TRUE(st.ok()) << st.to_string(); + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + } + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, WideTailPrefixAvoidsPerExpansionLookup) { + const Corpus corpus = BuildWideTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + metered.reset_metrics(); + const std::vector terms = {"lead", "aa_"}; + std::vector got; + const Status st = query::phrase_prefix_query(idx, terms, &got); + ASSERT_TRUE(st.ok()) << st.to_string(); + + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + EXPECT_LT(metered.metrics().read_at_calls, corpus.docs.size() / 3) + << "phrase_prefix_query must reuse PrefixHit entries, not lookup every tail term"; + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, RepeatedExactTermsMatchPositionOracle) { + const Corpus corpus = BuildRepeatedExactCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector> cases = { + {"x", "x", "br"}, {"x", "brown", "x"}, {"x", "x", "missing"}}; + for (const std::vector& terms : cases) { + std::vector got; + const Status st = query::phrase_prefix_query(idx, terms, &got); + ASSERT_TRUE(st.ok()) << st.to_string(); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + } + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, WideTailPrefixReusesExactTermPostingReads) { + const Corpus corpus = BuildSharedExactWideTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + metered.reset_metrics(); + const std::vector terms = {"lead", "aa_"}; + std::vector got; + const Status st = query::phrase_prefix_query(idx, terms, &got); + ASSERT_TRUE(st.ok()) << st.to_string(); + + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + EXPECT_LT(metered.metrics().read_at_calls, corpus.docs.size() / 4) + << "wide phrase-prefix must not re-read exact term postings for every tail hit"; + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, SegmentRelativeDfGatePrefiltersLeadingPositions) { + const Corpus corpus = BuildLeadingPrefilterCorpus(/*tail_docs=*/128); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + query::QueryProfile profile; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got, &profile).ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs({"lead", "aa_"})); + EXPECT_EQ(profile.phrase_query_stats.prefix_leading_candidate_docs, 128U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, HalfDfTailKeepsDirectLeadingDecode) { + const Corpus corpus = BuildLeadingPrefilterCorpus(/*tail_docs=*/512); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + query::QueryProfile profile; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got, &profile).ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs({"lead", "aa_"})); + EXPECT_EQ(profile.phrase_query_stats.prefix_leading_candidate_docs, 1024U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixQuery, NearEqualTailDfKeepsDirectLeadingDecode) { + const Corpus corpus = BuildNearEqualLeadingAndTailDfCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + query::QueryProfile profile; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got, &profile).ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs({"lead", "aa_"})); + EXPECT_EQ(profile.phrase_query_stats.prefix_leading_candidate_docs, 256U); + + std::remove(path.c_str()); +} + +// --------------------------------------------------------------------------- +// Merged multi-tail path (collect_merged_tail_matches): the per-tail verify+union +// loop was replaced by a single batched, forward-merge sweep. Every case below +// asserts the merged result equals the independent position oracle -- i.e. is +// identical to the old per-tail semantics -- across the 0/1/many-expansion +// trichotomy, max_expansions truncation, an empty expected set, cross-window +// positions and CJK/unicode terms. +// --------------------------------------------------------------------------- + +namespace qinternal = doris::snii::query::internal; + +// Zero expansions: the tail prefix matches no real term -> empty result. Also +// exercises the single-expansion (untouched) branch for comparison. +TEST(SniiPhrasePrefixMerge, ZeroAndSingleExpansionMatchOracle) { + const Corpus corpus = BuildSharedExactWideTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector zero_terms = {"lead", "zzz"}; + std::vector zero_got; + ASSERT_TRUE(query::phrase_prefix_query(idx, zero_terms, &zero_got).ok()); + EXPECT_TRUE(zero_got.empty()); + EXPECT_EQ(zero_got, corpus.phrase_prefix_docs(zero_terms)); + + // Prefix "aa_000" matches exactly one full term -> single-tail branch. + const std::vector single_terms = {"lead", "aa_000"}; + std::vector single_got; + ASSERT_TRUE(query::phrase_prefix_query(idx, single_terms, &single_got).ok()); + EXPECT_TRUE(std::ranges::is_sorted(single_got)); + EXPECT_EQ(single_got, corpus.phrase_prefix_docs(single_terms)); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector single_matches; + ASSERT_TRUE( + query::phrase_prefix_query_with_frequencies(idx, single_terms, &single_matches).ok()); + EXPECT_EQ(single_matches, (std::vector {{.docid = 0, .frequency = 1}})); + EXPECT_EQ(qinternal::query_test_counters().expected_docids_build, 0U); + + std::remove(path.c_str()); +} + +// Many expansions spanning several resident-capped groups (768 tails, batch 32). +// The merged/grouped path must equal the full oracle and stay sorted. +TEST(SniiPhrasePrefixMerge, ManyExpansionGroupsMatchOracle) { + const Corpus corpus = BuildSharedExactWideTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "aa_"}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + // The expected-docid projection must still be built exactly once per query + // (hoisted out of the per-tail loop), proving the multi-tail branch ran. + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_EQ(qinternal::query_test_counters().expected_docids_build, 1U); + EXPECT_EQ(qinternal::query_test_counters().resolved_term_entry_copies, 0U); + constexpr uint64_t baseline_group_visits = 768U * (768U / 32U); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, baseline_group_visits); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, NonAdjacentTailGroupsScanStableCandidateSet) { + const Corpus corpus = BuildWideNonAdjacentTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got).ok()); + EXPECT_TRUE(got.empty()); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, 256U * 3U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, DefaultExpansionCapKeepsSingleRemainingGroupDirect) { + const Corpus corpus = BuildDefaultExpansionPruneCorpus(/*adjacent_docs=*/256); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got, + /*max_expansions=*/50) + .ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs_capped({"lead", "aa_"}, 50)); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, 1024U * 2U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, SparseFirstGroupMatchScansStableCandidateSet) { + const Corpus corpus = BuildDefaultExpansionPruneCorpus(/*adjacent_docs=*/1); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got, + /*max_expansions=*/50) + .ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs_capped({"lead", "aa_"}, 50)); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, 1024U * 2U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, AllDocsMatchedBeforeFinalGroupStopsEarly) { + const Corpus corpus = BuildAllMatchedBeforeFinalGroupCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got).ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs({"lead", "aa_"})); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, 1024U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, AllDocsMatchedAfterPartialGroupEmitOnce) { + const Corpus corpus = BuildAllMatchedAfterPartialGroupCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + qinternal::query_test_counters() = qinternal::QueryTestCounters {}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"lead", "aa_"}, &got).ok()); + EXPECT_EQ(got, corpus.phrase_prefix_docs({"lead", "aa_"})); + EXPECT_EQ(qinternal::query_test_counters().prefix_expected_doc_visits, 1024U * 2U); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, CrossGroupDuplicateMatchesEmitOnce) { + const Corpus corpus = BuildCrossGroupDuplicateTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "aa_"}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_EQ(got, (std::vector {0})); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixFrequency, CountsDistinctPhraseStartsAcrossTailGroups) { + const Corpus corpus = BuildCrossGroupDuplicateTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "aa_"}; + for (const auto [cap, expected_frequency] : + {std::pair {32, 32U}, std::pair {33, 33U}, std::pair {65, 65U}}) { + std::vector matches; + ASSERT_TRUE(query::phrase_prefix_query_with_frequencies(idx, terms, &matches, nullptr, cap) + .ok()) + << "cap=" << cap; + EXPECT_EQ(matches, + (std::vector {{.docid = 0, .frequency = expected_frequency}})) + << "cap=" << cap; + } + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixFrequency, CountsStackedLeadingTermOccurrences) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.add_token("lead", 0, 0); + buffer.add_token("lead", 0, 0); + buffer.add_token("tail_a", 0, 1); + + const std::string path = TempPath(); + WritePostings(&buffer, 1, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + std::vector matches; + ASSERT_TRUE(query::phrase_prefix_query_with_frequencies(idx, {"lead", "tail"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 2}})); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixFrequency, CountsStackedLeadingTermWithSparseMiddleClause) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.add_token("lead", 0, 0); + buffer.add_token("lead", 0, 0); + buffer.add_token("middle", 0, 1); + buffer.add_token("tail_a", 0, 2); + for (uint32_t docid = 1; docid < 8; ++docid) { + buffer.add_token("lead", docid, 0); + } + + const std::string path = TempPath(); + WritePostings(&buffer, 8, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + std::vector matches; + ASSERT_TRUE( + query::phrase_prefix_query_with_frequencies(idx, {"lead", "middle", "tail"}, &matches) + .ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 2}})); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixFrequency, CountsOneOccurrenceAcrossSamePositionTailExpansions) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.add_token("lead", 0, 0); + buffer.add_token("tail_a", 0, 1); + buffer.add_token("tail_b", 0, 1); + + const std::string path = TempPath(); + WritePostings(&buffer, 1, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + std::vector matches; + ASSERT_TRUE(query::phrase_prefix_query_with_frequencies(idx, {"lead", "tail"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 1}})); + + std::remove(path.c_str()); +} + +TEST(SniiPhrasePrefixMerge, EmptyFinalGroupEmitsMatchesFromEarlierGroup) { + const Corpus corpus = BuildFinalGroupWithoutCandidatesCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "aa_"}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_EQ(got, (std::vector {0})); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + + std::vector matches; + ASSERT_TRUE(query::phrase_prefix_query_with_frequencies(idx, terms, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 1}})); + + std::remove(path.c_str()); +} + +// max_expansions truncation is byte-exact: lexicographic order, real unigrams +// only, and independent of how the tails are later grouped for the merge. +TEST(SniiPhrasePrefixMerge, MaxExpansionsTruncationMatchesOracle) { + const Corpus corpus = BuildSharedExactWideTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "aa_"}; + for (int32_t cap : {1, 10, 32, 50, 100, 768, 1000}) { + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got, cap).ok()) << "cap=" << cap; + EXPECT_TRUE(std::ranges::is_sorted(got)) << "cap=" << cap; + EXPECT_EQ(got, corpus.phrase_prefix_docs_capped(terms, cap)) << "cap=" << cap; + } + + std::remove(path.c_str()); +} + +// An empty expected set (two leading terms never adjacent) short-circuits the +// multi-tail branch to an empty result even though the tail prefix expands. +TEST(SniiPhrasePrefixMerge, EmptyExpectedSetReturnsEmpty) { + const Corpus corpus = BuildEmptyExpectedCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"alpha", "beta", "res_"}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_TRUE(got.empty()); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + + std::remove(path.c_str()); +} + +// Cross-window positions: high-df leading + tails span multiple windows, and the +// forward merge must accept adjacent matches and reject the non-adjacent / +// leading-absent injections while sweeping across window boundaries. +TEST(SniiPhrasePrefixMerge, CrossWindowPositionsMatchOracle) { + const Corpus corpus = BuildCrossWindowTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"lead", "res_"}; + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + + std::remove(path.c_str()); +} + +// CJK / multi-byte terms merge identically to the oracle. +TEST(SniiPhrasePrefixMerge, CjkUnicodeTailsMatchOracle) { + const Corpus corpus = BuildCjkTailCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + const std::vector terms = {"\xE8\xBF\x9E\xE6\x8E\xA5", // 连接 + "\xE7\xBB\x93\xE6\x9E\x9C"}; // 结果 + std::vector got; + ASSERT_TRUE(query::phrase_prefix_query(idx, terms, &got).ok()); + EXPECT_TRUE(std::ranges::is_sorted(got)); + EXPECT_EQ(got, corpus.phrase_prefix_docs(terms)); + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/phrase_skip_test.cpp b/be/test/storage/index/snii/query/phrase_skip_test.cpp new file mode 100644 index 00000000000000..b783aff2bf29b6 --- /dev/null +++ b/be/test/storage/index/snii/query/phrase_skip_test.cpp @@ -0,0 +1,659 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +// Differential test for phrase_query WINDOW SKIPPING (design spec section 6.2). +// +// Builds a corpus with a VERY high-df term ("aa_hi", df=2500) whose .frq +// posting spans MANY 256-doc windows, plus low/mid-df terms, and plants known +// phrases (including a 5-term phrase whose FIRST term is the high-df term). For +// every query the skipping phrase_query result must equal: +// (a) an in-memory brute-force ORACLE, and +// (b) an independent FULL-READ reference (decode every term's whole posting, +// intersect, positional check) -- the pre-skipping behavior. +// Finally it asserts the byte/round reduction: a phrase whose low-df lead term +// concentrates the candidates in a few windows reads FAR fewer bytes than the +// full-read path, and the high-df windows touched are NOT proportional to the +// term's total window count. +// +// NOTE on term naming: all real terms use an "aa_" prefix and a "zz_NNN" filler +// vocabulary fills the lexicographic tail, so every real term sorts within the +// SampledTermIndex's candidate range (the index samples per-block first terms). +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_phrase_skip_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +constexpr uint32_t kDocCount = 60000; +constexpr uint32_t kHiGap = 24; // aa_hi at 0, 24, 48 ... -> 2500 occurrences +constexpr uint32_t kHiCount = 2500; // df of aa_hi (>512 -> windowed, 10 windows) +// The 5-term phrase is planted only in the FIRST kPhraseSpan occurrences of +// aa_hi, concentrating its candidate docids into the first window(s). +constexpr uint32_t kPhraseSpan = 150; + +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; // docs[d] = ordered tokens + + bool phrase_in_doc(uint32_t d, const std::vector& phrase) const { + if (phrase.empty()) { + return false; + } + const auto& toks = docs[d]; + if (toks.size() < phrase.size()) { + return false; + } + for (size_t i = 0; i + phrase.size() <= toks.size(); ++i) { + bool match = true; + for (size_t k = 0; k < phrase.size(); ++k) { + if (toks[i + k] != phrase[k]) { + match = false; + break; + } + } + if (match) { + return true; + } + } + return false; + } + + std::vector oracle(const std::vector& phrase) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + if (phrase_in_doc(d, phrase)) { + out.push_back(d); + } + } + return out; + } +}; + +Corpus BuildCorpus() { + Corpus c; + c.doc_count = kDocCount; + c.docs.resize(c.doc_count); + const std::vector vocab = {"alpha", "bravo", "charlie", "delta"}; + + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + const bool is_hi = (d % kHiGap == 0) && (d / kHiGap < kHiCount); + if (is_hi) { + toks.emplace_back("aa_hi"); // high-df, position 0 + } + const uint32_t occ = d / kHiGap; // aa_hi occurrence ordinal + + if (is_hi && occ < kPhraseSpan) { + // 5-term phrase led by the high-df term, concentrated in the first + // windows. + toks.emplace_back("aa_quick"); + toks.emplace_back("aa_brown"); + toks.emplace_back("aa_fox"); + toks.emplace_back("aa_jumps"); + } else if (d % 13 == 0) { + // A phrase NOT containing the high-df term, scattered across the corpus. + toks.emplace_back("aa_lazy"); + toks.emplace_back("aa_dog"); + } + for (uint32_t k = 0; k < 2; ++k) { + toks.push_back(vocab[(d + k) % vocab.size()]); + } + if (d % 101 == 0) { + toks.emplace_back("aa_rare"); + } + if (d % 997 == 0) { + toks.emplace_back("aa_echo"); + toks.emplace_back("aa_echo"); + } + // Larger filler vocabulary to occupy the lexicographic tail blocks. + char nm[16]; + std::snprintf(nm, sizeof(nm), "zz_%03u", d % 500); + toks.emplace_back(nm); + } + return c; +} + +Corpus BuildDensePhraseCorpus() { + Corpus c; + c.doc_count = 4096; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + toks.emplace_back("aa_dense"); + if (d >= 16 && d < 96) { + toks.emplace_back("aa_rare_driver"); + } else { + toks.emplace_back("aa_gap"); + } + char nm[16]; + std::snprintf(nm, sizeof(nm), "zz_%03u", d % 257); + toks.emplace_back(nm); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.doc_count; ++d) { + const auto& toks = c.docs[d]; + for (uint32_t pos = 0; pos < toks.size(); ++pos) { + buf.add_token(toks[pos], d, pos); + } + } + std::vector terms = buf.finalize_sorted(); + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = c.doc_count; + in.terms = std::move(terms); + in.target_dict_block_bytes = 256; + + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(in).ok()); + ASSERT_TRUE(cw.finish().ok()); +} + +// Decodes one term's FULL posting (every window for windowed; the single window +// for slim/inline) -> sorted docids + aligned per-doc positions. This mirrors +// the pre-skipping reference path used to cross-check the skipping result. +struct FullPosting { + std::vector docids; + std::vector> positions; +}; + +bool DecodeFullTerm(const LogicalIndexReader& idx, const std::string& term, FullPosting* out) { + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + EXPECT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + if (!found) { + return false; + } + + if (entry.kind == DictEntryKind::kPodRef && entry.enc == DictEntryEnc::kWindowed) { + DecodedPosting dp; + EXPECT_TRUE(read_windowed_posting(idx, entry, frq_base, prx_base, + /*want_positions=*/true, + /*want_freq=*/false, &dp) + .ok()); + out->docids = std::move(dp.docids); + out->positions = std::move(dp.positions); + return true; + } + std::vector frq, prx; + if (entry.kind == DictEntryKind::kPodRef) { + uint64_t foff = 0, flen = 0, poff = 0, plen = 0; + EXPECT_TRUE(idx.resolve_frq_window(entry, frq_base, &foff, &flen).ok()); + EXPECT_TRUE(idx.resolve_prx_window(entry, prx_base, &poff, &plen).ok()); + EXPECT_TRUE(idx.reader()->read_at(foff, flen, &frq).ok()); + EXPECT_TRUE(idx.reader()->read_at(poff, plen, &prx).ok()); + } else { + frq = entry.frq_bytes; + prx = entry.prx_bytes; + } + // Slim/inline window = [dd_region][freq_region]; the dd region is the docs + // prefix. + WindowMeta meta; + meta.win_base = 0; + meta.doc_count = entry.df; + meta.dd_zstd = entry.dd_meta.zstd; + meta.dd_uncomp_len = entry.dd_meta.uncomp_len; + meta.dd_disk_len = entry.dd_meta.disk_len; + meta.crc_dd = entry.dd_meta.crc; + // INLINE entries (format v2) carry no per-region crc -- their bytes are + // covered by the dict block crc32c -- so decode must skip the region crc + // check. + meta.verify_crc = entry.dd_meta.verify_crc; + EXPECT_LE(entry.dd_meta.disk_len, frq.size()); + Slice dd_region(frq.data(), static_cast(entry.dd_meta.disk_len)); + std::vector freqs; + Status st = decode_window_slices(meta, dd_region, Slice(), Slice(prx), + /*want_positions=*/true, /*want_freq=*/false, &out->docids, + &freqs, &out->positions); + EXPECT_TRUE(st.ok()) << st.msg(); + return true; +} + +// Independent FULL-READ phrase reference: intersect full postings, positional +// check. Mirrors the pre-skipping algorithm. +std::vector FullReadPhrase(const LogicalIndexReader& idx, + const std::vector& phrase) { + std::vector posts(phrase.size()); + for (size_t t = 0; t < phrase.size(); ++t) { + if (!DecodeFullTerm(idx, phrase[t], &posts[t])) { + return {}; + } + } + std::vector cand = posts[0].docids; + for (size_t t = 1; t < posts.size(); ++t) { + std::vector next; + std::set_intersection(cand.begin(), cand.end(), posts[t].docids.begin(), + posts[t].docids.end(), std::back_inserter(next)); + cand.swap(next); + } + auto positions_for = [](const FullPosting& p, uint32_t d) -> const std::vector& { + const auto it = std::ranges::lower_bound(p.docids, d); + return p.positions[static_cast(it - p.docids.begin())]; + }; + std::vector out; + for (uint32_t d : cand) { + const auto& first = positions_for(posts[0], d); + bool hit = false; + for (uint32_t start : first) { + bool ok = true; + for (size_t t = 1; t < posts.size(); ++t) { + const auto& ps = positions_for(posts[t], d); + if (!std::ranges::binary_search(ps, start + static_cast(t))) { + ok = false; + break; + } + } + if (ok) { + hit = true; + break; + } + } + if (hit) { + out.push_back(d); + } + } + return out; +} + +// Looks up the high-df term and reports its full .frq/.prx length + window +// count. +void HighDfStats(const LogicalIndexReader& idx, const std::string& term, uint64_t* frq_len, + uint64_t* prx_len, uint32_t* windows) { + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + ASSERT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + ASSERT_EQ(entry.enc, DictEntryEnc::kWindowed); + FrqPreludeReader prelude; + ASSERT_TRUE(fetch_windowed_prelude(idx, entry, frq_base, &prelude).ok()); + *windows = prelude.n_windows(); + *frq_len = entry.frq_len; + *prx_len = entry.prx_len; +} + +} // namespace + +// boolean_and (MATCH all-terms) must equal the intersection of the per-term +// docid sets (term_query is the trusted oracle). Covers high+low df mix, +// all-present, an absent term (-> empty), and single-term (-> that term's +// docs). +TEST(SniiBooleanAnd, EqualsTermIntersection) { + Corpus c = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + auto intersect_terms = [&](const std::vector& terms) { + std::vector acc; + bool first = true; + for (const auto& t : terms) { + std::vector d; + EXPECT_TRUE(query::term_query(idx, t, &d).ok()); + if (first) { + acc = d; + first = false; + } else { + std::vector out; + std::ranges::set_intersection(acc, d, std::back_inserter(out)); + acc = std::move(out); + } + } + return acc; + }; + + const std::vector> cases = { + {"aa_hi", "aa_quick"}, // high + low df + {"aa_quick", "aa_brown", "aa_fox"}, + {"aa_hi", "nope_absent"}, // absent term -> empty + {"aa_hi"}, // single term -> its docs + }; + for (const auto& terms : cases) { + std::string label; + for (const auto& t : terms) { + label += t + " "; + } + std::vector got; + ASSERT_TRUE(query::boolean_and(idx, terms, &got).ok()) << label; + EXPECT_TRUE(std::ranges::is_sorted(got)) << label; + EXPECT_EQ(got, intersect_terms(terms)) + << "boolean_and != term-intersection: [" << label << "]"; + } + std::remove(path.c_str()); +} + +// prefix_terms ordered enumeration: the full enumeration (empty prefix) must be +// strictly sorted; any prefix scan must equal the full enumeration filtered by +// that prefix; and every enumerated term must lookup() to the SAME DictEntry +// (df). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPrefixTerms, OrderedEnumerationMatchesFilterAndLookup) { + Corpus c = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + std::vector all; + ASSERT_TRUE(idx.prefix_terms("", &all).ok()); + ASSERT_FALSE(all.empty()); + for (size_t i = 1; i < all.size(); ++i) { + EXPECT_LT(all[i - 1].term, all[i].term) << "at " << i; + } + for (const auto& h : all) { + bool found = false; + doris::snii::format::DictEntry e; + uint64_t fb = 0, pb = 0; + ASSERT_TRUE(idx.lookup(h.term, &found, &e, &fb, &pb).ok()); + EXPECT_TRUE(found) << h.term; + EXPECT_EQ(e.df, h.entry.df) << h.term; + } + auto filtered = [&](const std::string& pfx) { + std::vector v; + for (const auto& h : all) { + if (h.term.size() >= pfx.size() && h.term.starts_with(pfx)) { + v.push_back(h.term); + } + } + return v; + }; + auto scanned = [&](const std::string& pfx) { + std::vector hits; + EXPECT_TRUE(idx.prefix_terms(pfx, &hits).ok()); + std::vector v; + for (const auto& h : hits) { + v.push_back(h.term); + } + return v; + }; + for (const char* pfx : {"aa_", "aa_h", "zz_", "zz_0", "alpha"}) { + EXPECT_EQ(scanned(pfx), filtered(pfx)) << "prefix=" << pfx; + } + std::vector none; + ASSERT_TRUE(idx.prefix_terms("zzzznope", &none).ok()); + EXPECT_TRUE(none.empty()); + // Null output is rejected, not dereferenced. + EXPECT_FALSE(idx.prefix_terms("", nullptr).ok()); + std::remove(path.c_str()); +} + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPhraseSkip, SkippingEqualsOracleAndFullRead) { + Corpus c = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + // Small cache block so window-level skipping shows up in remote_bytes / GETs. + io::MeteredFileReader metered(&local, /*block_size=*/4096); + + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + ASSERT_EQ(idx.stats().doc_count, c.doc_count); + + uint64_t hi_frq_len = 0, hi_prx_len = 0; + uint32_t hi_windows = 0; + HighDfStats(idx, "aa_hi", &hi_frq_len, &hi_prx_len, &hi_windows); + ASSERT_GE(hi_windows, 8U) << "high-df term should span many windows"; + + // Phrases: present (incl. 5-term led by high-df term), absent, sub-phrases. + const std::vector> phrases = { + {"aa_hi", "aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // 5-term, hi lead + {"aa_hi", "aa_quick"}, // hi + mid + {"aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // no high-df term + {"aa_lazy", "aa_dog"}, // no high-df term + {"aa_brown", "aa_fox"}, // sub-phrase + {"aa_fox", "aa_brown"}, // reversed -> absent + {"aa_hi", "aa_lazy"}, // present terms, absent phrase + {"aa_quick", "aa_jumps"}, // non-consecutive -> absent + {"aa_echo", "aa_echo"}, // repeated term -> unique-term mapping + {"aa_hi"}, // single high-df term + {"nope", "missing"}, // absent terms -> empty + }; + + for (const auto& p : phrases) { + std::string label; + for (const auto& w : p) { + label += w + " "; + } + + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, p, &got).ok()) << label; + EXPECT_TRUE(std::ranges::is_sorted(got)) << label; + + const std::vector want = c.oracle(p); + EXPECT_EQ(got, want) << "oracle mismatch: [" << label << "]"; + + const std::vector full = FullReadPhrase(idx, p); + EXPECT_EQ(got, full) << "full-read mismatch: [" << label << "]"; + } + + // ---- Byte/round reduction for a phrase CONTAINING the high-df term. ---- + // The low-df lead "aa_quick" concentrates candidates into the first window(s) + // of "aa_hi", so the skip path reads only those windows + the prelude, NOT + // the whole posting. Compare against a full-read of the same phrase. + const std::vector hi_phrase = {"aa_hi", "aa_quick", "aa_brown", "aa_fox", + "aa_jumps"}; + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, hi_phrase, &got).ok()); + const io::IoMetrics skip = metered.metrics(); + + metered.reset_metrics(); + (void)FullReadPhrase(idx, hi_phrase); + const io::IoMetrics full = metered.metrics(); + + // 1. Skipping requests SUBSTANTIALLY fewer bytes than the full-read path: the + // high-df term contributes only its prelude + the candidate-covering + // windows (dd + prx sub-ranges), never its whole posting, so the 5-term + // skip query reads comfortably under 80% of the full-read bytes (a clear + // reduction even with the richer Phase-D prelude and 4 companion terms). + EXPECT_LT(skip.total_request_bytes * 10, full.total_request_bytes * 8) + << "skip=" << skip.total_request_bytes << " full=" << full.total_request_bytes; + // 2. Skipping requests strictly fewer bytes than the full-read path. + EXPECT_LT(skip.total_request_bytes, full.total_request_bytes) + << "skip=" << skip.total_request_bytes << " full=" << full.total_request_bytes; + // 3. Range GETs stay small and are NOT proportional to the term's window + // count + // (the high-df term contributes only its prelude + the covering windows). + EXPECT_LT(skip.range_gets, hi_windows) + << "range_gets=" << skip.range_gets << " windows=" << hi_windows; + // 4. Skipping issues no more range GETs than the full-read path, and stays in + // a small number of serial rounds (preludes + the selected-window batch). + EXPECT_LE(skip.range_gets, full.range_gets) + << "skip_gets=" << skip.range_gets << " full_gets=" << full.range_gets; + EXPECT_GE(skip.serial_rounds, 1U); + EXPECT_LE(skip.serial_rounds, 4U) << "skip did not stay in few serial rounds"; + + std::remove(path.c_str()); +} + +TEST(SniiPhraseSkip, DenseLeadingTermWithRareAnchorSkipsDocsButKeepsPositions) { + Corpus c = BuildDensePhraseCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + uint64_t dense_frq = 0; + uint64_t dense_prx = 0; + uint32_t dense_windows = 0; + HighDfStats(idx, "aa_dense", &dense_frq, &dense_prx, &dense_windows); + ASSERT_GE(dense_windows, 8U); + + const std::vector phrase = {"aa_dense", "aa_rare_driver"}; + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, phrase, &got).ok()); + const uint64_t skipped_bytes = metered.metrics().total_request_bytes; + + metered.reset_metrics(); + const std::vector full = FullReadPhrase(idx, phrase); + const uint64_t full_bytes = metered.metrics().total_request_bytes; + + EXPECT_EQ(got, c.oracle(phrase)); + EXPECT_EQ(got, full); + EXPECT_LT(skipped_bytes, full_bytes) + << "dense leading phrase should keep the rare-position anchor but avoid " + "reading dense docid regions"; + + std::remove(path.c_str()); +} + +TEST(SniiPhraseFrequency, CountsOverlappingAndRepeatedTermOccurrences) { + Corpus corpus; + corpus.doc_count = 4; + corpus.docs = {{"a", "b", "a", "b"}, + {"a", "a", "a", "a"}, + {"a", "b", "c", "a", "b", "c"}, + {"a", "x", "b"}}; + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + LogicalIndexReader index; + ASSERT_TRUE(segment.open_index(1, "body", &index).ok()); + + std::vector matches; + ASSERT_TRUE(query::phrase_query_with_frequencies(index, {"a", "b"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 2}, + {.docid = 2, .frequency = 2}})); + + ASSERT_TRUE(query::phrase_query_with_frequencies(index, {"a", "a"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 1, .frequency = 3}})); + + ASSERT_TRUE(query::phrase_query_with_frequencies(index, {"a", "a", "a"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 1, .frequency = 2}})); + + ASSERT_TRUE(query::phrase_query_with_frequencies(index, {"a", "b", "c"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 2, .frequency = 2}})); + + std::remove(path.c_str()); +} + +TEST(SniiPhraseFrequency, UsesFirstClauseMultiplicityWithStackedAnchorPositions) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.add_token("high", 0, 0); + buffer.add_token("rare", 0, 1); + buffer.add_token("rare", 0, 1); + buffer.add_token("tail", 0, 2); + buffer.add_token("high", 1, 0); + buffer.add_token("high", 1, 0); + buffer.add_token("rare", 1, 1); + buffer.add_token("rare", 1, 1); + buffer.add_token("tail", 1, 2); + for (uint32_t docid = 2; docid < 8; ++docid) { + buffer.add_token("high", docid, 0); + } + + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "body"; + input.config = IndexConfig::kDocsPositions; + input.doc_count = 8; + input.terms = buffer.finalize_sorted(); + input.target_dict_block_bytes = 256; + + const std::string path = TempPath(); + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(input).ok()); + ASSERT_TRUE(compound.finish().ok()); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + LogicalIndexReader index; + ASSERT_TRUE(segment.open_index(1, "body", &index).ok()); + + std::vector matches; + ASSERT_TRUE( + query::phrase_query_with_frequencies(index, {"high", "rare", "tail"}, &matches).ok()); + EXPECT_EQ(matches, (std::vector {{.docid = 0, .frequency = 1}, + {.docid = 1, .frequency = 2}})); + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/position_math_test.cpp b/be/test/storage/index/snii/query/position_math_test.cpp new file mode 100644 index 00000000000000..365adfb3f8a24c --- /dev/null +++ b/be/test/storage/index/snii/query/position_math_test.cpp @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "storage/index/snii/query/internal/position_math.h" + +#include + +#include +#include +#include + +#include "common/status.h" + +using doris::Status; // RETURN_IF_ERROR expands to bare Status // NOLINT(misc-unused-using-decls) + +TEST(SniiPositionMath, AddsOffsetWhenRepresentable) { + uint32_t out = 0; + EXPECT_TRUE(doris::snii::query::internal::add_position_offset(41, 1, &out)); + EXPECT_EQ(out, 42U); + EXPECT_TRUE(doris::snii::query::internal::add_position_offset( + std::numeric_limits::max() - 2, 2, &out)); + EXPECT_EQ(out, std::numeric_limits::max()); +} + +TEST(SniiPositionMath, RejectsWraparound) { + uint32_t out = 7; + EXPECT_FALSE(doris::snii::query::internal::add_position_offset( + std::numeric_limits::max(), 1, &out)); + EXPECT_EQ(out, 7U); +} + +TEST(SniiPositionMath, BuildsDenseOffsets) { + std::vector offsets; + EXPECT_TRUE(doris::snii::query::internal::build_position_offsets(4, &offsets)); + EXPECT_EQ(offsets, (std::vector {0U, 1U, 2U, 3U})); +} + +TEST(SniiPositionMath, RejectsUnrepresentableOffsetCount) { + std::vector offsets = {9}; + EXPECT_FALSE(doris::snii::query::internal::build_position_offsets( + std::numeric_limits::max(), &offsets)); + EXPECT_EQ(offsets, (std::vector {9U})); +} diff --git a/be/test/storage/index/snii/query/posting_grouping_test.cpp b/be/test/storage/index/snii/query/posting_grouping_test.cpp new file mode 100644 index 00000000000000..53466d1ef9d9e1 --- /dev/null +++ b/be/test/storage/index/snii/query/posting_grouping_test.cpp @@ -0,0 +1,514 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/batch_range_fetcher.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/scoring_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +// PHASE D differential + contiguity test (design 1.6: posting-level dd/freq +// grouping). The windowed .frq payload is laid out +// [prelude][dd-block][freq-block] so a docid-only / phrase reader fetches the +// docs-only data ([prelude][dd-block]) as ONE CONTIGUOUS run. Over a ~5000-doc +// kDocsPositionsScoring index with a high-df term spanning MANY adaptive +// windows plus mid/low terms and a planted 5-term phrase, this asserts: +// (a) term_query + phrase_query (incl the 5-term phrase) docids == a +// brute-force +// ORACLE; scoring top-K (exhaustive == wand == selective) unchanged. +// (b) Through a MeteredFileReader the docid-only posting reader for the +// high-df +// term reads the dd-block as a CONTIGUOUS region: read_at is SMALL (a +// couple of ranges, NOT one-per-window / not thousands), range_gets is +// small, AND request_bytes is strictly LESS than fetching the full +// posting (the freq-block is skipped). All three (read_at, range_gets, +// request_bytes) drop vs the full-posting read, with identical docids. +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::snii::query::Bm25Params; +using doris::snii::query::ScoredDoc; +using doris::snii::stats::SniiStatsProvider; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_posting_grouping_" + std::to_string(getpid()) + "_" + + std::to_string(counter++) + ".idx"; +} + +constexpr uint32_t kDocCount = 5000; +constexpr uint32_t kHiCount = 4800; // df of aa_hi (>> 256 -> many windows) +constexpr uint32_t kPhraseSpan = 150; // 5-term phrase planted in early aa_hi docs + +// aa_hi appears with varied, large frequency so its per-window freq region is +// big in absolute terms (the freq-block the docs-only path never fetches). +uint32_t HiFreq(uint32_t occ) { + return 50U + (occ * 17U + 11U) % 240U; +} + +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; + std::vector doc_len; + + bool phrase_in_doc(uint32_t d, const std::vector& phrase) const { + if (phrase.empty()) { + return false; + } + const auto& toks = docs[d]; + if (toks.size() < phrase.size()) { + return false; + } + for (size_t i = 0; i + phrase.size() <= toks.size(); ++i) { + bool match = true; + for (size_t k = 0; k < phrase.size(); ++k) { + if (toks[i + k] != phrase[k]) { + match = false; + break; + } + } + if (match) { + return true; + } + } + return false; + } + + std::vector term_oracle(const std::string& term) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + const auto& toks = docs[d]; + if (std::ranges::find(toks, term) != toks.end()) { + out.push_back(d); + } + } + return out; + } + + std::vector phrase_oracle(const std::vector& phrase) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + if (phrase_in_doc(d, phrase)) { + out.push_back(d); + } + } + return out; + } +}; + +Corpus BuildCorpus() { + Corpus c; + c.doc_count = kDocCount; + c.docs.resize(c.doc_count); + c.doc_len.assign(c.doc_count, 0); + const std::vector vocab = {"alpha", "bravo", "charlie", "delta"}; + + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + const bool is_hi = d < kHiCount; + if (is_hi) { + toks.emplace_back("aa_hi"); // high-df, position 0 + } + const uint32_t occ = d; + + if (is_hi && occ < kPhraseSpan) { + toks.emplace_back("aa_quick"); + toks.emplace_back("aa_brown"); + toks.emplace_back("aa_fox"); + toks.emplace_back("aa_jumps"); + } else if (d % 13 == 0) { + toks.emplace_back("aa_lazy"); + toks.emplace_back("aa_dog"); + } + if (d % 5 == 0) { + toks.emplace_back("aa_mid"); + if (d % 15 == 0) { + toks.emplace_back("aa_mid"); + } + } + for (uint32_t k = 0; k < 2; ++k) { + toks.push_back(vocab[(d + k) % vocab.size()]); + } + if (d % 101 == 0) { + toks.emplace_back("aa_rare"); + } + char nm[16]; + std::snprintf(nm, sizeof(nm), "zz_%03u", d % 400); + toks.emplace_back(nm); + if (is_hi) { + const uint32_t extra = HiFreq(occ) - 1U; // total freq = HiFreq(occ) + for (uint32_t k = 0; k < extra; ++k) { + toks.emplace_back("aa_hi"); + } + } + c.doc_len[d] = toks.size(); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.doc_count; ++d) { + const auto& toks = c.docs[d]; + for (uint32_t pos = 0; pos < toks.size(); ++pos) { + buf.add_token(toks[pos], d, pos); + } + } + std::vector terms = buf.finalize_sorted(); + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositionsScoring; + in.doc_count = c.doc_count; + in.terms = std::move(terms); + in.target_dict_block_bytes = 256; + in.encoded_norms.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + in.encoded_norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(in).ok()); + ASSERT_TRUE(cw.finish().ok()); +} + +// dd-block / freq-block byte sizes of a windowed term (the contiguous runs). +struct Blocks { + uint64_t prelude_len = 0; + uint64_t dd_block = 0; + uint64_t freq_block = 0; + uint32_t windows = 0; +}; + +Blocks MeasureBlocks(const LogicalIndexReader& idx, const std::string& term) { + Blocks b; + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + EXPECT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + EXPECT_TRUE(found); + EXPECT_EQ(entry.enc, DictEntryEnc::kWindowed); + b.prelude_len = entry.prelude_len; + FrqPreludeReader prelude; + EXPECT_TRUE(fetch_windowed_prelude(idx, entry, frq_base, &prelude).ok()); + b.windows = prelude.n_windows(); + b.dd_block = prelude.dd_block_len(); + b.freq_block = prelude.freq_block_len(); + return b; +} + +std::vector ReferenceRanking(const Corpus& c, const std::vector& terms, + uint32_t k, const Bm25Params& params) { + uint64_t sum_ttf = 0; + for (uint64_t dl : c.doc_len) { + sum_ttf += dl; + } + const double avgdl = static_cast(sum_ttf) / std::max(1, c.doc_count); + + std::unordered_map scores; + for (const auto& term : terms) { + std::map plist; + for (uint32_t d = 0; d < c.doc_count; ++d) { + uint32_t f = 0; + for (const auto& t : c.docs[d]) { + if (t == term) { + ++f; + } + } + if (f > 0) { + plist[d] = f; + } + } + if (plist.empty()) { + continue; + } + const uint64_t df = plist.size(); + const double idf = + std::log(1.0 + (static_cast(c.doc_count) - df + 0.5) / (df + 0.5)); + for (const auto& [docid, freq] : plist) { + const double dl = doris::snii::query::decode_norm( + doris::snii::query::encode_norm(c.doc_len[docid])); + const double denom = freq + params.k1 * (1.0 - params.b + params.b * dl / avgdl); + scores[docid] += idf * (freq * (params.k1 + 1.0)) / denom; + } + } + + std::vector all; + all.reserve(scores.size()); + for (const auto& [docid, s] : scores) { + all.push_back({docid, s}); + } + std::ranges::sort(all, [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) { + return a.score > b.score; + } + return a.docid < b.docid; + }); + if (all.size() > k) { + all.resize(k); + } + return all; +} + +void ExpectRankingEqual(const std::vector& got, const std::vector& ref, + const char* label) { + ASSERT_EQ(got.size(), ref.size()) << label; + for (size_t i = 0; i < ref.size(); ++i) { + EXPECT_EQ(got[i].docid, ref[i].docid) << label << " docid i=" << i; + EXPECT_NEAR(got[i].score, ref[i].score, 1e-9) << label << " score i=" << i; + } +} + +// PRE-GROUPING SIMULATION: decode a windowed term's docids by fetching EACH +// window's dd region in its OWN read round (one BatchRangeFetcher per window), +// as the prior per-window [dd][freq] layout forced -- each window's docs prefix +// had to be fetched on its own because the next window's dd was separated by +// the current window's (skipped) freq region. This reproduces the read_at / +// range_get explosion (one per window) that the grouped, contiguous dd-block +// eliminates. Same docids. +std::vector DecodePerWindow(const LogicalIndexReader& idx, const DictEntry& entry, + uint64_t frq_base, uint64_t prx_base) { + FrqPreludeReader prelude; + EXPECT_TRUE(fetch_windowed_prelude(idx, entry, frq_base, &prelude).ok()); + std::vector docids; + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowAbsRange r; + EXPECT_TRUE(windowed_window_range(idx, entry, frq_base, prx_base, prelude, w, + /*want_positions=*/false, /*want_freq=*/false, &r) + .ok()); + WindowMeta m; + EXPECT_TRUE(prelude.window(w, &m).ok()); + // One read round per window (the un-grouped reader could not coalesce + // across the interleaved freq regions of the old layout). + doris::snii::io::BatchRangeFetcher fetcher(idx.reader(), /*coalesce_gap=*/0); + const size_t h = fetcher.add(r.dd_off, r.dd_len); + EXPECT_TRUE(fetcher.fetch().ok()); + std::vector wd, wf; + std::vector> wp; + EXPECT_TRUE(decode_window_slices(m, fetcher.get(h), Slice(), Slice(), + /*want_positions=*/false, + /*want_freq=*/false, &wd, &wf, &wp) + .ok()); + docids.insert(docids.end(), wd.begin(), wd.end()); + } + return docids; +} + +} // namespace + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingGrouping, ContiguousDdBlockSavesAllThreeMetrics) { + Corpus c = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(c, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + // Fine-grained FileCache block (< a window's freq region) so the grouped + // docs-only run occupies strictly fewer cache blocks than (i) the full + // posting and (ii) a per-window dd fetch whose gaps span the skipped freq + // regions. + io::MeteredFileReader metered(&local, /*block_size=*/64); + + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + ASSERT_EQ(idx.stats().doc_count, c.doc_count); + + // The high-df term must span many windows (the contiguity win is meaningful). + const Blocks hi = MeasureBlocks(idx, "aa_hi"); + ASSERT_GE(hi.windows, 8U) << "high-df term should span many windows"; + ASSERT_GT(hi.freq_block, 0U) << "freq-block must carry real bytes"; + + // ---- (a) term_query / phrase_query docids == ORACLE ----------------------- + const std::vector terms_to_check = {"aa_hi", "aa_mid", "aa_quick", "aa_lazy", + "aa_rare"}; + for (const auto& term : terms_to_check) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, term, &got).ok()) << term; + EXPECT_TRUE(std::ranges::is_sorted(got)) << term; + EXPECT_EQ(got, c.term_oracle(term)) << "term_query oracle mismatch: " << term; + } + + const std::vector> phrases = { + {"aa_hi", "aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // 5-term, hi lead + {"aa_hi", "aa_quick"}, // hi + mid + {"aa_quick", "aa_brown", "aa_fox", "aa_jumps"}, // no high-df term + {"aa_lazy", "aa_dog"}, // scattered phrase + {"aa_brown", "aa_fox"}, // sub-phrase + {"aa_fox", "aa_brown"}, // reversed -> absent + {"aa_hi", "aa_lazy"}, // present terms, absent phrase + }; + for (const auto& p : phrases) { + std::string label; + for (const auto& w : p) { + label += w + " "; + } + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, p, &got).ok()) << label; + EXPECT_TRUE(std::ranges::is_sorted(got)) << label; + EXPECT_EQ(got, c.phrase_oracle(p)) << "phrase oracle mismatch: [" << label << "]"; + } + + // ---- scoring: exhaustive == wand == selective == reference ---------------- + SniiStatsProvider stats; + ASSERT_TRUE(SniiStatsProvider::open(&idx, &stats).ok()); + const Bm25Params params; + const std::vector score_terms = {"aa_hi", "aa_mid", "aa_rare"}; + for (uint32_t k : {1U, 10U, 100U}) { + std::vector ex, wa, sel; + ASSERT_TRUE(query::scoring_query_exhaustive(idx, stats, score_terms, k, params, &ex).ok()); + ASSERT_TRUE(query::scoring_query_wand(idx, stats, score_terms, k, params, &wa).ok()); + ASSERT_TRUE( + query::scoring_query_wand_selective(idx, stats, score_terms, k, params, &sel).ok()); + const std::vector ref = ReferenceRanking(c, score_terms, k, params); + ExpectRankingEqual(ex, ref, "exhaustive"); + ExpectRankingEqual(wa, ex, "wand"); + ExpectRankingEqual(sel, ex, "selective"); + } + + DictEntry hi_entry; + uint64_t hi_frq_base = 0, hi_prx_base = 0; + bool hi_found = false; + ASSERT_TRUE(idx.lookup("aa_hi", &hi_found, &hi_entry, &hi_frq_base, &hi_prx_base).ok()); + ASSERT_TRUE(hi_found); + ASSERT_EQ(hi_entry.frq_docs_len, hi.prelude_len + hi.dd_block); + ASSERT_LT(hi_entry.frq_docs_len, hi_entry.frq_len); + + DictEntry oversized_docs_prefix = hi_entry; + ++oversized_docs_prefix.frq_docs_len; + std::vector corrupt_docs; + // Integrated SNII reports posting corruption with the inverted-index-file + // corruption code (the "docs prefix length mismatch" guard), not the generic + // CORRUPTION code the standalone test used. + EXPECT_TRUE(query::internal::read_docid_posting(idx, oversized_docs_prefix, hi_frq_base, + hi_prx_base, &corrupt_docs) + .is()); + + // ---- (b) docid-only posting read fetches the dd-block CONTIGUOUSLY -------- + // GROUPED = the Phase-D grouped docid-only path: read [prelude][dd-block] as + // one contiguous prefix; the freq-block is skipped on the wire. The lookup is + // deliberately outside this metrics window so the assertion covers posting + // I/O. + metered.reset_metrics(); + std::vector grouped_docs; + ASSERT_TRUE(query::internal::read_docid_posting(idx, hi_entry, hi_frq_base, hi_prx_base, + &grouped_docs) + .ok()); + const io::IoMetrics grouped = metered.metrics(); + ASSERT_EQ(grouped_docs, c.term_oracle("aa_hi")); + EXPECT_EQ(grouped.serial_rounds, 1U) + << "docid-only windowed term_query should fetch [prelude][dd-block] in " + "one " + "batched round"; + EXPECT_EQ(grouped.range_gets, 1U) << "docid-only windowed term_query should " + "issue one contiguous prefix range"; + + // PER-WINDOW = the pre-grouping layout: fetch EACH window's dd region as its + // own physical range (gaps = the skipped freq regions). Same resolved term. + metered.reset_metrics(); + const std::vector perwin_docs = + DecodePerWindow(idx, hi_entry, hi_frq_base, hi_prx_base); + const io::IoMetrics perwin = metered.metrics(); + ASSERT_EQ(perwin_docs, c.term_oracle("aa_hi")); + + // FULL = fetch the whole posting (prelude + dd-block + freq-block). Same + // resolved term. + metered.reset_metrics(); + { + DecodedPosting full_posting; + ASSERT_TRUE(read_windowed_posting(idx, hi_entry, hi_frq_base, hi_prx_base, + /*want_positions=*/false, + /*want_freq=*/true, &full_posting) + .ok()); + EXPECT_EQ(full_posting.docids, c.term_oracle("aa_hi")); + EXPECT_EQ(full_posting.freqs.size(), full_posting.docids.size()); + } + const io::IoMetrics full = metered.metrics(); + + // The grouped docid-only path issues only a HANDFUL of logical reads (lookup + // + prelude + ONE dd-block range), NOT one-per-window and not thousands. + EXPECT_LE(grouped.read_at_calls, 4U) + << "grouped docid-only read_at not contiguous-small: " << grouped.read_at_calls; + EXPECT_LT(grouped.read_at_calls, hi.windows) + << "grouped must not read per-window: read_at=" << grouped.read_at_calls + << " windows=" << hi.windows; + + // (1) vs PER-WINDOW: all three drop because the dd-block is ONE contiguous + // range instead of one fetch per window separated by the skipped freq + // regions. + EXPECT_LT(grouped.read_at_calls, perwin.read_at_calls) + << "read_at vs per-window did not drop: grouped=" << grouped.read_at_calls + << " perwin=" << perwin.read_at_calls; + EXPECT_LT(grouped.range_gets, perwin.range_gets) + << "range_gets vs per-window did not drop: grouped=" << grouped.range_gets + << " perwin=" << perwin.range_gets; + EXPECT_LE(grouped.total_request_bytes, perwin.total_request_bytes) + << "request_bytes vs per-window grew: grouped=" << grouped.total_request_bytes + << " perwin=" << perwin.total_request_bytes; + + // (2) vs FULL posting: request_bytes and remote_bytes strictly drop + // (freq-block skipped on the wire and its cache blocks never touched). + EXPECT_LT(grouped.total_request_bytes, full.total_request_bytes) + << "request_bytes vs full did not drop: grouped=" << grouped.total_request_bytes + << " full=" << full.total_request_bytes; + EXPECT_EQ(full.total_request_bytes - grouped.total_request_bytes, hi.freq_block) + << "skipped bytes != freq-block size"; + EXPECT_LT(grouped.remote_bytes, full.remote_bytes) + << "remote_bytes vs full did not drop: grouped=" << grouped.remote_bytes + << " full=" << full.remote_bytes; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/prefix_query_test.cpp b/be/test/storage/index/snii/query/prefix_query_test.cpp new file mode 100644 index 00000000000000..3651c264cc7dfa --- /dev/null +++ b/be/test/storage/index/snii/query/prefix_query_test.cpp @@ -0,0 +1,185 @@ +// 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 "storage/index/snii/query/prefix_query.h" + +#include +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_prefix_query_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; + + std::vector prefix_docs(const std::string& prefix) const { + std::set ids; + for (uint32_t d = 0; d < docs.size(); ++d) { + for (const std::string& term : docs[d]) { + if (term.size() >= prefix.size() && term.starts_with(prefix)) { + ids.insert(d); + } + } + } + return {ids.begin(), ids.end()}; + } +}; + +Corpus BuildMixedCorpus() { + Corpus c; + c.doc_count = 120; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& terms = c.docs[d]; + if (d < 80) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + terms.emplace_back(term); + } + if (d < 80 && d % 2 == 0) { + terms.emplace_back("aa_even"); + } + if (d < 50) { + char term[16]; + std::snprintf(term, sizeof(term), "ab_%03u", d); + terms.emplace_back(term); + } + char filler[16]; + std::snprintf(filler, sizeof(filler), "zz_%03u", d); + terms.emplace_back(filler); + } + return c; +} + +Corpus BuildLowDfPrefixCorpus() { + Corpus c; + c.doc_count = 96; + c.docs.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + c.docs[d].emplace_back(term); + } + return c; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/false); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + const std::vector& terms = c.docs[d]; + for (uint32_t pos = 0; pos < terms.size(); ++pos) { + buf.add_token(terms[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = doris::snii::format::IndexConfig::kDocsOnly; + in.doc_count = c.doc_count; + in.terms = buf.finalize_sorted(); + in.target_dict_block_bytes = 2048; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +LogicalIndexReader OpenIndex(io::LocalFileReader* file, SniiSegmentReader* segment, + const std::string& path) { + EXPECT_TRUE(file->open(path).ok()); + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +LogicalIndexReader OpenMeteredIndex(io::MeteredFileReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +} // namespace + +TEST(SniiPrefixQuery, MatchesPrefixOracle) { + const Corpus corpus = BuildMixedCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader file; + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment, path); + + for (const char* prefix : {"aa_", "aa_e", "ab_0", "zz_", "missing"}) { + std::vector got; + ASSERT_TRUE(query::prefix_query(idx, prefix, &got).ok()) << prefix; + EXPECT_TRUE(std::ranges::is_sorted(got)) << prefix; + EXPECT_EQ(got, corpus.prefix_docs(prefix)) << prefix; + } + + std::remove(path.c_str()); +} + +TEST(SniiPrefixQuery, UsesEnumeratedEntriesWithoutPerTermLookup) { + const Corpus corpus = BuildLowDfPrefixCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/4096); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::prefix_query(idx, "aa_", &got).ok()); + + EXPECT_EQ(got, corpus.prefix_docs("aa_")); + EXPECT_LT(metered.metrics().read_at_calls, corpus.doc_count / 3) + << "prefix_query must reuse PrefixHit entries, not lookup every term again"; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/query_operator_error_test.cpp b/be/test/storage/index/snii/query/query_operator_error_test.cpp new file mode 100644 index 00000000000000..eb5833f986bd12 --- /dev/null +++ b/be/test/storage/index/snii/query/query_operator_error_test.cpp @@ -0,0 +1,286 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_query_operator_error_" + std::to_string(getpid()) + "_" + + std::to_string(counter++) + ".idx"; +} + +// In-memory FileReader (backed by a flat byte buffer) whose reads can be made +// to fail on demand, so queries can be driven over the I/O-error propagation +// paths after the index has been opened successfully. +class FaultInjectingReader : public io::FileReader { +public: + explicit FaultInjectingReader(std::vector bytes) : bytes_(std::move(bytes)) {} + + void set_fail_reads(bool v) { fail_reads_ = v; } + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (fail_reads_) { + return Status::IOError("injected read failure"); + } + if (out == nullptr) { + return Status::InvalidArgument("memory reader null out"); + } + if (offset > bytes_.size() || len > bytes_.size() - offset) { + return Status::Corruption("memory reader read past end"); + } + out->assign(bytes_.begin() + static_cast(offset), + bytes_.begin() + static_cast(offset + len)); + return Status::OK(); + } + + uint64_t size() const override { return bytes_.size(); } + +private: + std::vector bytes_; + bool fail_reads_ = false; +}; + +struct Corpus { + std::vector> docs = { + {"alpha", "beta", "gamma"}, {"alpha", "beta", "delta"}, {"alpha", "bravo", "gamma"}, + {"beta", "gamma"}, {"alphabeta", "beta"}, + }; +}; + +struct EnvGuard { + std::string name; + bool had = false; + std::string old; + + explicit EnvGuard(const char* env_name) : name(env_name) { + const char* value = std::getenv(env_name); + had = value != nullptr; + if (had) { + old = value; + } + } + + ~EnvGuard() { + if (had) { + ::setenv(name.c_str(), old.c_str(), 1); + } else { + ::unsetenv(name.c_str()); + } + } +}; + +void BuildIndexBytes(const Corpus& corpus, doris::snii::format::IndexConfig config, + std::vector* bytes) { + SpimiTermBuffer buf(/*has_positions=*/config != doris::snii::format::IndexConfig::kDocsOnly); + for (uint32_t d = 0; d < corpus.docs.size(); ++d) { + const std::vector& terms = corpus.docs[d]; + for (uint32_t pos = 0; pos < terms.size(); ++pos) { + buf.add_token(terms[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = config; + in.doc_count = static_cast(corpus.docs.size()); + if (config == doris::snii::format::IndexConfig::kDocsPositionsScoring) { + in.encoded_norms.assign(corpus.docs.size(), 1); + doris::segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = + doris::segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + metadata.scoring_coverage = doris::segment_v2::inverted_index::ScoringCoverage::kComplete; + metadata.scoring_stats_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + metadata.base_analyzer_fingerprint = "query-operator-error-test"; + metadata.scoring_doc_count = corpus.docs.size(); + for (const auto& terms : corpus.docs) { + metadata.scoring_token_count += terms.size(); + } + in.common_grams_metadata = std::move(metadata); + } + in.terms = buf.finalize_sorted(); + in.target_dict_block_bytes = 512; + + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); + } + { + io::LocalFileReader file; + ASSERT_TRUE(file.open(path).ok()); + ASSERT_TRUE(file.read_at(0, file.size(), bytes).ok()); + } + std::remove(path.c_str()); +} + +LogicalIndexReader OpenIndex(FaultInjectingReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +void ExpectIoError(const Status& status) { + EXPECT_TRUE(status.is()) << status.to_string(); +} + +} // namespace + +TEST(SniiQueryOperatorBoundaries, RejectNullOutputPointers) { + LogicalIndexReader idx; + // A typed null output pointer selects each operator's std::vector* + // overload unambiguously (the integrated API also exposes DocIdSink* + // overloads, so a bare nullptr would be ambiguous). + auto* null_docs = static_cast*>(nullptr); + EXPECT_TRUE( + query::term_query(idx, "alpha", null_docs).is()); + EXPECT_TRUE( + query::boolean_or(idx, {"alpha"}, null_docs).is()); + EXPECT_TRUE( + query::boolean_and(idx, {"alpha"}, null_docs).is()); + EXPECT_TRUE(query::prefix_query(idx, "al", null_docs).is()); + EXPECT_TRUE( + query::wildcard_query(idx, "a*", null_docs).is()); + EXPECT_TRUE( + query::regexp_query(idx, "a.*", null_docs).is()); + EXPECT_TRUE(query::phrase_query(idx, {"alpha"}, null_docs) + .is()); + EXPECT_TRUE(query::phrase_prefix_query(idx, {"alpha"}, null_docs) + .is()); +} + +TEST(SniiQueryOperatorBoundaries, EmptyMissingAndSingleTermCasesAreWellDefined) { + std::vector bytes; + BuildIndexBytes(Corpus {}, doris::snii::format::IndexConfig::kDocsPositionsScoring, &bytes); + FaultInjectingReader file(std::move(bytes)); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment); + + std::vector got = {99}; + ASSERT_TRUE(query::term_query(idx, "missing", &got).ok()); + EXPECT_TRUE(got.empty()); + + got = {99}; + ASSERT_TRUE(query::boolean_or(idx, {}, &got).ok()); + EXPECT_TRUE(got.empty()); + + got = {99}; + ASSERT_TRUE(query::boolean_and(idx, {}, &got).ok()); + EXPECT_TRUE(got.empty()); + + got = {99}; + ASSERT_TRUE(query::boolean_and(idx, {"alpha", "missing"}, &got).ok()); + EXPECT_TRUE(got.empty()); + + got = {99}; + ASSERT_TRUE(query::phrase_query(idx, {}, &got).ok()); + EXPECT_TRUE(got.empty()); + + got = {99}; + ASSERT_TRUE(query::phrase_prefix_query(idx, {}, &got).ok()); + EXPECT_TRUE(got.empty()); + + ASSERT_TRUE(query::prefix_query(idx, "", &got).ok()); + EXPECT_EQ(got, (std::vector {0, 1, 2, 3, 4})); + + ASSERT_TRUE(query::phrase_query(idx, {"alpha"}, &got).ok()); + EXPECT_EQ(got, (std::vector {0, 1, 2})); + + ASSERT_TRUE(query::phrase_prefix_query(idx, {"alpha"}, &got).ok()); + EXPECT_EQ(got, (std::vector {0, 1, 2, 4})); + + ASSERT_TRUE(query::wildcard_query(idx, "", &got).ok()); + EXPECT_TRUE(got.empty()); + + EXPECT_TRUE(query::regexp_query(idx, "[", &got).is()); +} + +TEST(SniiQueryOperatorBoundaries, PositionQueriesRejectDocsOnlyIndex) { + std::vector bytes; + BuildIndexBytes(Corpus {}, doris::snii::format::IndexConfig::kDocsOnly, &bytes); + FaultInjectingReader file(std::move(bytes)); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment); + + std::vector got; + EXPECT_TRUE(query::phrase_query(idx, {"alpha", "beta"}, &got) + .is()); + EXPECT_TRUE(query::phrase_prefix_query(idx, {"alpha", "b"}, &got) + .is()); +} + +TEST(SniiQueryOperatorIoErrors, PropagateUnderlyingReadFailures) { + EnvGuard dict_resident("SNII_DICT_RESIDENT_MAX"); + ::setenv("SNII_DICT_RESIDENT_MAX", "0", 1); + + std::vector bytes; + BuildIndexBytes(Corpus {}, doris::snii::format::IndexConfig::kDocsPositionsScoring, &bytes); + FaultInjectingReader file(std::move(bytes)); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenIndex(&file, &segment); + file.set_fail_reads(true); + + std::vector got; + ExpectIoError(query::term_query(idx, "alpha", &got)); + ExpectIoError(query::boolean_or(idx, {"alpha", "beta"}, &got)); + ExpectIoError(query::boolean_and(idx, {"alpha", "beta"}, &got)); + ExpectIoError(query::prefix_query(idx, "al", &got)); + ExpectIoError(query::prefix_query(idx, "", &got)); + ExpectIoError(query::wildcard_query(idx, "alpha*", &got)); + ExpectIoError(query::wildcard_query(idx, "*alpha", &got)); + ExpectIoError(query::regexp_query(idx, "alpha.*", &got)); + ExpectIoError(query::regexp_query(idx, ".*alpha", &got)); + ExpectIoError(query::phrase_query(idx, {"alpha", "beta"}, &got)); + ExpectIoError(query::phrase_prefix_query(idx, {"alpha", "b"}, &got)); +} diff --git a/be/test/storage/index/snii/query/query_profile_test.cpp b/be/test/storage/index/snii/query/query_profile_test.cpp new file mode 100644 index 00000000000000..6a0efa18861bd2 --- /dev/null +++ b/be/test/storage/index/snii/query/query_profile_test.cpp @@ -0,0 +1,985 @@ +// 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 "storage/index/snii/query/query_profile.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "runtime/runtime_profile.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/query/phrase_prx_validation.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/snii_prx_profile.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "util/defer_op.h" + +using namespace doris::snii; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +template +concept HasSniiPrxEncodedBytes = requires(T value) { value.snii_prx_encoded_bytes; }; +template +concept HasSniiPrxPayloadBytes = requires(T value) { value.snii_prx_payload_bytes; }; +template +concept HasSniiPrxCompressedBytes = requires(T value) { value.snii_prx_compressed_bytes; }; +template +concept HasSniiPrxChild128Touches = requires(T value) { value.snii_prx_child_128_touches; }; +template +concept HasSniiPrxChild256Touches = requires(T value) { value.snii_prx_child_256_touches; }; +template +concept HasSniiPrxCrcValidationNs = requires(T value) { value.snii_prx_crc_validation_ns; }; +template +concept HasSniiPrxDecompressNs = requires(T value) { value.snii_prx_decompress_ns; }; +template +concept HasSniiPrxScratchAllocationEvents = + requires(T value) { value.snii_prx_scratch_allocation_events; }; +template +concept HasSniiPrxContainerAllocationEvents = + requires(T value) { value.snii_prx_container_allocation_events; }; + +static_assert(!HasSniiPrxEncodedBytes); +static_assert(!HasSniiPrxPayloadBytes); +static_assert(!HasSniiPrxCompressedBytes); +static_assert(!HasSniiPrxChild128Touches); +static_assert(!HasSniiPrxChild256Touches); +static_assert(!HasSniiPrxCrcValidationNs); +static_assert(!HasSniiPrxDecompressNs); +static_assert(!HasSniiPrxScratchAllocationEvents); +static_assert(!HasSniiPrxContainerAllocationEvents); + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_query_profile_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +struct Corpus { + std::vector> docs; +}; + +Corpus BuildCorpus() { + Corpus c; + c.docs.resize(128); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + std::vector& doc = c.docs[d]; + doc.emplace_back("lead"); + doc.emplace_back("quick"); + doc.emplace_back(d % 2 == 0 ? "brown" : "bronze"); + char term[16]; + std::snprintf(term, sizeof(term), "aa_%03u", d); + doc.emplace_back(term); + } + return c; +} + +Corpus BuildHighTfCorpus() { + constexpr uint32_t kTailCount = 33; + constexpr uint32_t kRepetitions = 48; + Corpus corpus; + corpus.docs.resize(600); + for (uint32_t docid = 0; docid < corpus.docs.size(); ++docid) { + std::vector& doc = corpus.docs[docid]; + doc.reserve(kRepetitions * 5); + char tail[32]; + std::snprintf(tail, sizeof(tail), "epsilon_%02u", docid % kTailCount); + for (uint32_t repetition = 0; repetition < kRepetitions; ++repetition) { + doc.insert(doc.end(), {"alpha", "beta", "gamma", "delta", tail}); + } + } + return corpus; +} + +Corpus BuildSingleTailGroupCorpus() { + Corpus corpus; + corpus.docs.resize(96); + for (uint32_t docid = 0; docid < corpus.docs.size(); ++docid) { + char tail[32]; + std::snprintf(tail, sizeof(tail), "epsilon_%02u", docid % 3); + corpus.docs[docid] = {"alpha", "beta", tail}; + } + return corpus; +} + +Corpus BuildDisjointTailGroupCorpus() { + Corpus corpus; + corpus.docs.resize(99); + for (uint32_t docid = 0; docid < 96; ++docid) { + corpus.docs[docid] = {"alpha", "beta", "other"}; + } + for (uint32_t tail = 0; tail < 3; ++tail) { + char term[32]; + std::snprintf(term, sizeof(term), "epsilon_%02u", tail); + corpus.docs[96 + tail] = {term}; + } + return corpus; +} + +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.docs.size(); ++d) { + const std::vector& terms = c.docs[d]; + for (uint32_t pos = 0; pos < terms.size(); ++pos) { + buf.add_token(terms[pos], d, pos); + } + } + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = doris::snii::format::IndexConfig::kDocsPositionsScoring; + in.doc_count = static_cast(c.docs.size()); + in.encoded_norms.assign(c.docs.size(), 1); + doris::segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = + doris::segment_v2::inverted_index::PlainTermKeyVersion::kRawNoInternal; + metadata.scoring_coverage = doris::segment_v2::inverted_index::ScoringCoverage::kComplete; + metadata.scoring_stats_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + metadata.base_analyzer_fingerprint = "query-profile-test"; + metadata.scoring_doc_count = c.docs.size(); + for (const auto& terms : c.docs) { + metadata.scoring_token_count += terms.size(); + } + in.common_grams_metadata = std::move(metadata); + in.terms = buf.finalize_sorted(); + in.target_dict_block_bytes = 512; + + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + ASSERT_TRUE(compound.add_logical_index(in).ok()); + ASSERT_TRUE(compound.finish().ok()); +} + +LogicalIndexReader OpenMeteredIndex(io::MeteredFileReader* file, SniiSegmentReader* segment) { + EXPECT_TRUE(SniiSegmentReader::open(file, segment).ok()); + LogicalIndexReader idx; + EXPECT_TRUE(segment->open_index(1, "body", &idx).ok()); + return idx; +} + +void ExpectProfileMatchesMeteredDelta(io::MeteredFileReader* metered, + const std::function& run) { + metered->reset_metrics(); + query::QueryProfile profile; + const Status st = run(&profile); + ASSERT_TRUE(st.ok()) << st.to_string(); + + EXPECT_GT(profile.elapsed_ns, 0U); + ASSERT_TRUE(profile.has_io_metrics); + EXPECT_EQ(profile.io_delta.read_at_calls, metered->metrics().read_at_calls); + EXPECT_EQ(profile.io_delta.serial_rounds, metered->metrics().serial_rounds); + EXPECT_EQ(profile.io_delta.range_gets, metered->metrics().range_gets); + EXPECT_EQ(profile.io_delta.remote_bytes, metered->metrics().remote_bytes); + EXPECT_EQ(profile.io_delta.total_request_bytes, metered->metrics().total_request_bytes); +} + +template +concept CanCallPhraseQuery = requires(Args... args) { query::phrase_query(args...); }; + +template +concept CanCallPhrasePrefixQuery = requires(Args... args) { query::phrase_prefix_query(args...); }; + +} // namespace + +TEST(SniiQueryProfileTest, PhraseProfileInterfacesDoNotExposeTraceSinks) { + using PhraseProfileFunction = + Status (*)(const LogicalIndexReader&, const std::vector&, + std::vector*, query::QueryProfile*); + using PhrasePrefixProfileFunction = + Status (*)(const LogicalIndexReader&, const std::vector&, + std::vector*, query::QueryProfile*, int32_t); + + EXPECT_NE(static_cast(&query::phrase_query), nullptr); + EXPECT_NE(static_cast(&query::phrase_prefix_query), nullptr); + static_assert( + !CanCallPhraseQuery&, + std::vector*, query::QueryProfile*, std::nullptr_t>); + static_assert(!CanCallPhrasePrefixQuery&, std::vector*, + query::QueryProfile*, int32_t, std::nullptr_t>); +} + +TEST(SniiQueryProfileTest, SuccessfulDecodeKeepsStatsWhenCallerShapeValidationFails) { + const std::vector> positions = {{1}, {2, 4}, {3}}; + ByteSink sink; + ASSERT_TRUE(format::build_prx_window(positions, /*zstd_level=*/0, &sink).ok()); + + const std::vector selected_ordinals = {0, 2}; + std::vector flat; + std::vector offsets; + format::PrxDecodeStats stats; + format::PrxDecodeContext context {.stats = &stats}; + ByteSource source(sink.view()); + EXPECT_TRUE(query::internal::decode_and_validate_prx_frame( + &source, selected_ordinals, + /*decode_full=*/false, + /*all_docs_selected=*/false, + /*expected_total_docs=*/static_cast(positions.size() + 1), + /*expected_selected_docs=*/selected_ordinals.size(), &flat, &offsets, + &context) + .is()); + EXPECT_EQ(stats.raw_frames, 1U); + EXPECT_EQ(stats.total_docs, positions.size()); + EXPECT_EQ(stats.selected_docs, selected_ordinals.size()); + EXPECT_EQ(stats.total_positions, 4U); + EXPECT_EQ(stats.selected_positions, 2U); +} + +TEST(SniiQueryProfileTest, PhrasePrxValidationRejectsOnlyTotalDocMismatch) { + const std::vector flat = {1, 2, 3, 4}; + const std::vector selected_ordinals = {0, 2}; + EXPECT_TRUE(query::internal::validate_prx_frame( + flat, std::vector {0, 1, 4}, + /*actual_total_docs=*/3, /*expected_total_docs=*/4, + /*expected_selected_docs=*/selected_ordinals.size(), selected_ordinals, + /*offsets_by_prx_ordinal=*/false, + /*all_docs_selected=*/false) + .is()); +} + +TEST(SniiQueryProfileTest, PhrasePrxValidationRejectsOnlyOffsetCountMismatch) { + const std::vector flat = {1, 2, 3, 4}; + const std::vector selected_ordinals = {0, 2}; + EXPECT_TRUE(query::internal::validate_prx_frame( + flat, std::vector {0, 4}, + /*actual_total_docs=*/4, /*expected_total_docs=*/4, + /*expected_selected_docs=*/selected_ordinals.size(), selected_ordinals, + /*offsets_by_prx_ordinal=*/false, + /*all_docs_selected=*/false) + .is()); +} + +TEST(SniiQueryProfileTest, PhrasePrxValidationRejectsOnlyFinalOffsetMismatch) { + const std::vector flat = {1, 2, 3, 4}; + const std::vector selected_ordinals = {0, 2}; + EXPECT_TRUE(query::internal::validate_prx_frame( + flat, std::vector {0, 1, 3}, + /*actual_total_docs=*/4, /*expected_total_docs=*/4, + /*expected_selected_docs=*/selected_ordinals.size(), selected_ordinals, + /*offsets_by_prx_ordinal=*/false, + /*all_docs_selected=*/false) + .is()); +} + +TEST(SniiQueryProfileTest, PhrasePrxValidationRejectsOnlySelectionCountMismatch) { + const std::vector flat = {1, 2, 3, 4}; + const std::vector selected_ordinals = {0, 2}; + EXPECT_TRUE(query::internal::validate_prx_frame(flat, std::vector {0, 4}, + /*actual_total_docs=*/4, + /*expected_total_docs=*/4, + /*expected_selected_docs=*/1, selected_ordinals, + /*offsets_by_prx_ordinal=*/false, + /*all_docs_selected=*/false) + .is()); +} + +TEST(SniiQueryProfileTest, PhrasePrxValidationRejectsOnlyOrdinalRangeMismatch) { + const std::vector flat = {1, 2, 3, 4}; + const std::vector selected_ordinals = {0, 4}; + EXPECT_TRUE(query::internal::validate_prx_frame( + flat, std::vector {0, 1, 2, 3, 4}, + /*actual_total_docs=*/4, /*expected_total_docs=*/4, + /*expected_selected_docs=*/selected_ordinals.size(), selected_ordinals, + /*offsets_by_prx_ordinal=*/true, + /*all_docs_selected=*/false) + .is()); +} + +TEST(SniiQueryProfileTest, PhraseVerifyTimeExcludesDecodeDelta) { + EXPECT_EQ(query::internal::exclusive_phrase_verify_ns( + /*elapsed_ns=*/100, /*decode_ns_before=*/20, /*decode_ns_after=*/65), + 55U); + EXPECT_EQ(query::internal::exclusive_phrase_verify_ns( + /*elapsed_ns=*/100, /*decode_ns_before=*/65, /*decode_ns_after=*/65), + 100U); + EXPECT_EQ(query::internal::exclusive_phrase_verify_ns( + /*elapsed_ns=*/30, /*decode_ns_before=*/20, /*decode_ns_after=*/65), + 0U); +} + +TEST(SniiQueryProfileTest, UnprofiledSelectivePhraseKeepsInstrumentationDisabled) { + Corpus corpus; + corpus.docs.resize(600); + for (uint32_t docid = 0; docid < corpus.docs.size(); ++docid) { + corpus.docs[docid] = {"common", docid % 4 == 0 ? "rare" : "other"}; + } + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + std::vector docs; + format::testing::reset_prx_clock_read_count(); + query::internal::testing::reset_phrase_verify_clock_read_count(); + DEFER(format::testing::reset_prx_clock_read_count()); + DEFER(query::internal::testing::reset_phrase_verify_clock_read_count()); + ASSERT_TRUE(query::phrase_query(idx, {"common", "rare"}, &docs).ok()); + + EXPECT_EQ(docs.size(), 150U); + EXPECT_EQ(format::testing::prx_clock_read_count(), 0U); + EXPECT_EQ(query::internal::testing::phrase_verify_clock_read_count(), 0U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, PhraseVerifyClockIsDisabledWithoutProfile) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + std::vector docs; + query::internal::testing::reset_phrase_verify_clock_read_count(); + ASSERT_TRUE(query::phrase_query(idx, {"quick", "brown"}, &docs).ok()); + EXPECT_EQ(query::internal::testing::phrase_verify_clock_read_count(), 0U); + + query::QueryProfile profile; + ASSERT_TRUE(query::phrase_query(idx, {"quick", "brown"}, &docs, &profile).ok()); + EXPECT_GT(query::internal::testing::phrase_verify_clock_read_count(), 0U); + EXPECT_GT(profile.prx_decode_stats.decode_ns, 0U); + EXPECT_GT(profile.prx_decode_stats.phrase_verify_ns, 0U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, MultiWindowFiveTermQueriesAggregateEveryPrxFrame) { + Corpus corpus; + corpus.docs.resize(600); + for (auto& doc : corpus.docs) { + doc = {"alpha", "beta", "gamma", "delta", "epsilon"}; + } + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + const auto expect_totals = [](const query::QueryProfile& profile) { + EXPECT_EQ(profile.prx_decode_stats.frame_count(), 15U); + EXPECT_EQ(profile.prx_decode_stats.total_docs, 3000U); + EXPECT_EQ(profile.prx_decode_stats.selected_docs, 3000U); + EXPECT_EQ(profile.prx_decode_stats.total_positions, 3000U); + EXPECT_EQ(profile.prx_decode_stats.selected_positions, 3000U); + }; + + std::vector docs; + ASSERT_TRUE(query::phrase_query(idx, {"alpha", "beta"}, &docs).ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + + query::QueryProfile phrase_profile; + ASSERT_TRUE(query::phrase_query(idx, {"alpha", "beta", "gamma", "delta", "epsilon"}, &docs, + &phrase_profile) + .ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + expect_totals(phrase_profile); + EXPECT_EQ(phrase_profile.prx_decode_stats.zstd_frames, 10U); + + query::QueryProfile prefix_profile; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"alpha", "beta", "gamma", "delta", "eps"}, &docs, + &prefix_profile) + .ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + expect_totals(prefix_profile); + EXPECT_EQ(prefix_profile.prx_decode_stats.zstd_frames, 10U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, HighTfPhraseAndPrefixAggregateRetainedPrxStats) { + const Corpus corpus = BuildHighTfCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + std::vector docs; + ASSERT_TRUE(query::phrase_query(idx, {"alpha", "beta"}, &docs).ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + const std::vector matched_docs = docs; + + ASSERT_TRUE(query::phrase_query(idx, {"alpha", "gamma"}, &docs).ok()); + EXPECT_TRUE(docs.empty()); + + std::vector frequency_matches; + ASSERT_TRUE( + query::phrase_query_with_frequencies(idx, {"alpha", "beta"}, &frequency_matches).ok()); + std::vector frequency_docids; + frequency_docids.reserve(frequency_matches.size()); + for (const auto& match : frequency_matches) { + frequency_docids.push_back(match.docid); + } + EXPECT_EQ(matched_docs, frequency_docids); + + ASSERT_TRUE(query::phrase_query(idx, {"delta", "epsilon_00"}, &docs).ok()); + EXPECT_EQ(docs.size(), 19U); + + format::testing::reset_prx_clock_read_count(); + query::internal::testing::reset_phrase_verify_clock_read_count(); + DEFER(format::testing::reset_prx_clock_read_count()); + DEFER(query::internal::testing::reset_phrase_verify_clock_read_count()); + ASSERT_TRUE(query::phrase_query(idx, {"alpha", "beta", "gamma", "delta"}, &docs).ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + EXPECT_EQ(format::testing::prx_clock_read_count(), 0U); + EXPECT_EQ(query::internal::testing::phrase_verify_clock_read_count(), 0U); + + query::QueryProfile phrase_profile; + ASSERT_TRUE( + query::phrase_query(idx, {"alpha", "beta", "gamma", "delta"}, &docs, &phrase_profile) + .ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + EXPECT_EQ(phrase_profile.prx_decode_stats.zstd_frames, 12U); + EXPECT_EQ(phrase_profile.prx_decode_stats.frame_count(), 12U); + EXPECT_EQ(phrase_profile.prx_decode_stats.total_docs, 2400U); + EXPECT_EQ(phrase_profile.prx_decode_stats.selected_docs, 2400U); + EXPECT_EQ(phrase_profile.prx_decode_stats.total_positions, 115200U); + EXPECT_EQ(phrase_profile.prx_decode_stats.selected_positions, 115200U); + + query::internal::query_test_counters() = query::internal::QueryTestCounters {}; + DEFER(query::internal::query_test_counters() = query::internal::QueryTestCounters {}); + query::QueryProfile prefix_profile; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"alpha", "beta", "gamma", "delta", "eps"}, &docs, + &prefix_profile, /*max_expansions=*/0) + .ok()); + EXPECT_EQ(docs.size(), corpus.docs.size()); + EXPECT_EQ(prefix_profile.prx_decode_stats.zstd_frames, 45U); + EXPECT_EQ(prefix_profile.prx_decode_stats.frame_count(), 45U); + EXPECT_EQ(prefix_profile.prx_decode_stats.total_docs, 3000U); + EXPECT_EQ(prefix_profile.prx_decode_stats.selected_docs, 3000U); + EXPECT_EQ(prefix_profile.prx_decode_stats.total_positions, 144000U); + EXPECT_EQ(prefix_profile.prx_decode_stats.selected_positions, 144000U); + EXPECT_GT(query::internal::query_test_counters().monotonic_position_scans, 0U); + + format::testing::reset_prx_clock_read_count(); + query::internal::testing::reset_phrase_verify_clock_read_count(); + query::QueryProfile single_phrase_profile; + ASSERT_TRUE(query::phrase_query(idx, {"alpha"}, &docs, &single_phrase_profile).ok()); + EXPECT_EQ(single_phrase_profile.prx_decode_stats, format::PrxDecodeStats {}); + query::QueryProfile single_prefix_profile; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"eps"}, &docs, &single_prefix_profile).ok()); + EXPECT_EQ(single_prefix_profile.prx_decode_stats, format::PrxDecodeStats {}); + EXPECT_EQ(format::testing::prx_clock_read_count(), 0U); + EXPECT_EQ(query::internal::testing::phrase_verify_clock_read_count(), 0U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, SingleTailGroupVisitsExpectedDocsOnce) { + const Corpus corpus = BuildSingleTailGroupCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + query::internal::query_test_counters() = query::internal::QueryTestCounters {}; + DEFER(query::internal::query_test_counters() = query::internal::QueryTestCounters {}); + std::vector docs; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"alpha", "beta", "eps"}, &docs, + /*max_expansions=*/0) + .ok()); + + EXPECT_EQ(docs.size(), corpus.docs.size()); + EXPECT_EQ(query::internal::query_test_counters().prefix_expected_doc_visits, + corpus.docs.size()); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, DisjointTailGroupDoesNotVisitExpectedDocs) { + const Corpus corpus = BuildDisjointTailGroupCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + query::internal::query_test_counters() = query::internal::QueryTestCounters {}; + DEFER(query::internal::query_test_counters() = query::internal::QueryTestCounters {}); + std::vector docs; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"alpha", "beta", "eps"}, &docs, + /*max_expansions=*/0) + .ok()); + + EXPECT_TRUE(docs.empty()); + EXPECT_EQ(query::internal::query_test_counters().prefix_expected_doc_visits, 0U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, MultiTermPhraseResumesPairStartsAndEmitsEachDocOnce) { + Corpus corpus; + corpus.docs.resize(64, {"c"}); + corpus.docs[0] = {"a", "b", "x", "a", "b", "c"}; + corpus.docs[1] = {"a", "a", "a", "c"}; + corpus.docs[2] = {"a", "b", "c", "a", "b", "c"}; + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + std::vector docs; + ASSERT_TRUE(query::phrase_query(idx, {"a", "b", "c"}, &docs).ok()); + EXPECT_EQ(docs, (std::vector {0, 2})); + + ASSERT_TRUE(query::phrase_query(idx, {"a", "a", "a"}, &docs).ok()); + EXPECT_EQ(docs, (std::vector {1})); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, HalfDensePhraseUsesFullDecodeAndReportsOriginalSelection) { + Corpus corpus; + corpus.docs.resize(600); + for (uint32_t docid = 0; docid < corpus.docs.size(); ++docid) { + corpus.docs[docid] = {"common", docid % 2 == 0 ? "rare" : "other"}; + } + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader segment; + LogicalIndexReader idx; + ASSERT_TRUE(SniiSegmentReader::open(&local, &segment).ok()); + ASSERT_TRUE(segment.open_index(1, "body", &idx).ok()); + + query::QueryProfile profile; + std::vector docs; + ASSERT_TRUE(query::phrase_query(idx, {"common", "rare"}, &docs, &profile).ok()); + + ASSERT_EQ(docs.size(), 300U); + for (size_t i = 0; i < docs.size(); ++i) { + EXPECT_EQ(docs[i], i * 2); + } + EXPECT_EQ(profile.prx_decode_stats.frame_count(), 4U); + EXPECT_EQ(profile.prx_decode_stats.total_docs, 900U); + EXPECT_EQ(profile.prx_decode_stats.selected_docs, 600U); + EXPECT_EQ(profile.prx_decode_stats.total_positions, 900U); + EXPECT_EQ(profile.prx_decode_stats.selected_positions, 600U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfile, ReportsElapsedTimeAndMeteredIoForNativeOperators) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/512); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + std::vector docs; + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::term_query(idx, "lead", &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::boolean_or(idx, {"lead", "missing"}, &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::boolean_and(idx, {"quick", "brown"}, &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::prefix_query(idx, "aa_", &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::wildcard_query(idx, "aa_0??", &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::regexp_query(idx, "aa_00[0-9]", &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::phrase_query(idx, {"quick", "brown"}, &docs, profile); + }); + ExpectProfileMatchesMeteredDelta(&metered, [&](query::QueryProfile* profile) { + return query::phrase_prefix_query(idx, {"quick", "bro"}, &docs, profile); + }); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfile, ReportsElapsedTimeForOperatorErrorPath) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/512); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + metered.reset_metrics(); + std::vector docs; + query::QueryProfile profile; + const Status st = query::regexp_query(idx, "(", &docs, &profile); + + EXPECT_FALSE(st.ok()); + EXPECT_GT(profile.elapsed_ns, 0U); + ASSERT_TRUE(profile.has_io_metrics); + EXPECT_EQ(profile.io_delta.read_at_calls, 0U); + EXPECT_EQ(profile.io_delta.serial_rounds, 0U); + EXPECT_EQ(profile.io_delta.range_gets, 0U); + EXPECT_EQ(profile.io_delta.remote_bytes, 0U); + EXPECT_EQ(profile.io_delta.total_request_bytes, 0U); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfile, ScopeFinalizesProfileOnEarlyReturn) { + query::QueryProfile profile; + { query::QueryProfileScope scope(/*reader=*/nullptr, &profile); } + + EXPECT_GT(profile.elapsed_ns, 0U); + EXPECT_FALSE(profile.has_io_metrics); +} + +TEST(SniiQueryProfileTest, AggregatesMultiTermPrxAndKeepsSingleTermControlsAtZero) { + const Corpus corpus = BuildCorpus(); + const std::string path = TempPath(); + WriteCorpus(corpus, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local, /*block_size=*/512); + SniiSegmentReader segment; + LogicalIndexReader idx = OpenMeteredIndex(&metered, &segment); + + std::vector docs; + query::QueryProfile phrase_profile; + ASSERT_TRUE(query::phrase_query(idx, {"quick", "brown"}, &docs, &phrase_profile).ok()); + EXPECT_GT(phrase_profile.prx_decode_stats.frame_count(), 0U); + EXPECT_GT(phrase_profile.prx_decode_stats.total_docs, 0U); + EXPECT_LE(phrase_profile.prx_decode_stats.selected_docs, + phrase_profile.prx_decode_stats.total_docs); + EXPECT_GT(phrase_profile.prx_decode_stats.phrase_verify_ns, 0U); + + query::QueryProfile prefix_profile; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"quick", "bro"}, &docs, &prefix_profile).ok()); + EXPECT_GT(prefix_profile.prx_decode_stats.frame_count(), 0U); + EXPECT_TRUE(prefix_profile.prx_decode_stats.is_valid()); + + query::QueryProfile one_phrase_profile; + ASSERT_TRUE(query::phrase_query(idx, {"quick"}, &docs, &one_phrase_profile).ok()); + EXPECT_EQ(one_phrase_profile.prx_decode_stats, doris::snii::format::PrxDecodeStats {}); + + query::QueryProfile one_prefix_profile; + ASSERT_TRUE(query::phrase_prefix_query(idx, {"bro"}, &docs, &one_prefix_profile).ok()); + EXPECT_EQ(one_prefix_profile.prx_decode_stats, doris::snii::format::PrxDecodeStats {}); + + std::remove(path.c_str()); +} + +TEST(SniiQueryProfileTest, ReaderAndRuntimePrxTotalsAreAdditiveWithStableNames) { + doris::snii::format::PrxDecodeStats delta; + delta.raw_frames = 1; + delta.zstd_frames = 2; + delta.pfor_frames = 3; + delta.plaintext_bytes = 101; + delta.total_docs = 8; + delta.selected_docs = 7; + delta.total_positions = 9; + delta.selected_positions = 6; + delta.fetch_ns = 10; + delta.decode_ns = 11; + delta.phrase_verify_ns = 12; + + doris::OlapReaderStatistics reader_stats; + doris::snii::add_prx_decode_stats(&reader_stats, delta); + doris::snii::add_prx_decode_stats(&reader_stats, delta); + EXPECT_EQ(reader_stats.snii_stats.prx_raw_frames, 2); + EXPECT_EQ(reader_stats.snii_stats.prx_zstd_frames, 4); + EXPECT_EQ(reader_stats.snii_stats.prx_pfor_frames, 6); + EXPECT_EQ(reader_stats.snii_stats.prx_plaintext_bytes, 202); + EXPECT_EQ(reader_stats.snii_stats.prx_total_docs, 16); + EXPECT_EQ(reader_stats.snii_stats.prx_selected_docs, 14); + EXPECT_EQ(reader_stats.snii_stats.prx_total_positions, 18); + EXPECT_EQ(reader_stats.snii_stats.prx_selected_positions, 12); + EXPECT_EQ(reader_stats.snii_stats.prx_fetch_ns, 20); + EXPECT_EQ(reader_stats.snii_stats.prx_decode_ns, 22); + EXPECT_EQ(reader_stats.snii_stats.prx_phrase_verify_ns, 24); + + doris::RuntimeProfile runtime_profile("IndexFilter"); + doris::snii::SniiPrxRuntimeProfileCounters counters; + counters.initialize(&runtime_profile); + + std::vector zero_nodes; + runtime_profile.to_thrift(&zero_nodes); + ASSERT_EQ(zero_nodes.size(), 1U); + for (const char* name : doris::snii::SniiPrxRuntimeProfileCounters::counter_names()) { + EXPECT_TRUE(std::ranges::none_of( + zero_nodes.front().counters, + [name](const doris::TCounter& counter) { return counter.name == name; })) + << name; + } + + counters.update(reader_stats); + counters.update(reader_stats); + + EXPECT_EQ(doris::snii::SniiPrxRuntimeProfileCounters::counter_names().size(), 11U); + for (const char* name : doris::snii::SniiPrxRuntimeProfileCounters::counter_names()) { + auto* counter = runtime_profile.get_counter(name); + ASSERT_NE(counter, nullptr) << name; + EXPECT_NE(dynamic_cast(counter), nullptr) << name; + } + EXPECT_EQ(runtime_profile.get_counter("SniiPrxEncodedBytes"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxPayloadBytes"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxCompressedBytes"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxChild128Touches"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxChild256Touches"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxNestedCrcValidationTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxNestedDecompressTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxScratchAllocationEvents"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxContainerAllocationEvents"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxCountScanTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxSkipTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxTraceAllocationEvents"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxTraceRecords"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxDecodeTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxCrcValidationTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxDecompressTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxPhraseVerifyTime"), nullptr); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxRawFrames")->type(), doris::TUnit::UNIT); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxPlaintextBytes")->type(), doris::TUnit::BYTES); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxInclusiveDecodeTime")->type(), + doris::TUnit::TIME_NS); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxExclusivePhraseVerifyTime")->type(), + doris::TUnit::TIME_NS); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxRawFrames")->value(), 4); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxZstdFrames")->value(), 8); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxPforFrames")->value(), 12); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxPlaintextBytes")->value(), 404); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxTotalDocs")->value(), 32); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxSelectedDocs")->value(), 28); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxTotalPositions")->value(), 36); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxSelectedPositions")->value(), 24); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxFetchTime")->value(), 40); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxInclusiveDecodeTime")->value(), 44); + EXPECT_EQ(runtime_profile.get_counter("SniiPrxExclusivePhraseVerifyTime")->value(), 48); + + std::vector nonzero_nodes; + runtime_profile.to_thrift(&nonzero_nodes); + ASSERT_EQ(nonzero_nodes.size(), 1U); + for (const char* name : doris::snii::SniiPrxRuntimeProfileCounters::counter_names()) { + EXPECT_TRUE(std::ranges::any_of( + nonzero_nodes.front().counters, + [name](const doris::TCounter& counter) { return counter.name == name; })) + << name; + } +} + +TEST(SniiQueryProfileTest, PhraseReaderAndRuntimeTotalsAreAdditiveAndNonZero) { + doris::snii::format::PhraseQueryExecutionStats delta; + delta.exact_candidate_docs = 1; + delta.exact_candidate_visits = 2; + delta.prefix_leading_candidate_docs = 3; + delta.prefix_tail_candidate_visits = 4; + delta.common_grams_candidate_queries = 5; + delta.common_grams_plain_plans = 6; + delta.common_grams_gram_plans = 7; + delta.common_grams_fallback_no_gram = 8; + delta.common_grams_fallback_incompatible = 9; + delta.common_grams_fallback_kill_switch = 10; + delta.common_grams_fallback_cost = 11; + delta.common_grams_authoritative_empty = 12; + delta.common_grams_plain_posting_bytes = 13; + delta.common_grams_gram_posting_bytes = 14; + delta.common_grams_plain_estimated_candidate_df = 15; + delta.common_grams_gram_estimated_candidate_df = 16; + delta.common_grams_plain_estimated_cost = 17; + delta.common_grams_gram_estimated_cost = 18; + delta.common_grams_fallback_base_analyzer_mismatch = 19; + delta.common_grams_fallback_prefix_tail_empty = 20; + delta.common_grams_planning_ns = 23; + + doris::OlapReaderStatistics reader_stats; + doris::snii::add_phrase_query_stats(&reader_stats, delta); + doris::snii::add_phrase_query_stats(&reader_stats, delta); + EXPECT_EQ(reader_stats.snii_stats.phrase_candidate_docs, 2); + EXPECT_EQ(reader_stats.snii_stats.phrase_candidate_visits, 4); + EXPECT_EQ(reader_stats.snii_stats.phrase_prefix_leading_candidate_docs, 6); + EXPECT_EQ(reader_stats.snii_stats.phrase_prefix_tail_candidate_visits, 8); + EXPECT_EQ(reader_stats.snii_stats.common_grams_candidate_queries, 10); + EXPECT_EQ(reader_stats.snii_stats.common_grams_plain_plans, 12); + EXPECT_EQ(reader_stats.snii_stats.common_grams_gram_plans, 14); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_no_gram, 16); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_incompatible, 18); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_kill_switch, 20); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_cost, 22); + EXPECT_EQ(reader_stats.snii_stats.common_grams_authoritative_empty, 24); + EXPECT_EQ(reader_stats.snii_stats.common_grams_plain_posting_bytes, 26); + EXPECT_EQ(reader_stats.snii_stats.common_grams_gram_posting_bytes, 28); + EXPECT_EQ(reader_stats.snii_stats.common_grams_plain_estimated_candidate_df, 30); + EXPECT_EQ(reader_stats.snii_stats.common_grams_gram_estimated_candidate_df, 32); + EXPECT_EQ(reader_stats.snii_stats.common_grams_plain_estimated_cost, 34); + EXPECT_EQ(reader_stats.snii_stats.common_grams_gram_estimated_cost, 36); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_base_analyzer_mismatch, 38); + EXPECT_EQ(reader_stats.snii_stats.common_grams_fallback_prefix_tail_empty, 40); + EXPECT_EQ(reader_stats.snii_stats.common_grams_planning_ns, 46); + + doris::RuntimeProfile runtime_profile("IndexFilter"); + doris::snii::SniiPhraseRuntimeProfileCounters counters; + counters.initialize(&runtime_profile); + + struct ExpectedCounter { + const char* name; + int64_t value_after_two_updates; + doris::TUnit::type unit; + }; + const ExpectedCounter expected_counters[] = { + {"SniiPhraseCandidateDocs", 4, doris::TUnit::UNIT}, + {"SniiPhraseCandidateVisits", 8, doris::TUnit::UNIT}, + {"SniiPhrasePrefixLeadingCandidateDocs", 12, doris::TUnit::UNIT}, + {"SniiPhrasePrefixTailCandidateVisits", 16, doris::TUnit::UNIT}, + {"SniiCommonGramsCandidateQueries", 20, doris::TUnit::UNIT}, + {"SniiCommonGramsPlainPlans", 24, doris::TUnit::UNIT}, + {"SniiCommonGramsGramPlans", 28, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackNoGram", 32, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackIncompatible", 36, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackKillSwitch", 40, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackCost", 44, doris::TUnit::UNIT}, + {"SniiCommonGramsAuthoritativeEmpty", 48, doris::TUnit::UNIT}, + {"SniiCommonGramsPlainPostingBytes", 52, doris::TUnit::BYTES}, + {"SniiCommonGramsGramPostingBytes", 56, doris::TUnit::BYTES}, + {"SniiCommonGramsPlainEstimatedCandidateDf", 60, doris::TUnit::UNIT}, + {"SniiCommonGramsGramEstimatedCandidateDf", 64, doris::TUnit::UNIT}, + {"SniiCommonGramsPlainEstimatedCost", 68, doris::TUnit::UNIT}, + {"SniiCommonGramsGramEstimatedCost", 72, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackBaseAnalyzerMismatch", 76, doris::TUnit::UNIT}, + {"SniiCommonGramsFallbackPrefixTailEmpty", 80, doris::TUnit::UNIT}, + {"SniiCommonGramsPlanningTime", 92, doris::TUnit::TIME_NS}, + }; + + std::vector zero_nodes; + runtime_profile.to_thrift(&zero_nodes); + ASSERT_EQ(zero_nodes.size(), 1U); + for (const auto& expected : expected_counters) { + EXPECT_TRUE(std::ranges::none_of( + zero_nodes.front().counters, + [&](const doris::TCounter& counter) { return counter.name == expected.name; })) + << expected.name; + } + + counters.update(reader_stats); + counters.update(reader_stats); + + EXPECT_EQ(doris::snii::SniiPrxRuntimeProfileCounters::counter_names().size(), 11U); + for (const auto& expected : expected_counters) { + auto* counter = runtime_profile.get_counter(expected.name); + ASSERT_NE(counter, nullptr) << expected.name; + EXPECT_NE(dynamic_cast(counter), nullptr) + << expected.name; + EXPECT_EQ(counter->type(), expected.unit) << expected.name; + EXPECT_EQ(counter->value(), expected.value_after_two_updates) << expected.name; + } + + std::vector nonzero_nodes; + runtime_profile.to_thrift(&nonzero_nodes); + ASSERT_EQ(nonzero_nodes.size(), 1U); + for (const auto& expected : expected_counters) { + EXPECT_TRUE(std::ranges::any_of( + nonzero_nodes.front().counters, + [&](const doris::TCounter& counter) { return counter.name == expected.name; })) + << expected.name; + } +} + +TEST(SniiQueryProfileTest, ExecutionProfileScopeFlushesNormalAndErrorReturnsAdditively) { + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_move_constructible_v); + + doris::OlapReaderStatistics reader_stats; + const auto execute = [&](bool fail) -> Status { + doris::snii::SniiPrxExecutionProfileScope scope(reader_stats); + scope.profile()->prx_decode_stats.raw_frames = 1; + scope.profile()->prx_decode_stats.plaintext_bytes = 17; + if (fail) { + return Status::InternalError("injected execution failure"); + } + return Status::OK(); + }; + + ASSERT_TRUE(execute(false).ok()); + EXPECT_EQ(reader_stats.snii_stats.prx_raw_frames, 1); + EXPECT_EQ(reader_stats.snii_stats.prx_plaintext_bytes, 17); + ASSERT_FALSE(execute(true).ok()); + EXPECT_EQ(reader_stats.snii_stats.prx_raw_frames, 2); + EXPECT_EQ(reader_stats.snii_stats.prx_plaintext_bytes, 34); +} diff --git a/be/test/storage/index/snii/query/query_term_resolution_batch_test.cpp b/be/test/storage/index/snii/query/query_term_resolution_batch_test.cpp new file mode 100644 index 00000000000000..1ae14d9329a46b --- /dev/null +++ b/be/test/storage/index/snii/query/query_term_resolution_batch_test.cpp @@ -0,0 +1,304 @@ +// 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 + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/internal/docid_conjunction.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii::query { +namespace { + +using snii_test::MemoryFile; +using snii_test::ScopedEnv; +using snii_test::assert_ok; +using snii_test::make_term; + +constexpr uint64_t kIndexId = 1; +constexpr const char* kIndexSuffix = "body"; + +class CountingReader final : public io::FileReader { +public: + explicit CountingReader(io::FileReader* inner) : inner_(inner) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + ++read_at_calls_; + return inner_->read_at(offset, len, out); + } + + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + ++read_batch_calls_; + batch_range_counts_.push_back(ranges.size()); + return inner_->read_batch(ranges, outs); + } + + uint64_t size() const override { return inner_->size(); } + const io::IoMetrics* io_metrics() const override { return inner_->io_metrics(); } + + void reset_counts() { + read_at_calls_ = 0; + read_batch_calls_ = 0; + batch_range_counts_.clear(); + } + + uint64_t read_at_calls() const { return read_at_calls_; } + uint64_t read_batch_calls() const { return read_batch_calls_; } + const std::vector& batch_range_counts() const { return batch_range_counts_; } + +private: + io::FileReader* inner_; + uint64_t read_at_calls_ = 0; + uint64_t read_batch_calls_ = 0; + std::vector batch_range_counts_; +}; + +Status write_index(MemoryFile* file, const std::vector& terms, + uint32_t target_dict_block_bytes) { + writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = kIndexSuffix; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = static_cast(terms.size()); + input.target_dict_block_bytes = target_dict_block_bytes; + input.terms.reserve(terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + input.terms.push_back( + make_term(terms[i], {{.docid = static_cast(i), .positions = {0}}})); + } + + writer::SniiCompoundWriter compound_writer(file); + RETURN_IF_ERROR(compound_writer.add_logical_index(input)); + return compound_writer.finish(); +} + +std::vector numbered_terms(size_t count) { + std::vector terms; + terms.reserve(count); + for (size_t i = 0; i < count; ++i) { + std::string term = "term_"; + if (i < 10) { + term.push_back('0'); + } + term += std::to_string(i); + terms.push_back(std::move(term)); + } + return terms; +} + +TEST(SniiQueryTermResolutionBatch, ResolvesColdDictBlocksInOnePhysicalBatch) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + writer::SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = 6; + input.target_dict_block_bytes = 1; + input.terms = { + make_term("alpha", {{.docid = 1, .positions = {0}}}), + make_term("bravo", {{.docid = 2, .positions = {1}}}), + make_term("kappa", {{.docid = 3, .positions = {2}}}), + make_term("lambda", {{.docid = 4, .positions = {3}}}), + make_term("omega", {{.docid = 5, .positions = {4}}}), + }; + + writer::SniiCompoundWriter compound_writer(&file); + assert_ok(compound_writer.add_logical_index(input)); + assert_ok(compound_writer.finish()); + + io::MeteredFileReader metered(&file, /*block_size=*/1); + CountingReader counting(&metered); + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&counting, &segment_reader)); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &index_reader)); + ASSERT_EQ(index_reader.n_dict_blocks(), input.terms.size()); + metered.reset_metrics(); + counting.reset_counts(); + const std::vector terms {"alpha", "kappa", "omega"}; + std::vector resolved; + std::vector found; + assert_ok(internal::resolve_query_terms_batch(index_reader, terms, &resolved, &found)); + + ASSERT_EQ(resolved.size(), terms.size()); + ASSERT_EQ(found, (std::vector {1, 1, 1})); + for (size_t i = 0; i < terms.size(); ++i) { + EXPECT_EQ(resolved[i].entry.term, terms[i]); + EXPECT_EQ(resolved[i].entry.df, 1U); + } + EXPECT_EQ(counting.read_at_calls(), 0U); + EXPECT_EQ(counting.read_batch_calls(), 1U); + EXPECT_EQ(counting.batch_range_counts(), (std::vector {3})); + EXPECT_EQ(metered.metrics().serial_rounds, 1U) + << "independent cold DICT blocks must be fetched in one physical batch"; +} + +TEST(SniiQueryTermResolutionBatch, AlignsAbsentTermsAndReadsOneColdBlockSynchronously) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + ScopedEnv bsbf_resident_max("SNII_BSBF_RESIDENT_MAX", "0"); + + MemoryFile file; + assert_ok(write_index(&file, {"alpha", "kappa", "omega"}, + /*target_dict_block_bytes=*/4096)); + + io::MeteredFileReader metered(&file, /*block_size=*/1); + CountingReader counting(&metered); + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&counting, &segment_reader)); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(kIndexId, kIndexSuffix, &index_reader)); + ASSERT_EQ(index_reader.n_dict_blocks(), 1U); + metered.reset_metrics(); + counting.reset_counts(); + const std::vector terms {"aardvark", "alpha", "beta", "kappa", + "lambda", "omega", "zulu"}; + std::vector resolved; + std::vector found; + assert_ok(internal::resolve_query_terms_batch(index_reader, terms, &resolved, &found)); + + ASSERT_EQ(resolved.size(), terms.size()); + ASSERT_EQ(found, (std::vector {0, 1, 0, 1, 0, 1, 0})); + for (size_t i : {1U, 3U, 5U}) { + EXPECT_EQ(resolved[i].entry.term, terms[i]); + EXPECT_EQ(resolved[i].entry.df, 1U); + } + EXPECT_EQ(counting.read_at_calls(), 1U); + EXPECT_EQ(counting.read_batch_calls(), 0U); + EXPECT_EQ(metered.metrics().serial_rounds, 1U); +} + +TEST(SniiQueryTermResolutionBatch, ResolvesResidentBlocksWithoutQueryIo) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "1048576"); + + MemoryFile file; + const std::vector terms {"alpha", "kappa", "omega"}; + assert_ok(write_index(&file, terms, /*target_dict_block_bytes=*/1)); + + io::MeteredFileReader metered(&file, /*block_size=*/1); + CountingReader counting(&metered); + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&counting, &segment_reader)); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(kIndexId, kIndexSuffix, &index_reader)); + ASSERT_EQ(index_reader.n_dict_blocks(), terms.size()); + + metered.reset_metrics(); + counting.reset_counts(); + std::vector resolved; + std::vector found; + assert_ok(internal::resolve_query_terms_batch(index_reader, terms, &resolved, &found)); + + ASSERT_EQ(found, (std::vector {1, 1, 1})); + ASSERT_EQ(resolved.size(), terms.size()); + for (size_t i = 0; i < terms.size(); ++i) { + EXPECT_EQ(resolved[i].entry.term, terms[i]); + } + EXPECT_EQ(counting.read_at_calls(), 0U); + EXPECT_EQ(counting.read_batch_calls(), 0U); + EXPECT_EQ(metered.metrics().serial_rounds, 0U); +} + +TEST(SniiQueryTermResolutionBatch, ResolvesCompressedDictBlocksFromBatchBuffers) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + const std::vector terms { + "a" + std::string(4096, 'x'), + "b" + std::string(4096, 'y'), + "c" + std::string(4096, 'z'), + }; + assert_ok(write_index(&file, terms, /*target_dict_block_bytes=*/1)); + + io::MeteredFileReader metered(&file, /*block_size=*/1); + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&metered, &segment_reader)); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(kIndexId, kIndexSuffix, &index_reader)); + ASSERT_EQ(index_reader.n_dict_blocks(), terms.size()); + ASSERT_LT(index_reader.section_refs().dict_region.length, terms.size() * terms.front().size()); + + metered.reset_metrics(); + std::vector resolved; + std::vector found; + assert_ok(internal::resolve_query_terms_batch(index_reader, terms, &resolved, &found)); + + ASSERT_EQ(found, (std::vector {1, 1, 1})); + for (size_t i = 0; i < terms.size(); ++i) { + EXPECT_EQ(resolved[i].entry.term, terms[i]); + } + EXPECT_EQ(metered.metrics().serial_rounds, 1U); +} + +// GTest assertion macros inflate clang-tidy's branch count for this table-style I/O check. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiQueryTermResolutionBatch, ResolvesSeventeenDisjointBlocksInTwoBoundedWaves) { + ScopedEnv dict_resident_max("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + const std::vector indexed_terms = numbered_terms(33); + assert_ok(write_index(&file, indexed_terms, /*target_dict_block_bytes=*/1)); + + io::MeteredFileReader metered(&file, /*block_size=*/1); + CountingReader counting(&metered); + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&counting, &segment_reader)); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(kIndexId, kIndexSuffix, &index_reader)); + ASSERT_EQ(index_reader.n_dict_blocks(), indexed_terms.size()); + + std::vector query_terms; + query_terms.reserve(17); + for (size_t i = 0; i < indexed_terms.size(); i += 2) { + query_terms.push_back(indexed_terms[i]); + } + ASSERT_EQ(query_terms.size(), 17U); + + metered.reset_metrics(); + counting.reset_counts(); + std::vector resolved; + std::vector found; + assert_ok(internal::resolve_query_terms_batch(index_reader, query_terms, &resolved, &found)); + + ASSERT_EQ(found, std::vector(query_terms.size(), 1)); + ASSERT_EQ(resolved.size(), query_terms.size()); + for (size_t i = 0; i < query_terms.size(); ++i) { + EXPECT_EQ(resolved[i].entry.term, query_terms[i]); + EXPECT_EQ(resolved[i].entry.df, 1U); + } + EXPECT_EQ(resolved.back().entry.term, "term_32"); + EXPECT_EQ(counting.read_at_calls(), 0U); + EXPECT_EQ(counting.read_batch_calls(), 2U); + EXPECT_EQ(counting.batch_range_counts(), (std::vector {16, 1})); + EXPECT_EQ(metered.metrics().serial_rounds, 2U); + EXPECT_EQ(metered.metrics().range_gets, query_terms.size()); +} + +} // namespace +} // namespace doris::snii::query diff --git a/be/test/storage/index/snii/query/scoring_query_test.cpp b/be/test/storage/index/snii/query/scoring_query_test.cpp new file mode 100644 index 00000000000000..d454f279de04d1 --- /dev/null +++ b/be/test/storage/index/snii/query/scoring_query_test.cpp @@ -0,0 +1,756 @@ +// 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 "storage/index/snii/query/scoring_query.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "runtime/thread_context.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::writer; +using doris::snii::query::Bm25Params; +using doris::snii::query::ScoredDoc; +using doris::snii::stats::SniiStatsProvider; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_score_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +class ControllableFileReader final : public io::FileReader { +public: + explicit ControllableFileReader(io::FileReader* inner) : inner_(inner) {} + + doris::Status read_at(uint64_t offset, size_t len, std::vector* out) override { + read_at_calls_.fetch_add(1, std::memory_order_relaxed); + if (fail_next_read_.exchange(false, std::memory_order_acq_rel)) { + return doris::Status::IOError("injected norms read failure"); + } + + { + std::unique_lock lock(mutex_); + if (block_next_read_) { + block_next_read_ = false; + blocked_read_started_ = true; + condition_.notify_all(); + condition_.wait(lock, [this] { return blocked_read_released_; }); + } + } + return inner_->read_at(offset, len, out); + } + + uint64_t size() const override { return inner_->size(); } + + void reset_read_at_calls() { read_at_calls_.store(0, std::memory_order_relaxed); } + uint64_t read_at_calls() const { return read_at_calls_.load(std::memory_order_relaxed); } + void fail_next_read() { fail_next_read_.store(true, std::memory_order_release); } + + void block_next_read() { + std::lock_guard lock(mutex_); + block_next_read_ = true; + blocked_read_started_ = false; + blocked_read_released_ = false; + } + + void wait_for_blocked_read() { + std::unique_lock lock(mutex_); + condition_.wait(lock, [this] { return blocked_read_started_; }); + } + + void release_blocked_read() { + std::lock_guard lock(mutex_); + blocked_read_released_ = true; + condition_.notify_all(); + } + +private: + io::FileReader* inner_ = nullptr; + std::atomic read_at_calls_ = 0; + std::atomic fail_next_read_ = false; + std::mutex mutex_; + std::condition_variable condition_; + bool block_next_read_ = false; + bool blocked_read_started_ = false; + bool blocked_read_released_ = false; +}; + +// A small in-memory corpus: each doc is a bag of (term -> freq). Doc lengths vary +// so length normalization matters. "common" is a high-df term (~half the docs), +// "rare" is a low-df term. +struct Corpus { + uint32_t doc_count = 0; + // term -> (docid -> freq), docids ascending. + std::map> postings; + std::vector doc_len; // per-doc total token count +}; + +// Builds ~60 docs with varied lengths and a high-df + low-df term. +Corpus MakeCorpus() { + Corpus c; + c.doc_count = 60; + c.doc_len.assign(c.doc_count, 0); + + auto add = [&](const std::string& term, uint32_t doc, uint32_t freq) { + c.postings[term][doc] += freq; + c.doc_len[doc] += freq; + }; + + for (uint32_t d = 0; d < c.doc_count; ++d) { + // "common": appears in even docs, freq varies 1..4. + if (d % 2 == 0) { + add("common", d, 1 + (d % 4)); + } + // "rare": appears in only a few docs. + if (d == 3 || d == 17 || d == 42) { + add("rare", d, 2); + } + // "filler": gives docs varied lengths so dl differs widely. + add("filler", d, 1 + (d % 7) * 3); + // a unique padding token per doc to spread lengths further. + add("pad" + std::to_string(d % 11), d, (d % 5) + 1); + } + return c; +} + +// Converts the corpus into a sorted SniiIndexInput with encoded norms. +SniiIndexInput ToInput(const Corpus& c) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositionsScoring; + in.doc_count = c.doc_count; + in.target_dict_block_bytes = 1; // one block per term + + in.encoded_norms.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + in.encoded_norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + + for (const auto& [term, plist] : c.postings) { + TermPostings tp; + tp.term = term; + for (const auto& [docid, freq] : plist) { + tp.docids.push_back(docid); + tp.freqs.push_back(freq); + for (uint32_t k = 0; k < freq; ++k) { + tp.positions_flat.push_back(k); // flat + } + } + in.terms.push_back(std::move(tp)); + } + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + return in; +} + +// Reference BM25 ranking computed directly from the corpus (same encode/decode). +std::vector ReferenceRanking(const Corpus& c, const std::vector& norms, + const std::vector& terms, uint32_t k, + const Bm25Params& params) { + uint64_t sum_ttf = 0; + for (const auto& dl : c.doc_len) { + sum_ttf += dl; + } + const double avgdl = static_cast(sum_ttf) / std::max(1, c.doc_count); + + std::unordered_map scores; + for (const auto& term : terms) { + auto it = c.postings.find(term); + if (it == c.postings.end()) { + continue; + } + const uint64_t df = it->second.size(); + const double idf = + std::log(1.0 + (static_cast(c.doc_count) - df + 0.5) / (df + 0.5)); + for (const auto& [docid, freq] : it->second) { + const double dl = doris::snii::query::decode_norm(norms[docid]); + const double denom = freq + params.k1 * (1.0 - params.b + params.b * dl / avgdl); + scores[docid] += idf * (freq * (params.k1 + 1.0)) / denom; + } + } + + std::vector all; + all.reserve(scores.size()); + for (const auto& [docid, s] : scores) { + all.push_back({docid, s}); + } + std::ranges::sort(all, [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) { + return a.score > b.score; + } + return a.docid < b.docid; + }); + if (all.size() > k) { + all.resize(k); + } + return all; +} + +} // namespace + +// Helper that mirrors ToInput's encoding so the reference path can decode norms. +namespace { +std::vector EncodeNorms(const Corpus& c) { + std::vector v(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + v[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + return v; +} +} // namespace + +// Fixture-free test: build, open, and compare. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiScoringQuery, ReferenceOracleAndWandEqualsExhaustive) { + const Corpus corpus = MakeCorpus(); + const std::vector norms = EncodeNorms(corpus); + const std::string path = TempPath(); + + // --- build the scoring index --- + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + + // --- open via SniiSegmentReader over a MeteredFileReader --- + io::LocalFileReader inner; + ASSERT_TRUE(inner.open(path).ok()); + io::MeteredFileReader metered(&inner); + reader::SniiSegmentReader seg; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered, &seg).ok()); + reader::LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + SniiStatsProvider stats; + ASSERT_TRUE(SniiStatsProvider::open(&idx, &stats).ok()); + + // (c) SniiStatsProvider df / ttf / avgdl / encoded_norm match brute force. + uint64_t sum_ttf = 0; + for (const auto& dl : corpus.doc_len) { + sum_ttf += dl; + } + EXPECT_EQ(stats.indexed_doc_count(), corpus.doc_count); + EXPECT_EQ(stats.sum_total_term_freq(), sum_ttf); + EXPECT_NEAR(stats.avgdl(), static_cast(sum_ttf) / corpus.doc_count, 1e-9); + + for (const auto& [term, plist] : corpus.postings) { + uint64_t df = 0, ttf = 0; + ASSERT_TRUE(stats.doc_freq(term, &df).ok()); + ASSERT_TRUE(stats.total_term_freq(term, &ttf).ok()); + uint64_t exp_ttf = 0; + for (const auto& [d, f] : plist) { + exp_ttf += f; + } + EXPECT_EQ(df, plist.size()) << term; + EXPECT_EQ(ttf, exp_ttf) << term; + } + for (uint32_t d = 0; d < corpus.doc_count; ++d) { + uint8_t got = 0; + ASSERT_TRUE(stats.encoded_norm(d, &got).ok()); + EXPECT_EQ(got, norms[d]) << "docid " << d; + } + + // (a) single-term scoring_query top-K matches the reference. + const Bm25Params params; // defaults k1=1.2, b=0.75 + const uint32_t k = 10; + + auto run_and_check = [&](const std::vector& terms) { + std::vector reference = ReferenceRanking(corpus, norms, terms, k, params); + std::vector exhaustive; + ASSERT_TRUE(doris::snii::query::scoring_query_exhaustive(idx, stats, terms, k, params, + &exhaustive) + .ok()); + std::vector wand; + ASSERT_TRUE( + doris::snii::query::scoring_query_wand(idx, stats, terms, k, params, &wand).ok()); + + ASSERT_EQ(exhaustive.size(), reference.size()); + for (size_t i = 0; i < reference.size(); ++i) { + EXPECT_EQ(exhaustive[i].docid, reference[i].docid) << "rank " << i; + EXPECT_NEAR(exhaustive[i].score, reference[i].score, 1e-6) << "rank " << i; + } + // (b) WAND-pruned top-K equals the exhaustive top-K. + ASSERT_EQ(wand.size(), exhaustive.size()); + for (size_t i = 0; i < wand.size(); ++i) { + EXPECT_EQ(wand[i].docid, exhaustive[i].docid) << "wand rank " << i; + EXPECT_NEAR(wand[i].score, exhaustive[i].score, 1e-6) << "wand rank " << i; + } + }; + + run_and_check({"common"}); + run_and_check({"rare"}); + run_and_check({"common", "rare"}); + run_and_check({"common", "rare", "filler"}); + + std::remove(path.c_str()); +} + +TEST(SniiScoringQuery, StatsProviderSharesValidatedNormsAcrossQueries) { + const Corpus corpus = MakeCorpus(); + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound_writer(&writer); + ASSERT_TRUE(compound_writer.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + } + + io::LocalFileReader file_reader; + ASSERT_TRUE(file_reader.open(path).ok()); + io::MeteredFileReader metered_reader(&file_reader, /*block_size=*/64); + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered_reader, &segment_reader).ok()); + reader::LogicalIndexReader logical_reader; + ASSERT_TRUE(segment_reader.open_index(1, "body", &logical_reader).ok()); + + metered_reader.reset_metrics(); + const size_t memory_usage_before_load = logical_reader.memory_usage(); + SniiStatsProvider first; + ASSERT_TRUE(SniiStatsProvider::open(&logical_reader, &first).ok()); + const io::IoMetrics after_first = metered_reader.metrics(); + EXPECT_GT(after_first.total_request_bytes, 0U); + + SniiStatsProvider second; + ASSERT_TRUE(SniiStatsProvider::open(&logical_reader, &second).ok()); + EXPECT_EQ(metered_reader.metrics().read_at_calls, after_first.read_at_calls); + EXPECT_EQ(metered_reader.metrics().total_request_bytes, after_first.total_request_bytes); + EXPECT_EQ(logical_reader.memory_usage(), memory_usage_before_load); + + uint8_t first_norm = 0; + uint8_t second_norm = 0; + ASSERT_TRUE(first.encoded_norm(17, &first_norm).ok()); + ASSERT_TRUE(second.encoded_norm(17, &second_norm).ok()); + EXPECT_EQ(first_norm, second_norm); + + std::remove(path.c_str()); +} + +TEST(SniiScoringQuery, StatsProviderSharesOneConcurrentNormsLoad) { + const Corpus corpus = MakeCorpus(); + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound_writer(&writer); + ASSERT_TRUE(compound_writer.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + } + + io::LocalFileReader local_reader; + ASSERT_TRUE(local_reader.open(path).ok()); + ControllableFileReader controlled_reader(&local_reader); + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&controlled_reader, &segment_reader).ok()); + reader::LogicalIndexReader logical_reader; + ASSERT_TRUE(segment_reader.open_index(1, "body", &logical_reader).ok()); + + constexpr size_t kThreadCount = 16; + std::barrier start(static_cast(kThreadCount + 1)); + std::vector statuses(kThreadCount); + std::vector norms(kThreadCount); + std::vector threads; + threads.reserve(kThreadCount); + controlled_reader.reset_read_at_calls(); + controlled_reader.block_next_read(); + for (size_t i = 0; i < kThreadCount; ++i) { + threads.emplace_back([&, i] { + SCOPED_INIT_THREAD_CONTEXT(); + start.arrive_and_wait(); + SniiStatsProvider provider; + statuses[i] = SniiStatsProvider::open(&logical_reader, &provider); + if (statuses[i].ok()) { + statuses[i] = provider.encoded_norm(17, &norms[i]); + } + }); + } + start.arrive_and_wait(); + controlled_reader.wait_for_blocked_read(); + controlled_reader.release_blocked_read(); + for (auto& thread : threads) { + thread.join(); + } + + for (size_t i = 0; i < kThreadCount; ++i) { + EXPECT_TRUE(statuses[i].ok()) << statuses[i].to_string(); + EXPECT_EQ(norms[i], norms[0]); + } + EXPECT_EQ(controlled_reader.read_at_calls(), 1U); + + std::remove(path.c_str()); +} + +TEST(SniiScoringQuery, StatsProviderRetriesTransientNormsReadFailure) { + const Corpus corpus = MakeCorpus(); + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound_writer(&writer); + ASSERT_TRUE(compound_writer.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + } + + io::LocalFileReader local_reader; + ASSERT_TRUE(local_reader.open(path).ok()); + ControllableFileReader controlled_reader(&local_reader); + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&controlled_reader, &segment_reader).ok()); + reader::LogicalIndexReader logical_reader; + ASSERT_TRUE(segment_reader.open_index(1, "body", &logical_reader).ok()); + + controlled_reader.reset_read_at_calls(); + controlled_reader.fail_next_read(); + SniiStatsProvider failed; + const doris::Status first_status = SniiStatsProvider::open(&logical_reader, &failed); + EXPECT_TRUE(first_status.is()) << first_status.to_string(); + + SniiStatsProvider retry; + ASSERT_TRUE(SniiStatsProvider::open(&logical_reader, &retry).ok()); + uint8_t norm = 0; + ASSERT_TRUE(retry.encoded_norm(17, &norm).ok()); + EXPECT_EQ(controlled_reader.read_at_calls(), 2U); + + std::remove(path.c_str()); +} + +TEST(SniiScoringQuery, CandidatesUseCollectionStatisticsAndPreserveDuplicateClauses) { + const Corpus corpus = MakeCorpus(); + const std::vector norms = EncodeNorms(corpus); + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound_writer(&writer); + ASSERT_TRUE(compound_writer.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + } + + io::LocalFileReader file_reader; + ASSERT_TRUE(file_reader.open(path).ok()); + io::MeteredFileReader metered_reader(&file_reader); + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered_reader, &segment_reader).ok()); + reader::LogicalIndexReader logical_reader; + ASSERT_TRUE(segment_reader.open_index(1, "body", &logical_reader).ok()); + SniiStatsProvider segment_stats; + ASSERT_TRUE(SniiStatsProvider::open(&logical_reader, &segment_stats).ok()); + + constexpr double kCollectionAvgdl = 100.0; + const Bm25Params params; + const std::vector clauses { + {.physical_term = "common", .idf = 0.25}, + {.physical_term = "rare", .idf = 2.5}, + {.physical_term = "rare", .idf = 2.5}}; + roaring::Roaring final_candidates; + final_candidates.add(3); + final_candidates.add(42); + + std::vector scored; + ASSERT_TRUE(doris::snii::query::scoring_query_candidates(logical_reader, segment_stats, clauses, + final_candidates, kCollectionAvgdl, + params, &scored) + .ok()); + + ASSERT_EQ(scored.size(), 2u); + EXPECT_EQ(scored[0].docid, 3u); + EXPECT_EQ(scored[1].docid, 42u); + auto expected_score = [&](uint32_t docid) { + double score = 0.0; + for (const auto& clause : clauses) { + const auto& term_postings = corpus.postings.at(clause.physical_term); + const auto posting = term_postings.find(docid); + if (posting == term_postings.end()) { + continue; + } + const double tf = posting->second; + const double dl = doris::snii::query::decode_norm(norms[docid]); + const double denominator = + tf + params.k1 * (1.0 - params.b + params.b * dl / kCollectionAvgdl); + score += clause.idf * (tf * (params.k1 + 1.0)) / denominator; + } + return score; + }; + EXPECT_NEAR(scored[0].score, expected_score(3), 1e-9); + EXPECT_NEAR(scored[1].score, expected_score(42), 1e-9); + + std::remove(path.c_str()); +} + +TEST(SniiScoringQuery, CandidatesReadOnlyCoveredFrequencyWindows) { + constexpr uint32_t kDocCount = 32 * format::kAdaptiveWindowDocs; + constexpr double kCollectionAvgdl = 100.0; + const std::string path = TempPath(); + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "body"; + input.config = IndexConfig::kDocsPositionsScoring; + input.doc_count = kDocCount; + input.target_dict_block_bytes = 1; + input.encoded_norms.reserve(kDocCount); + TermPostings posting; + posting.term = "dense"; + posting.docids.reserve(kDocCount); + posting.freqs.reserve(kDocCount); + TermPostings leading_posting; + leading_posting.term = "leading"; + leading_posting.docids.reserve(kDocCount / 2); + leading_posting.freqs.reserve(kDocCount / 2); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + const uint32_t frequency = 1 + docid % 4; + input.encoded_norms.push_back(doris::snii::query::encode_norm(8 + docid % 17)); + posting.docids.push_back(docid); + posting.freqs.push_back(frequency); + for (uint32_t position = 0; position < frequency; ++position) { + posting.positions_flat.push_back(position); + } + if (docid < kDocCount / 2) { + leading_posting.docids.push_back(docid); + leading_posting.freqs.push_back(1); + leading_posting.positions_flat.push_back(0); + } + } + input.terms.push_back(std::move(posting)); + input.terms.push_back(std::move(leading_posting)); + uint64_t token_count = 0; + for (const auto& term : input.terms) { + token_count += term.positions_flat.size(); + } + input.common_grams_metadata = + snii_test::make_plain_scoring_metadata(input.doc_count, token_count); + { + io::LocalFileWriter writer; + ASSERT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound_writer(&writer); + ASSERT_TRUE(compound_writer.add_logical_index(input).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + } + + io::LocalFileReader file_reader; + ASSERT_TRUE(file_reader.open(path).ok()); + io::MeteredFileReader metered_reader(&file_reader, /*block_size=*/256); + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered_reader, &segment_reader).ok()); + reader::LogicalIndexReader logical_reader; + ASSERT_TRUE(segment_reader.open_index(1, "body", &logical_reader).ok()); + SniiStatsProvider segment_stats; + ASSERT_TRUE(SniiStatsProvider::open(&logical_reader, &segment_stats).ok()); + + const std::vector clauses { + {.physical_term = "dense", .idf = 0.75}, {.physical_term = "dense", .idf = 0.75}}; + roaring::Roaring sparse_candidates; + sparse_candidates.add(17); + sparse_candidates.add(kDocCount - 19); + metered_reader.reset_metrics(); + format::testing::reset_window_probe_count(); + std::vector sparse_scores; + ASSERT_TRUE(scoring_query_candidates(logical_reader, segment_stats, clauses, sparse_candidates, + kCollectionAvgdl, Bm25Params {}, &sparse_scores) + .ok()); + const io::IoMetrics sparse_metrics = metered_reader.metrics(); + const uint64_t sparse_window_probes = format::testing::window_probe_count(); + + roaring::Roaring clustered_candidates; + clustered_candidates.addRange(0, 129); + metered_reader.reset_metrics(); + format::testing::reset_window_probe_count(); + std::vector clustered_scores; + ASSERT_TRUE(scoring_query_candidates(logical_reader, segment_stats, clauses, + clustered_candidates, kCollectionAvgdl, Bm25Params {}, + &clustered_scores) + .ok()); + const io::IoMetrics clustered_metrics = metered_reader.metrics(); + const uint64_t clustered_window_probes = format::testing::window_probe_count(); + + roaring::Roaring dense_candidates; + dense_candidates.addRange(0, kDocCount); + metered_reader.reset_metrics(); + format::testing::reset_window_probe_count(); + std::vector dense_scores; + ASSERT_TRUE(scoring_query_candidates(logical_reader, segment_stats, clauses, dense_candidates, + kCollectionAvgdl, Bm25Params {}, &dense_scores) + .ok()); + const io::IoMetrics dense_metrics = metered_reader.metrics(); + const uint64_t dense_window_probes = format::testing::window_probe_count(); + + const std::vector leading_clause { + {.physical_term = "leading", .idf = 0.5}}; + roaring::Roaring leading_hit_candidates; + leading_hit_candidates.addRange(0, 129); + metered_reader.reset_metrics(); + std::vector leading_hit_scores; + ASSERT_TRUE(scoring_query_candidates(logical_reader, segment_stats, leading_clause, + leading_hit_candidates, kCollectionAvgdl, Bm25Params {}, + &leading_hit_scores) + .ok()); + const io::IoMetrics leading_hit_metrics = metered_reader.metrics(); + + roaring::Roaring disjoint_candidates; + disjoint_candidates.addRange(kDocCount / 2, kDocCount); + metered_reader.reset_metrics(); + std::vector disjoint_scores; + ASSERT_TRUE(scoring_query_candidates(logical_reader, segment_stats, leading_clause, + disjoint_candidates, kCollectionAvgdl, Bm25Params {}, + &disjoint_scores) + .ok()); + const io::IoMetrics disjoint_metrics = metered_reader.metrics(); + + ASSERT_EQ(sparse_scores.size(), 2U); + ASSERT_EQ(clustered_scores.size(), 129U); + ASSERT_EQ(dense_scores.size(), kDocCount); + const auto scorer = doris::snii::query::ScorerContext::from_idf(0.75); + for (const auto& scored_doc : sparse_scores) { + const uint32_t frequency = 1 + scored_doc.docid % 4; + const uint8_t norm = input.encoded_norms[scored_doc.docid]; + const double expected = 2 * scorer.score(frequency, norm, kCollectionAvgdl, Bm25Params {}); + EXPECT_NEAR(scored_doc.score, expected, 1e-9); + EXPECT_NEAR(dense_scores[scored_doc.docid].score, expected, 1e-9); + } + EXPECT_LT(sparse_metrics.total_request_bytes, dense_metrics.total_request_bytes); + EXPECT_LT(sparse_metrics.remote_bytes, dense_metrics.remote_bytes); + EXPECT_LT(clustered_metrics.total_request_bytes, dense_metrics.total_request_bytes); + EXPECT_LT(clustered_metrics.remote_bytes, dense_metrics.remote_bytes); + EXPECT_LE(sparse_metrics.serial_rounds, dense_metrics.serial_rounds); + EXPECT_GT(sparse_window_probes, 0U); + EXPECT_GT(clustered_window_probes, 0U); + EXPECT_EQ(dense_window_probes, 0U); + ASSERT_EQ(disjoint_scores.size(), kDocCount / 2); + for (const auto& scored_doc : disjoint_scores) { + EXPECT_DOUBLE_EQ(scored_doc.score, 0.0); + } + EXPECT_LT(disjoint_metrics.total_request_bytes, leading_hit_metrics.total_request_bytes); + EXPECT_LT(disjoint_metrics.remote_bytes, leading_hit_metrics.remote_bytes); + + std::remove(path.c_str()); +} + +namespace { + +// A corpus engineered to produce SCORE TIES at the top-k boundary and to drive +// the WINDOWED posting path: uniform doc length (so length-norm is constant) and +// high-df terms (df >= kSlimDfThreshold = 512 -> windowed pod_ref + frq_prelude). +// Every doc has the same length L=8, so docs sharing a term/freq score identically. +Corpus MakeWindowedTieCorpus() { + Corpus c; + c.doc_count = 700; // >= 512 so "anchor" becomes a windowed term + c.doc_len.assign(c.doc_count, 0); + auto add = [&](const std::string& term, uint32_t doc, uint32_t freq) { + c.postings[term][doc] += freq; + c.doc_len[doc] += freq; + }; + for (uint32_t d = 0; d < c.doc_count; ++d) { + add("anchor", d, 1); // df=700 (windowed), freq=1 everywhere -> ties + if (d % 2 == 0) { + add("evenz", d, 1); // df=350 (windowed), another high-df term + } + add("u" + std::to_string(d), d, 6); // unique pad: keeps every dl == 8 exactly + } + return c; +} + +} // namespace + +// Differential: WAND top-k MUST equal exhaustive top-k EVEN with boundary ties and +// windowed (block-max) terms, across many k. Strict-'>' pruning would drop ties. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiScoringQuery, WandEqualsExhaustiveWithTiesAndWindowedTerms) { + const Corpus corpus = MakeWindowedTieCorpus(); + const std::string path = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(ToInput(corpus)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + io::LocalFileReader inner; + ASSERT_TRUE(inner.open(path).ok()); + io::MeteredFileReader metered(&inner); + reader::SniiSegmentReader seg; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered, &seg).ok()); + reader::LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + SniiStatsProvider stats; + ASSERT_TRUE(SniiStatsProvider::open(&idx, &stats).ok()); + + const Bm25Params params; + const std::vector norms = EncodeNorms(corpus); + auto check = [&](const std::vector& terms, uint32_t k) { + std::vector ex, wa; + ASSERT_TRUE(scoring_query_exhaustive(idx, stats, terms, k, params, &ex).ok()); + ASSERT_TRUE(scoring_query_wand(idx, stats, terms, k, params, &wa).ok()); + const std::vector ref = ReferenceRanking(corpus, norms, terms, k, params); + ASSERT_EQ(wa.size(), ex.size()); + ASSERT_EQ(ex.size(), ref.size()); + for (size_t i = 0; i < ex.size(); ++i) { + EXPECT_EQ(wa[i].docid, ex[i].docid) + << "terms[0]=" << terms[0] << " k=" << k << " i=" << i; + EXPECT_EQ(ex[i].docid, ref[i].docid) << "ref k=" << k << " i=" << i; + EXPECT_NEAR(wa[i].score, ex[i].score, 1e-9); + } + }; + // Single high-df term: all 700 docs tie -> top-k must be the k smallest docids. + for (uint32_t k : {1U, 3U, 5U, 50U, 200U}) { + check({"anchor"}, k); + } + // Two windowed terms: even docs score higher (two terms) -> ties within each tier. + for (uint32_t k : {1U, 4U, 10U, 100U}) { + check({"anchor", "evenz"}, k); + } + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/query/scoring_wand_selective_test.cpp b/be/test/storage/index/snii/query/scoring_wand_selective_test.cpp new file mode 100644 index 00000000000000..ac2d436ff0dbac --- /dev/null +++ b/be/test/storage/index/snii/query/scoring_wand_selective_test.cpp @@ -0,0 +1,520 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/scoring_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +// Phase C differential test: scoring_query_wand_selective (block-max SELECTIVE +// FETCH) MUST return top-K (docid sequence AND scores within 1e-9) byte-identical +// to scoring_query_exhaustive AND scoring_query_wand AND an in-memory brute-force +// reference -- for MANY random queries, varied k (incl k=1), varied term sets +// (incl a high-df windowed term spanning many .frq windows), AND a corpus crafted +// to force SCORE TIES (uniform doc length). It additionally asserts the selective +// path fetches FEWER .frq windows / bytes than reading all windows for a small-k +// high-df query, WITHOUT ever changing the result. +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +using doris::snii::query::Bm25Params; +using doris::snii::query::ScoredDoc; +using doris::snii::stats::SniiStatsProvider; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_wand_sel_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +// term -> (docid -> freq); plus per-doc length (token count) for norms. +struct Corpus { + uint32_t doc_count = 0; + std::map> postings; + std::vector doc_len; +}; + +// A scoring corpus with a HIGH-DF term ("hi", df=doc_count -> windowed, spans +// many .frq windows), two MID-df windowed terms, and several low-df terms. Doc +// lengths VARY (so length-norm matters and scores spread out for the random +// queries), but every doc that carries "hi" carries exactly freq=1 of it. +Corpus MakeVariedCorpus(uint32_t doc_count) { + Corpus c; + c.doc_count = doc_count; + c.doc_len.assign(doc_count, 0); + auto add = [&](const std::string& t, uint32_t d, uint32_t f) { + c.postings[t][d] += f; + c.doc_len[d] += f; + }; + // All queryable terms use an "aa_" prefix and a "zz_NNN" filler vocabulary fills + // the lexicographic tail, so every real term sorts within the SampledTermIndex's + // candidate range (the index samples per-block first terms; without tail filler + // a term sorting last can fall outside the sampled range and miss on lookup). + for (uint32_t d = 0; d < doc_count; ++d) { + add("aa_hi", d, 1 + (d % 5)); // df=N, windowed, varied freq + if (d % 2 == 0) { + add("aa_evenmid", d, 1 + (d % 3)); + } + if (d % 3 == 0) { + add("aa_thirdmid", d, 1 + (d % 4)); + } + if (d % 37 == 0) { + add("aa_rarea", d, 2); + } + if (d % 53 == 0) { + add("aa_rareb", d, 3); + } + if (d % 101 == 0) { + add("aa_rarec", d, 1); + } + add("aa_pad" + std::to_string(d % 13), d, 1 + (d % 7)); // spreads dl widely + char nm[16]; + std::snprintf(nm, sizeof(nm), "zz_%03u", d % 257); // tail filler vocabulary + add(nm, d, 1); + } + return c; +} + +// A TIE corpus: uniform doc length L=8 for every doc (so equal-freq docs of a +// term score identically -> ties at the top-K boundary). "hi"/"evenz" are high-df +// windowed terms with freq=1 everywhere -> massive ties broken only by docid. +Corpus MakeTieCorpus(uint32_t doc_count) { + Corpus c; + c.doc_count = doc_count; + c.doc_len.assign(doc_count, 0); + auto add = [&](const std::string& t, uint32_t d, uint32_t f) { + c.postings[t][d] += f; + c.doc_len[d] += f; + }; + for (uint32_t d = 0; d < doc_count; ++d) { + add("aa_hi", d, 1); // df=N, freq=1 everywhere -> ties + if (d % 2 == 0) { + add("aa_evenz", d, 1); // another high-df windowed term + } + if (d % 41 == 0) { + add("aa_rarez", d, 1); // low-df, still freq=1 + } + // unique pad keeps every dl == 8 exactly; named in the "m_" mid range so the + // queryable "aa_" terms still sort within the sampled candidate range. + add("mm_u" + std::to_string(d), d, 6); + } + return c; +} + +// A DECAYING corpus engineered so block-max WAND genuinely SKIPS later windows of +// the high-df term "hi": "hi"'s in-doc freq decreases with docid (early docs have +// the highest tf) and doc length increases with docid (early docs are shortest, +// so least length-penalized). Thus EARLY windows hold the top scorers and fill the +// heap with a high theta, after which LATE windows' block-max (small tf, long dl) +// falls below theta and is provably skippable -- selective never fetches them. +Corpus MakeDecayCorpus(uint32_t doc_count) { + Corpus c; + c.doc_count = doc_count; + c.doc_len.assign(doc_count, 0); + auto add = [&](const std::string& t, uint32_t d, uint32_t f) { + c.postings[t][d] += f; + c.doc_len[d] += f; + }; + for (uint32_t d = 0; d < doc_count; ++d) { + // tf decays from ~20 (docid 0) down to 1 (last docids). + const uint32_t band = d / 256; // one step per window-sized band + const uint32_t tf = band < 20 ? (20 - band) : 1; + add("aa_hi", d, tf); + // Padding grows with docid so later docs are longer (higher length penalty). + add("mm_pad" + std::to_string(d % 7), d, 1 + band); + } + return c; +} + +SniiIndexInput ToInput(const Corpus& c) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositionsScoring; + in.doc_count = c.doc_count; + in.target_dict_block_bytes = 256; + in.encoded_norms.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + in.encoded_norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + for (const auto& [term, plist] : c.postings) { + TermPostings tp; + tp.term = term; + for (const auto& [docid, freq] : plist) { + tp.docids.push_back(docid); + tp.freqs.push_back(freq); + for (uint32_t k = 0; k < freq; ++k) { + tp.positions_flat.push_back(k); // flat + } + } + in.terms.push_back(std::move(tp)); + } + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + return in; +} + +std::vector EncodeNorms(const Corpus& c) { + std::vector v(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + v[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + return v; +} + +// Independent brute-force BM25 ranking straight from the corpus. +std::vector ReferenceRanking(const Corpus& c, const std::vector& norms, + const std::vector& terms, uint32_t k, + const Bm25Params& params) { + uint64_t sum_ttf = 0; + for (const auto& dl : c.doc_len) { + sum_ttf += dl; + } + const double avgdl = static_cast(sum_ttf) / std::max(1, c.doc_count); + std::unordered_map scores; + for (const auto& term : terms) { + auto it = c.postings.find(term); + if (it == c.postings.end()) { + continue; + } + const uint64_t df = it->second.size(); + const double idf = + std::log(1.0 + (static_cast(c.doc_count) - df + 0.5) / (df + 0.5)); + for (const auto& [docid, freq] : it->second) { + const double dl = doris::snii::query::decode_norm(norms[docid]); + const double denom = freq + params.k1 * (1.0 - params.b + params.b * dl / avgdl); + scores[docid] += idf * (freq * (params.k1 + 1.0)) / denom; + } + } + std::vector all; + all.reserve(scores.size()); + for (const auto& [docid, s] : scores) { + all.push_back({docid, s}); + } + std::ranges::sort(all, [](const ScoredDoc& a, const ScoredDoc& b) { + if (a.score != b.score) { + return a.score > b.score; + } + return a.docid < b.docid; + }); + if (all.size() > k) { + all.resize(k); + } + return all; +} + +// Opens a built index file over a metered reader. +struct OpenedIndex { + io::LocalFileReader inner; + io::MeteredFileReader metered; + SniiSegmentReader seg; + LogicalIndexReader idx; + SniiStatsProvider stats; + explicit OpenedIndex(const std::string& path, size_t block_size) : metered(&inner, block_size) { + EXPECT_TRUE(inner.open(path).ok()); + } +}; + +void BuildAndOpen(const Corpus& c, const std::string& path) { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(ToInput(c)).ok()); + ASSERT_TRUE(cw.finish().ok()); +} + +// Counts the .frq window count of a windowed term (its full posting window span). +uint32_t WindowCount(const LogicalIndexReader& idx, const std::string& term) { + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + EXPECT_TRUE(idx.lookup(term, &found, &entry, &frq_base, &prx_base).ok()); + if (!found || entry.enc != DictEntryEnc::kWindowed) { + return 0; + } + FrqPreludeReader prelude; + EXPECT_TRUE(fetch_windowed_prelude(idx, entry, frq_base, &prelude).ok()); + return prelude.n_windows(); +} + +// Asserts two top-K vectors are BYTE-IDENTICAL: same length, same docid sequence, +// scores equal within tol. This is the hard soundness gate used for selective vs +// the eager WAND (both implement the documented (score desc, docid asc) order and +// the >= theta tie rule, so they must agree exactly, ties included). +void ExpectIdentical(const std::vector& a, const std::vector& b, + const std::string& label, double tol) { + ASSERT_EQ(a.size(), b.size()) << label << " size"; + for (size_t i = 0; i < a.size(); ++i) { + EXPECT_EQ(a[i].docid, b[i].docid) << label << " docid rank " << i; + EXPECT_NEAR(a[i].score, b[i].score, tol) << label << " score rank " << i; + } +} + +// Asserts a is a VALID top-K relative to oracle b, tolerant ONLY of the floating- +// point ULP tie-order ambiguity inherent to BM25 (two docs whose summed scores +// differ by < tol may legitimately swap order across summation orders; the +// exhaustive path sums per doc in unordered_map order, the WAND path in cursor +// order). It groups each vector into maximal tie bands (consecutive scores within +// tol) and requires: +// - identical length and position-wise scores within tol (any real score / +// ranking error shows here -- positions only slide WITHIN an equal band), +// - the docid MULTISET of each tie band is identical (catches wrong / dropped +// docids, including at the top-K boundary when the boundary is not in a band), +// - a's docids are ascending within each band (a is the canonical winner). +void ExpectValidTopK(const std::vector& a, const std::vector& b, + const std::string& label, double tol) { + ASSERT_EQ(a.size(), b.size()) << label << " size"; + size_t i = 0; + while (i < a.size()) { + size_t j = i + 1; + while (j < a.size() && std::abs(a[j].score - a[i].score) <= tol) { + ++j; + } + std::vector ad, bd; + for (size_t t = i; t < j; ++t) { + EXPECT_NEAR(a[t].score, b[t].score, tol) << label << " band score " << t; + ad.push_back(a[t].docid); + bd.push_back(b[t].docid); + if (t > i) { + EXPECT_LT(a[t - 1].docid, a[t].docid) + << label << " a not docid-ascending in tie band at " << t; + } + } + std::ranges::sort(ad); + std::ranges::sort(bd); + EXPECT_EQ(ad, bd) << label << " tie-band docid set differs near rank " << i; + i = j; + } +} + +} // namespace + +// Selective vs exhaustive / eager-WAND / brute-force over MANY random queries on a +// varied-length scoring corpus with a high-df windowed term spanning many windows. +// The hard gate: selective is BYTE-IDENTICAL to the eager WAND (same docids, same +// scores, ties included) -- selective only changes the bytes read. Against the +// exhaustive oracle and the brute-force reference it must be a VALID top-K (same +// scores, same docid sets per tie band); those two paths sum per-doc in a +// different order, so they may legitimately swap docids only WITHIN an equal-score +// tie band (BM25 floating-point ULP), which ExpectValidTopK tolerates while still +// catching any real ranking divergence. +TEST(SniiScoringWandSelective, MatchesExhaustiveOverRandomQueries) { + const Corpus corpus = MakeVariedCorpus(4000); + const std::vector norms = EncodeNorms(corpus); + const std::string path = TempPath(); + BuildAndOpen(corpus, path); + + OpenedIndex oi(path, /*block_size=*/4096); + ASSERT_TRUE(SniiSegmentReader::open(&oi.metered, &oi.seg).ok()); + ASSERT_TRUE(oi.seg.open_index(1, "body", &oi.idx).ok()); + ASSERT_TRUE(SniiStatsProvider::open(&oi.idx, &oi.stats).ok()); + + ASSERT_GE(WindowCount(oi.idx, "aa_hi"), 8U) << "aa_hi must span many windows"; + + const Bm25Params params; + const std::vector vocab = {"aa_hi", "aa_evenmid", "aa_thirdmid", + "aa_rarea", "aa_rareb", "aa_rarec", + "aa_pad0", "aa_pad7", "aa_missing"}; + std::mt19937 rng(0xC0FFEEU); + std::uniform_int_distribution n_terms(1, 4); + std::uniform_int_distribution pick(0, vocab.size() - 1); + const std::vector ks = {1, 2, 3, 5, 10, 25, 100, 500}; + std::uniform_int_distribution pick_k(0, ks.size() - 1); + + for (int iter = 0; iter < 300; ++iter) { + // Distinct query terms (a query never repeats a term in practice). + std::vector terms; + const uint32_t nt = n_terms(rng); + for (uint32_t t = 0; t < nt; ++t) { + const std::string cand = vocab[pick(rng)]; + if (std::ranges::find(terms, cand) == terms.end()) { + terms.push_back(cand); + } + } + if (terms.empty()) { + continue; + } + const uint32_t k = ks[pick_k(rng)]; + + std::vector sel, ex, wa; + ASSERT_TRUE(doris::snii::query::scoring_query_wand_selective(oi.idx, oi.stats, terms, k, + params, &sel) + .ok()); + ASSERT_TRUE(doris::snii::query::scoring_query_exhaustive(oi.idx, oi.stats, terms, k, params, + &ex) + .ok()); + ASSERT_TRUE(doris::snii::query::scoring_query_wand(oi.idx, oi.stats, terms, k, params, &wa) + .ok()); + const std::vector ref = ReferenceRanking(corpus, norms, terms, k, params); + + std::string label = "iter " + std::to_string(iter) + " k=" + std::to_string(k) + " terms:"; + for (const auto& t : terms) { + label += " " + t; + } + + // HARD GATE: selective == eager WAND exactly (docids + scores, ties included). + ExpectIdentical(sel, wa, label + " [sel==wand]", 1e-9); + // Selective is a valid top-K vs the exhaustive oracle and brute-force ref. + ExpectValidTopK(sel, ex, label + " [sel~ex]", 1e-9); + ExpectValidTopK(sel, ref, label + " [sel~ref]", 1e-9); + } + std::remove(path.c_str()); +} + +// Selective == exhaustive == wand == brute-force WITH SCORE TIES (uniform doc +// length) across many k, including k=1. Strict pruning that dropped ties would +// fail here; the >= theta tie rule must be preserved end to end. +TEST(SniiScoringWandSelective, MatchesExhaustiveWithTies) { + const Corpus corpus = MakeTieCorpus(2400); + const std::vector norms = EncodeNorms(corpus); + const std::string path = TempPath(); + BuildAndOpen(corpus, path); + + OpenedIndex oi(path, /*block_size=*/4096); + ASSERT_TRUE(SniiSegmentReader::open(&oi.metered, &oi.seg).ok()); + ASSERT_TRUE(oi.seg.open_index(1, "body", &oi.idx).ok()); + ASSERT_TRUE(SniiStatsProvider::open(&oi.idx, &oi.stats).ok()); + ASSERT_GE(WindowCount(oi.idx, "aa_hi"), 4U); + + const Bm25Params params; + auto check = [&](const std::vector& terms, uint32_t k) { + std::vector sel, ex, wa; + ASSERT_TRUE(doris::snii::query::scoring_query_wand_selective(oi.idx, oi.stats, terms, k, + params, &sel) + .ok()); + ASSERT_TRUE(doris::snii::query::scoring_query_exhaustive(oi.idx, oi.stats, terms, k, params, + &ex) + .ok()); + ASSERT_TRUE(doris::snii::query::scoring_query_wand(oi.idx, oi.stats, terms, k, params, &wa) + .ok()); + const std::vector ref = ReferenceRanking(corpus, norms, terms, k, params); + std::string label = "ties terms[0]=" + terms[0] + " k=" + std::to_string(k); + // Uniform doc length => EXACT ties (bitwise-equal scores): every path must + // agree exactly, ordering equal-score docs by ascending docid (the >= theta + // tie rule). This is the strongest possible selective==exhaustive gate. + ExpectIdentical(sel, ex, label + " [sel==ex]", 1e-9); + ExpectIdentical(sel, wa, label + " [sel==wand]", 1e-9); + ExpectIdentical(sel, ref, label + " [sel==ref]", 1e-9); + }; + for (uint32_t k : {1U, 2U, 3U, 5U, 25U, 200U, 1000U}) { + check({"aa_hi"}, k); + } + for (uint32_t k : {1U, 4U, 10U, 100U, 800U}) { + check({"aa_hi", "aa_evenz"}, k); + } + for (uint32_t k : {1U, 5U, 50U}) { + check({"aa_hi", "aa_evenz", "aa_rarez"}, k); + } + std::remove(path.c_str()); +} + +// Selective fetch reads FEWER .frq windows / bytes than reading EVERY window of +// the high-df term, for a small-k single-high-df-term query -- with NO change in +// the result. The decay corpus puts the top scorers in the EARLY windows so the +// block-max of later windows provably falls below theta and they are skipped. We +// measure the selective path's remote_bytes / requested bytes against the eager +// full-posting read of "hi" through the same metered reader. +TEST(SniiScoringWandSelective, FetchesFewerWindowsForSmallKHighDf) { + const Corpus corpus = MakeDecayCorpus(8000); + const std::string path = TempPath(); + BuildAndOpen(corpus, path); + + // Fine-grained cache block: the high-df frq span is only a few KiB, so a coarse + // 4096-byte block can round BOTH the selective and full reads up to the same + // aligned block set even though their raw request sizes differ. A 512-byte block + // keeps the block-aligned remote_bytes measure faithful to the genuine per-window + // savings (independent of where the interleaved posting region happens to land). + OpenedIndex oi(path, /*block_size=*/512); + ASSERT_TRUE(SniiSegmentReader::open(&oi.metered, &oi.seg).ok()); + ASSERT_TRUE(oi.seg.open_index(1, "body", &oi.idx).ok()); + ASSERT_TRUE(SniiStatsProvider::open(&oi.idx, &oi.stats).ok()); + const uint32_t hi_windows = WindowCount(oi.idx, "aa_hi"); + ASSERT_GE(hi_windows, 16U) << "aa_hi must span many windows"; + + const Bm25Params params; + const std::vector terms = {"aa_hi"}; + const uint32_t k = 5; + + // Selective path metrics (prelude + only the surviving windows fetched). + oi.metered.reset_metrics(); + std::vector sel; + ASSERT_TRUE(doris::snii::query::scoring_query_wand_selective(oi.idx, oi.stats, terms, k, params, + &sel) + .ok()); + const io::IoMetrics sel_m = oi.metered.metrics(); + + // Full-posting read of the same term (every window's full .frq), as the + // "read all windows" baseline through the same metered reader. + oi.metered.reset_metrics(); + DictEntry entry; + uint64_t frq_base = 0, prx_base = 0; + bool found = false; + ASSERT_TRUE(oi.idx.lookup("aa_hi", &found, &entry, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + DecodedPosting dp; + ASSERT_TRUE(read_windowed_posting(oi.idx, entry, frq_base, prx_base, + /*want_positions=*/false, /*want_freq=*/true, &dp) + .ok()); + const io::IoMetrics full_m = oi.metered.metrics(); + + // Result must be unchanged vs exhaustive (byte savings never change the answer). + std::vector ex; + ASSERT_TRUE( + doris::snii::query::scoring_query_exhaustive(oi.idx, oi.stats, terms, k, params, &ex) + .ok()); + ExpectValidTopK(sel, ex, "decay small-k [sel~ex]", 1e-9); + + // Selective requests strictly fewer raw bytes than reading all windows: later + // windows are skipped, not merely cached. + EXPECT_LT(sel_m.total_request_bytes, full_m.total_request_bytes) + << "sel=" << sel_m.total_request_bytes << " full=" << full_m.total_request_bytes; + // And fewer remote bytes leave the wire. + EXPECT_LT(sel_m.remote_bytes, full_m.remote_bytes) + << "sel=" << sel_m.remote_bytes << " full=" << full_m.remote_bytes; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/reader/snii_end_to_end_test.cpp b/be/test/storage/index/snii/reader/snii_end_to_end_test.cpp new file mode 100644 index 00000000000000..e5f2a2c7576dd8 --- /dev/null +++ b/be/test/storage/index/snii/reader/snii_end_to_end_test.cpp @@ -0,0 +1,295 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_e2e_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +// Deterministic corpus oracle: term -> set, and doc -> token sequence. +struct Corpus { + uint32_t doc_count = 0; + std::vector> docs; // docs[d] = ordered token list + std::map> term_docs; + + // Does the phrase occur consecutively in doc d? + bool phrase_in_doc(uint32_t d, const std::vector& phrase) const { + if (phrase.empty()) { + return false; + } + const auto& toks = docs[d]; + if (toks.size() < phrase.size()) { + return false; + } + for (size_t i = 0; i + phrase.size() <= toks.size(); ++i) { + bool match = true; + for (size_t k = 0; k < phrase.size(); ++k) { + if (toks[i + k] != phrase[k]) { + match = false; + break; + } + } + if (match) { + return true; + } + } + return false; + } + + std::vector phrase_oracle(const std::vector& phrase) const { + std::vector out; + for (uint32_t d = 0; d < doc_count; ++d) { + if (phrase_in_doc(d, phrase)) { + out.push_back(d); + } + } + return out; + } +}; + +// Builds a deterministic corpus of 60 docs. Plants a HIGH-frequency term +// "the" in every doc (df=60, but we also force a high-df term via padding to +// exercise the windowed path), several low-df terms, and known phrases. +Corpus BuildCorpus() { + Corpus c; + c.doc_count = 60; + c.docs.resize(c.doc_count); + + // Common base vocabulary cycled deterministically. + const std::vector vocab = {"alpha", "bravo", "charlie", "delta", + "echo", "foxtrot", "golf", "hotel"}; + + for (uint32_t d = 0; d < c.doc_count; ++d) { + std::vector& toks = c.docs[d]; + // "the" in every doc -> high df. + toks.emplace_back("the"); + // A planted 5-term phrase in docs where d % 7 == 0. + if (d % 7 == 0) { + toks.emplace_back("quick"); + toks.emplace_back("brown"); + toks.emplace_back("fox"); + toks.emplace_back("jumps"); + toks.emplace_back("over"); + } + // A planted 2-term phrase "lazy dog" in docs where d % 5 == 0. + if (d % 5 == 0) { + toks.emplace_back("lazy"); + toks.emplace_back("dog"); + } + // Repeated-term phrase to verify execution de-duplicates reads without + // weakening positional semantics. + if (d % 13 == 0) { + toks.emplace_back("repeat"); + toks.emplace_back("repeat"); + } + // Cycle through vocab to give terms varied dfs. + for (uint32_t k = 0; k < 4; ++k) { + toks.push_back(vocab[(d + k) % vocab.size()]); + } + // A unique low-df marker per doc bucket. + if (d % 11 == 0) { + toks.emplace_back("rare"); + } + } + + for (uint32_t d = 0; d < c.doc_count; ++d) { + for (const auto& t : c.docs[d]) { + c.term_docs[t].insert(d); + } + } + return c; +} + +// Adds a synthetic high-df term spanning > kSlimDfThreshold docs to force the +// windowed pod_ref path. Uses an expanded doc space; the term appears in docs +// [0, df) at position 0. Updates the oracle. +void PlantHighDfTerm(Corpus* c, const std::string& term, uint32_t df) { + // Expand the corpus to hold df docs if needed. + if (df > c->doc_count) { + uint32_t old = c->doc_count; + c->doc_count = df; + c->docs.resize(df); + (void)old; + } + for (uint32_t d = 0; d < df; ++d) { + c->docs[d].insert(c->docs[d].begin(), term); // position 0 + c->term_docs[term].insert(d); + } + // Recompute term_docs fully to keep positions/oracle consistent. + c->term_docs.clear(); + for (uint32_t d = 0; d < c->doc_count; ++d) { + for (const auto& t : c->docs[d]) { + c->term_docs[t].insert(d); + } + } +} + +// Feeds the corpus into a SpimiTermBuffer and writes a single-index container. +void WriteCorpus(const Corpus& c, const std::string& path) { + SpimiTermBuffer buf(/*has_positions=*/true); + for (uint32_t d = 0; d < c.doc_count; ++d) { + const auto& toks = c.docs[d]; + for (uint32_t pos = 0; pos < toks.size(); ++pos) { + buf.add_token(toks[pos], d, pos); + } + } + std::vector terms = buf.finalize_sorted(); + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = c.doc_count; + in.terms = std::move(terms); + // Small block target so we get multiple DICT blocks (exercises sampling). + in.target_dict_block_bytes = 256; + + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(in).ok()); + ASSERT_TRUE(cw.finish().ok()); +} + +std::vector SetToVec(const std::set& s) { + // NOLINTNEXTLINE(modernize-return-braced-init-list) + return std::vector(s.begin(), s.end()); +} + +} // namespace + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiEndToEnd, TermAndPhraseAgainstOracle) { + Corpus c = BuildCorpus(); + // Force a windowed (df >= 512) term to exercise the windowed pod_ref path. + PlantHighDfTerm(&c, "ubiquitous", 600); + + const std::string path = TempPath(); + WriteCorpus(c, path); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local); + + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&metered, &seg).ok()); + EXPECT_EQ(seg.n_logical_indexes(), 1U); + + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + EXPECT_EQ(idx.stats().doc_count, c.doc_count); + + // Small DICT blocks are resident after open: exact lookup should not add + // query path I/O once the index is open. + metered.reset_metrics(); + { + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + bool found = false; + ASSERT_TRUE(idx.lookup("quick", &found, &entry, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(metered.metrics().read_at_calls, 0U); + } + + // ---- term_query: present terms match the oracle docid set exactly. ---- + std::vector present_terms = {"the", "quick", "brown", "fox", + "jumps", "over", "lazy", "dog", + "rare", "alpha", "repeat", "ubiquitous"}; + for (const auto& t : present_terms) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, t, &got).ok()) << t; + std::vector want = SetToVec(c.term_docs[t]); + EXPECT_EQ(got, want) << "term=" << t; + } + + // ---- term_query: absent terms -> empty. ---- + std::vector absent_terms = {"nonexistent", "zzzz", "qqqq-absent"}; + for (const auto& t : absent_terms) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, t, &got).ok()) << t; + EXPECT_TRUE(got.empty()) << "term=" << t; + } + + // ---- phrase_query: present and absent phrases vs oracle. ---- + std::vector> phrases = { + {"quick", "brown", "fox", "jumps", "over"}, // 5-term planted phrase + {"lazy", "dog"}, // 2-term planted phrase + {"repeat", "repeat"}, // repeated-term phrase + {"brown", "fox"}, // sub-phrase + {"fox", "brown"}, // reversed -> absent + {"quick", "fox"}, // non-consecutive -> absent + {"the"}, // single-term phrase + {"nope", "missing"}, // absent terms -> empty + }; + for (const auto& p : phrases) { + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, p, &got).ok()); + std::vector want = c.phrase_oracle(p); + std::string label; + for (const auto& w : p) { + label += w + " "; + } + EXPECT_EQ(got, want) << "phrase=[" << label << "]"; + } + + // ---- metrics sanity: a fresh phrase query over windowed postings + // plans/batches its reads. Small DICT blocks are resident and small postings + // may be inline, so use the synthetic high-df term to force external posting + // I/O. + metered.reset_metrics(); + std::vector got; + ASSERT_TRUE(query::phrase_query(idx, {"ubiquitous", "the"}, &got).ok()); + const auto& m = metered.metrics(); + EXPECT_GE(m.serial_rounds, 1U); + EXPECT_GE(m.range_gets, 1U); + // The planned batch should keep rounds small (resident DICT + batched + // postings). + EXPECT_LE(m.serial_rounds, 8U) << "phrase query did not batch its I/O"; + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/reader/snii_segment_reader_open_test.cpp b/be/test/storage/index/snii/reader/snii_segment_reader_open_test.cpp new file mode 100644 index 00000000000000..7d238c1eb665ed --- /dev/null +++ b/be/test/storage/index/snii/reader/snii_segment_reader_open_test.cpp @@ -0,0 +1,383 @@ +// 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. + +// Focused open()-path tests for SniiSegmentReader. These pin two guarantees of +// the offset-0 bootstrap-read removal: +// 1. open() issues NO read intersecting the bootstrap header region +// [0, kBootstrapHeaderSize) -- the redundant offset-0 cache block / remote +// round-trip is gone. +// 2. The container version gate is preserved by the tail pointer: a corrupt +// offset-0 bootstrap header no longer fails open(), but a corrupt tail +// pointer format_version still does. + +#include +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/bootstrap_header.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/metadata_directory.h" +#include "storage/index/snii/format/tail_pointer.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::reader; +using namespace doris::snii::writer; +namespace ErrorCode = doris::ErrorCode; +using doris::Status; + +namespace { + +constexpr size_t kV1TailPointerSize = 31; + +struct TailFields { + uint64_t directory_offset = 0; + uint64_t directory_length = 0; +}; + +TailFields ReadTailFields(const std::vector& bytes) { + EXPECT_GE(bytes.size(), kV1TailPointerSize); + ByteSource source(Slice(bytes.data() + bytes.size() - kV1TailPointerSize, kV1TailPointerSize)); + uint32_t magic = 0; + uint16_t version = 0; + TailFields tail; + uint32_t directory_crc = 0; + uint8_t encoded_size = 0; + uint32_t tail_crc = 0; + EXPECT_TRUE(source.get_fixed32(&magic).ok()); + EXPECT_TRUE(source.get_fixed16(&version).ok()); + EXPECT_TRUE(source.get_fixed64(&tail.directory_offset).ok()); + EXPECT_TRUE(source.get_fixed64(&tail.directory_length).ok()); + EXPECT_TRUE(source.get_fixed32(&directory_crc).ok()); + EXPECT_TRUE(source.get_u8(&encoded_size).ok()); + EXPECT_TRUE(source.get_fixed32(&tail_crc).ok()); + EXPECT_EQ(kTailMagic, magic); + EXPECT_EQ(kFormatVersion, version); + EXPECT_EQ(kV1TailPointerSize, encoded_size); + return tail; +} + +MetadataDirectory ReadDirectory(const std::vector& bytes, const TailFields& tail) { + MetadataDirectory directory; + EXPECT_LE(tail.directory_offset, bytes.size()); + EXPECT_LE(tail.directory_length, bytes.size() - tail.directory_offset); + EXPECT_TRUE(MetadataDirectory::decode(Slice(bytes.data() + tail.directory_offset, + static_cast(tail.directory_length)), + &directory) + .ok()); + return directory; +} + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_seg_open_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +// An in-memory FileReader over an owned byte buffer that RECORDS every read +// range. The buffer is mutable so a test can corrupt specific on-disk bytes +// before re-opening. read_batch is overridden so batched reads are recorded too +// (open() currently uses only read_at, but recording both keeps the assertion +// honest if that ever changes). +class RecordingFileReader : public io::FileReader { +public: + explicit RecordingFileReader(std::vector bytes) : bytes_(std::move(bytes)) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + reads_.push_back(io::Range {offset, len}); + if (offset > bytes_.size() || len > bytes_.size() - offset) { + return Status::Error( + "recording reader: read past EOF"); + } + out->assign(bytes_.begin() + static_cast(offset), + bytes_.begin() + static_cast(offset + len)); + return Status::OK(); + } + + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + outs->resize(ranges.size()); + for (size_t i = 0; i < ranges.size(); ++i) { + RETURN_IF_ERROR(read_at(ranges[i].offset, ranges[i].len, &(*outs)[i])); + } + return Status::OK(); + } + + uint64_t size() const override { return bytes_.size(); } + + const std::vector& reads() const { return reads_; } + std::vector& bytes() { return bytes_; } + + // True iff any recorded read overlaps [lo, hi). + bool any_read_intersects(uint64_t lo, uint64_t hi) const { + for (const auto& r : reads_) { + const uint64_t r_lo = r.offset; + const uint64_t r_hi = r.offset + r.len; + if (r_lo < hi && lo < r_hi) { + return true; + } + } + return false; + } + +private: + std::vector bytes_; + std::vector reads_; +}; + +class TailOnlyDeclaredSizeReader final : public io::FileReader { +public: + explicit TailOnlyDeclaredSizeReader(uint64_t directory_length) + : size_(directory_length + tail_pointer_size()) { + TailPointer tail; + tail.directory_offset = 0; + tail.directory_length = directory_length; + ByteSink sink; + EXPECT_TRUE(encode_tail_pointer(tail, &sink).ok()); + tail_bytes_ = sink.buffer(); + } + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (offset == size_ - tail_bytes_.size() && len == tail_bytes_.size()) { + *out = tail_bytes_; + return Status::OK(); + } + ++non_footer_reads_; + return Status::Error("unexpected non-footer read"); + } + + uint64_t size() const override { return size_; } + size_t non_footer_reads() const { return non_footer_reads_; } + +private: + uint64_t size_; + size_t non_footer_reads_ = 0; + std::vector tail_bytes_; +}; + +// Writes a minimal single-index docs+positions container and returns its bytes. +std::vector BuildContainerBytes() { + SpimiTermBuffer buf(/*has_positions=*/true); + // A tiny deterministic corpus: a couple of terms across a few docs. + const char* docs[] = {"alpha bravo", "bravo charlie", "alpha charlie delta"}; + for (uint32_t d = 0; d < 3; ++d) { + std::string s = docs[d]; + uint32_t pos = 0; + size_t start = 0; + while (start <= s.size()) { + size_t sp = s.find(' ', start); + std::string tok = + s.substr(start, sp == std::string::npos ? std::string::npos : sp - start); + if (!tok.empty()) { + buf.add_token(tok, d, pos++); + } + if (sp == std::string::npos) { + break; + } + start = sp + 1; + } + } + std::vector terms = buf.finalize_sorted(); + + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = 3; + in.terms = std::move(terms); + in.target_dict_block_bytes = 256; + + const std::string path = TempPath(); + { + io::LocalFileWriter w; + EXPECT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + EXPECT_TRUE(cw.add_logical_index(in).ok()); + EXPECT_TRUE(cw.finish().ok()); + } + + io::LocalFileReader r; + EXPECT_TRUE(r.open(path).ok()); + std::vector bytes; + EXPECT_TRUE(r.read_at(0, r.size(), &bytes).ok()); + std::remove(path.c_str()); + return bytes; +} + +} // namespace + +// Catches a metadata-layout or eager-open regression before the PB migration: +// the compact custom baseline must not grow, take extra reads, or undercharge +// the searcher cache. The literals are independently measured from the +// pre-PB container built above, not derived from reader implementation details. +TEST(SniiSegmentReaderOpen, PreservesPrePbMetadataBounds) { + constexpr size_t kCompleteImageBytesUpperBound = 454; + constexpr uint64_t kTailMetadataBytesUpperBound = 165; + constexpr size_t kSegmentOpenReadCount = 2; + constexpr size_t kLogicalIndexOpenAdditionalReadCount = 3; + constexpr size_t kCoreOnlyAdditionalReadCount = 1; + constexpr size_t kLogicalReaderMemoryUsageUpperBound = 1380; + + const std::vector bytes = BuildContainerBytes(); + EXPECT_LE(bytes.size(), kCompleteImageBytesUpperBound); + ASSERT_EQ(kV1TailPointerSize, tail_pointer_size()); + const TailFields tail = ReadTailFields(bytes); + const MetadataDirectory directory = ReadDirectory(bytes, tail); + ASSERT_EQ(1U, directory.size()); + const auto& entry = directory.entries().front(); + ASSERT_LE(entry.core_metadata.offset, tail.directory_offset); + const uint64_t tail_metadata_bytes = + tail.directory_offset + tail.directory_length - entry.core_metadata.offset; + EXPECT_LE(tail_metadata_bytes, kTailMetadataBytesUpperBound); + + RecordingFileReader query_reader(bytes); + SniiSegmentReader query_seg; + ASSERT_TRUE(SniiSegmentReader::open(&query_reader, &query_seg).ok()); + const size_t segment_open_read_count = query_reader.reads().size(); + EXPECT_EQ(segment_open_read_count, kSegmentOpenReadCount); + LogicalIndexReader query_index; + ASSERT_TRUE(query_seg.open_index(1, "body", &query_index).ok()); + const size_t logical_index_open_additional_read_count = + query_reader.reads().size() - segment_open_read_count; + EXPECT_EQ(logical_index_open_additional_read_count, kLogicalIndexOpenAdditionalReadCount); + ASSERT_GE(query_reader.reads().size(), segment_open_read_count + 1); + const auto& group_read = query_reader.reads()[segment_open_read_count]; + EXPECT_EQ(entry.core_metadata.offset, group_read.offset); + EXPECT_EQ(entry.core_metadata.length + entry.sampled_term_index.length + + entry.dict_block_directory.length, + group_read.len); + EXPECT_LE(query_index.memory_usage(), kLogicalReaderMemoryUsageUpperBound); + + RecordingFileReader core_reader(bytes); + SniiSegmentReader core_seg; + ASSERT_TRUE(SniiSegmentReader::open(&core_reader, &core_seg).ok()); + const size_t core_segment_open_read_count = core_reader.reads().size(); + SectionRefs section_refs; + ASSERT_TRUE(core_seg.section_refs_for_index(1, "body", §ion_refs).ok()); + const size_t core_only_additional_read_count = + core_reader.reads().size() - core_segment_open_read_count; + EXPECT_EQ(core_only_additional_read_count, kCoreOnlyAdditionalReadCount); + ASSERT_GE(core_reader.reads().size(), core_segment_open_read_count + 1); + const auto& core_read = core_reader.reads()[core_segment_open_read_count]; + EXPECT_EQ(entry.core_metadata.offset, core_read.offset); + EXPECT_EQ(entry.core_metadata.length, core_read.len); + EXPECT_GT(section_refs.dict_region.length, 0U); +} + +TEST(SniiSegmentReaderOpen, ReadsExactlyFooterThenRawDirectory) { + const std::vector bytes = BuildContainerBytes(); + const TailFields tail = ReadTailFields(bytes); + RecordingFileReader reader(bytes); + SniiSegmentReader segment; + ASSERT_TRUE(SniiSegmentReader::open(&reader, &segment).ok()); + + ASSERT_EQ(2U, reader.reads().size()); + EXPECT_EQ(bytes.size() - kV1TailPointerSize, reader.reads()[0].offset); + EXPECT_EQ(kV1TailPointerSize, reader.reads()[0].len); + EXPECT_EQ(tail.directory_offset, reader.reads()[1].offset); + EXPECT_EQ(tail.directory_length, reader.reads()[1].len); +} + +TEST(SniiSegmentReaderOpen, RejectsDirectoryLargerThanProtobufIntLimitBeforeReadingIt) { + const uint64_t oversized_directory = static_cast(std::numeric_limits::max()) + 1; + TailOnlyDeclaredSizeReader reader(oversized_directory); + SniiSegmentReader segment; + const Status status = SniiSegmentReader::open(&reader, &segment); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(0U, reader.non_footer_reads()); +} + +// open() must succeed and must NOT read any byte in [0, kBootstrapHeaderSize). +TEST(SniiSegmentReaderOpen, IssuesNoReadAtBootstrapRegion) { + std::vector bytes = BuildContainerBytes(); + ASSERT_GT(bytes.size(), kBootstrapHeaderSize + tail_pointer_size()); + + RecordingFileReader reader(std::move(bytes)); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&reader, &seg).ok()); + EXPECT_EQ(seg.n_logical_indexes(), 1U); + + // The container must still be usable (real, not vacuous): the logical index + // opens and reports the corpus doc count. + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + EXPECT_EQ(idx.stats().doc_count, 3U); + + // Core assertion: segment and logical-index open touched the tail, never the + // bootstrap header at the front of the file. + EXPECT_FALSE(reader.any_read_intersects(0, kBootstrapHeaderSize)) + << "open path read the offset-0 bootstrap region"; + // And it issued at least one read (otherwise the assertion above is vacuous). + EXPECT_GE(reader.reads().size(), 1U); +} + +// A corrupt offset-0 bootstrap header no longer fails open(): nothing reads it. +TEST(SniiSegmentReaderOpen, IgnoresCorruptBootstrapHeader) { + std::vector bytes = BuildContainerBytes(); + ASSERT_GE(bytes.size(), kBootstrapHeaderSize); + + RecordingFileReader reader(std::move(bytes)); + // Smash the entire bootstrap header region. + for (uint32_t i = 0; i < kBootstrapHeaderSize; ++i) { + reader.bytes()[i] = 0xFFU; + } + + SniiSegmentReader seg; + EXPECT_TRUE(SniiSegmentReader::open(&reader, &seg).ok()) + << "open() must not depend on the offset-0 bootstrap header"; +} + +// The container version gate is preserved by the tail pointer: corrupting the +// tail pointer's format_version makes open() fail. +TEST(SniiSegmentReaderOpen, RejectsCorruptTailFormatVersion) { + std::vector good = BuildContainerBytes(); + ASSERT_GE(good.size(), tail_pointer_size()); + + // Sanity: the unmodified container opens. + { + RecordingFileReader ok_reader(good); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&ok_reader, &seg).ok()); + } + + // The tail pointer is the last tail_pointer_size() bytes. Its layout is + // u32 magic, u16 format_version, ... so format_version sits at + // (size - tail_pointer_size) + 4. + RecordingFileReader bad_reader(good); + const uint64_t tp_start = bad_reader.size() - tail_pointer_size(); + const uint64_t fv_off = tp_start + 4; // skip the u32 magic + // Write a wrong, never-valid format_version (kFormatVersion is small). + bad_reader.bytes()[fv_off] = 0xFFU; + bad_reader.bytes()[fv_off + 1] = 0xFFU; + + SniiSegmentReader seg; + EXPECT_FALSE(SniiSegmentReader::open(&bad_reader, &seg).ok()) + << "open() must reject a container whose tail format_version is wrong"; +} diff --git a/be/test/storage/index/snii/snii_index_reader_count_fallback_test.cpp b/be/test/storage/index/snii/snii_index_reader_count_fallback_test.cpp new file mode 100644 index 00000000000000..48c39c0081f731 --- /dev/null +++ b/be/test/storage/index/snii/snii_index_reader_count_fallback_test.cpp @@ -0,0 +1,1628 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/check.h" +#include "common/config.h" +#include "core/column/column_const.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_string.h" +#include "cpp/sync_point.h" +#include "exprs/function/function_multi_match.h" +#include "exprs/function/match.h" +#include "io/fs/local_file_system.h" +#include "runtime/exec_env.h" +#include "runtime/runtime_state.h" +#include "runtime/thread_context.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/inverted/analyzer/custom_analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/inverted_index_cache.h" +#include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/phrase_verify_timer.h" +#include "storage/index/snii/query/query_profile.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/index/snii/snii_prx_profile.h" +// Exercise the reader router without acquiring process-global query-cache ownership. +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wkeyword-macro" +#define private public +#include "storage/index/snii/snii_index_reader.h" +#undef private +#pragma clang diagnostic pop +#include "storage/index/snii_query_test_util.h" +#include "storage/olap_common.h" +#include "storage/tablet/tablet_schema.h" +#include "util/debug_points.h" +#include "util/defer_op.h" + +namespace doris::segment_v2 { +namespace { + +using namespace doris::snii::snii_test; + +constexpr int64_t kIndexId = 31; +constexpr const char* kTestDir = "./ut_dir/snii_index_reader_count_fallback_test"; +constexpr const char* kIndexPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/positional_segment"; + +using PrxStatsSnapshot = std::array; + +struct SingleFlightFollowerCounts { + std::atomic count_consumers {0}; + std::atomic row_consumers {0}; + std::binary_semaphore* row_follower_joined = nullptr; +}; + +struct SingleFlightLeaderGate { + std::binary_semaphore leader_entered {0}; + std::binary_semaphore release_leader {0}; +}; + +thread_local bool is_count_consumer = false; + +void record_single_flight_follower(void* opaque) noexcept { + auto* counts = static_cast(opaque); + if (is_count_consumer) { + counts->count_consumers.fetch_add(1, std::memory_order_relaxed); + } else { + counts->row_consumers.fetch_add(1, std::memory_order_relaxed); + if (counts->row_follower_joined != nullptr) { + counts->row_follower_joined->release(); + } + } +} + +void block_single_flight_leader_before_compute(void* opaque) noexcept { + auto* gate = static_cast(opaque); + gate->leader_entered.release(); + gate->release_leader.acquire(); +} + +void record_single_flight_leader(void* opaque) noexcept { + static_cast*>(opaque)->fetch_add(1, std::memory_order_relaxed); +} + +void record_searcher_open(void* opaque) noexcept { + static_cast*>(opaque)->fetch_add(1, std::memory_order_relaxed); +} + +PrxStatsSnapshot prx_stats_snapshot(const OlapReaderStatistics& stats) { + return {stats.snii_stats.prx_raw_frames, stats.snii_stats.prx_zstd_frames, + stats.snii_stats.prx_pfor_frames, stats.snii_stats.prx_plaintext_bytes, + stats.snii_stats.prx_total_docs, stats.snii_stats.prx_selected_docs, + stats.snii_stats.prx_total_positions, stats.snii_stats.prx_selected_positions, + stats.snii_stats.prx_fetch_ns, stats.snii_stats.prx_decode_ns, + stats.snii_stats.prx_phrase_verify_ns}; +} + +void set_prx_stats_sentinel(OlapReaderStatistics* stats) { + stats->snii_stats.prx_raw_frames = 101; + stats->snii_stats.prx_zstd_frames = 102; + stats->snii_stats.prx_pfor_frames = 103; + stats->snii_stats.prx_plaintext_bytes = 104; + stats->snii_stats.prx_total_docs = 105; + stats->snii_stats.prx_selected_docs = 106; + stats->snii_stats.prx_total_positions = 107; + stats->snii_stats.prx_selected_positions = 108; + stats->snii_stats.prx_fetch_ns = 109; + stats->snii_stats.prx_decode_ns = 110; + stats->snii_stats.prx_phrase_verify_ns = 111; +} + +std::vector bitmap_docids(const roaring::Roaring& bitmap) { + return {bitmap.begin(), bitmap.end()}; +} + +struct QueryExecutionContext { + explicit QueryExecutionContext(bool enable_query_cache, bool count_on_index_fastpath = false) { + TQueryOptions query_options; + query_options.query_type = TQueryType::SELECT; + query_options.enable_inverted_index_query_cache = enable_query_cache; + query_options.enable_inverted_index_searcher_cache = false; + runtime_state.set_query_options(query_options); + context->io_ctx = &io_ctx; + context->stats = &stats; + context->runtime_state = &runtime_state; + context->count_on_index_fastpath = count_on_index_fastpath; + } + + OlapReaderStatistics stats; + io::IOContext io_ctx; + RuntimeState runtime_state; + IndexQueryContextPtr context = std::make_shared(); +}; + +struct CommonGramsCounterSnapshot { + int64_t candidate_queries = 0; + int64_t plain_plans = 0; + int64_t fallback_no_gram = 0; + int64_t fallback_incompatible = 0; + int64_t fallback_cost = 0; +}; + +CommonGramsCounterSnapshot common_grams_counter_snapshot(const OlapReaderStatistics& stats) { + return {.candidate_queries = stats.snii_stats.common_grams_candidate_queries, + .plain_plans = stats.snii_stats.common_grams_plain_plans, + .fallback_no_gram = stats.snii_stats.common_grams_fallback_no_gram, + .fallback_incompatible = stats.snii_stats.common_grams_fallback_incompatible, + .fallback_cost = stats.snii_stats.common_grams_fallback_cost}; +} + +void init_index_meta(TabletIndex* meta, int64_t index_id = kIndexId, + std::string parser = "english") { + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(index_id); + pb.set_index_name("count_fallback_idx"); + pb.add_col_unique_id(0); + pb.mutable_properties()->insert({"parser", std::move(parser)}); + pb.mutable_properties()->insert({"lower_case", "true"}); + pb.mutable_properties()->insert({"support_phrase", "true"}); + meta->init_from_pb(pb); +} + +void write_positional_segment() { + std::vector terms { + make_term("failed", {{.docid = 0, .positions = {0}}, + {.docid = 1, .positions = {0}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {0}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {0}}}), + make_term("order", {{.docid = 0, .positions = {1}}, + {.docid = 1, .positions = {0}}, + {.docid = 3, .positions = {1}}, + {.docid = 4, .positions = {0}}, + {.docid = 5, .positions = {1}}}), + make_term("ordered", {{.docid = 2, .positions = {1}}}), + }; + + doris::snii::writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = ""; + input.config = doris::snii::format::IndexConfig::kDocsPositions; + input.doc_count = 6; + input.terms = std::move(terms); + input.target_dict_block_bytes = 64; + + MemoryFile memory_file; + doris::snii::writer::SniiCompoundWriter compound(&memory_file); + assert_ok(compound.add_logical_index(input)); + assert_ok(compound.finish()); + + const std::string path = InvertedIndexDescriptor::get_index_file_path_v2(kIndexPathPrefix); + doris::snii::io::LocalFileWriter file; + assert_ok(file.open(path)); + assert_ok( + file.append(doris::snii::Slice(memory_file.data().data(), memory_file.data().size()))); + assert_ok(file.finalize()); +} + +std::shared_ptr make_common_grams_provider() { + inverted_index::Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + inverted_index::CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", tokenizer_settings); + builder.add_token_filter_config("lowercase", {}); + builder.add_token_filter_config("common_grams", {}); + return std::make_shared(builder.build()); +} + +std::shared_ptr make_plain_provider() { + inverted_index::Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + inverted_index::CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", tokenizer_settings); + builder.add_token_filter_config("lowercase", {}); + return std::make_shared(builder.build()); +} + +std::string encode_plain_test_term(std::string_view term) { + auto encoded = inverted_index::encode_plain_term( + term, inverted_index::PlainTermKeyVersion::kEscapedV1); + DORIS_CHECK(encoded.has_value()); + return std::move(*encoded); +} + +std::string encode_gram_test_term(std::string_view left, std::string_view right) { + auto encoded = inverted_index::encode_common_gram(left, right); + DORIS_CHECK(encoded.has_value()); + return std::move(*encoded); +} + +Status write_common_grams_segment(std::string_view index_path_prefix, + const inverted_index::CommonGramsQueryIdentity& analyzer_identity, + inverted_index::CommonGramsCoverage coverage, bool include_gram, + uint32_t dense_doc_count = 0, bool dense_common_pair = false) { + std::vector alpha_docs {{.docid = 1, .positions = {0}}}; + std::vector beta_docs {{.docid = 1, .positions = {1}}}; + std::vector the_docs {{.docid = 0, .positions = {0}}}; + std::vector wolf_docs {{.docid = 0, .positions = {1}}}; + if (dense_doc_count != 0) { + alpha_docs.clear(); + beta_docs.clear(); + alpha_docs.reserve(dense_doc_count); + beta_docs.reserve(dense_doc_count); + for (uint32_t docid = 0; docid < dense_doc_count; ++docid) { + alpha_docs.push_back({.docid = docid, .positions = {0}}); + beta_docs.push_back({.docid = docid, .positions = {1}}); + } + if (dense_common_pair) { + // Make "the wolf" an adjacent phrase in EVERY doc with dense postings + // for both terms, mirroring the dense alpha/beta shape whose plain + // execution provably reads postings through the adapter. This is the + // planning-entry variant: "the" is a common word, so the wordset + // pre-proof cannot rule gram usage out. Callers must combine this + // with include_gram=true: omitting the gram under kComplete coverage + // would let the planner prove the phrase authoritatively empty from + // the resident dictionary alone and skip posting IO entirely. + the_docs.clear(); + wolf_docs.clear(); + the_docs.reserve(dense_doc_count); + wolf_docs.reserve(dense_doc_count); + for (uint32_t docid = 0; docid < dense_doc_count; ++docid) { + the_docs.push_back({.docid = docid, .positions = {0}}); + wolf_docs.push_back({.docid = docid, .positions = {1}}); + } + } + } + std::vector terms { + make_term(encode_plain_test_term("alpha"), std::move(alpha_docs)), + make_term(encode_plain_test_term("beta"), std::move(beta_docs)), + make_term(encode_plain_test_term("the"), std::move(the_docs)), + make_term(encode_plain_test_term("wolf"), std::move(wolf_docs)), + }; + if (include_gram) { + std::vector gram_docs {{.docid = 0, .positions = {0}}}; + if (dense_doc_count != 0 && dense_common_pair) { + // Keep the gram postings consistent with the dense the@0/wolf@1 + // docs above: the_wolf occurs at position 0 in every doc. + gram_docs.clear(); + gram_docs.reserve(dense_doc_count); + for (uint32_t docid = 0; docid < dense_doc_count; ++docid) { + gram_docs.push_back({.docid = docid, .positions = {0}}); + } + } + terms.push_back(make_term(encode_gram_test_term("the", "wolf"), std::move(gram_docs))); + } + std::ranges::sort(terms, [](const auto& lhs, const auto& rhs) { return lhs.term < rhs.term; }); + + inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = inverted_index::PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = coverage; + metadata.common_grams_semantics_version = inverted_index::COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = inverted_index::COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = analyzer_identity.common_grams_dictionary_identity; + metadata.base_analyzer_fingerprint = analyzer_identity.base_analyzer_fingerprint; + metadata.common_grams_fingerprint = analyzer_identity.common_grams_fingerprint; + + doris::snii::writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = ""; + input.config = doris::snii::format::IndexConfig::kDocsPositions; + input.doc_count = dense_doc_count == 0 ? 2 : dense_doc_count * 4; + input.terms = std::move(terms); + input.target_dict_block_bytes = 64; + input.common_grams_metadata = std::move(metadata); + + MemoryFile memory_file; + doris::snii::writer::SniiCompoundWriter compound(&memory_file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + + doris::snii::io::LocalFileWriter local_file; + RETURN_IF_ERROR(local_file.open( + InvertedIndexDescriptor::get_index_file_path_v2(std::string(index_path_prefix)))); + RETURN_IF_ERROR(local_file.append( + doris::snii::Slice(memory_file.data().data(), memory_file.data().size()))); + return local_file.finalize(); +} + +Status write_scoring_segment(std::string_view index_path_prefix, + const inverted_index::CommonGramsQueryIdentity& analyzer_identity, + bool corrupt_norms = false) { + std::vector terms { + make_term(encode_plain_test_term("alpha"), + {{.docid = 0, .positions = {0, 1, 2, 3}}, {.docid = 1, .positions = {0, 2}}}), + make_term(encode_plain_test_term("beta"), + {{.docid = 1, .positions = {1, 3}}, {.docid = 2, .positions = {0, 1}}}), + }; + std::ranges::sort(terms, [](const auto& lhs, const auto& rhs) { return lhs.term < rhs.term; }); + + auto metadata = inverted_index::make_common_grams_segment_metadata(analyzer_identity); + metadata.scoring_doc_count = 3; + metadata.scoring_token_count = 10; + + doris::snii::writer::SniiIndexInput input; + input.index_id = kIndexId; + input.index_suffix = ""; + input.config = doris::snii::format::IndexConfig::kDocsPositionsScoring; + input.doc_count = 3; + input.encoded_norms = {doris::snii::query::encode_norm(4), doris::snii::query::encode_norm(4), + doris::snii::query::encode_norm(2)}; + input.terms = std::move(terms); + input.target_dict_block_bytes = 64; + input.common_grams_metadata = std::move(metadata); + + MemoryFile memory_file; + doris::snii::writer::SniiCompoundWriter compound(&memory_file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + + std::vector bytes = memory_file.data(); + if (corrupt_norms) { + doris::snii::reader::SniiSegmentReader segment_reader; + RETURN_IF_ERROR( + doris::snii::reader::SniiSegmentReader::open(&memory_file, &segment_reader)); + doris::snii::reader::LogicalIndexReader logical_reader; + RETURN_IF_ERROR(segment_reader.open_index(kIndexId, "", &logical_reader)); + const auto norms = logical_reader.section_refs().norms; + DORIS_CHECK_GT(norms.length, 0); + DORIS_CHECK_LE(norms.offset + norms.length, bytes.size()); + bytes[norms.offset + norms.length - 1] ^= 0x01; + } + + doris::snii::io::LocalFileWriter local_file; + RETURN_IF_ERROR(local_file.open( + InvertedIndexDescriptor::get_index_file_path_v2(std::string(index_path_prefix)))); + RETURN_IF_ERROR(local_file.append(doris::snii::Slice(bytes.data(), bytes.size()))); + return local_file.finalize(); +} + +Status write_nullable_phrase_segment(std::string_view index_path_prefix, bool has_null, + int64_t index_id = kIndexId, bool corrupt_null = false) { + std::vector terms { + make_term("alpha", {{.docid = 0, .positions = {0}}, + {.docid = 2, .positions = {0}}, + {.docid = 3, .positions = {1}}}), + make_term("beta", {{.docid = 0, .positions = {1}}, {.docid = 3, .positions = {0}}}), + make_term("betamax", {{.docid = 2, .positions = {1}}}), + }; + + doris::snii::writer::SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = ""; + input.config = doris::snii::format::IndexConfig::kDocsPositions; + input.doc_count = 4; + input.terms = std::move(terms); + input.target_dict_block_bytes = 64; + if (has_null) { + input.null_docids = {1}; + } + + MemoryFile memory_file; + doris::snii::writer::SniiCompoundWriter compound(&memory_file); + RETURN_IF_ERROR(compound.add_logical_index(input)); + RETURN_IF_ERROR(compound.finish()); + + std::vector bytes = memory_file.data(); + if (corrupt_null) { + doris::snii::reader::SniiSegmentReader segment_reader; + RETURN_IF_ERROR( + doris::snii::reader::SniiSegmentReader::open(&memory_file, &segment_reader)); + doris::snii::reader::LogicalIndexReader logical_reader; + RETURN_IF_ERROR(segment_reader.open_index(index_id, "", &logical_reader)); + const auto null_bitmap = logical_reader.section_refs().null_bitmap; + DORIS_CHECK_GT(null_bitmap.length, 0); + DORIS_CHECK_LE(null_bitmap.offset + null_bitmap.length, bytes.size()); + bytes[null_bitmap.offset + null_bitmap.length - 1] ^= 0x01; + } + + doris::snii::io::LocalFileWriter local_file; + RETURN_IF_ERROR(local_file.open( + InvertedIndexDescriptor::get_index_file_path_v2(std::string(index_path_prefix)))); + RETURN_IF_ERROR(local_file.append(doris::snii::Slice(bytes.data(), bytes.size()))); + return local_file.finalize(); +} + +class FixedCollectionStatistics final : public CollectionStatistics { +public: + float get_or_calculate_idf(const std::wstring& field_name, const std::wstring& term) override { + fields.push_back(field_name); + idf_terms.push_back(term); + return idfs.at(term); + } + + float get_or_calculate_avg_dl(const std::wstring& field_name) override { + fields.push_back(field_name); + return avgdl; + } + + float avgdl = 3.0F; + std::unordered_map idfs {{L"alpha", 1.0F}, {L"beta", 2.0F}}; + std::vector fields; + std::vector idf_terms; +}; + +std::vector positive_score_docids(const CollectionSimilarity& similarity, + std::initializer_list candidates) { + roaring::Roaring bitmap; + for (uint32_t docid : candidates) { + bitmap.add(docid); + } + IColumn::MutablePtr scores; + auto row_ids = std::make_unique>(); + auto positive = std::make_shared( + ScoreRangeFilter {.op = TExprOpcode::GT, .threshold = 0.0}); + similarity.get_bm25_scores(&bitmap, scores, row_ids, positive); + return bitmap_docids(bitmap); +} + +float score_for_doc(const CollectionSimilarity& similarity, uint32_t docid) { + roaring::Roaring bitmap; + bitmap.add(docid); + IColumn::MutablePtr scores; + auto row_ids = std::make_unique>(); + similarity.get_bm25_scores(&bitmap, scores, row_ids); + DORIS_CHECK_EQ(row_ids->size(), 1); + const auto& nullable = assert_cast(*scores); + const auto& values = assert_cast(nullable.get_nested_column()); + return values.get_data()[0]; +} + +struct OpenedSniiIndex { + std::shared_ptr file_reader; + std::shared_ptr index_reader; +}; + +Status open_snii_index(const TabletIndex* meta, std::string index_path_prefix, + OpenedSniiIndex* opened) { + opened->file_reader = std::make_shared(io::global_local_filesystem(), + std::move(index_path_prefix), + InvertedIndexStorageFormatPB::SNII); + RETURN_IF_ERROR(opened->file_reader->init()); + opened->index_reader = SniiIndexReader::create_shared(meta, opened->file_reader, + InvertedIndexReaderType::FULLTEXT); + return Status::OK(); +} + +bool has_cached_null_bitmap(const TabletIndex& meta, const OpenedSniiIndex& opened) { + const auto index_file_key = opened.file_reader->get_index_file_cache_key(&meta); + InvertedIndexQueryCache::CacheKey cache_key { + index_file_key, "", InvertedIndexQueryType::UNKNOWN_QUERY, "null_bitmap"}; + InvertedIndexQueryCacheHandle cache_handle; + return InvertedIndexQueryCache::instance()->lookup(cache_key, &cache_handle); +} + +uint32_t lookup_df(const doris::snii::reader::LogicalIndexReader& index, const std::string& term) { + bool found = false; + doris::snii::format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup(term, &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + return entry.df; +} + +class SniiIndexReaderCountFallback : public testing::Test { +protected: + void SetUp() override { + assert_ok(io::global_local_filesystem()->delete_directory(kTestDir)); + assert_ok(io::global_local_filesystem()->create_directory(kTestDir)); + init_index_meta(&_meta); + write_positional_segment(); + _file_reader = + std::make_shared(io::global_local_filesystem(), kIndexPathPrefix, + InvertedIndexStorageFormatPB::SNII); + assert_ok(_file_reader->init()); + _index_reader = SniiIndexReader::create_shared(&_meta, _file_reader, + InvertedIndexReaderType::FULLTEXT); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); + _query_cache.reset(InvertedIndexQueryCache::create_global_cache(1024 * 1024, 1)); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_query_cache.get()); + } + + void TearDown() override { + _index_reader.reset(); + _file_reader.reset(); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); + _query_cache.reset(); + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); + } + + TabletIndex _meta; + std::shared_ptr _file_reader; + std::shared_ptr _index_reader; + InvertedIndexQueryCache* _previous_query_cache = nullptr; + std::unique_ptr _query_cache; +}; + +template +void expect_function_match_reuses_reader(const TabletIndex& meta, std::string_view path_prefix, + std::string_view query, bool has_null, + std::vector expected_docids) { + assert_ok(write_nullable_phrase_segment(path_prefix, has_null)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&meta, std::string(path_prefix), &opened)); + + std::atomic searcher_opens {0}; + opened.index_reader->set_searcher_open_observer_for_test(record_searcher_open, &searcher_opens); + + QueryExecutionContext execution(/*enable_query_cache=*/false); + std::unique_ptr iterator; + assert_ok(opened.index_reader->new_iterator(&iterator)); + iterator->set_context(execution.context); + + auto query_column = ColumnString::create(); + query_column->insert_data(query.data(), query.size()); + ColumnsWithTypeAndName arguments {{ColumnConst::create(std::move(query_column), 1), + std::make_shared(), "query"}}; + std::vector data_type_with_names { + {"content", std::make_shared()}}; + std::vector iterators {iterator.get()}; + + MatchFunction function; + InvertedIndexResultBitmap result; + assert_ok(function.evaluate_inverted_index(arguments, data_type_with_names, iterators, + /*num_rows=*/4, nullptr, result)); + + ASSERT_NE(result.get_data_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_data_bitmap()), expected_docids); + ASSERT_NE(result.get_null_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_null_bitmap()), + has_null ? std::vector {1} : std::vector {}); + EXPECT_EQ(searcher_opens.load(std::memory_order_relaxed), 1); +} + +TEST_F(SniiIndexReaderCountFallback, PhraseQueryReusesReaderForNullBitmap) { + expect_function_match_reuses_reader( + _meta, std::string(kTestDir) + "/phrase_nullable", "alpha beta", /*has_null=*/true, + {0}); + expect_function_match_reuses_reader( + _meta, std::string(kTestDir) + "/phrase_null_free", "alpha beta", /*has_null=*/false, + {0}); +} + +TEST_F(SniiIndexReaderCountFallback, PhrasePrefixQueryReusesReaderForNullBitmap) { + expect_function_match_reuses_reader( + _meta, std::string(kTestDir) + "/prefix_nullable", "alpha be", /*has_null=*/true, + {0, 2}); + expect_function_match_reuses_reader( + _meta, std::string(kTestDir) + "/prefix_null_free", "alpha be", /*has_null=*/false, + {0, 2}); +} + +TEST_F(SniiIndexReaderCountFallback, FunctionMatchUsesQuerySelectedReaderForNullBitmap) { + constexpr int64_t kSelectedIndexId = kIndexId + 1; + const std::string decoy_path = std::string(kTestDir) + "/affinity_decoy"; + const std::string selected_path = std::string(kTestDir) + "/affinity_selected"; + TabletIndex decoy_meta; + TabletIndex selected_meta; + init_index_meta(&decoy_meta, kIndexId, "standard"); + init_index_meta(&selected_meta, kSelectedIndexId, "english"); + assert_ok(write_nullable_phrase_segment(decoy_path, /*has_null=*/true, kIndexId)); + assert_ok(write_nullable_phrase_segment(selected_path, /*has_null=*/false, kSelectedIndexId)); + + OpenedSniiIndex decoy; + OpenedSniiIndex selected; + assert_ok(open_snii_index(&decoy_meta, decoy_path, &decoy)); + assert_ok(open_snii_index(&selected_meta, selected_path, &selected)); + std::atomic decoy_opens {0}; + std::atomic selected_opens {0}; + decoy.index_reader->set_searcher_open_observer_for_test(record_searcher_open, &decoy_opens); + selected.index_reader->set_searcher_open_observer_for_test(record_searcher_open, + &selected_opens); + + QueryExecutionContext execution(/*enable_query_cache=*/false); + std::unique_ptr iterator; + assert_ok(decoy.index_reader->new_iterator(&iterator)); + assert_ok(selected.index_reader->new_iterator(&iterator)); + iterator->set_context(execution.context); + + auto query_column = ColumnString::create(); + query_column->insert_data("alpha beta", 10); + ColumnsWithTypeAndName arguments {{ColumnConst::create(std::move(query_column), 1), + std::make_shared(), "query"}}; + std::vector data_type_with_names { + {"content", std::make_shared()}}; + std::vector iterators {iterator.get()}; + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.analyzer_name = "english"; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + + FunctionMatchPhrase function; + InvertedIndexResultBitmap result; + assert_ok(function.evaluate_inverted_index(arguments, data_type_with_names, iterators, + /*num_rows=*/4, &analyzer_ctx, result)); + + ASSERT_NE(result.get_data_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_data_bitmap()), (std::vector {0})); + ASSERT_NE(result.get_null_bitmap(), nullptr); + EXPECT_TRUE(result.get_null_bitmap()->isEmpty()); + EXPECT_EQ(selected_opens.load(std::memory_order_relaxed), 1); + EXPECT_EQ(decoy_opens.load(std::memory_order_relaxed), 0); + EXPECT_TRUE(has_cached_null_bitmap(selected_meta, selected)); + EXPECT_FALSE(has_cached_null_bitmap(decoy_meta, decoy)); +} + +TEST_F(SniiIndexReaderCountFallback, FunctionMultiMatchDoesNotCacheSniiNullBitmap) { + const std::string path = std::string(kTestDir) + "/multi_match_no_null"; + assert_ok(write_nullable_phrase_segment(path, /*has_null=*/true)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, path, &opened)); + std::atomic searcher_opens {0}; + opened.index_reader->set_searcher_open_observer_for_test(record_searcher_open, &searcher_opens); + + QueryExecutionContext execution(/*enable_query_cache=*/false); + std::unique_ptr iterator; + assert_ok(opened.index_reader->new_iterator(&iterator)); + iterator->set_context(execution.context); + + auto type_column = ColumnString::create(); + type_column->insert_data("phrase", 6); + auto query_column = ColumnString::create(); + query_column->insert_data("alpha beta", 10); + ColumnsWithTypeAndName arguments {{ColumnConst::create(std::move(type_column), 1), + std::make_shared(), "type"}, + {ColumnConst::create(std::move(query_column), 1), + std::make_shared(), "query"}}; + std::vector data_type_with_names { + {"content", std::make_shared()}}; + std::vector iterators {iterator.get()}; + + FunctionMultiMatch function; + InvertedIndexResultBitmap result; + assert_ok(function.evaluate_inverted_index(arguments, data_type_with_names, iterators, + /*num_rows=*/4, nullptr, result)); + + ASSERT_NE(result.get_data_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_data_bitmap()), (std::vector {0})); + ASSERT_NE(result.get_null_bitmap(), nullptr); + EXPECT_TRUE(result.get_null_bitmap()->isEmpty()); + EXPECT_EQ(searcher_opens.load(std::memory_order_relaxed), 1); + EXPECT_FALSE(has_cached_null_bitmap(_meta, opened)); + EXPECT_EQ(execution.stats.inverted_index_query_null_bitmap_timer, 0); +} + +TEST_F(SniiIndexReaderCountFallback, FunctionMatchCacheHitLoadsNullFromSelectedReaderOnce) { + const std::string path = std::string(kTestDir) + "/function_match_cache_hit"; + assert_ok(write_nullable_phrase_segment(path, /*has_null=*/true)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, path, &opened)); + std::atomic searcher_opens {0}; + opened.index_reader->set_searcher_open_observer_for_test(record_searcher_open, &searcher_opens); + + const Field query_value = Field::create_field(std::string("alpha beta")); + QueryExecutionContext populate(/*enable_query_cache=*/true); + std::shared_ptr populated_bitmap; + assert_ok(opened.index_reader->query(populate.context, "content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + populated_bitmap)); + EXPECT_FALSE(has_cached_null_bitmap(_meta, opened)); + searcher_opens.store(0, std::memory_order_relaxed); + + QueryExecutionContext cache_hit(/*enable_query_cache=*/true); + std::unique_ptr iterator; + assert_ok(opened.index_reader->new_iterator(&iterator)); + iterator->set_context(cache_hit.context); + auto query_column = ColumnString::create(); + query_column->insert_data("alpha beta", 10); + ColumnsWithTypeAndName arguments {{ColumnConst::create(std::move(query_column), 1), + std::make_shared(), "query"}}; + std::vector data_type_with_names { + {"content", std::make_shared()}}; + std::vector iterators {iterator.get()}; + + FunctionMatchPhrase function; + InvertedIndexResultBitmap result; + assert_ok(function.evaluate_inverted_index(arguments, data_type_with_names, iterators, + /*num_rows=*/4, nullptr, result)); + + ASSERT_NE(result.get_data_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_data_bitmap()), (std::vector {0})); + ASSERT_NE(result.get_null_bitmap(), nullptr); + EXPECT_EQ(bitmap_docids(*result.get_null_bitmap()), (std::vector {1})); + EXPECT_EQ(cache_hit.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(searcher_opens.load(std::memory_order_relaxed), 1); + EXPECT_TRUE(has_cached_null_bitmap(_meta, opened)); +} + +TEST_F(SniiIndexReaderCountFallback, PublicPhraseQueryLeaderRecordsPrxWork) { + QueryExecutionContext execution(/*enable_query_cache=*/false); + Field query_value = Field::create_field(std::string("failed order")); + std::shared_ptr bitmap; + + assert_ok(_index_reader->query(execution.context, "content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap)); + + ASSERT_NE(bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*bitmap), (std::vector {0, 3, 5})); + EXPECT_GT(execution.stats.snii_stats.prx_raw_frames + + execution.stats.snii_stats.prx_zstd_frames + + execution.stats.snii_stats.prx_pfor_frames, + 0); + EXPECT_GT(execution.stats.snii_stats.prx_plaintext_bytes, 0); + EXPECT_GT(execution.stats.snii_stats.prx_total_docs, 0); +} + +TEST_F(SniiIndexReaderCountFallback, PublicBooleanQueriesScoreOnlyFinalBitmapWithPlainTerms) { + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/scoring_segment"; + assert_ok(write_scoring_segment(kPathPrefix, *provider->common_grams_identity())); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + const Field query_value = Field::create_field(std::string("alpha beta")); + + auto all_stats = std::make_shared(); + QueryExecutionContext match_all(/*enable_query_cache=*/true); + match_all.context->collection_statistics = all_stats; + match_all.context->collection_similarity = std::make_shared(); + std::shared_ptr all_bitmap; + assert_ok(opened.index_reader->query(match_all.context, "content", query_value, + InvertedIndexQueryType::MATCH_ALL_QUERY, all_bitmap, + &analyzer_ctx)); + ASSERT_NE(all_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*all_bitmap), (std::vector {1})); + EXPECT_EQ(positive_score_docids(*match_all.context->collection_similarity, {0, 1, 2}), + (std::vector {1})); + EXPECT_EQ(all_stats->idf_terms, (std::vector {L"alpha", L"beta"})); + EXPECT_EQ(all_stats->fields, (std::vector {L"content", L"content", L"content"})); + EXPECT_EQ(match_all.stats.inverted_index_query_cache_lookup, 0); + + auto any_stats = std::make_shared(); + QueryExecutionContext match_any(/*enable_query_cache=*/true); + match_any.context->collection_statistics = any_stats; + match_any.context->collection_similarity = std::make_shared(); + std::shared_ptr any_bitmap; + assert_ok(opened.index_reader->query(match_any.context, "content", query_value, + InvertedIndexQueryType::MATCH_ANY_QUERY, any_bitmap, + &analyzer_ctx)); + ASSERT_NE(any_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*any_bitmap), (std::vector {0, 1, 2})); + EXPECT_EQ(positive_score_docids(*match_any.context->collection_similarity, {0, 1, 2}), + (std::vector {0, 1, 2})); + EXPECT_EQ(any_stats->idf_terms, (std::vector {L"alpha", L"beta"})); + EXPECT_EQ(any_stats->fields, (std::vector {L"content", L"content", L"content"})); + EXPECT_EQ(match_any.stats.inverted_index_query_cache_lookup, 0); +} + +TEST_F(SniiIndexReaderCountFallback, PublicScoringZeroHitDoesNotLoadNorms) { + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/scoring_zero_hit"; + assert_ok(write_scoring_segment(kPathPrefix, *provider->common_grams_identity(), + /*corrupt_norms=*/true)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + QueryExecutionContext execution(/*enable_query_cache=*/true); + execution.context->collection_statistics = std::make_shared(); + execution.context->collection_similarity = std::make_shared(); + const Field query_value = Field::create_field(std::string("missing")); + std::shared_ptr bitmap; + + const Status status = opened.index_reader->query( + execution.context, "scoring_zero_hit_content", query_value, + InvertedIndexQueryType::MATCH_ANY_QUERY, bitmap, &analyzer_ctx); + + ASSERT_TRUE(status.ok()) << status; + ASSERT_NE(bitmap, nullptr); + EXPECT_TRUE(bitmap->isEmpty()); +} + +TEST_F(SniiIndexReaderCountFallback, PublicPhraseQueriesScoreOccurrenceFrequencyWithPlainTerms) { + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/scoring_phrase"; + assert_ok(write_scoring_segment(kPathPrefix, *provider->common_grams_identity())); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + + QueryExecutionContext exact(/*enable_query_cache=*/true); + auto exact_stats = std::make_shared(); + exact.context->collection_statistics = exact_stats; + exact.context->collection_similarity = std::make_shared(); + const Field exact_value = Field::create_field(std::string("alpha beta")); + std::shared_ptr exact_bitmap; + assert_ok(opened.index_reader->query(exact.context, "scoring_phrase_content", exact_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, exact_bitmap, + &analyzer_ctx)); + ASSERT_NE(exact_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*exact_bitmap), (std::vector {1})); + EXPECT_EQ(exact_stats->idf_terms, (std::vector {L"alpha", L"beta"})); + EXPECT_EQ(exact.stats.inverted_index_query_cache_lookup, 0); + const double expected_exact = doris::snii::query::ScorerContext::from_idf(3.0).score( + 2, doris::snii::query::encode_norm(4), 3.0, doris::snii::query::Bm25Params {}); + EXPECT_FLOAT_EQ(score_for_doc(*exact.context->collection_similarity, 1), + static_cast(expected_exact)); + + QueryExecutionContext prefix(/*enable_query_cache=*/true); + auto prefix_stats = std::make_shared(); + prefix.context->collection_statistics = prefix_stats; + prefix.context->collection_similarity = std::make_shared(); + const Field prefix_value = Field::create_field(std::string("alpha be")); + std::shared_ptr prefix_bitmap; + assert_ok(opened.index_reader->query(prefix.context, "scoring_phrase_content", prefix_value, + InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, + prefix_bitmap, &analyzer_ctx)); + ASSERT_NE(prefix_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*prefix_bitmap), (std::vector {1})); + EXPECT_EQ(prefix_stats->idf_terms, (std::vector {L"alpha"})); + EXPECT_EQ(prefix.stats.inverted_index_query_cache_lookup, 0); + const double expected_prefix = doris::snii::query::ScorerContext::from_idf(1.0).score( + 2, doris::snii::query::encode_norm(4), 3.0, doris::snii::query::Bm25Params {}); + EXPECT_FLOAT_EQ(score_for_doc(*prefix.context->collection_similarity, 1), + static_cast(expected_prefix)); + + QueryExecutionContext single_prefix(/*enable_query_cache=*/true); + single_prefix.context->collection_statistics = std::make_shared(); + single_prefix.context->collection_similarity = std::make_shared(); + const Field single_prefix_value = Field::create_field(std::string("be")); + std::shared_ptr single_prefix_bitmap; + const Status single_prefix_status = opened.index_reader->query( + single_prefix.context, "scoring_phrase_content", single_prefix_value, + InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, single_prefix_bitmap, &analyzer_ctx); + EXPECT_EQ(single_prefix_status.code(), ErrorCode::INVERTED_INDEX_NOT_SUPPORTED) + << single_prefix_status; +} + +TEST_F(SniiIndexReaderCountFallback, PublicPhraseQueryCacheHitLeavesPrxStatsUnchanged) { + QueryExecutionContext leader(/*enable_query_cache=*/true); + Field query_value = Field::create_field(std::string("failed order")); + std::shared_ptr leader_bitmap; + assert_ok(_index_reader->query(leader.context, "cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, leader_bitmap)); + EXPECT_GT(leader.stats.snii_stats.prx_plaintext_bytes, 0); + + QueryExecutionContext cache_hit(/*enable_query_cache=*/true); + set_prx_stats_sentinel(&cache_hit.stats); + const PrxStatsSnapshot before = prx_stats_snapshot(cache_hit.stats); + std::shared_ptr cached_bitmap; + assert_ok(_index_reader->query(cache_hit.context, "cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, cached_bitmap)); + + ASSERT_NE(cached_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*cached_bitmap), bitmap_docids(*leader_bitmap)); + EXPECT_EQ(cache_hit.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(prx_stats_snapshot(cache_hit.stats), before); +} + +TEST_F(SniiIndexReaderCountFallback, NonCommonGramsAnalyzerRetainsAnalyzedTermCache) { + inverted_index::Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + inverted_index::CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", tokenizer_settings); + builder.add_token_filter_config("lowercase", {}); + auto provider = std::make_shared(builder.build()); + ASSERT_FALSE(provider->uses_common_grams()); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = std::move(provider); + const Field query_value = Field::create_field(std::string("FAILED ORDER")); + + QueryExecutionContext first(/*enable_query_cache=*/true); + std::shared_ptr first_bitmap; + assert_ok(_index_reader->query(first.context, "non_cg_cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, first_bitmap, + &analyzer_ctx)); + ASSERT_NE(first_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*first_bitmap), (std::vector {0, 3, 5})); + EXPECT_EQ(first.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(first.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(first.stats.inverted_index_query_cache_insert, 1); + + QueryExecutionContext second(/*enable_query_cache=*/true); + std::shared_ptr second_bitmap; + assert_ok(_index_reader->query(second.context, "non_cg_cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, second_bitmap, + &analyzer_ctx)); + ASSERT_NE(second_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*second_bitmap), bitmap_docids(*first_bitmap)); + EXPECT_EQ(second.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(second.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(second.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(second.stats.inverted_index_query_cache_insert, 0); +} + +TEST_F(SniiIndexReaderCountFallback, PublicQueryWithCacheDisabledDoesNotLookupOrInsert) { + Field query_value = Field::create_field(std::string("failed order")); + + QueryExecutionContext disabled(/*enable_query_cache=*/false); + std::shared_ptr disabled_bitmap; + assert_ok(_index_reader->query(disabled.context, "disabled_cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, disabled_bitmap)); + + ASSERT_NE(disabled_bitmap, nullptr); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_lookup, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_insert, 0); + + QueryExecutionContext enabled_miss(/*enable_query_cache=*/true); + std::shared_ptr enabled_bitmap; + assert_ok(_index_reader->query(enabled_miss.context, "disabled_cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, enabled_bitmap)); + + ASSERT_NE(enabled_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*enabled_bitmap), bitmap_docids(*disabled_bitmap)); + EXPECT_EQ(enabled_miss.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(enabled_miss.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(enabled_miss.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(enabled_miss.stats.inverted_index_query_cache_insert, 1); + + QueryExecutionContext enabled_hit(/*enable_query_cache=*/true); + std::shared_ptr cached_bitmap; + assert_ok(_index_reader->query(enabled_hit.context, "disabled_cache_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, cached_bitmap)); + + ASSERT_NE(cached_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*cached_bitmap), bitmap_docids(*disabled_bitmap)); + EXPECT_EQ(enabled_hit.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(enabled_hit.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(enabled_hit.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(enabled_hit.stats.inverted_index_query_cache_insert, 0); +} + +TEST_F(SniiIndexReaderCountFallback, + AuthoritativeEmptyCacheIsBypassedWhenKillSwitchForcesPlainPlan) { + const bool original_enabled = config::enable_common_grams_query_plan; + Defer restore_config([original_enabled] { + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original_enabled ? "true" : "false", + /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/authoritative_empty"; + assert_ok(write_common_grams_segment(kPathPrefix, *provider->common_grams_identity(), + inverted_index::CommonGramsCoverage::kComplete, + /*include_gram=*/false)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + const Field query_value = Field::create_field(std::string("the wolf")); + + QueryExecutionContext enabled_miss(/*enable_query_cache=*/true); + std::shared_ptr empty_bitmap; + assert_ok(opened.index_reader->query(enabled_miss.context, "authoritative_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, empty_bitmap, + &analyzer_ctx)); + ASSERT_NE(empty_bitmap, nullptr); + EXPECT_TRUE(empty_bitmap->isEmpty()); + EXPECT_EQ(enabled_miss.stats.snii_stats.common_grams_candidate_queries, 1); + EXPECT_EQ(enabled_miss.stats.snii_stats.common_grams_gram_plans, 1); + EXPECT_EQ(enabled_miss.stats.snii_stats.common_grams_authoritative_empty, 1); + EXPECT_EQ(enabled_miss.stats.inverted_index_query_cache_insert, 1); + + QueryExecutionContext enabled_hit(/*enable_query_cache=*/true); + std::shared_ptr cached_empty_bitmap; + assert_ok(opened.index_reader->query(enabled_hit.context, "authoritative_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + cached_empty_bitmap, &analyzer_ctx)); + ASSERT_NE(cached_empty_bitmap, nullptr); + EXPECT_TRUE(cached_empty_bitmap->isEmpty()); + EXPECT_EQ(enabled_hit.stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(enabled_hit.stats.snii_stats.common_grams_candidate_queries, 0); + + const auto enabled_snapshot = config::common_grams_query_plan_config_snapshot(); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "false", + /*need_persist=*/false) + .ok()); + const auto disabled_snapshot = config::common_grams_query_plan_config_snapshot(); + EXPECT_GT(disabled_snapshot.cache_generation, enabled_snapshot.cache_generation); + QueryExecutionContext disabled(/*enable_query_cache=*/true); + std::shared_ptr plain_bitmap; + assert_ok(opened.index_reader->query(disabled.context, "authoritative_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, plain_bitmap, + &analyzer_ctx)); + ASSERT_NE(plain_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*plain_bitmap), (std::vector {0})); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_lookup, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_miss, 0); + EXPECT_EQ(disabled.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(disabled.stats.snii_stats.common_grams_candidate_queries, 1); + EXPECT_EQ(disabled.stats.snii_stats.common_grams_plain_plans, 1); + EXPECT_EQ(disabled.stats.snii_stats.common_grams_fallback_kill_switch, 1); +} + +TEST_F(SniiIndexReaderCountFallback, KillSwitchBypassesCachedGramResultBeforeSegmentAnalysis) { + const bool original_enabled = config::enable_common_grams_query_plan; + Defer restore_config([original_enabled] { + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original_enabled ? "true" : "false", + /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/query_safety_fence"; + assert_ok(write_common_grams_segment(kPathPrefix, *provider->common_grams_identity(), + inverted_index::CommonGramsCoverage::kComplete, + /*include_gram=*/false)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + std::atomic single_flight_leader_calls {0}; + opened.index_reader->set_single_flight_leader_before_compute_observer_for_test( + record_single_flight_leader, &single_flight_leader_calls); + + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + const Field query_value = Field::create_field(std::string("the wolf")); + + QueryExecutionContext admitted(/*enable_query_cache=*/true); + std::shared_ptr authoritative_empty; + assert_ok(opened.index_reader->query(admitted.context, "safety_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + authoritative_empty, &analyzer_ctx)); + ASSERT_NE(authoritative_empty, nullptr); + EXPECT_TRUE(authoritative_empty->isEmpty()); + EXPECT_EQ(admitted.stats.inverted_index_query_cache_insert, 1); + EXPECT_EQ(admitted.stats.snii_stats.common_grams_gram_plans, 1); + + const auto plain_provider = make_plain_provider(); + ASSERT_EQ(plain_provider->base_analyzer_fingerprint(), provider->base_analyzer_fingerprint()); + ASSERT_FALSE(plain_provider->uses_common_grams()); + InvertedIndexAnalyzerCtx plain_analyzer_ctx = analyzer_ctx; + plain_analyzer_ctx.analyzer_provider = plain_provider; + // Disabling the BE switch is the only forced-plain path now; the flip itself advances the + // local cache generation, so gram-plan cache entries become unreachable. + const auto before_disable = config::common_grams_query_plan_config_snapshot().cache_generation; + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "false", + /*need_persist=*/false) + .ok()); + EXPECT_GT(config::common_grams_query_plan_config_snapshot().cache_generation, before_disable); + + QueryExecutionContext analyzer_mismatch_disabled(/*enable_query_cache=*/true); + std::shared_ptr analyzer_mismatch_bitmap; + assert_ok(opened.index_reader->query(analyzer_mismatch_disabled.context, "safety_content", + query_value, InvertedIndexQueryType::MATCH_PHRASE_QUERY, + analyzer_mismatch_bitmap, &plain_analyzer_ctx)); + ASSERT_NE(analyzer_mismatch_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*analyzer_mismatch_bitmap), (std::vector {0})); + EXPECT_EQ(analyzer_mismatch_disabled.stats.inverted_index_query_cache_lookup, 0); + EXPECT_EQ(analyzer_mismatch_disabled.stats.inverted_index_query_cache_insert, 0); + + QueryExecutionContext forced_plain(/*enable_query_cache=*/true); + std::shared_ptr plain_bitmap; + assert_ok(opened.index_reader->query(forced_plain.context, "safety_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, plain_bitmap, + &analyzer_ctx)); + ASSERT_NE(plain_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*plain_bitmap), (std::vector {0})); + EXPECT_EQ(forced_plain.stats.inverted_index_query_cache_lookup, 0); + EXPECT_EQ(forced_plain.stats.inverted_index_query_cache_insert, 0); + EXPECT_EQ(forced_plain.stats.snii_stats.common_grams_plain_plans, 1); + EXPECT_EQ(forced_plain.stats.snii_stats.common_grams_fallback_kill_switch, 1); + EXPECT_EQ(single_flight_leader_calls.load(std::memory_order_relaxed), 1); + + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + QueryExecutionContext readmitted(/*enable_query_cache=*/true); + std::shared_ptr readmitted_bitmap; + assert_ok(opened.index_reader->query(readmitted.context, "safety_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + readmitted_bitmap, &analyzer_ctx)); + ASSERT_NE(readmitted_bitmap, nullptr); + EXPECT_TRUE(readmitted_bitmap->isEmpty()); + EXPECT_EQ(readmitted.stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(readmitted.stats.inverted_index_query_cache_hit, 0); + EXPECT_EQ(readmitted.stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(readmitted.stats.inverted_index_query_cache_insert, 1); + EXPECT_EQ(single_flight_leader_calls.load(std::memory_order_relaxed), 2); + opened.index_reader->set_single_flight_leader_before_compute_observer_for_test(nullptr, + nullptr); +} + +TEST_F(SniiIndexReaderCountFallback, PublicPlannerRecordsEveryPlainFallbackReason) { + const bool original_enabled = config::enable_common_grams_query_plan; + const int32_t original_ratio = config::common_grams_plan_cost_ratio_percent; + Defer restore_config([original_enabled, original_ratio] { + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original_enabled ? "true" : "false", + /*need_persist=*/false) + .ok()); + EXPECT_TRUE(config::set_config("common_grams_plan_cost_ratio_percent", + std::to_string(original_ratio), + /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + + const auto run_query = [&](std::string_view path_suffix, + inverted_index::CommonGramsCoverage coverage, bool include_gram, + std::string query, CommonGramsCounterSnapshot* counters, + std::vector* docids) -> Status { + const std::string path_prefix = std::string(kTestDir) + "/" + std::string(path_suffix); + RETURN_IF_ERROR(write_common_grams_segment(path_prefix, *provider->common_grams_identity(), + coverage, include_gram)); + OpenedSniiIndex opened; + RETURN_IF_ERROR(open_snii_index(&_meta, path_prefix, &opened)); + QueryExecutionContext execution(/*enable_query_cache=*/false); + std::shared_ptr bitmap; + const Field query_value = Field::create_field(std::move(query)); + RETURN_IF_ERROR(opened.index_reader->query( + execution.context, "fallback_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap, &analyzer_ctx)); + DORIS_CHECK(bitmap != nullptr); + *counters = common_grams_counter_snapshot(execution.stats); + *docids = bitmap_docids(*bitmap); + return Status::OK(); + }; + + CommonGramsCounterSnapshot no_gram; + std::vector no_gram_docids; + assert_ok(run_query("no_gram", inverted_index::CommonGramsCoverage::kComplete, + /*include_gram=*/false, "alpha beta", &no_gram, &no_gram_docids)); + EXPECT_EQ(no_gram_docids, (std::vector {1})); + EXPECT_EQ(no_gram.candidate_queries, 1); + EXPECT_EQ(no_gram.plain_plans, 1); + EXPECT_EQ(no_gram.fallback_no_gram, 1); + EXPECT_EQ(no_gram.fallback_incompatible, 0); + EXPECT_EQ(no_gram.fallback_cost, 0); + + CommonGramsCounterSnapshot incompatible; + std::vector incompatible_docids; + assert_ok(run_query("incompatible", inverted_index::CommonGramsCoverage::kMixed, + /*include_gram=*/false, "the wolf", &incompatible, &incompatible_docids)); + EXPECT_EQ(incompatible_docids, (std::vector {0})); + EXPECT_EQ(incompatible.candidate_queries, 1); + EXPECT_EQ(incompatible.plain_plans, 1); + EXPECT_EQ(incompatible.fallback_no_gram, 0); + EXPECT_EQ(incompatible.fallback_incompatible, 1); + EXPECT_EQ(incompatible.fallback_cost, 0); + + ASSERT_TRUE(config::set_config("common_grams_plan_cost_ratio_percent", "0", + /*need_persist=*/false) + .ok()); + CommonGramsCounterSnapshot cost; + std::vector cost_docids; + assert_ok(run_query("cost", inverted_index::CommonGramsCoverage::kComplete, + /*include_gram=*/true, "the wolf", &cost, &cost_docids)); + EXPECT_EQ(cost_docids, (std::vector {0})); + EXPECT_EQ(cost.candidate_queries, 1); + EXPECT_EQ(cost.plain_plans, 1); + EXPECT_EQ(cost.fallback_no_gram, 0); + EXPECT_EQ(cost.fallback_incompatible, 0); + EXPECT_EQ(cost.fallback_cost, 1); +} + +TEST_F(SniiIndexReaderCountFallback, DebugForceGramCannotBypassProcessKillSwitch) { + const bool original_enabled = config::enable_common_grams_query_plan; + const bool original_enable_debug_points = config::enable_debug_points; + Defer restore_state([original_enabled, original_enable_debug_points] { + DebugPoints::instance()->remove("snii.common_grams.force_gram_plan"); + config::enable_debug_points = original_enable_debug_points; + EXPECT_TRUE(config::set_config("enable_common_grams_query_plan", + original_enabled ? "true" : "false", + /*need_persist=*/false) + .ok()); + }); + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "true", + /*need_persist=*/false) + .ok()); + + const auto provider = make_common_grams_provider(); + ASSERT_NE(provider->common_grams_identity(), nullptr); + constexpr std::string_view kPathPrefix = + "./ut_dir/snii_index_reader_count_fallback_test/debug_plan_safety"; + assert_ok(write_common_grams_segment(kPathPrefix, *provider->common_grams_identity(), + inverted_index::CommonGramsCoverage::kComplete, + /*include_gram=*/true)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, std::string(kPathPrefix), &opened)); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_ctx.analyzer_provider = provider; + + config::enable_debug_points = true; + DebugPoints::instance()->add("snii.common_grams.force_gram_plan"); + + const Field exact_query = Field::create_field(std::string("the wolf")); + const Field prefix_query = Field::create_field(std::string("the wo")); + const auto assert_forced_plain = [&]() { + QueryExecutionContext exact_execution( + /*enable_query_cache=*/false, /*count_on_index_fastpath=*/false); + std::shared_ptr exact_bitmap; + assert_ok(opened.index_reader->query( + exact_execution.context, "debug_safety_content", exact_query, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, exact_bitmap, &analyzer_ctx)); + ASSERT_NE(exact_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*exact_bitmap), (std::vector {0})); + EXPECT_EQ(exact_execution.stats.snii_stats.common_grams_plain_plans, 1); + EXPECT_EQ(exact_execution.stats.snii_stats.common_grams_gram_plans, 0); + EXPECT_EQ(exact_execution.stats.snii_stats.common_grams_fallback_kill_switch, 1); + + QueryExecutionContext prefix_execution( + /*enable_query_cache=*/false, /*count_on_index_fastpath=*/false); + std::shared_ptr prefix_bitmap; + assert_ok(opened.index_reader->query( + prefix_execution.context, "debug_safety_content", prefix_query, + InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, prefix_bitmap, &analyzer_ctx)); + ASSERT_NE(prefix_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*prefix_bitmap), (std::vector {0})); + EXPECT_EQ(prefix_execution.stats.snii_stats.common_grams_plain_plans, 1); + EXPECT_EQ(prefix_execution.stats.snii_stats.common_grams_gram_plans, 0); + EXPECT_EQ(prefix_execution.stats.snii_stats.common_grams_fallback_kill_switch, 1); + }; + + ASSERT_TRUE(config::set_config("enable_common_grams_query_plan", "false", + /*need_persist=*/false) + .ok()); + assert_forced_plain(); +} + +// A plain-only pair is proven gram-free by the wordset pre-proof, so the plan +// decision cache is legitimately never consulted (miss stays 0) -- but the plain +// posting IO is NOT bypassed: the injected adapter failure must surface as the +// query error and nothing may be published. +// GTest assertion macros inflate clang-tidy's branch count; both scenarios below are linear. +// NOLINTBEGIN(readability-function-cognitive-complexity) +TEST_F(SniiIndexReaderCountFallback, PublicSingleTermCountFastPathLeavesPrxStatsUnchanged) { + QueryExecutionContext execution(/*enable_query_cache=*/false, + /*count_on_index_fastpath=*/true); + set_prx_stats_sentinel(&execution.stats); + const PrxStatsSnapshot before = prx_stats_snapshot(execution.stats); + Field query_value = Field::create_field(std::string("failed")); + std::shared_ptr bitmap; + + assert_ok(_index_reader->query(execution.context, "count_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap)); + + ASSERT_NE(bitmap, nullptr); + EXPECT_EQ(bitmap->cardinality(), 6U); + EXPECT_TRUE(execution.context->count_on_index_fastpath_hit); + EXPECT_EQ(prx_stats_snapshot(execution.stats), before); +} + +TEST_F(SniiIndexReaderCountFallback, CountFastPathPublishesHitAfterRequestedNullBitmap) { + const std::string path = std::string(kTestDir) + "/count_null_failure"; + assert_ok(write_nullable_phrase_segment(path, /*has_null=*/true, kIndexId, + /*corrupt_null=*/true)); + OpenedSniiIndex opened; + assert_ok(open_snii_index(&_meta, path, &opened)); + QueryExecutionContext execution(/*enable_query_cache=*/false, + /*count_on_index_fastpath=*/true); + const Field query_value = Field::create_field(std::string("missing")); + std::shared_ptr bitmap; + InvertedIndexQueryCacheHandle null_bitmap_cache_handle; + + const Status status = opened.index_reader->query_with_null_bitmap( + execution.context, "count_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap, &null_bitmap_cache_handle); + + EXPECT_FALSE(status.ok()) << status.to_string(); + EXPECT_FALSE(execution.context->count_on_index_fastpath_hit); +} + +TEST_F(SniiIndexReaderCountFallback, ConcurrentCountShapedResultsNeverJoinRowAccurateSingleFlight) { + Field query_value = Field::create_field(std::string("order")); + + QueryExecutionContext baseline_row(/*enable_query_cache=*/false); + std::shared_ptr row_bitmap; + assert_ok(_index_reader->query(baseline_row.context, "g02_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, row_bitmap)); + ASSERT_NE(row_bitmap, nullptr); + const auto expected_row_docids = bitmap_docids(*row_bitmap); + + QueryExecutionContext baseline_count(/*enable_query_cache=*/false, + /*count_on_index_fastpath=*/true); + std::shared_ptr count_bitmap; + assert_ok(_index_reader->query(baseline_count.context, "g02_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, count_bitmap)); + ASSERT_NE(count_bitmap, nullptr); + ASSERT_TRUE(baseline_count.context->count_on_index_fastpath_hit); + const auto expected_count_docids = bitmap_docids(*count_bitmap); + + constexpr size_t kLeader = 0; + constexpr size_t kCountConsumer = 1; + constexpr size_t kRowFollower = 2; + constexpr size_t kConsumerCount = 3; + std::array status_ok {}; + std::array count_fastpath_hit {}; + std::array, kConsumerCount> docids; + std::binary_semaphore row_follower_joined {0}; + SingleFlightFollowerCounts follower_counts; + follower_counts.row_follower_joined = &row_follower_joined; + SingleFlightLeaderGate leader_gate; + _index_reader->set_single_flight_follower_joined_observer_for_test( + record_single_flight_follower, &follower_counts); + _index_reader->set_single_flight_leader_before_compute_observer_for_test( + block_single_flight_leader_before_compute, &leader_gate); + + auto run_query = [&](size_t consumer, bool count_consumer) { + SCOPED_INIT_THREAD_CONTEXT(); + is_count_consumer = count_consumer; + QueryExecutionContext execution(/*enable_query_cache=*/false, + /*count_on_index_fastpath=*/count_consumer); + std::shared_ptr bitmap; + const Status status = + _index_reader->query(execution.context, "g02_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap); + status_ok[consumer] = status.ok(); + count_fastpath_hit[consumer] = execution.context->count_on_index_fastpath_hit; + if (bitmap != nullptr) { + docids[consumer] = bitmap_docids(*bitmap); + } + is_count_consumer = false; + }; + + std::vector consumers; + consumers.emplace_back(run_query, kLeader, false); + leader_gate.leader_entered.acquire(); + consumers.emplace_back(run_query, kCountConsumer, true); + consumers.emplace_back(run_query, kRowFollower, false); + const bool row_joined_while_leader_blocked = + row_follower_joined.try_acquire_for(std::chrono::seconds(5)); + leader_gate.release_leader.release(); + for (auto& consumer : consumers) { + consumer.join(); + } + _index_reader->set_single_flight_follower_joined_observer_for_test(nullptr, nullptr); + _index_reader->set_single_flight_leader_before_compute_observer_for_test(nullptr, nullptr); + + EXPECT_TRUE(row_joined_while_leader_blocked); + ASSERT_TRUE(status_ok[kLeader]); + EXPECT_FALSE(count_fastpath_hit[kLeader]); + EXPECT_EQ(docids[kLeader], expected_row_docids); + ASSERT_TRUE(status_ok[kCountConsumer]); + EXPECT_TRUE(count_fastpath_hit[kCountConsumer]); + EXPECT_EQ(docids[kCountConsumer], expected_count_docids); + ASSERT_TRUE(status_ok[kRowFollower]); + EXPECT_FALSE(count_fastpath_hit[kRowFollower]); + EXPECT_EQ(docids[kRowFollower], expected_row_docids); + EXPECT_EQ(follower_counts.count_consumers.load(std::memory_order_relaxed), 0U); + EXPECT_EQ(follower_counts.row_consumers.load(std::memory_order_relaxed), 1U); +} + +TEST_F(SniiIndexReaderCountFallback, DistinctRawQueriesDoNotShareSingleFlight) { + const std::string lower_query = "failed order"; + const std::string upper_query = "FAILED ORDER"; + const auto index_file_key = _file_reader->get_index_file_cache_key(&_meta); + const InvertedIndexRawQuerySemantic lower_semantic { + .raw_query_bytes = lower_query, + .query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY}; + const InvertedIndexRawQuerySemantic upper_semantic { + .raw_query_bytes = upper_query, + .query_type = InvertedIndexQueryType::MATCH_PHRASE_QUERY}; + const InvertedIndexQueryCache::CacheKey lower_key {index_file_key, "raw_query_content", + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + lower_semantic.encode()}; + const InvertedIndexQueryCache::CacheKey upper_key {index_file_key, "raw_query_content", + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + upper_semantic.encode()}; + ASSERT_NE(lower_key.encode(), upper_key.encode()); + + std::latch ready(2); + std::latch start(1); + std::array status_ok {}; + std::array, 2> docids; + SingleFlightFollowerCounts follower_counts; + _index_reader->set_single_flight_follower_joined_observer_for_test( + record_single_flight_follower, &follower_counts); + std::array consumers { + std::thread([&] { + SCOPED_INIT_THREAD_CONTEXT(); + QueryExecutionContext execution(/*enable_query_cache=*/false); + Field query_value = Field::create_field(lower_query); + std::shared_ptr bitmap; + ready.count_down(); + start.wait(); + const Status status = + _index_reader->query(execution.context, "raw_query_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap); + status_ok[0] = status.ok(); + if (bitmap != nullptr) { + docids[0] = bitmap_docids(*bitmap); + } + }), + std::thread([&] { + SCOPED_INIT_THREAD_CONTEXT(); + QueryExecutionContext execution(/*enable_query_cache=*/false); + Field query_value = Field::create_field(upper_query); + std::shared_ptr bitmap; + ready.count_down(); + start.wait(); + const Status status = + _index_reader->query(execution.context, "raw_query_content", query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, bitmap); + status_ok[1] = status.ok(); + if (bitmap != nullptr) { + docids[1] = bitmap_docids(*bitmap); + } + })}; + ready.wait(); + start.count_down(); + for (auto& consumer : consumers) { + consumer.join(); + } + _index_reader->set_single_flight_follower_joined_observer_for_test(nullptr, nullptr); + + ASSERT_TRUE(status_ok[0]); + ASSERT_TRUE(status_ok[1]); + EXPECT_EQ(docids[0], docids[1]); + EXPECT_EQ(follower_counts.count_consumers.load(std::memory_order_relaxed), 0U); + EXPECT_EQ(follower_counts.row_consumers.load(std::memory_order_relaxed), 0U); +} + +TEST_F(SniiIndexReaderCountFallback, PrxProfileScopeOnlyExistsForMultiTermPhraseQueries) { + struct ControlQuery { + InvertedIndexQueryType type; + std::string search_str; + std::vector terms; + }; + const std::array controls {{ + {InvertedIndexQueryType::EQUAL_QUERY, "failed", {"failed"}}, + {InvertedIndexQueryType::MATCH_ANY_QUERY, "failed order", {"failed", "order"}}, + {InvertedIndexQueryType::MATCH_ALL_QUERY, "failed order", {"failed", "order"}}, + {InvertedIndexQueryType::MATCH_PHRASE_QUERY, "failed", {"failed"}}, + {InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "fail", {"fail"}}, + {InvertedIndexQueryType::MATCH_REGEXP_QUERY, "fail.*", {}}, + {InvertedIndexQueryType::WILDCARD_QUERY, "fail*", {}}, + }}; + QueryExecutionContext execution(/*enable_query_cache=*/false); + InvertedIndexQueryInfo query_info; + const auto set_plain_terms = [&query_info](const std::vector& terms) { + query_info.term_infos.clear(); + query_info.term_infos.reserve(terms.size()); + int32_t position = 0; + for (const auto& term : terms) { + query_info.term_infos.emplace_back(term, position++); + } + }; + + doris::snii::testing::reset_prx_execution_profile_scope_counters(); + doris::snii::query::testing::reset_query_profile_clock_read_count(); + doris::snii::format::testing::reset_prx_clock_read_count(); + doris::snii::query::internal::testing::reset_phrase_verify_clock_read_count(); + for (const auto& control : controls) { + SCOPED_TRACE(query_type_to_string(control.type)); + set_plain_terms(control.terms); + auto terms = control.terms; + std::shared_ptr bitmap; + assert_ok(_index_reader->_compute_query_bitmap(execution.context, control.type, query_info, + control.search_str, &terms, 50, &bitmap)); + ASSERT_NE(bitmap, nullptr); + } + set_plain_terms({"failed", "order"}); + query_info.slop = 1; + std::vector sloppy_phrase_terms {"failed", "order"}; + std::shared_ptr sloppy_phrase_bitmap; + const Status sloppy_phrase_status = _index_reader->_compute_query_bitmap( + execution.context, InvertedIndexQueryType::MATCH_PHRASE_QUERY, query_info, + "failed order", &sloppy_phrase_terms, 50, &sloppy_phrase_bitmap); + EXPECT_TRUE(sloppy_phrase_status.is()); + query_info.slop = 0; + + EXPECT_EQ(doris::snii::testing::prx_execution_profile_scope_construction_count(), 0U); + EXPECT_EQ(doris::snii::testing::prx_execution_profile_scope_flush_count(), 0U); + EXPECT_EQ(doris::snii::query::testing::query_profile_clock_read_count(), 0U); + EXPECT_EQ(doris::snii::format::testing::prx_clock_read_count(), 0U); + EXPECT_EQ(doris::snii::query::internal::testing::phrase_verify_clock_read_count(), 0U); + + const auto run_profiled_query = [&](InvertedIndexQueryType type, std::string_view search_str, + std::vector terms) { + set_plain_terms(terms); + std::shared_ptr bitmap; + assert_ok(_index_reader->_compute_query_bitmap(execution.context, type, query_info, + search_str, &terms, 50, &bitmap)); + ASSERT_NE(bitmap, nullptr); + }; + run_profiled_query(InvertedIndexQueryType::MATCH_PHRASE_QUERY, "failed order", + {"failed", "order"}); + run_profiled_query(InvertedIndexQueryType::MATCH_PHRASE_PREFIX_QUERY, "failed ord", + {"failed", "ord"}); + + EXPECT_EQ(doris::snii::testing::prx_execution_profile_scope_construction_count(), 2U); + EXPECT_EQ(doris::snii::testing::prx_execution_profile_scope_flush_count(), 2U); + EXPECT_GT(doris::snii::query::testing::query_profile_clock_read_count(), 0U); + EXPECT_GT(doris::snii::format::testing::prx_clock_read_count(), 0U); + EXPECT_GT(doris::snii::query::internal::testing::phrase_verify_clock_read_count(), 0U); +} + +TEST_F(SniiIndexReaderCountFallback, MultiTermPhraseUsesNormalPositionalQuery) { + auto logical_reader = _file_reader->open_snii_index(&_meta); + ASSERT_TRUE(logical_reader.has_value()) << logical_reader.error(); + const std::vector terms {"failed", "order"}; + EXPECT_EQ(lookup_df(*logical_reader.value(), "failed"), 6U); + + OlapReaderStatistics stats; + io::IOContext io_ctx; + RuntimeState runtime_state; + TQueryOptions query_options; + query_options.enable_inverted_index_query_cache = false; + query_options.enable_inverted_index_searcher_cache = false; + runtime_state.set_query_options(query_options); + auto context = std::make_shared(); + context->io_ctx = &io_ctx; + context->stats = &stats; + context->runtime_state = &runtime_state; + context->count_on_index_fastpath = true; + + InvertedIndexQueryInfo query_info; + bool handled = false; + std::shared_ptr count_bitmap; + assert_ok(_index_reader->_try_count_only_fastpath(context, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + query_info, terms, &handled, &count_bitmap)); + + EXPECT_FALSE(handled); + EXPECT_EQ(count_bitmap, nullptr); + EXPECT_EQ(count_bitmap == nullptr ? 0 : count_bitmap->cardinality(), 0U); + + snii_doris::DorisSniiFileReader::ScopedIOContext io_context_scope(&io_ctx); + std::vector phrase_docids; + assert_ok(doris::snii::query::phrase_query(*logical_reader.value(), terms, &phrase_docids)); + EXPECT_EQ(phrase_docids, (std::vector {0, 3, 5})); +} + +} // namespace +} // namespace doris::segment_v2 diff --git a/be/test/storage/index/snii/snii_test_env.cpp b/be/test/storage/index/snii/snii_test_env.cpp new file mode 100644 index 00000000000000..c1b8de6728f762 --- /dev/null +++ b/be/test/storage/index/snii/snii_test_env.cpp @@ -0,0 +1,57 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include +#include +#include + +#include "io/fs/local_file_system.h" +#include "runtime/exec_env.h" +#include "storage/index/index_writer.h" // segment_v2::TmpFileDirs +#include "storage/options.h" // StorePath + +namespace doris { +namespace { + +// Initializes ExecEnv's tmp_file_dirs ONCE for the whole doris_be_test binary so the +// SNII writer's resolve_temp_dir() (ExecEnv::get_tmp_file_dirs()->get_tmp_file_dir()) +// works in unit tests, which otherwise leave tmp_file_dirs null. Mirrors the pattern +// used by ann_index_writer_test / segcompaction_test. +class SniiTmpDirEnvironment : public ::testing::Environment { +public: + void SetUp() override { + if (ExecEnv::GetInstance()->get_tmp_file_dirs() != nullptr) { + return; // another test environment already configured it + } + const std::string tmp_dir = "./ut_dir/snii_tmp"; + static_cast(io::global_local_filesystem()->delete_directory(tmp_dir)); + static_cast(io::global_local_filesystem()->create_directory(tmp_dir)); + std::vector paths; + paths.emplace_back(tmp_dir, -1); + auto dirs = std::make_unique(paths); + static_cast(dirs->init()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(dirs)); + } +}; + +[[maybe_unused]] ::testing::Environment* const g_snii_tmp_dir_env = + ::testing::AddGlobalTestEnvironment(new SniiTmpDirEnvironment); + +} // namespace +} // namespace doris diff --git a/be/test/storage/index/snii/spimi_common_gram_chain_encoding_test.cpp b/be/test/storage/index/snii/spimi_common_gram_chain_encoding_test.cpp new file mode 100644 index 00000000000000..d9b30d606d5ba7 --- /dev/null +++ b/be/test/storage/index/snii/spimi_common_gram_chain_encoding_test.cpp @@ -0,0 +1,137 @@ +// 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 + +#include +#include + +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::writer { +namespace { + +TEST(SpimiCommonGramChainEncodingTest, NativeDocsOnlyPairUsesStatlessEncoding) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + const PlainTermId left = buffer.intern_plain_term("left"); + const PlainTermId right = buffer.intern_plain_term("right"); + + buffer.add_common_gram(left, right, /*docid=*/1, /*pos=*/0, + /*retain_positions=*/false); + buffer.add_common_gram(left, right, /*docid=*/1, /*pos=*/1, + /*retain_positions=*/false); + buffer.add_common_gram(left, right, /*docid=*/4, /*pos=*/0, + /*retain_positions=*/false); + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {1, 4})); + EXPECT_TRUE(terms[0].freqs.empty()); + EXPECT_TRUE(terms[0].positions_flat.empty()); + EXPECT_FALSE(terms[0].retain_positions); +} + +TEST(SpimiCommonGramChainEncodingTest, NativeDocsOnlySingletonDefersPostingArena) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + const PlainTermId left = buffer.intern_plain_term("left"); + const PlainTermId right = buffer.intern_plain_term("right"); + const uint64_t bytes_before = buffer.resident_bytes_for_test(); + + buffer.add_common_gram(left, right, /*docid=*/7, /*pos=*/0, + /*retain_positions=*/false); + + EXPECT_LT(buffer.resident_bytes_for_test() - bytes_before, 32U << 10); + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {7})); + EXPECT_TRUE(terms[0].freqs.empty()); + EXPECT_TRUE(terms[0].positions_flat.empty()); +} + +TEST(SpimiCommonGramChainEncodingTest, NativeDocsOnlyBackfillPreservesOutOfOrderDocids) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + const PlainTermId left = buffer.intern_plain_term("left"); + const PlainTermId right = buffer.intern_plain_term("right"); + + for (const uint32_t docid : {9U, 7U, 9U, UINT32_MAX, 0U}) { + buffer.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {0U, 7U, 9U, UINT32_MAX})); +} + +TEST(SpimiCommonGramChainEncodingTest, NativeDocsOnlySingletonMergesAcrossSpill) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + buffer.set_forced_spill_min_arena_bytes(0); + const PlainTermId left = buffer.intern_plain_term("left"); + const PlainTermId right = buffer.intern_plain_term("right"); + + buffer.request_global_spill_for_test(); + buffer.add_common_gram_and_plain(left, right, /*docid=*/7, /*gram_pos=*/0, + /*plain_pos=*/1, /*retain_gram_positions=*/false); + ASSERT_EQ(buffer.run_count_for_test(), 1U); + buffer.add_common_gram(left, right, /*docid=*/7, /*pos=*/2, + /*retain_positions=*/false); + buffer.add_common_gram(left, right, /*docid=*/9, /*pos=*/0, + /*retain_positions=*/false); + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + const TermPostings* gram = nullptr; + for (const TermPostings& term : terms) { + if (!term.retain_positions) { + gram = &term; + } + } + ASSERT_NE(gram, nullptr); + EXPECT_EQ(gram->docids, (std::vector {7U, 9U})); +} + +TEST(SpimiCommonGramChainEncodingTest, PositionedNativePairRemainsTagged) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + const PlainTermId left = buffer.intern_plain_term("left"); + const PlainTermId right = buffer.intern_plain_term("right"); + + buffer.add_common_gram(left, right, /*docid=*/1, /*pos=*/2, + /*retain_positions=*/true); + buffer.add_common_gram(left, right, /*docid=*/1, /*pos=*/5, + /*retain_positions=*/true); + buffer.add_common_gram(left, right, /*docid=*/3, /*pos=*/1, + /*retain_positions=*/true); + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {1, 3})); + EXPECT_EQ(terms[0].freqs, (std::vector {2, 1})); + EXPECT_EQ(terms[0].positions_flat, (std::vector {2, 5, 1})); + EXPECT_TRUE(terms[0].retain_positions); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/spimi_dense_rank_order_test.cpp b/be/test/storage/index/snii/spimi_dense_rank_order_test.cpp new file mode 100644 index 00000000000000..39e98703f7556a --- /dev/null +++ b/be/test/storage/index/snii/spimi_dense_rank_order_test.cpp @@ -0,0 +1,69 @@ +// 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 + +#include +#include +#include + +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::writer { +namespace { + +TEST(SpimiDenseRankOrderTest, FullCommonGramVocabularyUsesDenseRankInverse) { + testing::reset_rank_ordering_counts(); + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + + const PlainTermId z = buffer.intern_plain_term("z"); + const PlainTermId a = buffer.intern_plain_term("a"); + const PlainTermId middle = buffer.intern_plain_term("middle"); + buffer.add_plain_token(z, /*docid=*/0, /*pos=*/0); + buffer.add_plain_token(a, /*docid=*/0, /*pos=*/1); + buffer.add_plain_token(middle, /*docid=*/0, /*pos=*/2); + buffer.add_common_gram(z, a, /*docid=*/0, /*pos=*/0, /*retain_positions=*/false); + buffer.add_common_gram(a, middle, /*docid=*/0, /*pos=*/1, + /*retain_positions=*/false); + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 5U); + EXPECT_TRUE(std::ranges::is_sorted(terms, {}, &TermPostings::term)); + EXPECT_EQ(testing::dense_rank_inversions(), 1U); + EXPECT_EQ(testing::rank_comparison_sorts(), 0U); +} + +TEST(SpimiDenseRankOrderTest, PartialBorrowedVocabularyUsesComparisonSort) { + testing::reset_rank_ordering_counts(); + const std::vector vocabulary = {"z", "a", "middle", "b"}; + SpimiTermBuffer buffer(&vocabulary, /*has_positions=*/false); + buffer.add_token(/*term_id=*/0, /*docid=*/0, /*pos=*/0); + buffer.add_token(/*term_id=*/2, /*docid=*/0, /*pos=*/1); + + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].term, "middle"); + EXPECT_EQ(terms[1].term, "z"); + EXPECT_EQ(testing::dense_rank_inversions(), 0U); + EXPECT_EQ(testing::rank_comparison_sorts(), 1U); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/spimi_spill_rank_test.cpp b/be/test/storage/index/snii/spimi_spill_rank_test.cpp new file mode 100644 index 00000000000000..62ae72cf6e8c74 --- /dev/null +++ b/be/test/storage/index/snii/spimi_spill_rank_test.cpp @@ -0,0 +1,221 @@ +// 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 + +#include +#include +#include + +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::writer { +namespace { + +TEST(SpimiSpillRankTest, RebuildsGrowingVocabularyRankOnlyForFinalMerge) { + testing::reset_string_rank_rebuilds(); + SpimiTermBuffer buffer(/*has_positions=*/false, /*spill_threshold_bytes=*/0); + buffer.set_max_run_files(/*cap=*/0); + buffer.set_forced_spill_min_arena_bytes(0); + + buffer.add_token("b", /*docid=*/0, /*pos=*/0); + buffer.request_global_spill_for_test(); + buffer.add_token("c", /*docid=*/1, /*pos=*/0); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + + buffer.add_token("z", /*docid=*/2, /*pos=*/0); + buffer.add_token("a", /*docid=*/3, /*pos=*/0); + buffer.add_token("b", /*docid=*/4, /*pos=*/0); + buffer.request_global_spill_for_test(); + buffer.add_token("m", /*docid=*/5, /*pos=*/0); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + EXPECT_GE(buffer.string_rank_capacity_for_test(), 5); + EXPECT_GE(buffer.resident_bytes_for_test(), + buffer.string_rank_capacity_for_test() * sizeof(uint32_t)); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(buffer.run_count_for_test(), 2); + + const std::vector terms = buffer.finalize_sorted(); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + EXPECT_EQ(testing::string_rank_rebuilds(), 2); + ASSERT_EQ(terms.size(), 5); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[0].docids, std::vector({3})); + EXPECT_EQ(terms[0].freqs, std::vector({1})); + EXPECT_EQ(terms[1].term, "b"); + EXPECT_EQ(terms[1].docids, std::vector({0, 4})); + EXPECT_EQ(terms[1].freqs, std::vector({1, 1})); + EXPECT_EQ(terms[2].term, "c"); + EXPECT_EQ(terms[2].docids, std::vector({1})); + EXPECT_EQ(terms[2].freqs, std::vector({1})); + EXPECT_EQ(terms[3].term, "m"); + EXPECT_EQ(terms[3].docids, std::vector({5})); + EXPECT_EQ(terms[3].freqs, std::vector({1})); + EXPECT_EQ(terms[4].term, "z"); + EXPECT_EQ(terms[4].docids, std::vector({2})); + EXPECT_EQ(terms[4].freqs, std::vector({1})); +} + +TEST(SpimiSpillRankTest, ReusesRankAcrossFixedVocabularySpillsAndFinalMerge) { + testing::reset_string_rank_rebuilds(); + const std::vector vocabulary = {"z", "a", "m"}; + SpimiTermBuffer buffer(&vocabulary, /*has_positions=*/false, /*spill_threshold_bytes=*/1); + buffer.set_max_run_files(/*cap=*/0); + + buffer.add_token(/*term_id=*/0, /*docid=*/0, /*pos=*/0); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + buffer.add_token(/*term_id=*/1, /*docid=*/1, /*pos=*/0); + buffer.add_token(/*term_id=*/2, /*docid=*/2, /*pos=*/0); + buffer.add_token(/*term_id=*/0, /*docid=*/3, /*pos=*/0); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(buffer.run_count_for_test(), 4); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + + const std::vector terms = buffer.finalize_sorted(); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + ASSERT_EQ(terms.size(), 3); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[0].docids, std::vector({1})); + EXPECT_EQ(terms[1].term, "m"); + EXPECT_EQ(terms[1].docids, std::vector({2})); + EXPECT_EQ(terms[2].term, "z"); + EXPECT_EQ(terms[2].docids, std::vector({0, 3})); +} + +TEST(SpimiSpillRankTest, RefreshesStaleRankBeforeRunCompaction) { + testing::reset_run_compactions(); + testing::reset_string_rank_rebuilds(); + SpimiTermBuffer buffer(/*has_positions=*/false, /*spill_threshold_bytes=*/0); + buffer.set_max_run_files(/*cap=*/1); + buffer.set_forced_spill_min_arena_bytes(0); + + buffer.add_token("h", /*docid=*/0, /*pos=*/0); + buffer.request_global_spill_for_test(); + buffer.add_token("i", /*docid=*/1, /*pos=*/0); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + + buffer.add_token("z", /*docid=*/2, /*pos=*/0); + buffer.add_token("h", /*docid=*/3, /*pos=*/0); + buffer.request_global_spill_for_test(); + buffer.add_token("a", /*docid=*/4, /*pos=*/0); + EXPECT_EQ(testing::string_rank_rebuilds(), 1); + + buffer.add_token("m", /*docid=*/5, /*pos=*/0); + buffer.request_global_spill_for_test(); + buffer.add_token("b", /*docid=*/6, /*pos=*/0); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + EXPECT_EQ(testing::run_compactions(), 1); + EXPECT_EQ(testing::string_rank_rebuilds(), 2); + + const std::vector terms = buffer.finalize_sorted(); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + EXPECT_EQ(testing::string_rank_rebuilds(), 2); + ASSERT_EQ(terms.size(), 6); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[0].docids, std::vector({4})); + EXPECT_EQ(terms[1].term, "b"); + EXPECT_EQ(terms[1].docids, std::vector({6})); + EXPECT_EQ(terms[2].term, "h"); + EXPECT_EQ(terms[2].docids, std::vector({0, 3})); + EXPECT_EQ(terms[3].term, "i"); + EXPECT_EQ(terms[3].docids, std::vector({1})); + EXPECT_EQ(terms[4].term, "m"); + EXPECT_EQ(terms[4].docids, std::vector({5})); + EXPECT_EQ(terms[5].term, "z"); + EXPECT_EQ(terms[5].docids, std::vector({2})); +} + +TEST(SpimiSpillRankTest, PhysicalCommonGramsSurviveSpillAndRunCompaction) { + namespace inverted_index = doris::segment_v2::inverted_index; + const std::string first = inverted_index::encode_common_gram("of", "the").value(); + const std::string second = inverted_index::encode_common_gram("the", "world").value(); + + auto feed = [&](SpimiTermBuffer* buffer) { + buffer->add_token(second, 0, 1); + buffer->add_token("plain", 0, 0); + buffer->request_global_spill_for_test(); + buffer->add_token(first, 0, 0); + buffer->add_token(first, 1, 2); + buffer->request_global_spill_for_test(); + buffer->add_token(second, 2, 3); + buffer->request_global_spill_for_test(); + buffer->add_token("plain", 3, 4); + }; + + SpimiTermBuffer unspilled(/*has_positions=*/true); + feed(&unspilled); + + SpimiTermBuffer spilled(/*has_positions=*/true); + spilled.set_max_run_files(/*cap=*/1); + spilled.set_forced_spill_min_arena_bytes(0); + feed(&spilled); + + ASSERT_TRUE(unspilled.status().ok()) << unspilled.status(); + ASSERT_TRUE(spilled.status().ok()) << spilled.status(); + EXPECT_GE(spilled.run_count_for_test(), 1U); + + const std::vector expected = unspilled.finalize_sorted(); + const std::vector actual = spilled.finalize_sorted(); + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(actual[i].term, expected[i].term); + EXPECT_EQ(actual[i].docids, expected[i].docids); + EXPECT_EQ(actual[i].freqs, expected[i].freqs); + EXPECT_EQ(actual[i].positions_flat, expected[i].positions_flat); + } +} + +TEST(SpimiSpillRankTest, PerTermDocsOnlyDeduplicatesAcrossSpillAndCompaction) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.set_max_run_files(/*cap=*/1); + buffer.set_forced_spill_min_arena_bytes(0); + + buffer.add_token("docs", /*docid=*/7, /*pos=*/1, /*retain_positions=*/false); + buffer.request_global_spill_for_test(); + buffer.add_token("docs", /*docid=*/7, /*pos=*/2, /*retain_positions=*/false); + buffer.request_global_spill_for_test(); + buffer.add_token("positioned", /*docid=*/7, /*pos=*/3, /*retain_positions=*/true); + buffer.request_global_spill_for_test(); + buffer.add_token("positioned", /*docid=*/7, /*pos=*/4, /*retain_positions=*/true); + buffer.add_token("docs", /*docid=*/8, /*pos=*/5, /*retain_positions=*/false); + + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + EXPECT_GE(buffer.run_count_for_test(), 1U); + const std::vector terms = buffer.finalize_sorted(); + ASSERT_TRUE(buffer.status().ok()) << buffer.status(); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].term, "docs"); + EXPECT_FALSE(terms[0].retain_positions); + EXPECT_EQ(terms[0].docids, (std::vector {7U, 8U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U, 1U})); + EXPECT_TRUE(terms[0].positions_flat.empty()); + EXPECT_EQ(terms[1].term, "positioned"); + EXPECT_TRUE(terms[1].retain_positions); + EXPECT_EQ(terms[1].docids, (std::vector {7U})); + EXPECT_EQ(terms[1].freqs, (std::vector {2U})); + EXPECT_EQ(terms[1].positions_flat, (std::vector {3U, 4U})); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/compact_posting_pool_test.cpp b/be/test/storage/index/snii/writer/compact_posting_pool_test.cpp new file mode 100644 index 00000000000000..5a05412dce9da8 --- /dev/null +++ b/be/test/storage/index/snii/writer/compact_posting_pool_test.cpp @@ -0,0 +1,602 @@ +// 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 "storage/index/snii/writer/compact_posting_pool.h" + +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using doris::snii::writer::CompactPostingPool; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::TermPostings; + +namespace { + +// Test helper bundling a chain's append handle, its level, and its head -- the +// same per-term state the real accumulator keeps -- so tests can append by value. +struct Chain { + CompactPostingPool::SliceWriter w; + uint8_t level = 0; + uint32_t head = 0; + void start(CompactPostingPool* pool) { head = pool->start_chain(&w, &level); } + void put(CompactPostingPool* pool, uint8_t b) { pool->append_byte(&w, &level, b); } +}; + +// Reads back the whole chain into a vector for comparison. +std::vector ReadChain(const CompactPostingPool& pool, uint32_t head, uint64_t len) { + std::vector out; + out.reserve(len); + CompactPostingPool::Cursor c = pool.cursor(head, len); + while (c.has_next()) { + out.push_back(c.next()); + } + return out; +} + +} // namespace + +// A single chain shorter than one slice round-trips exactly. +TEST(SniiCompactPostingPool, TinyChainRoundTrips) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + const std::vector data = {7, 0, 255, 42}; + for (uint8_t b : data) { + ch.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + EXPECT_EQ(pool.payload_bytes(), data.size()); +} + +// A chain that spans many slice levels round-trips exactly (exercises forward +// pointers across several geometric slice sizes). +TEST(SniiCompactPostingPool, MultiSliceChainRoundTrips) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + std::vector data; + for (uint32_t i = 0; i < 5000; ++i) { + data.push_back(static_cast(i * 31 + 7)); + } + for (uint8_t b : data) { + ch.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + EXPECT_EQ(pool.payload_bytes(), data.size()); +} + +// Many INTERLEAVED chains (the real SPIMI access pattern) stay independent: a byte +// written to chain A never appears in chain B's read-back. +TEST(SniiCompactPostingPool, InterleavedChainsIndependent) { + CompactPostingPool pool; + constexpr int kChains = 64; + std::vector chains(kChains); + std::vector> expect(kChains); + for (auto& ch : chains) { + ch.start(&pool); + } + + std::mt19937 rng(12345); + // Append bytes to chains in a randomized interleaving so slices for different + // chains land in the same blocks intermixed. + for (int round = 0; round < 20000; ++round) { + const int c = static_cast(rng() % kChains); + const auto b = static_cast(rng()); + chains[c].put(&pool, b); + expect[c].push_back(b); + } + for (int i = 0; i < kChains; ++i) { + EXPECT_EQ(ReadChain(pool, chains[i].head, expect[i].size()), expect[i]) << "chain " << i; + } +} + +// MANY chains + MANY bytes force the arena across several 32 KiB block +// boundaries. This is the regression for a block-boundary bump bug: a run that +// exactly fills a block must allocate the next block before handing out the +// boundary offset, never returning an offset into a not-yet-allocated block. +TEST(SniiCompactPostingPool, ManyChainsAcrossBlockBoundaries) { + CompactPostingPool pool; + constexpr int kChains = 2000; + std::vector chains(kChains); + std::vector> expect(kChains); + for (auto& ch : chains) { + ch.start(&pool); + } + + std::mt19937 rng(98765); + for (int round = 0; round < 1'000'000; ++round) { + const int c = static_cast(rng() % kChains); + const auto b = static_cast(rng()); + chains[c].put(&pool, b); + expect[c].push_back(b); + } + for (int i = 0; i < kChains; ++i) { + EXPECT_EQ(ReadChain(pool, chains[i].head, expect[i].size()), expect[i]) << "chain " << i; + } +} + +// An empty chain (started but never written) reads back as zero bytes. +TEST(SniiCompactPostingPool, EmptyChain) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + EXPECT_TRUE(ReadChain(pool, ch.head, 0).empty()); +} + +// A chain that exactly fills a slice boundary (no extra byte) reads back exactly, +// without dereferencing the (still-zero) forward pointer. +TEST(SniiCompactPostingPool, ExactSliceBoundary) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + // Exactly fill the level-0 slice (kSliceSizes[0] payload bytes) and stop. + std::vector data; + for (uint32_t i = 0; i < CompactPostingPool::kSliceSizes_level0(); ++i) { + data.push_back(static_cast(i + 1)); + } + for (uint8_t b : data) { + ch.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + // Now extend by one byte (forces the forward link) and re-read fully. + ch.put(&pool, 99); + data.push_back(99); + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); +} + +// Cursor CONTRACT: `budget` is an UPPER BOUND on bytes yielded, NOT a required exact +// length. The cursor is SELF-TERMINATING -- it stops at the chain tail (a zero forward +// pointer) no matter how large the budget. This pins both halves of the single contract: +// (1) an exact-length budget round-trips the written bytes, and +// (2) an OVER-LARGE budget (as the production caller passes -- the write-head offset) +// stays MEMORY-SAFE: next() never follows the tail's zero forward pointer off the +// chain into block 0 (UB); it yields at most the tail-slice's payload region (the +// written bytes plus the slice's zero-initialized unwritten tail) and then stops. +// +// WITHOUT the tail check (next_head == 0 -> stop), looping on has_next() with an +// over-large budget would hit the slice boundary, read the still-zero tail forward +// pointer, jump cur_ to offset 0, and re-read block 0's live bytes (an ALIAS) -- exactly +// the misuse the old "remaining" contract left latent. This test fails without the fix. +TEST(SniiCompactPostingPool, CursorOverLargeBudgetSelfTerminates) { + CompactPostingPool pool; + // Lay down a DISTINCT first chain so block 0 offset 0 holds recognizable bytes; if the + // cursor ever aliased offset 0 it would yield these, which we assert it never does. + Chain decoy; + decoy.start(&pool); + const std::vector decoy_bytes = {0xDE, 0xAD, 0xBE, 0xEF}; + for (uint8_t b : decoy_bytes) { + decoy.put(&pool, b); + } + ASSERT_EQ(decoy.head, 0U) << "first chain must own pool offset 0 (the alias target)"; + + // A second short chain whose tail forward pointer is still zero (single level-0 slice). + Chain ch; + ch.start(&pool); + const std::vector data = {11, 22, 33}; + for (uint8_t b : data) { + ch.put(&pool, b); + } + + // (1) exact-length budget round-trips the written bytes. + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + + // (2) MISUSE: a budget far larger than the payload. The cursor must self-terminate at + // the chain tail, yielding at most the level-0 slice's payload region and stopping. + const uint32_t slice0 = CompactPostingPool::kSliceSizes_level0(); + const uint32_t over_budget = 100U * CompactPostingPool::kBlockSize; // absurdly large + CompactPostingPool::Cursor c = pool.cursor(ch.head, over_budget); + std::vector pulled; + uint32_t pulls = 0; + while (c.has_next()) { + pulled.push_back(c.next()); + ASSERT_LT(++pulls, over_budget) << "cursor failed to self-terminate at the chain tail"; + } + // It stops at the tail: exactly the level-0 slice's payload region (written bytes plus + // the slice's zero-initialized unwritten tail), and NOTHING beyond. + EXPECT_EQ(pulled.size(), slice0) + << "over-large budget must stop at the tail slice's end, not run on"; + ASSERT_GE(pulled.size(), data.size()); + EXPECT_EQ(std::vector(pulled.begin(), pulled.begin() + data.size()), data); + for (size_t i = data.size(); i < pulled.size(); ++i) { + EXPECT_EQ(pulled[i], 0U) << "unwritten tail byte " << i << " must read as zero"; + } + // CRITICAL: the cursor must NEVER have aliased the decoy at offset 0. + for (uint8_t b : pulled) { + EXPECT_NE(b, 0xDEU) << "cursor aliased block 0 -- tail check failed"; + } +} + +// Computes an EXACT one-block fill layout analytically from the public slice schedule: +// `pad` level-0 chains (each consuming one level-0 slice allocation) followed by ONE +// chain grown until it has just allocated the slice at `grow_to_level`, such that the +// total arena bytes consumed equal exactly kBlockSize. Returns false if no such layout +// exists for the current schedule. No magic numbers -- everything derives from the +// kSliceSize_at / kNextLevel_at accessors, so a schedule change is handled automatically. +static bool ExactBlockFillLayout(uint32_t* pad, int* grow_to_level) { + const uint32_t block = CompactPostingPool::kBlockSize; + const uint32_t l0_alloc = + CompactPostingPool::kSliceSizes_level0() + CompactPostingPool::kPtrBytes; + // cum = bytes a single growing chain consumes after allocating slices for levels + // 0..L inclusive (each slice allocation is payload-cap + kPtrBytes). + uint32_t cum = 0; + for (int level = 0; level < CompactPostingPool::kLevelCount; ++level) { + cum += CompactPostingPool::kSliceSize_at(level) + CompactPostingPool::kPtrBytes; + if (cum > block) { + break; + } + const uint32_t rem = block - cum; + if (rem % l0_alloc == 0) { + *pad = rem / l0_alloc; + *grow_to_level = level; + return true; + } + } + return false; +} + +// DETERMINISTIC coverage of alloc_run case (c): a previous allocation EXACTLY fills a +// 32 KiB block, leaving next_offset_ on the block boundary (in_block == 0) over a block +// that is already fully consumed. The next allocation MUST detect this (via the +// tail_exists guard) and start a FRESH block -- it must NOT mistake in_block == 0 for an +// empty fresh block and hand back offset 0, which would alias block 0's live bytes. +// +// Today this branch is only hit probabilistically by the RNG-seeded interleave tests. +// Here we drive it EXACTLY: compute a layout (analytically, from the public schedule) +// whose allocations fill block 0 to its LAST byte (next_offset_ == kBlockSize), realize +// it on a real pool, then start one more chain. Its head MUST be exactly kBlockSize +// (block 1, byte 0) -- never 0 -- and every chain in block 0 MUST still round-trip, +// proving the case-(c) allocation did not return offset 0 and clobber block 0. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiCompactPostingPool, AllocRunExactBlockFillStartsFreshBlock) { + const uint32_t block = CompactPostingPool::kBlockSize; + + uint32_t pad = 0; + int grow_to_level = 0; + ASSERT_TRUE(ExactBlockFillLayout(&pad, &grow_to_level)) + << "no exact one-block fill exists for the current slice schedule"; + + CompactPostingPool pool; + std::vector chains; + std::vector> data; + + // Lay down `pad` level-0 chains, each with one identifiable payload byte. None may + // spill out of block 0 (the layout was computed so they all fit). + for (uint32_t i = 0; i < pad; ++i) { + Chain c; + c.start(&pool); + const auto b = static_cast(0x40 + (i & 0x3F)); + c.put(&pool, b); + chains.push_back(c); + data.push_back({b}); + ASSERT_EQ(pool.arena_bytes(), block) << "padding chain " << i << " spilled early"; + } + + // Grow ONE chain until it has just allocated the slice at `grow_to_level`, writing no + // byte into that final slice so its tail ends exactly on the block boundary. + Chain g; + g.start(&pool); + std::vector gbytes; + uint32_t guard = 0; + while (std::cmp_less(g.level, grow_to_level)) { + const auto b = static_cast(0x80 + (gbytes.size() & 0x3F)); + g.put(&pool, b); + gbytes.push_back(b); + ASSERT_EQ(pool.arena_bytes(), block) << "grown chain spilled before the boundary"; + ASSERT_LT(++guard, block) << "grow loop did not converge"; + } + chains.push_back(g); + data.push_back(gbytes); + + // Block 0 is now full to its LAST byte: next_offset_ == kBlockSize, still one block. + ASSERT_EQ(pool.arena_bytes(), block) << "block 0 should be exactly full, one block"; + + // The case-(c) allocation: starting a new chain must open a FRESH block at offset + // kBlockSize, NOT alias offset 0. + Chain probe; + probe.start(&pool); + EXPECT_EQ(probe.head, block) << "case (c) must hand out block 1 byte 0, not offset 0"; + EXPECT_NE(probe.head, 0U) << "case (c) must not alias block 0"; + EXPECT_EQ(pool.arena_bytes(), 2U * block) << "exactly two blocks after the probe"; + + // The probe chain owns block 1 byte 0; write+read it to confirm it is live. + probe.put(&pool, 0xEE); + EXPECT_EQ(ReadChain(pool, probe.head, 1U), (std::vector {0xEE})); + + // Every chain laid down in block 0 must still round-trip -- proving the case-(c) + // allocation did NOT return offset 0 and overwrite block 0's live bytes. + for (size_t i = 0; i < chains.size(); ++i) { + EXPECT_EQ(ReadChain(pool, chains[i].head, data[i].size()), data[i]) + << "block-0 chain " << i << " corrupted across the exact block boundary"; + } +} + +// reset() drops all blocks; a fresh chain after reset starts clean. +TEST(SniiCompactPostingPool, ResetClears) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + for (int i = 0; i < 1000; ++i) { + ch.put(&pool, static_cast(i)); + } + EXPECT_GT(pool.payload_bytes(), 0U); + pool.reset(); + EXPECT_EQ(pool.payload_bytes(), 0U); + Chain ch2; + ch2.start(&pool); + std::vector data = {5, 6, 7}; + for (uint8_t b : data) { + ch2.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch2.head, data.size()), data); +} + +// ====================================================================================== +// T13: regression net for the append_byte / Cursor::has_next / Cursor::next INLINE move. +// +// The move is a pure code relocation (.cpp out-of-line bodies -> .h inline), so every +// assertion below must hold byte-for-byte BOTH before and after the change: the golden +// values are the contract. F1-F8 drive the arena encode (append_byte) and decode +// (Cursor) micro-loops directly; F9 covers the full SpimiTermBuffer encode+decode path +// (put_byte->append_byte ... decode_chain_varint->Cursor::next) end to end; F10 pins the +// out-of-vocab error path. F4/F5/F9/F10 add coverage the pre-existing suite lacked +// (tail no-phantom via an over-large budget, budget truncation, end-to-end postings, +// latched InvalidArgument). +// ====================================================================================== + +// T13-F1: a chain that fits in one level-0 slice round-trips and never overflows. +TEST(SniiCompactPostingPoolTest, RoundTripsSingleSlice) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + std::vector data; + for (uint32_t i = 0; i < CompactPostingPool::kSliceSizes_level0(); ++i) { + data.push_back(static_cast(i * 9 + 1)); + } + for (uint8_t b : data) { + ch.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + EXPECT_EQ(pool.payload_bytes(), data.size()); + EXPECT_EQ(ch.level, uint8_t {0}) << "a single-slice fill must stay at level 0"; +} + +// T13-F2: 5000 pseudo-random bytes (fixed seed) cross many slice levels and round-trip +// byte-identically -- the encode/decode forward-pointer chain golden. +TEST(SniiCompactPostingPoolTest, RoundTripsAcrossSliceLevels) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + std::mt19937 rng(0xC0FFEEU); + std::vector data; + data.reserve(5000); + for (uint32_t i = 0; i < 5000; ++i) { + const auto b = static_cast(rng()); + data.push_back(b); + ch.put(&pool, b); + } + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + EXPECT_EQ(pool.payload_bytes(), data.size()); + EXPECT_GT(ch.level, uint8_t {0}) << "5000 bytes must advance past level 0"; +} + +// T13-F3: a single chain longer than one 32 KiB block spans >= 2 blocks and round-trips +// byte-identically (exercises at()'s two-level block index across alloc_run new blocks). +TEST(SniiCompactPostingPoolTest, RoundTripsAcrossBlockBoundary) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + std::mt19937 rng(0xB10CU); + const uint32_t n = 3U * CompactPostingPool::kBlockSize; // payload alone exceeds 2 blocks + std::vector data; + data.reserve(n); + for (uint32_t i = 0; i < n; ++i) { + const auto b = static_cast(rng()); + data.push_back(b); + ch.put(&pool, b); + } + EXPECT_GE(pool.arena_bytes(), 2ULL * CompactPostingPool::kBlockSize) + << "payload over one block must occupy >= 2 blocks"; + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); + EXPECT_EQ(pool.payload_bytes(), data.size()); +} + +// T13-F4: with a budget LARGER than the payload, has_next() must stop at the tail slice's +// zero forward pointer and report no phantom trailing byte (the stop is the tail, not the +// budget). next() at the tail yields 0. +TEST(SniiCompactPostingPoolTest, HasNextStopsAtTailNoPhantom) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + const uint32_t cap = CompactPostingPool::kSliceSize_at(0); + std::vector data; + for (uint32_t i = 0; i < cap; ++i) { + const auto b = static_cast(i + 1); + data.push_back(b); + ch.put(&pool, b); + } + // Budget far exceeds the payload, so a false has_next() can only come from the tail + // zero pointer (a phantom-byte bug would instead read the zero pointer as a byte). + CompactPostingPool::Cursor c = pool.cursor(ch.head, CompactPostingPool::kBlockSize); + std::vector out; + while (c.has_next()) { + out.push_back(c.next()); + } + EXPECT_EQ(out, data) << "exactly the written payload region, no phantom tail byte"; + EXPECT_EQ(out.size(), cap); + EXPECT_FALSE(c.has_next()) << "tail zero pointer must stop has_next()"; + EXPECT_EQ(c.next(), 0U) << "next() past the tail yields 0"; +} + +// T13-F5: a budget SMALLER than the payload truncates the cursor to exactly `budget` +// bytes (the first ones, in write order), then has_next() goes false. +TEST(SniiCompactPostingPoolTest, BudgetCapsYieldedBytes) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + std::vector data; + for (uint32_t i = 0; i < 100; ++i) { + const auto b = static_cast(i * 7 + 1); + data.push_back(b); + ch.put(&pool, b); + } + constexpr uint64_t kBudget = 10; + CompactPostingPool::Cursor c = pool.cursor(ch.head, kBudget); + std::vector out; + while (c.has_next()) { + out.push_back(c.next()); + } + EXPECT_EQ(out.size(), kBudget); + EXPECT_EQ(out, std::vector(data.begin(), data.begin() + kBudget)); + EXPECT_FALSE(c.has_next()) << "budget spent"; + EXPECT_EQ(c.next(), 0U); +} + +// T13-F6: an absurd budget over a short multi-slice chain self-terminates at the tail +// slice's payload end -- yielding the written bytes plus the slice's zero-initialized +// unwritten tail, and NEVER aliasing block 0 (offset 0 is owned by a decoy chain). +TEST(SniiCompactPostingPoolTest, OverLargeBudgetSelfTerminates) { + CompactPostingPool pool; + // Decoy owns offset 0 so an erroneous alias to block 0 is detectable. + Chain decoy; + decoy.start(&pool); + for (uint8_t b : {static_cast(0xDE), static_cast(0xAD)}) { + decoy.put(&pool, b); + } + ASSERT_EQ(decoy.head, 0U) << "first chain must own pool offset 0"; + + Chain ch; + ch.start(&pool); + std::vector data; + for (uint32_t i = 0; i < 7; ++i) { + const auto b = static_cast(0x11 * (i + 1)); // 0x11..0x77, never 0 or 0xDE + data.push_back(b); + ch.put(&pool, b); + } + + const uint64_t over = 1ULL << 20; // far beyond the 7-byte payload + CompactPostingPool::Cursor c = pool.cursor(ch.head, over); + std::vector out; + uint64_t guard = 0; + while (c.has_next()) { + out.push_back(c.next()); + ASSERT_LT(++guard, over) << "cursor failed to self-terminate at the chain tail"; + } + // 7 bytes fill level 0 (cap kSliceSize_at(0)) then spill into the next level; the + // cursor stops at that tail slice's payload end. + const size_t expect_len = + CompactPostingPool::kSliceSize_at(0) + + CompactPostingPool::kSliceSize_at(CompactPostingPool::kNextLevel_at(0)); + EXPECT_EQ(out.size(), expect_len) << "must stop at the tail slice's end, not run on"; + ASSERT_GE(out.size(), data.size()); + EXPECT_EQ(std::vector(out.begin(), out.begin() + data.size()), data); + for (size_t i = data.size(); i < out.size(); ++i) { + EXPECT_EQ(out[i], 0U) << "unwritten tail byte " << i << " must read as zero"; + } + for (uint8_t b : out) { + EXPECT_NE(b, static_cast(0xDE)) << "cursor aliased block 0"; + } +} + +// T13-F7: a started-but-never-written chain with a zero budget yields nothing. +TEST(SniiCompactPostingPoolTest, EmptyChainYieldsNothing) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + CompactPostingPool::Cursor c = pool.cursor(ch.head, 0); + EXPECT_FALSE(c.has_next()); + EXPECT_EQ(c.next(), 0U); +} + +// T13-F8: exactly filling a level-0 slice then writing ONE more byte advances the chain +// to the scheduled next level and links the slices; the whole chain round-trips. +TEST(SniiCompactPostingPoolTest, SliceOverflowLinksCorrectly) { + CompactPostingPool pool; + Chain ch; + ch.start(&pool); + const uint32_t cap0 = CompactPostingPool::kSliceSize_at(0); + std::vector data; + for (uint32_t i = 0; i < cap0; ++i) { + const auto b = static_cast(0xA0 + i); + data.push_back(b); + ch.put(&pool, b); + } + EXPECT_EQ(ch.level, uint8_t {0}) << "an exact fill has not overflowed yet"; + // The boundary+1 byte forces the overflow + forward link. + ch.put(&pool, 0x5A); + data.push_back(0x5A); + EXPECT_EQ(ch.level, CompactPostingPool::kNextLevel_at(0)) + << "overflow must advance to the scheduled next level"; + EXPECT_EQ(ReadChain(pool, ch.head, data.size()), data); +} + +// T13-F9 (equivalence/golden): the full SpimiTermBuffer encode+decode path. The token +// feed exercises lexicographic term ordering (NOT first-seen), a multi-doc term, a freq>1 +// doc (banana@10 twice), and an out-of-order docid (cherry 100 then 50) that drives the +// finalize sort_by_docid + position reorder. The golden TermPostings were derived by hand +// from the documented tagged-varint contract and MUST stay identical across the inline +// move -- this is the core "new path == old path" check for T13. +TEST(SniiCompactPostingPoolTest, EndToEndPostingsEquivalence) { + SpimiTermBuffer buf(/*has_positions=*/true); // owned-vocab (string-keyed add_token) + buf.add_token("banana", 10, 0); + buf.add_token("apple", 5, 3); + buf.add_token("cherry", 100, 0); + buf.add_token("banana", 10, 5); // same doc -> freq 2, positions {0,5} + buf.add_token("apple", 7, 2); + buf.add_token("cherry", 50, 9); // docid < previous -> triggers sort_by_docid + reorder + buf.add_token("banana", 20, 1); + + std::vector got = buf.finalize_sorted(); + EXPECT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_EQ(got.size(), 3U); + + EXPECT_EQ(got[0].term, "apple"); + EXPECT_EQ(got[0].docids, (std::vector {5, 7})); + EXPECT_EQ(got[0].freqs, (std::vector {1, 1})); + EXPECT_EQ(got[0].positions_flat, (std::vector {3, 2})); + + EXPECT_EQ(got[1].term, "banana"); + EXPECT_EQ(got[1].docids, (std::vector {10, 20})); + EXPECT_EQ(got[1].freqs, (std::vector {2, 1})); + EXPECT_EQ(got[1].positions_flat, (std::vector {0, 5, 1})); + + EXPECT_EQ(got[2].term, "cherry"); + EXPECT_EQ(got[2].docids, (std::vector {50, 100})); + EXPECT_EQ(got[2].freqs, (std::vector {1, 1})); + EXPECT_EQ(got[2].positions_flat, (std::vector {9, 0})); +} + +// T13-F10 (error path): a borrowed-vocab buffer fed an out-of-range term-id latches an +// InvalidArgument into status(), ignores the token, and finalize_sorted() yields nothing +// (no spill, no crash). +TEST(SniiCompactPostingPoolTest, OutOfVocabTokenLatchesError) { + const std::vector vocab = {"alpha", "beta"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true); + buf.add_token(/*term_id=*/5U, /*docid=*/1U, /*pos=*/0U); // 5 >= vocab.size() == 2 + + std::vector got = buf.finalize_sorted(); + EXPECT_TRUE(got.empty()); + EXPECT_FALSE(buf.status().ok()); + EXPECT_TRUE(buf.status().is()) << buf.status().to_string(); +} diff --git a/be/test/storage/index/snii/writer/logical_index_writer_freq_stats_test.cpp b/be/test/storage/index/snii/writer/logical_index_writer_freq_stats_test.cpp new file mode 100644 index 00000000000000..8bdb50466e2ab2 --- /dev/null +++ b/be/test/storage/index/snii/writer/logical_index_writer_freq_stats_test.cpp @@ -0,0 +1,562 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/query/internal/docid_posting_reader.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +// T12 -- writer fused freqs statistics (single pass total + max), reused for the +// has_prx position-count check, stats_.sum_total_term_freq, and the DictEntry +// ttf_delta/max_freq. This suite guards: +// * the deterministic op-count: exactly ONE term-level freqs scan per term +// (was 3N docs-only / 4N with positions) via the term_freq_scans() seam; +// * value bit-identity: ttf_delta / max_freq / sum_total_term_freq read back +// equal to an independent reference, across windowed + slim(pod_ref/inline); +// * the fused pure helper on boundary inputs (empty/single/zeros/equal/u32max/ +// large random) vs a naive reference; +// * the validate_term error paths preserved after dropping its internal +// freqs-sum loop (length / position-count / strict-ascending), and that the +// has_prx position-count check now consumes the FUSED total. +namespace doris::snii::writer { +namespace { + +using doris::Status; // RETURN_IF_ERROR / Status::OK() expand to a bare Status. +using snii_test::assert_ok; +using snii_test::make_term; +using snii_test::MemoryFile; +using snii_test::PostingDoc; + +// One term whose postings are known up front, so the reference total/max can be +// recomputed in the test independently of the writer. +struct KnownTerm { + std::string term; + std::vector docs; + + uint64_t ref_sum() const { + uint64_t sum = 0; + for (const PostingDoc& doc : docs) { + sum += doc.positions.size(); + } + return sum; + } + uint32_t ref_max() const { + uint32_t max = 0; + for (const PostingDoc& doc : docs) { + max = std::max(max, static_cast(doc.positions.size())); + } + return max; + } + uint32_t df() const { return static_cast(docs.size()); } +}; + +// 480 docids with irregular gaps (deterministic LCG): the slim docs region PFOR +// exceeds the 256B inline threshold so the term becomes a pod_ref, while df < 512 +// keeps it slim (not windowed). Mirrors the proven generator in snii_query_test. +std::vector slim_pod_ref_docids() { + std::vector ids; + ids.reserve(480); + uint32_t cur = 0; + uint32_t state = 0x9e3779b9U; + for (int i = 0; i < 480; ++i) { + ids.push_back(cur); + state = state * 1664525U + 1013904223U; + cur += 1U + (state >> 23) % 250U; // gap in [1, 250] + } + return ids; +} + +// A corpus covering every term-level encoding branch with non-trivial freqs: +// "wide" df=600 >= kSlimDfThreshold(512) -> windowed (3 base-unit windows), +// freqs cycle 1,2,3 (max=3) so the per-window MaxOf is also exercised. +// "slimref" df=480 < 512, irregular gaps -> slim pod_ref, freq 1 (max=1). +// "tiny" df=4 -> slim inline, one doc freq 2 (max=2). +// "mid" df=300 consecutive -> slim, one doc freq 2 (max=2). +std::vector make_known_corpus() { + std::vector corpus; + + KnownTerm wide; + wide.term = "wide"; + for (uint32_t i = 0; i < 600; ++i) { + const uint32_t freq = (i % 3) + 1; + std::vector positions; + positions.reserve(freq); + for (uint32_t p = 0; p < freq; ++p) { + positions.push_back(p); + } + wide.docs.push_back({i, std::move(positions)}); + } + corpus.push_back(std::move(wide)); + + KnownTerm slimref; + slimref.term = "slimref"; + for (uint32_t docid : slim_pod_ref_docids()) { + slimref.docs.push_back({docid, {0}}); + } + corpus.push_back(std::move(slimref)); + + KnownTerm tiny; + tiny.term = "tiny"; + tiny.docs = {{.docid = 10, .positions = {0, 1}}, + {.docid = 20, .positions = {0}}, + {.docid = 30, .positions = {0}}, + {.docid = 40, .positions = {0}}}; + corpus.push_back(std::move(tiny)); + + KnownTerm mid; + mid.term = "mid"; + for (uint32_t i = 0; i < 300; ++i) { + mid.docs.push_back({i, i == 0 ? std::vector {0, 1} : std::vector {0}}); + } + corpus.push_back(std::move(mid)); + + return corpus; +} + +const KnownTerm& find_term(const std::vector& corpus, std::string_view term) { + for (const KnownTerm& kt : corpus) { + if (std::string_view(kt.term) == term) { + return kt; + } + } + ADD_FAILURE() << "term not in corpus: " << term; + return corpus.front(); +} + +// Builds the corpus into `file` and opens a reader over it. The corpus is left +// untouched so the caller can recompute references. +Status build_corpus_reader(MemoryFile* file, reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader, + const std::vector& corpus, format::IndexConfig config, + bool write_freq = true) { + SniiIndexInput input; + input.index_id = 7; + input.index_suffix = "Body"; + input.config = config; + input.write_freq = write_freq; + uint32_t max_docid = 0; + input.terms.reserve(corpus.size()); + for (const KnownTerm& kt : corpus) { + for (const PostingDoc& doc : kt.docs) { + max_docid = std::max(max_docid, doc.docid); + } + input.terms.push_back(make_term(kt.term, kt.docs)); + } + input.doc_count = max_docid + 1; + std::ranges::sort(input.terms, [](const TermPostings& lhs, const TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + + RETURN_IF_ERROR(reader::SniiSegmentReader::open(file, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +format::DictEntry lookup_entry(const reader::LogicalIndexReader& index_reader, + std::string_view term) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup(term, &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found) << term; + return entry; +} + +// ------------------------------------------------------------------------- +// Deterministic perf: exactly one term-level freqs scan per term. +// ------------------------------------------------------------------------- + +// The op-count seam (testing::term_freq_scans) increments once per fuse_freq_stats +// call, and process_term calls it exactly once per term -> N for an N-term build. +// Before the fusion the instrumented build scanned freqs 4x per term with +// positions (validate-sum + stats-sum + ttf-sum + max), i.e. 4 * N, so this +// assertion fails on the un-fused code and passes after the single-pass fuse. +TEST(SniiWriterTest, ProcessTermScansFreqsOncePerTerm) { + const std::vector corpus = make_known_corpus(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + + testing::reset_term_freq_scans(); + assert_ok(build_corpus_reader(&file, &segment_reader, &index_reader, corpus, + format::IndexConfig::kDocsPositions)); + + EXPECT_EQ(testing::term_freq_scans(), corpus.size()); +} + +// ------------------------------------------------------------------------- +// Value bit-identity: ttf_delta / max_freq / sum_total_term_freq. +// ------------------------------------------------------------------------- + +// FW-EQ-1 / FW-EQ-2: every term's read-back ttf_delta == sum(freqs) and +// max_freq == max(freqs), and the index-level sum_total_term_freq == the sum of +// every term's sum(freqs). The reference is computed independently from the known +// corpus. Holds before AND after the fusion (guards the CSE did not change values). +TEST(SniiWriterTest, FusedFreqStatsPreserveTtfMaxAndSum) { + const std::vector corpus = make_known_corpus(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_corpus_reader(&file, &segment_reader, &index_reader, corpus, + format::IndexConfig::kDocsPositions)); + + uint64_t expected_total = 0; + for (const KnownTerm& kt : corpus) { + const format::DictEntry entry = lookup_entry(index_reader, kt.term); + EXPECT_EQ(entry.df, kt.df()) << kt.term; + EXPECT_EQ(entry.ttf_delta, kt.ref_sum()) << kt.term; + EXPECT_EQ(entry.max_freq, kt.ref_max()) << kt.term; + expected_total += kt.ref_sum(); + } + EXPECT_EQ(index_reader.stats().sum_total_term_freq, expected_total); +} + +// FW-DF-windowed / FW-DF-slim-inline (+ slim pod_ref): the fused ttf/max are +// correct on each encoding branch. Also guards that the windowed per-window MaxOf +// / sum (NOT collapsed by this task) still produce a term whose term-level fused +// values match the reference. +TEST(SniiWriterTest, FusedFreqStatsCorrectAcrossWindowedAndSlimEncodings) { + const std::vector corpus = make_known_corpus(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_corpus_reader(&file, &segment_reader, &index_reader, corpus, + format::IndexConfig::kDocsPositions)); + + const format::DictEntry wide = lookup_entry(index_reader, "wide"); + EXPECT_EQ(wide.enc, format::DictEntryEnc::kWindowed); + EXPECT_EQ(wide.ttf_delta, find_term(corpus, "wide").ref_sum()); + EXPECT_EQ(wide.max_freq, find_term(corpus, "wide").ref_max()); + EXPECT_EQ(wide.max_freq, 3U); // sanity: cycling 1,2,3 + + const format::DictEntry slimref = lookup_entry(index_reader, "slimref"); + EXPECT_EQ(slimref.enc, format::DictEntryEnc::kSlim); + EXPECT_EQ(slimref.kind, format::DictEntryKind::kPodRef); + EXPECT_EQ(slimref.ttf_delta, find_term(corpus, "slimref").ref_sum()); + EXPECT_EQ(slimref.max_freq, find_term(corpus, "slimref").ref_max()); + + const format::DictEntry tiny = lookup_entry(index_reader, "tiny"); + EXPECT_EQ(tiny.enc, format::DictEntryEnc::kSlim); + EXPECT_EQ(tiny.kind, format::DictEntryKind::kInline); + EXPECT_EQ(tiny.ttf_delta, find_term(corpus, "tiny").ref_sum()); + EXPECT_EQ(tiny.max_freq, find_term(corpus, "tiny").ref_max()); +} + +// Frequency layout is an index-wide capability. Disabling it preserves ordinary +// position payloads while omitting both posting-level freqs and per-entry term +// frequency statistics across windowed and slim encodings. +TEST(SniiWriterTest, WriteFreqOffPreservesOrdinaryPositions) { + const std::vector corpus = make_known_corpus(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_corpus_reader(&file, &segment_reader, &index_reader, corpus, + format::IndexConfig::kDocsPositions, + /*write_freq=*/false)); + + bool found = false; + format::DictEntry wide; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup("wide", &found, &wide, &frq_base, &prx_base)); + ASSERT_TRUE(found); + ASSERT_EQ(wide.enc, format::DictEntryEnc::kWindowed); + EXPECT_EQ(wide.frq_docs_len, wide.frq_len); + EXPECT_GT(wide.prx_len, 0U); + EXPECT_FALSE(wide.term_stats_present); + EXPECT_EQ(wide.ttf_delta, 0U); + EXPECT_EQ(wide.max_freq, 0U); + EXPECT_EQ(wide.df, 600U); + + reader::DecodedPosting decoded; + assert_ok(reader::read_windowed_posting(index_reader, wide, frq_base, prx_base, + /*want_positions=*/true, + /*want_freq=*/false, &decoded)); + const KnownTerm& wide_term = find_term(corpus, "wide"); + ASSERT_EQ(decoded.docids.size(), wide_term.docs.size()); + ASSERT_EQ(decoded.positions.size(), wide_term.docs.size()); + uint64_t decoded_position_count = 0; + for (size_t i = 0; i < wide_term.docs.size(); ++i) { + EXPECT_EQ(decoded.docids[i], wide_term.docs[i].docid); + EXPECT_EQ(decoded.positions[i], wide_term.docs[i].positions); + decoded_position_count += decoded.positions[i].size(); + } + EXPECT_EQ(decoded_position_count, wide_term.ref_sum()); + EXPECT_TRUE(decoded.freqs.empty()); + + const format::DictEntry tiny = lookup_entry(index_reader, "tiny"); + ASSERT_EQ(tiny.enc, format::DictEntryEnc::kSlim); + ASSERT_EQ(tiny.kind, format::DictEntryKind::kInline); + EXPECT_EQ(tiny.frq_bytes.size(), tiny.inline_dd_disk_len); + EXPECT_FALSE(tiny.prx_bytes.empty()); + EXPECT_FALSE(tiny.term_stats_present); + EXPECT_EQ(tiny.ttf_delta, 0U); + EXPECT_EQ(tiny.max_freq, 0U); + + std::vector tiny_docids; + assert_ok(query::internal::read_docid_posting(index_reader, tiny, 0, 0, &tiny_docids)); + const KnownTerm& tiny_term = find_term(corpus, "tiny"); + ASSERT_EQ(tiny_docids.size(), tiny_term.docs.size()); + + ByteSource tiny_prx_source(Slice(tiny.prx_bytes)); + std::vector> tiny_positions; + assert_ok(format::read_prx_window(&tiny_prx_source, &tiny_positions)); + EXPECT_TRUE(tiny_prx_source.eof()); + ASSERT_EQ(tiny_positions.size(), tiny_term.docs.size()); + for (size_t i = 0; i < tiny_term.docs.size(); ++i) { + EXPECT_EQ(tiny_docids[i], tiny_term.docs[i].docid); + EXPECT_EQ(tiny_positions[i], tiny_term.docs[i].positions); + } +} + +// ------------------------------------------------------------------------- +// Pure helper boundaries (real production helper via the testing seam). +// ------------------------------------------------------------------------- + +// FW-PURE-*: fuse_freq_stats on degenerate/boundary inputs equals the hand +// computed reference. max stays 0 for an all-zero input (a freq of 0 never lowers +// the running max), and the u32 sum promotes into the u64 total. +TEST(SniiWriterTest, FuseFreqStatsMatchesReferenceOnEdgeInputs) { + { + const FreqStats fs = testing::fuse_freq_stats_for_test({}); + EXPECT_EQ(fs.total_freq, 0U); + EXPECT_EQ(fs.max_freq, 0U); + } + { + const FreqStats fs = testing::fuse_freq_stats_for_test({7}); + EXPECT_EQ(fs.total_freq, 7U); + EXPECT_EQ(fs.max_freq, 7U); + } + { + const FreqStats fs = testing::fuse_freq_stats_for_test({0, 0, 0}); + EXPECT_EQ(fs.total_freq, 0U); + EXPECT_EQ(fs.max_freq, 0U); + } + { + const FreqStats fs = testing::fuse_freq_stats_for_test({5, 5, 5, 5}); + EXPECT_EQ(fs.total_freq, 20U); + EXPECT_EQ(fs.max_freq, 5U); + } + { + const FreqStats fs = testing::fuse_freq_stats_for_test({1, UINT32_MAX, 2}); + EXPECT_EQ(fs.total_freq, static_cast(UINT32_MAX) + 3U); + EXPECT_EQ(fs.max_freq, UINT32_MAX); + } +} + +// FW-PURE-rand: a large fixed-seed random input agrees with an independent +// std::accumulate (u64) / std::max_element reference. +TEST(SniiWriterTest, FuseFreqStatsMatchesNaiveOnLargeRandomInput) { + std::mt19937 rng(0xC0FFEEU); + std::vector freqs(4096); + for (uint32_t& f : freqs) { + f = static_cast(rng()); + } + + const FreqStats fs = testing::fuse_freq_stats_for_test(freqs); + EXPECT_EQ(fs.total_freq, std::accumulate(freqs.begin(), freqs.end(), uint64_t {0})); + EXPECT_EQ(fs.max_freq, *std::ranges::max_element(freqs)); +} + +// ------------------------------------------------------------------------- +// Error paths preserved after dropping validate_term's internal sum loop. +// ------------------------------------------------------------------------- + +// FW-VAL-len-mismatch: freqs.size() != docids.size() is still rejected (the +// length check is independent of the freqs sum and must remain). +TEST(SniiWriterTest, ValidateTermRejectsFreqDocidLengthMismatch) { + TermPostings tp; + tp.term = "x"; + tp.docids = {0, 1}; + tp.freqs = {1}; // length 1 != 2 + tp.retain_positions = false; + + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = 2; + input.terms = {std::move(tp)}; + + MemoryFile file; + SniiCompoundWriter writer(&file); + const Status status = writer.add_logical_index(input); + EXPECT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find("span posting source: freqs length must equal docids"), + std::string::npos) + << status.to_string(); +} + +TEST(SniiWriterTest, ValidateTermAcceptsMissingFreqsForOrdinaryDocsOnlyTerm) { + TermPostings tp; + tp.term = "x"; + tp.docids = {0, 1}; + tp.retain_positions = false; + + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = 2; + input.terms = {std::move(tp)}; + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &index_reader)); + EXPECT_EQ(index_reader.stats().sum_total_term_freq, 2U); +} + +// FW-VAL-nonasc: non-strictly-ascending docids still rejected. +TEST(SniiWriterTest, ValidateTermRejectsNonAscendingDocids) { + TermPostings tp; + tp.term = "x"; + tp.docids = {5, 5}; // equal -> not strictly ascending + tp.freqs = {1, 1}; + + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = 6; + input.terms = {std::move(tp)}; + + MemoryFile file; + SniiCompoundWriter writer(&file); + const Status status = writer.add_logical_index(input); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +// FW-VAL-prx: the has_prx position-count check now consumes the FUSED total. +// A mismatch (positions != sum(freqs)) is rejected; a match is accepted AND the +// ttf/max read back from the fused stats. The accepted sub-case is the guard that +// validate_term receives the correct total (a mis-wired/zero total would wrongly +// reject this valid term). +TEST(SniiWriterTest, ValidateTermUsesFusedTotalForPositionCount) { + SniiIndexInput base; + base.index_id = 1; + base.index_suffix = "Body"; + base.config = format::IndexConfig::kDocsPositions; + base.doc_count = 1; + + { + TermPostings tp; + tp.term = "x"; + tp.docids = {0}; + tp.freqs = {2}; // sum(freqs) = 2 + tp.positions_flat = {7}; // but only 1 position + SniiIndexInput input = base; + input.terms = {std::move(tp)}; + + MemoryFile file; + SniiCompoundWriter writer(&file); + const Status status = writer.add_logical_index(input); + EXPECT_TRUE(status.is()) << status.to_string(); + } + { + TermPostings tp; + tp.term = "x"; + tp.docids = {0}; + tp.freqs = {2}; // sum(freqs) = 2 + tp.positions_flat = {7, 8}; // matches + SniiIndexInput input = base; + input.terms = {std::move(tp)}; + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &index_reader)); + + const format::DictEntry entry = lookup_entry(index_reader, "x"); + EXPECT_EQ(entry.ttf_delta, 2U); + EXPECT_EQ(entry.max_freq, 2U); + } +} + +// FW-docs-only-spimi: the real import path for an untokenized (keyword) or +// support_phrase=false index streams docs-only terms out of SpimiTermBuffer, +// whose to_postings derives retain_positions=false from the chain shape. The +// positioned-index invariant "a positionless term must be a declared common +// gram" must not fire for a docs-only config -- it aborted the BE on the first +// keyword-index memtable flush before the fix. Direct-vector inputs cannot +// cover this: TermPostings defaults retain_positions=true, so only the +// streamed source produces the false flag the invariant sees in production. +TEST(SniiWriterTest, DocsOnlyStreamedSpimiBuildSucceeds) { + SpimiTermBuffer buffer(/*has_positions=*/false); + buffer.add_token(std::string_view("alpha"), /*docid=*/0, /*pos=*/0); + buffer.add_token(std::string_view("beta"), /*docid=*/0, /*pos=*/1); + buffer.add_token(std::string_view("alpha"), /*docid=*/1, /*pos=*/0); + assert_ok(buffer.status()); + + SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = 2; + input.term_source = &buffer; + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &index_reader)); + + const format::DictEntry alpha = lookup_entry(index_reader, "alpha"); + EXPECT_EQ(alpha.df, 2U); + const format::DictEntry beta = lookup_entry(index_reader, "beta"); + EXPECT_EQ(beta.df, 1U); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/memory_reporter_test.cpp b/be/test/storage/index/snii/writer/memory_reporter_test.cpp new file mode 100644 index 00000000000000..f78975e353d16d --- /dev/null +++ b/be/test/storage/index/snii/writer/memory_reporter_test.cpp @@ -0,0 +1,85 @@ +// 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 "storage/index/snii/writer/memory_reporter.h" + +#include + +#include +#include + +#include "common/status.h" + +using doris::snii::writer::MemoryReporter; + +TEST(SniiMemoryReporter, StartsAtZero) { + MemoryReporter reporter; // null callback + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporter, TracksPositiveAndNegativeDeltas) { + MemoryReporter reporter; + reporter.report(+100); + reporter.report(+50); + EXPECT_EQ(reporter.current_bytes(), 150); + reporter.report(-30); + EXPECT_EQ(reporter.current_bytes(), 120); + reporter.report(-120); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporter, NullCallbackIsSafe) { + MemoryReporter reporter(nullptr); // explicit null ConsumeReleaseFn + reporter.report(+10); + reporter.report(-10); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporter, CallbackFiresWithSameDelta) { + std::vector sink; + MemoryReporter reporter([&sink](int64_t delta) { sink.push_back(delta); }); + reporter.report(+100); + reporter.report(-40); + ASSERT_EQ(sink.size(), 2U); + EXPECT_EQ(sink[0], 100); + EXPECT_EQ(sink[1], -40); + EXPECT_EQ(reporter.current_bytes(), 60); +} + +TEST(SniiMemoryReporter, CallbackSumMirrorsCurrentBytes) { + int64_t external_total = 0; + MemoryReporter reporter([&external_total](int64_t delta) { external_total += delta; }); + reporter.report(+100); + reporter.report(+250); + reporter.report(-75); + reporter.report(+12); + reporter.report(-200); + EXPECT_EQ(external_total, reporter.current_bytes()); +} + +TEST(SniiMemoryReporter, ZeroDeltaIsDebouncedNoOp) { + int fire_count = 0; + MemoryReporter reporter([&fire_count](int64_t) { ++fire_count; }); + reporter.report(+100); + EXPECT_EQ(reporter.current_bytes(), 100); + reporter.report(0); + EXPECT_EQ(reporter.current_bytes(), 100); + // report(0) is debounced: it neither moves the total nor invokes the + // callback (spill paths emit frequent zero deltas; forwarding them would + // spam the shared limiter for no state change). + EXPECT_EQ(fire_count, 1); +} diff --git a/be/test/storage/index/snii/writer/memory_reporter_wiring_test.cpp b/be/test/storage/index/snii/writer/memory_reporter_wiring_test.cpp new file mode 100644 index 00000000000000..2524f270c43f68 --- /dev/null +++ b/be/test/storage/index/snii/writer/memory_reporter_wiring_test.cpp @@ -0,0 +1,296 @@ +// 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. + +// P1 wiring + P3 seam: verifies that the writer-level MemoryReporter sees the REAL +// resident-byte deltas reported by SpillableByteBuffer (dict side: ram_bytes_) and +// SpimiTermBuffer (SPIMI side: arena_bytes() + slot_of_.capacity()*4), positive on +// grow and negative on every free, exactly once -- so the counter returns to a +// known measurement after a build and to 0 after a full drain (a missed negative +// would be a Doris LOAD-tracker leak). The reporter must NEVER count live_bytes_. + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/spillable_byte_buffer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +namespace doris::snii::writer { +using doris::Status; +namespace { + +// Discards appended bytes; only counts how many were written (SpillableByteBuffer +// stream_into / spill targets need a sink, but these tests check the reporter). +class NullWriter : public doris::snii::io::FileWriter { +public: + Status append(Slice s) override { + written_ += s.size(); + return Status::OK(); + } + Status finalize() override { return Status::OK(); } + uint64_t bytes_written() const override { return written_; } + +private: + uint64_t written_ = 0; +}; + +std::vector Block(size_t n, uint8_t seed) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = static_cast((i + seed) & 0xFF); + } + return v; +} + +// ---- dict side (SpillableByteBuffer) --------------------------------------- + +// Every RAM append reports exactly its byte count; the reporter equals the +// RAM-resident size, which (un-spilled) equals buf.size(). +TEST(SniiMemoryReporterWiring, DictRamDeltasMatchSize) { + MemoryReporter reporter; + SpillableByteBuffer buf(/*cap=*/UINT64_MAX, "dict", &reporter); + int64_t total = 0; + for (int b = 0; b < 5; ++b) { + auto chunk = Block(1000 + b, static_cast(b)); + total += static_cast(chunk.size()); + ASSERT_TRUE(buf.append(Slice(chunk)).ok()); + EXPECT_EQ(reporter.current_bytes(), total); + } + // Move-append path reports identically. + auto moved = Block(777, 9); + total += 777; + ASSERT_TRUE(buf.append_move(std::move(moved)).ok()); + EXPECT_EQ(reporter.current_bytes(), total); + EXPECT_EQ(reporter.current_bytes(), static_cast(buf.size())); +} + +// Destroying an UN-spilled buffer must release its resident RAM (the common path: +// most dict buffers never spill). A missed negative here would leak into Doris's +// LOAD MemTracker. After the buffer's scope, the reporter must be back to 0. +TEST(SniiMemoryReporterWiring, DictDtorReleasesUnspilledRam) { + MemoryReporter reporter; + { + SpillableByteBuffer buf(/*cap=*/UINT64_MAX, "dict", &reporter); + ASSERT_TRUE(buf.append(Slice(Block(5000, 1))).ok()); + ASSERT_TRUE(buf.append_move(Block(3000, 2)).ok()); + EXPECT_EQ(reporter.current_bytes(), 8000); + EXPECT_FALSE(buf.spilled()); + } // ~SpillableByteBuffer must report -8000 + EXPECT_EQ(reporter.current_bytes(), 0); +} + +// UNIFIED gate-2: a small cap on the shared reporter triggers the dict spill BEFORE +// its (huge) local cap_bytes_ is reached -- proving the spill is driven by the writer's +// total RAM (reporter->over_cap()), not a per-buffer threshold. +TEST(SniiMemoryReporterWiring, UnifiedCapDrivesDictSpill) { + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/8000); + SpillableByteBuffer buf(/*local cap=*/UINT64_MAX, "dict", &reporter); + ASSERT_TRUE(buf.append(Slice(Block(5000, 1))).ok()); + EXPECT_FALSE(buf.spilled()); // 5000 < unified cap 8000 + ASSERT_TRUE(buf.append(Slice(Block(4000, 2))).ok()); // 9000 >= 8000 -> spill + EXPECT_TRUE(buf.spilled()); // unified cap fired, not local + EXPECT_EQ(reporter.current_bytes(), 0); // resident handed to disk + ASSERT_TRUE(buf.seal().ok()); +} + +// A spill drops the resident tier: the reporter returns to 0 (one negative == +// prior ram_bytes_) while buf.size() still counts the on-disk bytes. +TEST(SniiMemoryReporterWiring, DictSpillFreesRam) { + MemoryReporter reporter; + SpillableByteBuffer buf(/*cap=*/8192, "dict", &reporter); + uint64_t total = 0; + for (int b = 0; b < 8; ++b) { // 8 x 4 KiB through an 8 KiB cap -> spills + auto chunk = Block(4096, static_cast(b)); + total += chunk.size(); + ASSERT_TRUE(buf.append(Slice(chunk)).ok()); + } + EXPECT_TRUE(buf.spilled()); + // RAM tier was handed to disk -> reporter back to 0; size() still totals bytes. + EXPECT_EQ(reporter.current_bytes(), 0); + EXPECT_EQ(buf.size(), total); + ASSERT_TRUE(buf.seal().ok()); +} + +// ---- SPIMI side (SpimiTermBuffer) ------------------------------------------ + +// A null reporter must not affect behavior (default off-Doris path). +TEST(SniiMemoryReporterWiring, SpimiNullReporterIsSafe) { + std::vector vocab = {"a", "b", "c"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, /*reporter=*/nullptr); + buf.add_token(0, 0, 0); + buf.add_token(1, 0, 1); + auto terms = buf.finalize_sorted(); + EXPECT_EQ(terms.size(), 2U); +} + +// Borrowed-vocab construction reports the vocab-sized slot index up front +// (slot_of_.capacity()*4). assign(N,0) makes capacity == N, so the initial report +// is exactly N*4 -- a directly-measurable accuracy check. +TEST(SniiMemoryReporterWiring, SpimiInitialSlotIndexIsExact) { + std::vector vocab(64, "x"); + MemoryReporter reporter; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, &reporter); + EXPECT_EQ(reporter.current_bytes(), static_cast(vocab.size()) * 4); +} + +// Adding tokens grows the reporter ABOVE the initial slot-index figure (the arena +// is now non-empty), and a full drain returns it to 0 -- every negative reported, +// no leak. Also proves the counter tracks arena_bytes(), not the (default-mode) +// live_bytes_ estimate, which stays 0 when spill_threshold_bytes_ == 0. +TEST(SniiMemoryReporterWiring, SpimiGrowThenDrainNetZero) { + std::vector vocab = {"alpha", "beta", "gamma", "delta"}; + MemoryReporter reporter; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, &reporter); + const int64_t base = reporter.current_bytes(); // slot index only + EXPECT_EQ(base, static_cast(vocab.size()) * 4); + for (uint32_t doc = 0; doc < 200; ++doc) { + for (uint32_t t = 0; t < 4; ++t) { + buf.add_token(t, doc, /*pos=*/0); + } + } + // Arena is now resident -> strictly above the slot-index-only baseline. Because + // spill_threshold_bytes_ == 0, live_bytes_ is never accumulated; a counter that + // mistakenly reported live_bytes_ would still be at `base` here. + EXPECT_GT(reporter.current_bytes(), base); + auto terms = buf.finalize_sorted(); + EXPECT_EQ(terms.size(), 4U); + // Drain freed the arena + slot index: every negative reported -> exactly 0. + EXPECT_EQ(reporter.current_bytes(), 0); +} + +// Streaming drain (for_each_term_sorted) frees identically -> net zero. +TEST(SniiMemoryReporterWiring, SpimiStreamingDrainNetZero) { + std::vector vocab = {"k0", "k1", "k2"}; + MemoryReporter reporter; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &reporter); + for (uint32_t doc = 0; doc < 100; ++doc) { + for (uint32_t t = 0; t < 3; ++t) { + buf.add_token(t, doc, 0); + } + } + EXPECT_GT(reporter.current_bytes(), 0); + size_t seen = 0; + // Integrated for_each_term_sorted returns a [[nodiscard]] doris::Status (drift + // from the standalone void); consume it and assert the happy-path drain succeeds. + ASSERT_TRUE(buf.for_each_term_sorted([&seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR(consume_streamed_term(std::move(source))); + ++seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(seen, 3U); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +// The spilled finalize path (merge_runs) must also balance: a small spill threshold +// forces at least one spill mid-build; after the full drain the reporter is 0 (the +// spill's arena-reset negative AND merge_runs' slot_of_ free are both reported). +TEST(SniiMemoryReporterWiring, SpimiSpillPathNetZero) { + std::vector vocab = {"w0", "w1", "w2", "w3", "w4", "w5", "w6", "w7"}; + // Tiny UNIFIED cap on the reporter so the REAL resident size (arena + slot index) + // crosses it early and forces spills, exercising the drain_to_writer / merge_runs + // free sites. (The local spill_threshold arg is unused once a reporter is attached.) + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/256, + MemoryReporter::CapPolicy::kSpillThreshold); + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, &reporter); + for (uint32_t doc = 0; doc < 500; ++doc) { + for (uint32_t t = 0; t < 8; ++t) { + buf.add_token(t, doc, doc & 7); + } + } + size_t seen = 0; + // for_each_term_sorted now surfaces spill/merge I/O via a [[nodiscard]] Status. + ASSERT_TRUE(buf.for_each_term_sorted([&seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR(consume_streamed_term(std::move(source))); + ++seen; + return Status::OK(); + }).ok()); + EXPECT_TRUE(buf.status().ok()); + EXPECT_EQ(seen, 8U); + // All spill/merge frees reported -> back to 0 (no MemTracker leak). + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiMemoryReporterWiring, SpimiSpillMaterializationIsAccountedThroughRunWrite) { + int64_t mirrored_bytes = 0; + int64_t peak_bytes = 0; + MemoryReporter reporter( + [&](int64_t delta) { + mirrored_bytes += delta; + peak_bytes = std::max(peak_bytes, mirrored_bytes); + }, + /*cap_bytes=*/0, MemoryReporter::CapPolicy::kSpillThreshold); + std::vector vocab = {"wide"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, &reporter); + buf.set_forced_spill_min_arena_bytes(1); + + constexpr uint32_t kDocumentsBeforeSpill = 8192; + for (uint32_t docid = 0; docid < kDocumentsBeforeSpill; ++docid) { + buf.add_token(/*term_id=*/0, docid, /*pos=*/docid); + } + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + const int64_t before_spill = reporter.current_bytes(); + ASSERT_EQ(mirrored_bytes, before_spill); + peak_bytes = before_spill; + + buf.request_global_spill_for_test(); + buf.add_token(/*term_id=*/0, kDocumentsBeforeSpill, /*pos=*/kDocumentsBeforeSpill); + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_EQ(buf.run_count_for_test(), 1U); + + constexpr int64_t kPostingArrayBytes = + static_cast(kDocumentsBeforeSpill + 1) * 3 * sizeof(uint32_t); + EXPECT_GE(peak_bytes - before_spill, 2 * kPostingArrayBytes); + + size_t terms_seen = 0; + ASSERT_TRUE(buf.for_each_term_sorted([&terms_seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR(consume_streamed_term(std::move(source))); + ++terms_seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(terms_seen, 1U); + EXPECT_EQ(reporter.current_bytes(), 0); + EXPECT_EQ(mirrored_bytes, 0); +} + +// Destructor balance: a buffer destroyed WITHOUT a drain (e.g. an aborted build) +// must return its reported resident bytes so nothing leaks in the tracker. +TEST(SniiMemoryReporterWiring, SpimiDestructorBalancesOnAbort) { + std::vector vocab = {"p", "q", "r"}; + MemoryReporter reporter; + { + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill=*/0, &reporter); + for (uint32_t doc = 0; doc < 50; ++doc) { + for (uint32_t t = 0; t < 3; ++t) { + buf.add_token(t, doc, 0); + } + } + EXPECT_GT(reporter.current_bytes(), 0); + // No drain: leave the build incomplete and let the destructor run. + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/phase_a_readback_test.cpp b/be/test/storage/index/snii/writer/phase_a_readback_test.cpp new file mode 100644 index 00000000000000..27b910f95d187a --- /dev/null +++ b/be/test/storage/index/snii/writer/phase_a_readback_test.cpp @@ -0,0 +1,303 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/scoring_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/stats/snii_stats_provider.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +// PHASE A read-back self-validation (design 2026-06-18-snii-read-byte-optimizations +// sections 1 + 3 + 4). Builds a scoring index with a high-df term that spans +// multiple ADAPTIVE-size windows plus low-df terms, then asserts: +// (a) for every window, decoding ONLY the docs-only prefix slice +// [frq window start, +frq_docs_len) via read_frq_window_docs yields exactly +// the same docids as decoding the full window; +// (b) sum of window doc_counts == df and the windows tile the posting in order; +// (c) max_norm is non-zero for windows whose docs have non-default norms and +// equals the tightest (smallest encoded) norm in the window; +// (d) term_query / phrase_query / scoring (exhaustive vs WAND) agree with the +// oracle docids. +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::format; // NOLINT +using namespace doris::snii::writer; // NOLINT +using doris::snii::query::Bm25Params; +using doris::snii::query::ScoredDoc; +using doris::snii::stats::SniiStatsProvider; + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_phaseA_readback_" + std::to_string(getpid()) + "_" + + std::to_string(counter++) + ".idx"; +} + +// Corpus: every doc has "hot" (very high df -> adaptive 1024-doc windows). "spark" +// follows "hot" consecutively in docs where d % 7 == 0. "rare" is a tiny term. +// Doc lengths vary by docid so encoded norms differ across windows and within a +// window, exercising real max_norm. +struct Corpus { + uint32_t doc_count = 12000; // df(hot)=12000 >= kAdaptiveWindowDfThreshold(8192) + std::vector hot_docs; // every doc + std::vector spark_docs; // d % 7 == 0 + std::vector rare_docs = {5, 91, 4096, 11999}; + std::vector phrase_oracle; // docs with consecutive "hot spark" + std::vector doc_len; // per-doc length (drives encoded norm) +}; + +// Per-doc length: varies so encoded norms are non-default (>1) and differ; the +// minimum within each 1024-doc window is deterministic. +uint64_t DocLen(uint32_t d) { + return 2 + (d % 250); +} + +Corpus MakeCorpus() { + Corpus c; + c.doc_len.assign(c.doc_count, 0); + for (uint32_t d = 0; d < c.doc_count; ++d) { + c.hot_docs.push_back(d); + c.doc_len[d] = DocLen(d); + if (d % 7 == 0) { + c.spark_docs.push_back(d); + c.phrase_oracle.push_back(d); + } + } + return c; +} + +TermPostings MakePosTerm(const std::string& term, const std::vector& docs, uint32_t pos) { + TermPostings tp; + tp.term = term; + tp.docids = docs; + tp.freqs.assign(docs.size(), 1); + tp.positions_flat.assign(docs.size(), pos); // one position per doc, flat + return tp; +} + +SniiIndexInput MakeIndex(const Corpus& c) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositionsScoring; // scoring -> norms available + in.doc_count = c.doc_count; + in.target_dict_block_bytes = 1; // one block per term + in.encoded_norms.resize(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + in.encoded_norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + // Lexicographically sorted: hot < rare < spark. + in.terms.push_back(MakePosTerm("hot", c.hot_docs, /*pos=*/0)); + in.terms.push_back(MakePosTerm("rare", c.rare_docs, /*pos=*/0)); + in.terms.push_back(MakePosTerm("spark", c.spark_docs, /*pos=*/1)); + uint64_t token_count = 0; + for (const auto& term : in.terms) { + token_count += term.positions_flat.size(); + } + in.common_grams_metadata = snii_test::make_plain_scoring_metadata(in.doc_count, token_count); + return in; +} + +// Smallest encoded norm across [start, end) docids (the tightest WAND norm). +uint8_t WindowMinNorm(const std::vector& norms, uint32_t start, uint32_t end) { + uint8_t best = std::numeric_limits::max(); + for (uint32_t d = start; d < end; ++d) { + best = std::min(best, norms[d]); + } + return best; +} + +} // namespace + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPhaseAReadBack, DocsPrefixTilesAndMaxNormTightAndQueriesAgree) { + const Corpus c = MakeCorpus(); + std::vector norms(c.doc_count); + for (uint32_t d = 0; d < c.doc_count; ++d) { + norms[d] = doris::snii::query::encode_norm(c.doc_len[d]); + } + + const std::string path = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(MakeIndex(c)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local); + reader::SniiSegmentReader seg; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered, &seg).ok()); + reader::LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + // Resolve "hot" -> windowed DictEntry + the block's frq_base. + bool found = false; + DictEntry hot; + uint64_t frq_base = 0, prx_base = 0; + ASSERT_TRUE(idx.lookup("hot", &found, &hot, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(hot.df, c.doc_count); + EXPECT_EQ(hot.enc, DictEntryEnc::kWindowed); + EXPECT_TRUE(hot.has_sb); + + const uint64_t prelude_abs = + idx.section_refs().posting_region.offset + frq_base + hot.frq_off_delta; + std::vector prelude_bytes; + ASSERT_TRUE(local.read_at(prelude_abs, hot.prelude_len, &prelude_bytes).ok()); + FrqPreludeReader prelude; + ASSERT_TRUE(FrqPreludeReader::open(Slice(prelude_bytes), &prelude).ok()); + + // Adaptive sizing: df=12000 -> 1024-doc windows -> >= 2 windows. + ASSERT_GT(prelude.n_windows(), 1U); + // Most windows are the adaptive size (the last may be a remainder). + WindowMeta w0; + ASSERT_TRUE(prelude.window(0, &w0).ok()); + EXPECT_EQ(w0.doc_count, doris::snii::format::kAdaptiveWindowDocs); + + // The grouped layout puts all dd regions in the dd-block right after the prelude. + const uint64_t dd_block_start = prelude_abs + hot.prelude_len; + std::vector tiled; + uint64_t summed = 0; + uint64_t expect_win_base = 0; + uint64_t expect_dd_off = 0; + uint32_t doc_cursor = 0; + bool any_nonzero_max_norm = false; + + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta m; + ASSERT_TRUE(prelude.window(w, &m).ok()); + EXPECT_EQ(m.win_base, expect_win_base) << "win_base w=" << w; + ASSERT_GT(m.dd_disk_len, 0U) << "w=" << w; + // (a) dd regions are contiguous in the dd-block (the docs-only run is one range). + EXPECT_EQ(m.dd_off, expect_dd_off) << "dd_off contiguity w=" << w; + + // Decode the window's dd region ALONE (the freq region bytes are NOT fetched). + FrqRegionMeta dd_meta; + dd_meta.zstd = m.dd_zstd; + dd_meta.uncomp_len = m.dd_uncomp_len; + dd_meta.disk_len = m.dd_disk_len; + dd_meta.crc = m.crc_dd; + std::vector dd_bytes; + ASSERT_TRUE(local.read_at(dd_block_start + m.dd_off, m.dd_disk_len, &dd_bytes).ok()); + std::vector dd_docs; + ASSERT_TRUE(decode_dd_region(Slice(dd_bytes), dd_meta, m.win_base, &dd_docs).ok()) + << "dd decode w=" << w; + ASSERT_EQ(dd_docs.size(), m.doc_count) << "w=" << w; + + // (c) max_norm equals the tightest (smallest encoded) norm in the window. + const uint32_t exp_min_norm = WindowMinNorm(norms, doc_cursor, doc_cursor + m.doc_count); + EXPECT_EQ(m.max_norm, exp_min_norm) << "max_norm w=" << w; + if (m.max_norm != 0) { + any_nonzero_max_norm = true; + } + + // (b) tiling: doc_counts sum and concatenated docids stay ascending. + summed += m.doc_count; + expect_win_base = m.last_docid; + expect_dd_off += m.dd_disk_len; + doc_cursor += m.doc_count; + tiled.insert(tiled.end(), dd_docs.begin(), dd_docs.end()); + } + // The dd-block length equals the sum of per-window dd region lengths. + EXPECT_EQ(prelude.dd_block_len(), expect_dd_off); + + // (b) windows tile the whole posting [0, doc_count). + EXPECT_EQ(summed, c.doc_count); + ASSERT_EQ(tiled.size(), c.doc_count); + EXPECT_EQ(tiled.front(), 0U); + EXPECT_EQ(tiled.back(), c.doc_count - 1); + for (size_t i = 1; i < tiled.size(); ++i) { + ASSERT_LT(tiled[i - 1], tiled[i]) << "non-ascending at " << i; + } + // (c) at least one window's docs have non-default norms -> non-zero max_norm. + EXPECT_TRUE(any_nonzero_max_norm); + + // --- slim/inline term frq_docs_len plumbing ('rare' is tiny -> inline) --- + bool rfound = false; + DictEntry rare; + uint64_t rf = 0, rp = 0; + ASSERT_TRUE(idx.lookup("rare", &rfound, &rare, &rf, &rp).ok()); + ASSERT_TRUE(rfound); + EXPECT_EQ(rare.df, c.rare_docs.size()); + + // (d) term_query returns exactly the oracle docid set for each term. + std::vector hot_q; + ASSERT_TRUE(doris::snii::query::term_query(idx, "hot", &hot_q).ok()); + EXPECT_EQ(hot_q, c.hot_docs); + + std::vector rare_q; + ASSERT_TRUE(doris::snii::query::term_query(idx, "rare", &rare_q).ok()); + EXPECT_EQ(rare_q, c.rare_docs); + + std::vector spark_q; + ASSERT_TRUE(doris::snii::query::term_query(idx, "spark", &spark_q).ok()); + EXPECT_EQ(spark_q, c.spark_docs); + + // (d) phrase_query "hot spark" returns exactly the consecutive-occurrence docs. + std::vector phrase_q; + ASSERT_TRUE(doris::snii::query::phrase_query(idx, {"hot", "spark"}, &phrase_q).ok()); + EXPECT_EQ(phrase_q, c.phrase_oracle); + + // (d) scoring: WAND top-K == exhaustive top-K (uses the real per-window max_norm). + SniiStatsProvider stats; + ASSERT_TRUE(SniiStatsProvider::open(&idx, &stats).ok()); + const Bm25Params params; + for (uint32_t k : {1U, 5U, 50U}) { + std::vector ex, wa; + ASSERT_TRUE(doris::snii::query::scoring_query_exhaustive(idx, stats, {"hot", "rare"}, k, + params, &ex) + .ok()); + ASSERT_TRUE( + doris::snii::query::scoring_query_wand(idx, stats, {"hot", "rare"}, k, params, &wa) + .ok()); + ASSERT_EQ(wa.size(), ex.size()) << "k=" << k; + for (size_t i = 0; i < ex.size(); ++i) { + EXPECT_EQ(wa[i].docid, ex[i].docid) << "k=" << k << " i=" << i; + EXPECT_NEAR(wa[i].score, ex[i].score, 1e-9) << "k=" << k << " i=" << i; + } + } + + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/writer/posting_interleave_test.cpp b/be/test/storage/index/snii/writer/posting_interleave_test.cpp new file mode 100644 index 00000000000000..05f015236bd90d --- /dev/null +++ b/be/test/storage/index/snii/writer/posting_interleave_test.cpp @@ -0,0 +1,560 @@ +// 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 +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/snii_compound_writer.h" + +// Interleaved posting-region read-back validation (docs/design/frqprx-interleave- +// design.md section 10.2). The former separate .frq POD and .prx POD are merged +// into ONE posting region in which each pod_ref term writes [prx span][frq span] +// contiguously, in term order. These tests assert the writer/reader contract for +// that layout: per-term contiguity (frq_off_delta == prx_off_delta + prx_len when +// has_prx), independent delta resolution, byte-correct sub-ranges, docs-only-prefix +// containment inside the frq span, multi-index isolation, empty-index tier recovery +// from the persisted flag, INLINE-between-pod_ref gaplessness, and the R1 docs-only- +// with-pod_ref tier-recovery regression guard. +namespace { + +using namespace doris::snii; // NOLINT +using namespace doris::snii::format; // NOLINT +using namespace doris::snii::writer; // NOLINT +using doris::snii::reader::DecodedPosting; +using doris::snii::reader::LogicalIndexReader; +using doris::snii::reader::SniiSegmentReader; +using doris::snii::reader::read_windowed_posting; + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_posting_interleave_" + std::to_string(getpid()) + "_" + + std::to_string(counter++) + ".idx"; +} + +TermPostings MakeTerm(const std::string& term, const std::vector& docs, + bool with_positions, uint32_t positions_per_doc) { + TermPostings tp; + tp.term = term; + tp.docids = docs; + tp.freqs.assign(docs.size(), positions_per_doc); + if (with_positions) { + for (size_t i = 0; i < docs.size(); ++i) { + for (uint32_t p = 0; p < positions_per_doc; ++p) { + tp.positions_flat.push_back(p); // deterministic ascending positions + } + } + } + tp.retain_positions = with_positions; + return tp; +} + +// Writes one container with one logical index, returns its path. Keeps the +// SniiIndexInput's terms (the writer reads but does not retain them). +std::string WriteOne(const SniiIndexInput& in) { + const std::string path = TempPath(); + io::LocalFileWriter w; + EXPECT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + EXPECT_TRUE(cw.add_logical_index(in).ok()); + EXPECT_TRUE(cw.finish().ok()); + return path; +} + +// Enumerates every term's DictEntry + owning block frq/prx bases via prefix_terms +// (empty prefix = all terms, in lexicographic order). +std::vector AllEntries(const LogicalIndexReader& idx) { + std::vector out; + EXPECT_TRUE(idx.prefix_terms("", &out).ok()); + return out; +} + +SniiIndexInput BaseInput(uint64_t id, std::string suffix, IndexConfig cfg) { + SniiIndexInput in; + in.index_id = id; + in.index_suffix = std::move(suffix); + in.config = cfg; + in.target_dict_block_bytes = 256; // many small blocks -> exercise multiple bases + return in; +} + +} // namespace + +// Test #1: round-trip with positions (interleave correctness). Slim-pod_ref and +// windowed terms (df below and above kSlimDfThreshold). For each pod_ref term: +// - frq_off_delta == prx_off_delta + prx_len (gated on has_prx); +// - resolve_frq_window / resolve_prx_window land inside posting_region; +// - decoded docids/freqs/positions equal the input. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, RoundTripWithPositionsContiguous) { + SniiIndexInput in = BaseInput(1, "body", IndexConfig::kDocsPositions); + const uint32_t doc_count = 4096; + in.doc_count = doc_count; + + // windowed term (df=2000 >= kSlimDfThreshold=512), slim pod_ref (df=600 forced + // large enough to exceed the inline threshold via several positions/doc), and an + // INLINE tiny term -- all positions-bearing. + std::vector wide_docs, slim_docs; + for (uint32_t d = 0; d < 2000; ++d) { + wide_docs.push_back(d * 2); + } + for (uint32_t d = 0; d < 600; ++d) { + slim_docs.push_back(d * 3 + 1); + } + in.terms.push_back(MakeTerm("aa_wide", wide_docs, /*with_positions=*/true, 3)); + in.terms.push_back(MakeTerm("bb_slim", slim_docs, /*with_positions=*/true, 4)); + in.terms.push_back(MakeTerm("cc_tiny", {7, 19, 33}, /*with_positions=*/true, 1)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + EXPECT_TRUE(idx.has_positions()); + + const RegionRef& region = idx.section_refs().posting_region; + for (const auto& hit : AllEntries(idx)) { + const DictEntry& e = hit.entry; + if (e.kind != DictEntryKind::kPodRef) { + continue; // INLINE term: nothing in region + } + // Contiguity property (gated on has_prx): frq span immediately follows prx span. + EXPECT_EQ(e.frq_off_delta, e.prx_off_delta + e.prx_len) << e.term; + + uint64_t foff = 0, flen = 0, poff = 0, plen = 0; + ASSERT_TRUE(idx.resolve_frq_window(e, hit.frq_base, &foff, &flen).ok()) << e.term; + ASSERT_TRUE(idx.resolve_prx_window(e, hit.prx_base, &poff, &plen).ok()) << e.term; + // Both windows land inside the single posting region. + EXPECT_GE(foff, region.offset) << e.term; + EXPECT_LE(foff + flen, region.offset + region.length) << e.term; + EXPECT_GE(poff, region.offset) << e.term; + EXPECT_LE(poff + plen, region.offset + region.length) << e.term; + } + + // Docids round-trip for every term. + for (const auto& [term, expect] : std::vector>> { + {"aa_wide", wide_docs}, {"bb_slim", slim_docs}, {"cc_tiny", {7, 19, 33}}}) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, term, &got).ok()) << term; + EXPECT_EQ(got, expect) << term; + } + + // Positions round-trip for the windowed term via the full decode path. + bool found = false; + DictEntry wide_e; + uint64_t fb = 0, pb = 0; + ASSERT_TRUE(idx.lookup("aa_wide", &found, &wide_e, &fb, &pb).ok()); + ASSERT_TRUE(found); + DecodedPosting dp; + ASSERT_TRUE(read_windowed_posting(idx, wide_e, fb, pb, /*want_positions=*/true, + /*want_freq=*/true, &dp) + .ok()); + ASSERT_EQ(dp.docids.size(), wide_docs.size()); + EXPECT_EQ(dp.docids, wide_docs); + ASSERT_EQ(dp.positions.size(), wide_docs.size()); + for (const auto& pv : dp.positions) { + ASSERT_EQ(pv.size(), 3U); // positions_per_doc + EXPECT_EQ(pv[0], 0U); + EXPECT_EQ(pv[1], 1U); + EXPECT_EQ(pv[2], 2U); + } + std::remove(path.c_str()); +} + +// Test #2 / #3: byte-correctness of the combined region. Read the raw posting +// region bytes; assert each term's [prx_off_delta, +prx_len) and +// [frq_off_delta, +frq_len) sub-ranges, and that the docs-only prefix +// [frq_off_delta, +frq_docs_len) stays INSIDE the frq span (never reads prx). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, ByteCorrectnessAndDocsPrefixStaysInFrqSpan) { + SniiIndexInput in = BaseInput(1, "body", IndexConfig::kDocsPositions); + in.doc_count = 4096; + std::vector wide_docs; + for (uint32_t d = 0; d < 1500; ++d) { + wide_docs.push_back(d); + } + in.terms.push_back(MakeTerm("aa_wide", wide_docs, /*with_positions=*/true, 2)); + std::vector slim_docs; + for (uint32_t d = 0; d < 600; ++d) { + slim_docs.push_back(d + 5); + } + in.terms.push_back(MakeTerm("bb_slim", slim_docs, /*with_positions=*/true, 3)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + const RegionRef& region = idx.section_refs().posting_region; + std::vector region_bytes; + ASSERT_TRUE(local.read_at(region.offset, region.length, ®ion_bytes).ok()); + + for (const auto& hit : AllEntries(idx)) { + const DictEntry& e = hit.entry; + if (e.kind != DictEntryKind::kPodRef) { + continue; + } + // prx span precedes frq span; both lie inside the region (relative to base). + const uint64_t prx_in = hit.prx_base + e.prx_off_delta; + const uint64_t frq_in = hit.frq_base + e.frq_off_delta; + ASSERT_LE(prx_in + e.prx_len, region.length) << e.term; + ASSERT_LE(frq_in + e.frq_len, region.length) << e.term; + // prx span ends exactly where frq span begins (contiguous, prx first). + EXPECT_EQ(prx_in + e.prx_len, frq_in) << e.term; + + // The docs-only prefix [frq_off_delta, +frq_docs_len) stays inside the frq span + // (begins after the prx span and fits within frq_len) -- it never reads prx. + EXPECT_GE(frq_in, prx_in + e.prx_len) << e.term; // prefix begins after prx + EXPECT_LE(e.frq_docs_len, e.frq_len) << e.term; // prefix fits in frq span + + // The absolute frq span stays inside the posting region (reader reads exactly + // what the writer wrote). + const uint64_t frq_span_abs = region.offset + frq_in; + EXPECT_LE(frq_span_abs + e.frq_len, region.offset + region.length) << e.term; + } + std::remove(path.c_str()); +} + +// Test #4: multi-index. One docs-positions index and one docs-only index in one +// container; each posting_region placement is independent (no offset bleed) and +// both resolve correctly. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, MultiIndexIndependentPlacements) { + SniiIndexInput a = BaseInput(1, "body", IndexConfig::kDocsPositions); + a.doc_count = 2048; + std::vector a_docs; + for (uint32_t d = 0; d < 1000; ++d) { + a_docs.push_back(d); + } + a.terms.push_back(MakeTerm("aa_wide", a_docs, /*with_positions=*/true, 2)); + + SniiIndexInput b = BaseInput(2, "tag", IndexConfig::kDocsOnly); + b.doc_count = 2048; + std::vector b_docs; + for (uint32_t d = 0; d < 800; ++d) { + b_docs.push_back(d * 2); + } + b.terms.push_back(MakeTerm("kk_keyword", b_docs, /*with_positions=*/false, 1)); + + const std::string path = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(a).ok()); + ASSERT_TRUE(cw.add_logical_index(b).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + + LogicalIndexReader ia, ib; + ASSERT_TRUE(seg.open_index(1, "body", &ia).ok()); + ASSERT_TRUE(seg.open_index(2, "tag", &ib).ok()); + EXPECT_TRUE(ia.has_positions()); + EXPECT_FALSE(ib.has_positions()); // docs-only recovered from the flag + EXPECT_EQ(ia.tier(), IndexTier::kT2); + EXPECT_EQ(ib.tier(), IndexTier::kT1); + + const RegionRef ra = ia.section_refs().posting_region; + const RegionRef rb = ib.section_refs().posting_region; + // Non-overlapping regions (independent placement; no cross-index bleed). + const bool disjoint = + (ra.offset + ra.length <= rb.offset) || (rb.offset + rb.length <= ra.offset); + EXPECT_TRUE(disjoint) << "ra=[" << ra.offset << "," << ra.length << "] rb=[" << rb.offset << "," + << rb.length << "]"; + + std::vector got_a, got_b; + ASSERT_TRUE(query::term_query(ia, "aa_wide", &got_a).ok()); + ASSERT_TRUE(query::term_query(ib, "kk_keyword", &got_b).ok()); + EXPECT_EQ(got_a, a_docs); + EXPECT_EQ(got_b, b_docs); + std::remove(path.c_str()); +} + +// Test #5: no-positions index -- the posting region holds only frq spans; +// prx_off_delta / prx_len are unset; the bytes are identical to today's .frq POD +// for the same input (a docs-only index's posting_region IS the frq concatenation, +// since the prx span is always empty). We verify prx spans are absent and the +// posting region tiles exactly the per-term frq spans with no gaps/prx. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, NoPositionsRegionIsFrqOnly) { + SniiIndexInput in = BaseInput(1, "tag", IndexConfig::kDocsOnly); + in.doc_count = 4096; + std::vector wide_docs, slim_docs; + for (uint32_t d = 0; d < 1500; ++d) { + wide_docs.push_back(d); + } + for (uint32_t d = 0; d < 600; ++d) { + slim_docs.push_back(d + 2); + } + in.terms.push_back(MakeTerm("aa_wide", wide_docs, /*with_positions=*/false, 1)); + in.terms.push_back(MakeTerm("bb_slim", slim_docs, /*with_positions=*/false, 1)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "tag", &idx).ok()); + EXPECT_FALSE(idx.has_positions()); + EXPECT_EQ(idx.tier(), IndexTier::kT1); + + const RegionRef region = idx.section_refs().posting_region; + uint64_t frq_bytes_sum = 0; + for (const auto& hit : AllEntries(idx)) { + const DictEntry& e = hit.entry; + if (e.kind != DictEntryKind::kPodRef) { + continue; + } + EXPECT_EQ(e.prx_len, 0U) << e.term; // no prx span at all + EXPECT_EQ(e.prx_off_delta, 0U) << e.term; + // frq span is the whole per-term posting (no preceding prx span). + EXPECT_EQ(hit.frq_base + e.frq_off_delta, frq_bytes_sum) << e.term; + frq_bytes_sum += e.frq_len; + } + // The frq spans tile the region exactly: no prx bytes, no gaps. + EXPECT_EQ(frq_bytes_sum, region.length); + + std::vector got; + ASSERT_TRUE(query::term_query(idx, "aa_wide", &got).ok()); + EXPECT_EQ(got, wide_docs); + std::remove(path.c_str()); +} + +// Test #6: empty term set. posting_region.length == 0 and dict_region.length == 0; +// the reader opens cleanly and lookups miss. Built BOTH docs-only and docs-positions +// to assert tier is recovered correctly from the flag despite both having a +// zero-length posting region. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, EmptyTermSetTierFromFlag) { + for (IndexConfig cfg : {IndexConfig::kDocsOnly, IndexConfig::kDocsPositions}) { + SniiIndexInput in = BaseInput(1, "body", cfg); + in.doc_count = 16; // docs exist, but no terms + const std::string path = WriteOne(in); + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + EXPECT_EQ(idx.section_refs().posting_region.length, 0U); + EXPECT_EQ(idx.section_refs().dict_region.length, 0U); + const bool want_positions = (cfg != IndexConfig::kDocsOnly); + EXPECT_EQ(idx.has_positions(), want_positions); + EXPECT_EQ(idx.tier(), want_positions ? IndexTier::kT2 : IndexTier::kT1); + + std::vector got; + ASSERT_TRUE(query::term_query(idx, "anything", &got).ok()); + EXPECT_TRUE(got.empty()); + std::remove(path.c_str()); + } +} + +// Test #8: INLINE-between-pod_ref. Mix tiny (INLINE) and large (pod_ref) terms; +// INLINE terms append nothing to the region (region length == sum of pod_ref +// spans) and pod_ref deltas stay contiguous across the interleaving. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, InlineBetweenPodRefIsGapless) { + SniiIndexInput in = BaseInput(1, "body", IndexConfig::kDocsPositions); + in.doc_count = 4096; + std::vector wide1, wide2; + for (uint32_t d = 0; d < 1500; ++d) { + wide1.push_back(d); + } + for (uint32_t d = 0; d < 1200; ++d) { + wide2.push_back(d + 1); + } + // Lexicographic order interleaves tiny INLINE terms between the two windowed + // pod_ref terms: aa_big < bb_tiny < cc_big < dd_tiny. + in.terms.push_back(MakeTerm("aa_big", wide1, /*with_positions=*/true, 2)); + in.terms.push_back(MakeTerm("bb_tiny", {3, 8}, /*with_positions=*/true, 1)); + in.terms.push_back(MakeTerm("cc_big", wide2, /*with_positions=*/true, 2)); + in.terms.push_back(MakeTerm("dd_tiny", {1, 99}, /*with_positions=*/true, 1)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + const RegionRef region = idx.section_refs().posting_region; + uint64_t pod_span_sum = 0; + uint64_t cursor = 0; // running region offset of the next pod_ref term's span + for (const auto& hit : AllEntries(idx)) { + const DictEntry& e = hit.entry; + if (e.kind == DictEntryKind::kInline) { + EXPECT_EQ(e.prx_len, 0U) << e.term; // inline carries bytes in the entry + EXPECT_EQ(e.frq_len, 0U) << e.term; + continue; // appends nothing to the region + } + // pod_ref term's prx span begins exactly where the previous span ended (no gap + // introduced by the interleaved INLINE terms). + const uint64_t prx_in = hit.prx_base + e.prx_off_delta; + EXPECT_EQ(prx_in, cursor) << e.term; + const uint64_t term_span = e.prx_len + e.frq_len; + cursor += term_span; + pod_span_sum += term_span; + } + // Region length is exactly the sum of pod_ref spans (INLINE adds nothing). + EXPECT_EQ(pod_span_sum, region.length); + + // All terms still resolve. + for (const auto& [term, expect] : std::vector>> { + {"aa_big", wide1}, {"bb_tiny", {3, 8}}, {"cc_big", wide2}, {"dd_tiny", {1, 99}}}) { + std::vector got; + ASSERT_TRUE(query::term_query(idx, term, &got).ok()) << term; + EXPECT_EQ(got, expect) << term; + } + std::remove(path.c_str()); +} + +// Test #9: tier recovery for docs-only-with-pod_ref (R1 regression guard). A +// docs-only (T1) index containing a windowed pod_ref term so posting_region.length +// > 0. open_index must recover tier == kT1 / has_positions == false (NOT inferred +// from the non-zero region length), and a lookup + docid decode must succeed with +// NO spurious positions-flag parse (DictBlockReader::open does not InvalidArgument). +TEST(SniiPostingInterleave, DocsOnlyWithPodRefTierRecovery) { + SniiIndexInput in = BaseInput(1, "tag", IndexConfig::kDocsOnly); + in.doc_count = 4096; + std::vector wide_docs; + for (uint32_t d = 0; d < 2000; ++d) { + wide_docs.push_back(d); // windowed pod_ref + } + in.terms.push_back(MakeTerm("kk_keyword", wide_docs, /*with_positions=*/false, 1)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + // This open would FAIL under the naive posting_region.length > 0 heuristic: + // has_positions would be wrongly true, DictBlockReader::open -> check_flags would + // hard-fail with InvalidArgument. With the persisted flag it opens cleanly. + ASSERT_TRUE(seg.open_index(1, "tag", &idx).ok()); + EXPECT_EQ(idx.tier(), IndexTier::kT1); + EXPECT_FALSE(idx.has_positions()); + EXPECT_GT(idx.section_refs().posting_region.length, 0U); // non-empty region + + // Lookup + docid decode succeeds (no spurious positions parse). + bool found = false; + DictEntry e; + uint64_t fb = 0, pb = 0; + ASSERT_TRUE(idx.lookup("kk_keyword", &found, &e, &fb, &pb).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(e.kind, DictEntryKind::kPodRef); + EXPECT_EQ(e.prx_len, 0U); // docs-only: no prx span + std::vector got; + ASSERT_TRUE(query::term_query(idx, "kk_keyword", &got).ok()); + EXPECT_EQ(got, wide_docs); + std::remove(path.c_str()); +} + +// M1 (code-review blocker): the SLIM pod_ref path is the ONLY write path whose +// physical byte order flipped to [prx][frq], yet no prior test exercised it -- every +// "slim" term elsewhere is df>=512 (windowed) or tiny (inline). Build a genuine +// kSlim/kPodRef term (df<512 so build_slim_entry runs; docids spread so the PFOR'd +// frq exceeds the 256B inline threshold so it is kPodRef not kInline), lock the +// routing, assert the interleave invariant, and round-trip docids (frq span) AND +// positions (prx span, decoded via the slim phrase path). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiPostingInterleave, SlimPodRefInterleaveAndPositionsRoundTrip) { + SniiIndexInput in = BaseInput(1, "body", IndexConfig::kDocsPositions); + in.doc_count = 30000; + std::vector docs; + for (uint32_t d = 0; d < 500; ++d) { + docs.push_back(d * 50 + 1); // df=500 (<512), spread + } + in.terms.push_back(MakeTerm("ss_aaa", docs, /*with_positions=*/true, 4)); + in.terms.push_back(MakeTerm("ss_bbb", docs, /*with_positions=*/true, 4)); + + const std::string path = WriteOne(in); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + SniiSegmentReader seg; + ASSERT_TRUE(SniiSegmentReader::open(&local, &seg).ok()); + LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + ASSERT_TRUE(idx.has_positions()); + + bool found = false; + DictEntry e; + uint64_t fb = 0, pb = 0; + ASSERT_TRUE(idx.lookup("ss_aaa", &found, &e, &fb, &pb).ok()); + ASSERT_TRUE(found); + // Lock the routing onto the slim pod_ref path (the one whose order flipped). + ASSERT_EQ(e.kind, DictEntryKind::kPodRef) << "frq must exceed the inline threshold"; + ASSERT_EQ(e.enc, DictEntryEnc::kSlim) << "df must be below kSlimDfThreshold"; + + // Interleave invariant + adjacency: the prx span ends exactly where the frq begins. + EXPECT_EQ(e.frq_off_delta, e.prx_off_delta + e.prx_len); + const RegionRef& region = idx.section_refs().posting_region; + uint64_t foff = 0, flen = 0, poff = 0, plen = 0; + ASSERT_TRUE(idx.resolve_frq_window(e, fb, &foff, &flen).ok()); + ASSERT_TRUE(idx.resolve_prx_window(e, pb, &poff, &plen).ok()); + EXPECT_GE(poff, region.offset); + EXPECT_LE(foff + flen, region.offset + region.length); + EXPECT_EQ(poff + plen, foff) << "prx span must abut the frq span ([prx][frq])"; + + // Docids round-trip via term_query: decodes dd from the FRQ span -> proves the frq + // offset lands after the prx in the [prx][frq] layout. + std::vector got; + ASSERT_TRUE(query::term_query(idx, "ss_aaa", &got).ok()); + EXPECT_EQ(got, docs); + + // Positions round-trip via the SLIM read path: both terms sit at positions {0,1,2,3} + // in every shared doc, so "ss_aaa ss_bbb" matches (ss_aaa@0, ss_bbb@1) iff each + // term's prx window decodes correctly from its prx span. + std::vector phr; + ASSERT_TRUE(query::phrase_query(idx, {"ss_aaa", "ss_bbb"}, &phr).ok()); + EXPECT_EQ(phr, docs); + std::remove(path.c_str()); +} diff --git a/be/test/storage/index/snii/writer/posting_window_emitter_test.cpp b/be/test/storage/index/snii/writer/posting_window_emitter_test.cpp new file mode 100644 index 00000000000000..c97bfcc8069c27 --- /dev/null +++ b/be/test/storage/index/snii/writer/posting_window_emitter_test.cpp @@ -0,0 +1,215 @@ +// 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 "storage/index/snii/writer/posting_window_emitter.h" + +#include + +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii::writer { +namespace { + +using format::DictEntry; +using format::DictEntryEnc; +using format::DictEntryKind; +using format::FrqPreludeReader; +using format::WindowMeta; +using snii_test::MemoryFile; + +void assert_ok(const Status& status) { + ASSERT_TRUE(status.ok()) << status.to_string(); +} + +std::vector read_windows(const MemoryFile& file, const DictEntry& entry) { + EXPECT_EQ(entry.enc, DictEntryEnc::kWindowed); + EXPECT_LE(entry.frq_off_delta + entry.prelude_len, file.data().size()); + if (entry.frq_off_delta + entry.prelude_len > file.data().size()) { + return {}; + } + FrqPreludeReader prelude; + assert_ok(FrqPreludeReader::open( + Slice(file.data().data() + entry.frq_off_delta, entry.prelude_len), &prelude)); + std::vector windows; + windows.reserve(prelude.n_windows()); + for (uint32_t i = 0; i < prelude.n_windows(); ++i) { + WindowMeta window; + assert_ok(prelude.window(i, &window)); + windows.push_back(window); + } + return windows; +} + +TEST(PostingWindowEmitterTest, EmitsPositionedRecutAndDocsOnlyMetadata) { + constexpr uint32_t kDocs = 512; + constexpr uint32_t kFarPosition = 1U << 28; + std::vector docids(kDocs); + std::iota(docids.begin(), docids.end(), 0); + std::vector freqs(kDocs, 2); + std::vector position_offsets(kDocs + 1); + std::vector positions; + positions.reserve(kDocs * 2); + for (uint32_t doc = 0; doc < kDocs; ++doc) { + position_offsets[doc] = static_cast(positions.size()); + positions.push_back(0); + positions.push_back(kFarPosition); + } + position_offsets.back() = static_cast(positions.size()); + + testing::reset_window_emitter_counters(); + MemoryFile positioned_file; + WindowEmitterOptions positioned_options; + positioned_options.posting_out = &positioned_file; + positioned_options.has_freq = true; + positioned_options.has_prx = true; + positioned_options.prx_window_limits = { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 2048, + }; + WindowEmitter positioned(positioned_options); + assert_ok(positioned.emit_window(PostingRunView { + .docids = docids, + .freqs = freqs, + .position_offsets = position_offsets, + .positions_flat = positions, + })); + DictEntry positioned_entry; + TermAggregateStats positioned_stats; + assert_ok(positioned.finish_term(&positioned_entry, &positioned_stats)); + + EXPECT_EQ(positioned_entry.kind, DictEntryKind::kPodRef); + EXPECT_EQ(positioned_entry.enc, DictEntryEnc::kWindowed); + EXPECT_EQ(positioned_stats.df, kDocs); + EXPECT_EQ(positioned_stats.total_freq, kDocs * 2); + EXPECT_EQ(positioned_stats.max_freq, 2); + const std::vector positioned_windows = + read_windows(positioned_file, positioned_entry); + ASSERT_GT(positioned_windows.size(), 1U); + EXPECT_LT(positioned_windows.front().doc_count, kDocs); + EXPECT_EQ(std::accumulate(positioned_windows.begin(), positioned_windows.end(), 0U, + [](uint32_t total, const WindowMeta& window) { + return total + window.doc_count; + }), + kDocs); + + MemoryFile docs_only_file; + WindowEmitterOptions docs_only_options; + docs_only_options.posting_out = &docs_only_file; + WindowEmitter docs_only(docs_only_options); + PostingRunView first_docs_only; + first_docs_only.docids = std::span(docids).first(256); + assert_ok(docs_only.emit_window(first_docs_only)); + PostingRunView second_docs_only; + second_docs_only.docids = std::span(docids).subspan(256); + assert_ok(docs_only.emit_window(second_docs_only)); + DictEntry docs_only_entry; + TermAggregateStats docs_only_stats; + assert_ok(docs_only.finish_term(&docs_only_entry, &docs_only_stats)); + + EXPECT_EQ(docs_only_stats.df, kDocs); + EXPECT_EQ(docs_only_stats.total_freq, kDocs); + EXPECT_EQ(docs_only_stats.max_freq, 0); + const std::vector docs_only_windows = read_windows(docs_only_file, docs_only_entry); + ASSERT_EQ(docs_only_windows.size(), 2U); + EXPECT_EQ(docs_only_windows[0].doc_count, 256U); + EXPECT_EQ(docs_only_windows[1].doc_count, 256U); + EXPECT_EQ(testing::window_emitter_finished_terms(), 2U); + EXPECT_EQ(testing::window_emitter_physical_windows(), + positioned_windows.size() + docs_only_windows.size()); +} + +TEST(PostingWindowEmitterTest, FailurePoisonsTheEmitter) { + std::vector docids {1, 2}; + std::vector freqs {1, 1}; + std::vector malformed_offsets {0, 1}; + std::vector positions {0, 0}; + MemoryFile file; + WindowEmitterOptions options; + options.posting_out = &file; + options.has_freq = true; + options.has_prx = true; + WindowEmitter emitter(options); + + EXPECT_FALSE(emitter.emit_window(PostingRunView { + .docids = docids, + .freqs = freqs, + .position_offsets = malformed_offsets, + .positions_flat = positions, + }) + .ok()); + DictEntry entry; + TermAggregateStats stats; + EXPECT_FALSE(emitter.finish_term(&entry, &stats).ok()); + EXPECT_TRUE(file.data().empty()); +} + +TEST(PostingWindowEmitterTest, RejectsPositionStatsWithoutPrxOffsets) { + std::vector docids {1, 2}; + MemoryFile file; + WindowEmitterOptions options; + options.posting_out = &file; + options.term_frequency_source = TermFrequencySource::kPositions; + WindowEmitter emitter(options); + + PostingRunView run; + run.docids = docids; + const Status status = emitter.emit_window(run); + EXPECT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find("position-derived statistics require PRX offsets"), + std::string::npos); + DictEntry entry; + TermAggregateStats stats; + EXPECT_FALSE(emitter.finish_term(&entry, &stats).ok()); + EXPECT_TRUE(file.data().empty()); +} + +TEST(PostingWindowEmitterTest, AcceptsBaseRelativePositionRunWithNonzeroOffsets) { + std::vector docids {1, 4}; + std::vector freqs {2, 1}; + std::vector position_offsets {10, 12, 13}; + std::vector positions {3, 7, 9}; + MemoryFile file; + WindowEmitterOptions options; + options.posting_out = &file; + options.has_freq = true; + options.has_prx = true; + WindowEmitter emitter(options); + + assert_ok(emitter.emit_window(PostingRunView { + .docids = docids, + .freqs = freqs, + .position_offsets = position_offsets, + .positions_flat = positions, + })); + DictEntry entry; + TermAggregateStats stats; + assert_ok(emitter.finish_term(&entry, &stats)); + EXPECT_EQ(stats.df, 2U); + EXPECT_EQ(stats.total_freq, 3U); + EXPECT_EQ(stats.max_freq, 2U); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/snii_compound_writer_test.cpp b/be/test/storage/index/snii/writer/snii_compound_writer_test.cpp new file mode 100644 index 00000000000000..0be2da9f676367 --- /dev/null +++ b/be/test/storage/index/snii/writer/snii_compound_writer_test.cpp @@ -0,0 +1,1264 @@ +// 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 "storage/index/snii/writer/snii_compound_writer.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "gen_cpp/snii.pb.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/section_framer.h" +#include "storage/index/snii/format/bootstrap_header.h" +#include "storage/index/snii/format/bsbf.h" +#include "storage/index/snii/format/core_metadata.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/format/metadata_blob.h" +#include "storage/index/snii/format/metadata_directory.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/format/tail_pointer.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/logical_index_writer.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +// Temp file path helper (process-unique). +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_cw_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +// Reads the whole file into a buffer. +std::vector ReadAll(const std::string& path) { + io::LocalFileReader r; + EXPECT_TRUE(r.open(path).ok()); + std::vector out; + EXPECT_TRUE(r.read_at(0, r.size(), &out).ok()); + return out; +} + +// Writes bytes to a fresh temp file and returns its path. +std::string WriteTemp(const std::vector& bytes) { + const std::string p = TempPath(); + io::LocalFileWriter w; + EXPECT_TRUE(w.open(p).ok()); + EXPECT_TRUE(w.append(Slice(bytes)).ok()); + EXPECT_TRUE(w.finalize().ok()); + return p; +} + +constexpr size_t kV1TailPointerSize = 31; +constexpr size_t kV1TailCrcOffset = 27; + +struct TailFields { + uint64_t directory_offset = 0; + uint64_t directory_length = 0; + uint32_t directory_crc32c = 0; +}; + +Status ParseTail(const std::vector& file, TailFields* out) { + if (file.size() < kV1TailPointerSize) { + return Status::Error( + "test: file shorter than v1 tail"); + } + const Slice bytes(file.data() + file.size() - kV1TailPointerSize, kV1TailPointerSize); + ByteSource source(bytes); + uint32_t magic = 0; + uint16_t version = 0; + uint8_t encoded_size = 0; + uint32_t tail_crc = 0; + RETURN_IF_ERROR(source.get_fixed32(&magic)); + RETURN_IF_ERROR(source.get_fixed16(&version)); + RETURN_IF_ERROR(source.get_fixed64(&out->directory_offset)); + RETURN_IF_ERROR(source.get_fixed64(&out->directory_length)); + RETURN_IF_ERROR(source.get_fixed32(&out->directory_crc32c)); + RETURN_IF_ERROR(source.get_u8(&encoded_size)); + RETURN_IF_ERROR(source.get_fixed32(&tail_crc)); + if (magic != kTailMagic || version != 1 || encoded_size != kV1TailPointerSize || + tail_crc != doris::snii::crc32c(bytes.subslice(0, kV1TailCrcOffset))) { + return Status::Error( + "test: invalid v1 tail"); + } + return Status::OK(); +} + +Status ParseDirectory(const std::vector& file, TailFields* tail, + MetadataDirectory* directory) { + RETURN_IF_ERROR(ParseTail(file, tail)); + if (tail->directory_offset > file.size() || + tail->directory_length > file.size() - tail->directory_offset) { + return Status::Error( + "test: directory range out of file"); + } + const Slice bytes(file.data() + tail->directory_offset, + static_cast(tail->directory_length)); + if (doris::snii::crc32c(bytes) != tail->directory_crc32c) { + return Status::Error( + "test: directory crc mismatch"); + } + return MetadataDirectory::decode(bytes, directory); +} + +Status SliceBlob(const std::vector& file, const MetadataBlobRef& ref, Slice* out) { + if (ref.offset > file.size() || ref.length > file.size() - ref.offset) { + return Status::Error( + "test: metadata blob out of file"); + } + *out = Slice(file.data() + ref.offset, static_cast(ref.length)); + return Status::OK(); +} + +Status ParseMetadataGroup(const std::vector& file, const LogicalIndexMetadataRef& ref, + CoreMetadata* core, SampledTermIndexReader* sti, + DictBlockDirectoryReader* dbd) { + Slice core_frame; + RETURN_IF_ERROR(SliceBlob(file, ref.core_metadata, &core_frame)); + RETURN_IF_ERROR(decode_core_metadata(core_frame, core)); + + Slice stored_sti; + RETURN_IF_ERROR(SliceBlob(file, ref.sampled_term_index, &stored_sti)); + std::vector sti_scratch; + Slice sti_frame; + RETURN_IF_ERROR(materialize_metadata_blob(stored_sti, SectionType::kSampledTermIndex, + SectionType::kSampledTermIndexZstd, &sti_scratch, + &sti_frame)); + RETURN_IF_ERROR(SampledTermIndexReader::open(sti_frame, sti)); + + Slice stored_dbd; + RETURN_IF_ERROR(SliceBlob(file, ref.dict_block_directory, &stored_dbd)); + std::vector dbd_scratch; + Slice dbd_frame; + RETURN_IF_ERROR(materialize_metadata_blob(stored_dbd, SectionType::kDictBlockDirectory, + SectionType::kDictBlockDirectoryZstd, &dbd_scratch, + &dbd_frame)); + return DictBlockDirectoryReader::open(dbd_frame, dbd); +} + +std::vector EncodeV1Tail(uint64_t directory_offset, Slice directory) { + ByteSink covered; + covered.put_fixed32(kTailMagic); + covered.put_fixed16(1); + covered.put_fixed64(directory_offset); + covered.put_fixed64(directory.size()); + covered.put_fixed32(doris::snii::crc32c(directory)); + covered.put_u8(kV1TailPointerSize); + EXPECT_EQ(kV1TailCrcOffset, covered.size()); + ByteSink tail; + tail.put_bytes(covered.view()); + tail.put_fixed32(doris::snii::crc32c(covered.view())); + return tail.buffer(); +} + +doris::snii::SniiMetadataDirectoryPB DecodeDirectoryPb(const std::vector& file, + const TailFields& tail) { + doris::snii::SniiMetadataDirectoryPB directory; + EXPECT_TRUE(directory.ParseFromArray(file.data() + tail.directory_offset, + static_cast(tail.directory_length))); + return directory; +} + +std::vector RebuildWithDirectory(const std::vector& file, + const doris::snii::SniiMetadataDirectoryPB& directory) { + TailFields old_tail; + EXPECT_TRUE(ParseTail(file, &old_tail).ok()); + std::string directory_bytes; + EXPECT_TRUE(directory.SerializeToString(&directory_bytes)); + std::vector rebuilt( + file.begin(), file.begin() + static_cast(old_tail.directory_offset)); + rebuilt.insert(rebuilt.end(), directory_bytes.begin(), directory_bytes.end()); + const auto tail = EncodeV1Tail(old_tail.directory_offset, Slice(directory_bytes)); + rebuilt.insert(rebuilt.end(), tail.begin(), tail.end()); + return rebuilt; +} + +std::vector RebuildSingleMetadataGroup(const std::vector& file, Slice new_sti, + Slice new_dbd) { + TailFields old_tail; + EXPECT_TRUE(ParseTail(file, &old_tail).ok()); + auto directory = DecodeDirectoryPb(file, old_tail); + EXPECT_EQ(1, directory.indexes_size()); + auto* index = directory.mutable_indexes(0); + const uint64_t core_end = index->core_metadata().offset() + index->core_metadata().length(); + std::vector rebuilt(file.begin(), + file.begin() + static_cast(core_end)); + index->mutable_sampled_term_index()->set_offset(rebuilt.size()); + index->mutable_sampled_term_index()->set_length(new_sti.size()); + rebuilt.insert(rebuilt.end(), new_sti.data(), new_sti.data() + new_sti.size()); + index->mutable_dict_block_directory()->set_offset(rebuilt.size()); + index->mutable_dict_block_directory()->set_length(new_dbd.size()); + rebuilt.insert(rebuilt.end(), new_dbd.data(), new_dbd.data() + new_dbd.size()); + + std::string directory_bytes; + EXPECT_TRUE(directory.SerializeToString(&directory_bytes)); + const uint64_t directory_offset = rebuilt.size(); + rebuilt.insert(rebuilt.end(), directory_bytes.begin(), directory_bytes.end()); + const auto tail = EncodeV1Tail(directory_offset, Slice(directory_bytes)); + rebuilt.insert(rebuilt.end(), tail.begin(), tail.end()); + return rebuilt; +} + +std::vector BuildIndexes(const std::vector& indexes) { + const std::string path = TempPath(); + { + io::LocalFileWriter writer; + EXPECT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + for (const auto& index : indexes) { + EXPECT_TRUE(compound.add_logical_index(index).ok()); + } + EXPECT_TRUE(compound.finish().ok()); + } + auto bytes = ReadAll(path); + std::remove(path.c_str()); + return bytes; +} + +SniiIndexInput EmptyIndex(uint64_t index_id, std::string suffix) { + SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = std::move(suffix); + input.config = IndexConfig::kDocsPositions; + return input; +} + +doris::segment_v2::inverted_index::CommonGramsSegmentMetadata CompleteCommonGramsMetadata() { + using doris::segment_v2::inverted_index::CommonGramsCoverage; + using doris::segment_v2::inverted_index::PlainTermKeyVersion; + doris::segment_v2::inverted_index::CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kEscapedV1; + metadata.common_grams_coverage = CommonGramsCoverage::kComplete; + metadata.common_grams_semantics_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_SEMANTICS_VERSION_V1; + metadata.common_grams_key_version = + doris::segment_v2::inverted_index::COMMON_GRAMS_KEY_VERSION_V1; + metadata.common_grams_dictionary_identity = "test:dictionary:v1"; + metadata.base_analyzer_fingerprint = "test:base:v1"; + metadata.common_grams_fingerprint = "test:grams:v1"; + return metadata; +} + +Status OpenBytes(const std::vector& bytes, uint64_t index_id, std::string_view suffix, + bool open_index) { + const std::string path = WriteTemp(bytes); + io::LocalFileReader local; + Status status = local.open(path); + reader::SniiSegmentReader segment; + if (status.ok()) { + status = reader::SniiSegmentReader::open(&local, &segment); + } + if (status.ok() && open_index) { + reader::LogicalIndexReader index; + status = segment.open_index(index_id, suffix, &index); + } + std::remove(path.c_str()); + return status; +} + +class FailOnAppendWriter final : public io::FileWriter { +public: + explicit FailOnAppendWriter(size_t fail_on_append) : fail_on_append_(fail_on_append) {} + + Status append(Slice data) override { + ++append_calls_; + if (append_calls_ == fail_on_append_) { + return Status::Error("injected append failure"); + } + bytes_.insert(bytes_.end(), data.data(), data.data() + data.size()); + if (data.size() == kV1TailPointerSize) { + TailPointer decoded; + if (decode_tail_pointer(data, &decoded).ok()) { + ++valid_footer_appends_; + } + } + return Status::OK(); + } + + Status finalize() override { + ++finalize_calls_; + return Status::OK(); + } + + uint64_t bytes_written() const override { return bytes_.size(); } + size_t append_calls() const { return append_calls_; } + size_t finalize_calls() const { return finalize_calls_; } + size_t valid_footer_appends() const { return valid_footer_appends_; } + +private: + size_t fail_on_append_; + size_t append_calls_ = 0; + size_t finalize_calls_ = 0; + size_t valid_footer_appends_ = 0; + std::vector bytes_; +}; + +void VerifyAppendFailurePoisonsWriter(size_t fail_on_append) { + FailOnAppendWriter file(fail_on_append); + SniiCompoundWriter writer(&file); + ASSERT_TRUE(writer.add_logical_index(EmptyIndex(7, std::string("body"))).ok()); + const Status first = writer.finish(); + EXPECT_FALSE(first.ok()) << "fail_on_append=" << fail_on_append; + EXPECT_EQ(fail_on_append, file.append_calls()); + EXPECT_EQ(0U, file.valid_footer_appends()); + EXPECT_EQ(0U, file.finalize_calls()); + + const Status retry = writer.finish(); + EXPECT_EQ(retry, first); + EXPECT_EQ(fail_on_append, file.append_calls()); + EXPECT_EQ(0U, file.valid_footer_appends()); + EXPECT_EQ(0U, file.finalize_calls()); +} + +// A FileReader decorator that counts how many physical reads (single or batched) +// touch a given byte window. Used to assert that the BSBF section is NOT read at +// all on the non-resident (L1) path: with the P1 cold-read fix, open must not +// read the 28B bloom header and lookup must not issue a 32B bloom probe. +class WindowTouchCountingReader : public io::FileReader { +public: + WindowTouchCountingReader(io::FileReader* inner, uint64_t win_off, uint64_t win_len) + : inner_(inner), win_off_(win_off), win_len_(win_len) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + account(offset, len); + return inner_->read_at(offset, len, out); + } + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + for (const auto& r : ranges) { + account(r.offset, r.len); + } + return inner_->read_batch(ranges, outs); + } + uint64_t size() const override { return inner_->size(); } + + uint64_t window_touches() const { return window_touches_; } + +private: + void account(uint64_t offset, size_t len) { + if (win_len_ == 0 || len == 0) { + return; + } + const uint64_t end = offset + len; + const uint64_t win_end = win_off_ + win_len_; + if (offset < win_end && win_off_ < end) { + ++window_touches_; + } + } + io::FileReader* inner_; + uint64_t win_off_; + uint64_t win_len_; + uint64_t window_touches_ = 0; +}; + +class FailNextReadReader final : public io::FileReader { +public: + explicit FailNextReadReader(io::FileReader* inner) : inner_(inner) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (fail_next_read_) { + fail_next_read_ = false; + return Status::Error("injected read failure"); + } + return inner_->read_at(offset, len, out); + } + + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + return inner_->read_batch(ranges, outs); + } + + uint64_t size() const override { return inner_->size(); } + void fail_next_read() { fail_next_read_ = true; } + +private: + io::FileReader* inner_; + bool fail_next_read_ = false; +}; + +void ExpectClosedLogicalReader(const reader::LogicalIndexReader& index) { + EXPECT_EQ(nullptr, index.reader()); + EXPECT_EQ(0U, index.stats().doc_count); + EXPECT_EQ(0U, index.n_dict_blocks()); + EXPECT_EQ(nullptr, index.common_grams_metadata()); +} + +// Builds a TermPostings with constant freq per doc and (optionally) positions. +TermPostings MakeTerm(const std::string& term, const std::vector& docids, + bool with_positions) { + TermPostings tp; + tp.term = term; + tp.docids = docids; + tp.freqs.assign(docids.size(), 0); + for (size_t i = 0; i < docids.size(); ++i) { + tp.freqs[i] = static_cast((i % 3) + 1); + } + if (with_positions) { + for (size_t i = 0; i < docids.size(); ++i) { + for (uint32_t k = 0; k < tp.freqs[i]; ++k) { + tp.positions_flat.push_back(k * 2); // ascending positions (flat, by freq) + } + } + } + return tp; +} + +// Builds a FrqRegionMeta for a window's dd region from its WindowMeta. +FrqRegionMeta DdMetaOf(const WindowMeta& m) { + FrqRegionMeta r; + r.zstd = m.dd_zstd; + r.uncomp_len = m.dd_uncomp_len; + r.disk_len = m.dd_disk_len; + r.crc = m.crc_dd; + return r; +} + +// Index 0: ~30 docs. Vocab includes a HIGH-df term (df=600 > 512 -> windowed) +// and several LOW-df terms (-> slim/inline). Terms must be lexicographically +// sorted. +SniiIndexInput MakeIndex(uint64_t index_id, const std::string& suffix, uint32_t doc_count) { + SniiIndexInput in; + in.index_id = index_id; + in.index_suffix = suffix; + in.config = IndexConfig::kDocsPositions; + // The writer validates docid < doc_count; the fixed "common" term below + // spans docids 0..599, so the declared doc count must cover it regardless + // of the caller's (smaller) nominal value. + in.doc_count = std::max(doc_count, 600); + // Force one DICT block per term so the sampled-term index covers every term + // (its max_term == the lexicographically largest term). + in.target_dict_block_bytes = 1; + + // low-df "apple" in a handful of docs. + in.terms.push_back(MakeTerm("apple", {0, 5, 12, 20}, true)); + // mid-low "banana". + in.terms.push_back(MakeTerm("banana", {1, 2, 3, 4, 9, 15}, true)); + // HIGH-df "common": df=600 (>=512) -> windowed pod_ref. + std::vector common_docs; + for (uint32_t d = 0; d < 600; ++d) { + common_docs.push_back(d); + } + in.terms.push_back(MakeTerm("common", common_docs, true)); + // low-df "zebra". + in.terms.push_back(MakeTerm("zebra", {7, 21, 29}, true)); + return in; +} + +// Locate a term through the full reader walk and return its DictEntry. +Status LocateEntry(const std::vector& file, const SampledTermIndexReader& sti, + const DictBlockDirectoryReader& dbd, const std::string& term, bool* found, + DictEntry* out) { + bool maybe = false; + uint32_t ordinal = 0; + RETURN_IF_ERROR(sti.locate(term, &maybe, &ordinal)); + if (!maybe) { + *found = false; + return Status::OK(); + } + BlockRef ref {}; + RETURN_IF_ERROR(dbd.get(ordinal, &ref)); + Slice block(file.data() + ref.offset, ref.length); + DictBlockReader br; + RETURN_IF_ERROR(DictBlockReader::open(block, IndexTier::kT2, true, &br)); + return br.find_term(term, found, out); +} + +} // namespace + +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiCompoundWriter, ReadBackSelfValidation) { + const std::string path = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(MakeIndex(10, "title", 30)).ok()); + ASSERT_TRUE(cw.add_logical_index(MakeIndex(11, "body", 30)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + + std::vector file = ReadAll(path); + ASSERT_GT(file.size(), kBootstrapHeaderSize + tail_pointer_size()); + + // --- bootstrap header --- + BootstrapHeader bh; + ASSERT_TRUE(decode_bootstrap_header(Slice(file.data(), kBootstrapHeaderSize), &bh).ok()); + EXPECT_EQ(bh.magic, kContainerMagic); + EXPECT_EQ(bh.format_version, kFormatVersion); + EXPECT_EQ(bh.tail_pointer_size, static_cast(tail_pointer_size())); + + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(file, &tail, &directory).ok()); + ASSERT_EQ(directory.size(), 2U); + + struct Expect { + uint64_t id; + std::string suffix; + }; + std::vector expects = {{.id = 10, .suffix = "title"}, {.id = 11, .suffix = "body"}}; + for (const auto& e : expects) { + const LogicalIndexMetadataRef* metadata_ref = directory.find(e.id, e.suffix); + ASSERT_NE(metadata_ref, nullptr) << "index " << e.id; + EXPECT_EQ(metadata_ref->sampled_term_index.offset, + metadata_ref->core_metadata.offset + metadata_ref->core_metadata.length); + EXPECT_EQ( + metadata_ref->dict_block_directory.offset, + metadata_ref->sampled_term_index.offset + metadata_ref->sampled_term_index.length); + EXPECT_LE(metadata_ref->dict_block_directory.offset + + metadata_ref->dict_block_directory.length, + tail.directory_offset); + + CoreMetadata core; + SampledTermIndexReader sti; + DictBlockDirectoryReader dbd; + ASSERT_TRUE(ParseMetadataGroup(file, *metadata_ref, &core, &sti, &dbd).ok()); + EXPECT_EQ(core.stats.doc_count, 600U); // MakeIndex widens to cover df=600 "common" + EXPECT_EQ(core.stats.term_count, 4U); + EXPECT_EQ(sti.n_blocks(), dbd.n_blocks()); + + const SectionRefs& refs = core.section_refs; + // posting_region / dict_region must be within file bounds. With the order flip + // (posting region first, then DICT trailer), the posting region precedes the + // DICT region: posting_off < dict_off and posting_off + posting_len == dict_off. + ASSERT_GT(refs.posting_region.length, 0U); + ASSERT_LE(refs.posting_region.offset + refs.posting_region.length, file.size()); + ASSERT_GT(refs.dict_region.length, 0U); + ASSERT_LE(refs.dict_region.offset + refs.dict_region.length, file.size()); + EXPECT_LT(refs.posting_region.offset, refs.dict_region.offset); + EXPECT_EQ(refs.posting_region.offset + refs.posting_region.length, refs.dict_region.offset); + // norms absent for docs-positions (no scoring). + EXPECT_EQ(refs.norms.offset, 0U); + EXPECT_EQ(refs.norms.length, 0U); + EXPECT_EQ(refs.null_bitmap.offset, 0U); + EXPECT_EQ(refs.null_bitmap.length, 0U); + + // --- XFilter (block-split bloom, physical section): present true, absent false --- + // Probe the on-disk filter directly: one 32-byte block at a self-computed offset. + ASSERT_GT(refs.bsbf.length, doris::snii::format::kBsbfHeaderSize); + ASSERT_LE(refs.bsbf.offset + refs.bsbf.length, file.size()); + const uint64_t bsbf_bitset = refs.bsbf.offset + doris::snii::format::kBsbfHeaderSize; + const auto bsbf_nblocks = + static_cast((refs.bsbf.length - doris::snii::format::kBsbfHeaderSize) / + doris::snii::format::kBsbfBytesPerBlock); + auto bsbf_present = [&](std::string_view term) { + const uint64_t h = doris::snii::format::bsbf_hash(term); + const uint64_t off = + bsbf_bitset + + static_cast(doris::snii::format::bsbf_block_index(h, bsbf_nblocks)) * + doris::snii::format::kBsbfBytesPerBlock; + return doris::snii::format::bsbf_block_contains(h, file.data() + off); + }; + EXPECT_TRUE(bsbf_present("apple")); + EXPECT_TRUE(bsbf_present("common")); + EXPECT_FALSE(bsbf_present("nonexistent-term-xyzzy-12345")); + + // --- windowed high-df term "common": read its .frq window --- + bool found_common = false; + DictEntry common_entry; + ASSERT_TRUE(LocateEntry(file, sti, dbd, "common", &found_common, &common_entry).ok()); + ASSERT_TRUE(found_common); + EXPECT_EQ(common_entry.df, 600U); + EXPECT_EQ(common_entry.kind, DictEntryKind::kPodRef); + EXPECT_EQ(common_entry.enc, DictEntryEnc::kWindowed); + + // The DICT block carrying "common" supplies frq_base via DictBlockReader. + // Recompute frq_base by re-locating the block. + bool maybe = false; + uint32_t ord = 0; + ASSERT_TRUE(sti.locate("common", &maybe, &ord).ok()); + ASSERT_TRUE(maybe); + BlockRef bref {}; + ASSERT_TRUE(dbd.get(ord, &bref).ok()); + Slice block(file.data() + bref.offset, bref.length); + DictBlockReader br; + ASSERT_TRUE(DictBlockReader::open(block, IndexTier::kT2, true, &br).ok()); + + // Absolute .frq offset = posting_region.offset + frq_base + frq_off_delta. The + // windowed payload is [prelude][dd-block][freq-block]; parse the two-level + // prelude, then decode every window's dd region from the dd-block (each with + // its prelude win_base) to reconstruct the full posting (docs-only path). + uint64_t frq_abs = refs.posting_region.offset + br.frq_base() + common_entry.frq_off_delta; + Slice prelude_bytes(file.data() + frq_abs, common_entry.prelude_len); + FrqPreludeReader prelude; + ASSERT_TRUE(FrqPreludeReader::open(prelude_bytes, &prelude).ok()); + const uint64_t dd_block_start = frq_abs + common_entry.prelude_len; + std::vector got_docs; + uint32_t summed = 0; + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta m; + ASSERT_TRUE(prelude.window(w, &m).ok()); + Slice dd(file.data() + dd_block_start + m.dd_off, m.dd_disk_len); + std::vector wdocs; + ASSERT_TRUE(decode_dd_region(dd, DdMetaOf(m), m.win_base, &wdocs).ok()); + ASSERT_EQ(wdocs.size(), m.doc_count); + summed += m.doc_count; + got_docs.insert(got_docs.end(), wdocs.begin(), wdocs.end()); + } + EXPECT_EQ(summed, 600U); // window doc_counts sum to df. + ASSERT_EQ(got_docs.size(), 600U); + EXPECT_EQ(got_docs.front(), 0U); + EXPECT_EQ(got_docs.back(), 599U); + + // --- inline low-df term "apple": decode its inline .frq bytes --- + bool found_apple = false; + DictEntry apple_entry; + ASSERT_TRUE(LocateEntry(file, sti, dbd, "apple", &found_apple, &apple_entry).ok()); + ASSERT_TRUE(found_apple); + EXPECT_EQ(apple_entry.df, 4U); + EXPECT_EQ(apple_entry.kind, DictEntryKind::kInline); + Slice apple_dd(apple_entry.frq_bytes.data(), + static_cast(apple_entry.dd_meta.disk_len)); + std::vector apple_docs; + ASSERT_TRUE(decode_dd_region(apple_dd, apple_entry.dd_meta, 0, &apple_docs).ok()); + std::vector expected_apple = {0, 5, 12, 20}; + EXPECT_EQ(apple_docs, expected_apple); + + // inline .prx bytes decode to per-doc positions. + ByteSource psrc(Slice(apple_entry.prx_bytes)); + std::vector> apple_pos; + ASSERT_TRUE(read_prx_window(&psrc, &apple_pos).ok()); + EXPECT_EQ(apple_pos.size(), 4U); + + // absent term -> not found through the dict walk. + bool found_absent = true; + DictEntry absent_entry; + ASSERT_TRUE(LocateEntry(file, sti, dbd, "zzz-not-here", &found_absent, &absent_entry).ok()); + EXPECT_FALSE(found_absent); + } + + std::remove(path.c_str()); +} + +// NOLINTNEXTLINE(readability-function-cognitive-complexity): end-to-end layout validation is clearer as one test. +TEST(SniiCompoundWriter, WritesAdjacentMandatoryGroupsForEmptyAndBinarySuffixIndexes) { + const std::string binary_suffix("binary\0suffix", 13); + const auto file = + BuildIndexes({EmptyIndex(7, binary_suffix), EmptyIndex(8, std::string("second"))}); + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(file, &tail, &directory).ok()); + ASSERT_EQ(2U, directory.size()); + EXPECT_NE(nullptr, directory.find(7, binary_suffix)); + EXPECT_NE(nullptr, directory.find(8, "second")); + + uint64_t previous_end = 0; + for (const auto& entry : directory.entries()) { + if (previous_end != 0) { + EXPECT_EQ(previous_end, entry.core_metadata.offset); + } + EXPECT_EQ(entry.core_metadata.offset + entry.core_metadata.length, + entry.sampled_term_index.offset); + EXPECT_EQ(entry.sampled_term_index.offset + entry.sampled_term_index.length, + entry.dict_block_directory.offset); + previous_end = entry.dict_block_directory.offset + entry.dict_block_directory.length; + + CoreMetadata core; + SampledTermIndexReader sti; + DictBlockDirectoryReader dbd; + ASSERT_TRUE(ParseMetadataGroup(file, entry, &core, &sti, &dbd).ok()); + EXPECT_EQ(0U, sti.n_blocks()); + EXPECT_EQ(0U, dbd.n_blocks()); + EXPECT_EQ(IndexConfig::kDocsPositions, core.index_config); + } + EXPECT_EQ(previous_end, tail.directory_offset); +} + +TEST(SniiCompoundWriter, DetectsDirectoryAndCoreCrcCorruptionAtTheirReadBoundaries) { + const auto file = BuildIndexes({EmptyIndex(7, std::string("body"))}); + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(file, &tail, &directory).ok()); + ASSERT_EQ(1U, directory.size()); + + auto bad_directory = file; + bad_directory[tail.directory_offset] ^= 1; + const Status directory_status = OpenBytes(bad_directory, 7, "body", false); + EXPECT_TRUE(directory_status.is()) + << directory_status; + + auto bad_core = file; + const auto& core = directory.entries().front().core_metadata; + bad_core[core.offset + core.length - 1] ^= 1; + EXPECT_TRUE(OpenBytes(bad_core, 7, "body", false).ok()); + const Status core_status = OpenBytes(bad_core, 7, "body", true); + EXPECT_TRUE(core_status.is()) << core_status; +} + +TEST(SniiCompoundWriter, RejectsGapsOverlapsWrongOrderAndOutOfRangeMetadataReferences) { + const auto file = BuildIndexes({EmptyIndex(7, std::string("body"))}); + TailFields tail; + ASSERT_TRUE(ParseTail(file, &tail).ok()); + + using Mutation = + std::function; + const std::vector mutations = { + [](auto* index, const auto&) { + index->mutable_sampled_term_index()->set_offset( + index->sampled_term_index().offset() + 1); + }, + [](auto* index, const auto&) { + index->mutable_sampled_term_index()->set_offset( + index->sampled_term_index().offset() - 1); + }, + [](auto* index, const auto&) { + const uint64_t sti_offset = index->sampled_term_index().offset(); + index->mutable_sampled_term_index()->set_offset( + index->dict_block_directory().offset()); + index->mutable_dict_block_directory()->set_offset(sti_offset); + }, + [](auto* index, const auto& fields) { + index->mutable_dict_block_directory()->set_length( + fields.directory_offset - index->dict_block_directory().offset() + 1); + }, + [](auto* index, const auto&) { + index->mutable_dict_block_directory()->set_offset(UINT64_MAX); + }, + }; + + for (size_t i = 0; i < mutations.size(); ++i) { + auto directory = DecodeDirectoryPb(file, tail); + ASSERT_EQ(1, directory.indexes_size()); + mutations[i](directory.mutable_indexes(0), tail); + const Status status = OpenBytes(RebuildWithDirectory(file, directory), 7, "body", false); + EXPECT_TRUE(status.is()) + << "mutation=" << i << " status=" << status; + } +} + +TEST(SniiCompoundWriter, RejectsWrongRecordTypesAndStiDbdCountMismatch) { + const auto file = BuildIndexes({EmptyIndex(7, std::string("body"))}); + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(file, &tail, &directory).ok()); + ASSERT_EQ(1U, directory.size()); + const auto& entry = directory.entries().front(); + Slice original_sti; + Slice original_dbd; + ASSERT_TRUE(SliceBlob(file, entry.sampled_term_index, &original_sti).ok()); + ASSERT_TRUE(SliceBlob(file, entry.dict_block_directory, &original_dbd).ok()); + + const auto wrong_order = RebuildSingleMetadataGroup(file, original_dbd, original_sti); + EXPECT_TRUE(OpenBytes(wrong_order, 7, "body", false).ok()); + const Status wrong_order_status = OpenBytes(wrong_order, 7, "body", true); + EXPECT_TRUE(wrong_order_status.is()) + << wrong_order_status; + + SampledTermIndexBuilder sampled_builder; + sampled_builder.add_block_first_term("term"); + ByteSink sampled; + sampled_builder.finish(&sampled); + const auto mismatch = RebuildSingleMetadataGroup(file, sampled.view(), original_dbd); + EXPECT_TRUE(OpenBytes(mismatch, 7, "body", false).ok()); + const Status mismatch_status = OpenBytes(mismatch, 7, "body", true); + EXPECT_TRUE(mismatch_status.is()) + << mismatch_status; +} + +// NOLINTNEXTLINE(readability-function-cognitive-complexity): one test validates every reopen failure stage. +TEST(SniiCompoundWriter, FailedLogicalIndexReopenLeavesReaderClosedAtEveryOpenStage) { + const auto good_file = BuildIndexes({MakeIndex(7, "body", 30)}); + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(good_file, &tail, &directory).ok()); + ASSERT_EQ(1U, directory.size()); + const auto& entry = directory.entries().front(); + + CoreMetadata core; + SampledTermIndexReader sti; + DictBlockDirectoryReader dbd; + ASSERT_TRUE(ParseMetadataGroup(good_file, entry, &core, &sti, &dbd).ok()); + ASSERT_GT(dbd.n_blocks(), 0U); + BlockRef dict_block; + ASSERT_TRUE(dbd.get(0, &dict_block).ok()); + ASSERT_GT(dict_block.length, 0U); + ASSERT_GT(core.section_refs.bsbf.length, kBsbfHeaderSize); + + Slice original_dbd; + ASSERT_TRUE(SliceBlob(good_file, entry.dict_block_directory, &original_dbd).ok()); + SampledTermIndexBuilder mismatched_sti_builder; + mismatched_sti_builder.add_block_first_term("only-one-block"); + ByteSink mismatched_sti; + mismatched_sti_builder.finish(&mismatched_sti); + + std::vector>> corrupt_files; + auto bad_core = good_file; + bad_core[entry.core_metadata.offset + entry.core_metadata.length - 1] ^= 1; + corrupt_files.emplace_back("core", std::move(bad_core)); + auto bad_sti = good_file; + bad_sti[entry.sampled_term_index.offset + entry.sampled_term_index.length - 1] ^= 1; + corrupt_files.emplace_back("sti", std::move(bad_sti)); + auto bad_dbd = good_file; + bad_dbd[entry.dict_block_directory.offset + entry.dict_block_directory.length - 1] ^= 1; + corrupt_files.emplace_back("dbd", std::move(bad_dbd)); + corrupt_files.emplace_back( + "count", RebuildSingleMetadataGroup(good_file, mismatched_sti.view(), original_dbd)); + auto bad_dict = good_file; + bad_dict[dict_block.offset + dict_block.length - 1] ^= 1; + corrupt_files.emplace_back("resident dict", std::move(bad_dict)); + auto bad_bsbf = good_file; + bad_bsbf[core.section_refs.bsbf.offset + kBsbfHeaderSize] ^= 1; + corrupt_files.emplace_back("resident bsbf", std::move(bad_bsbf)); + + const std::string good_path = WriteTemp(good_file); + io::LocalFileReader good_local; + ASSERT_TRUE(good_local.open(good_path).ok()); + reader::SniiSegmentReader good_segment; + ASSERT_TRUE(reader::SniiSegmentReader::open(&good_local, &good_segment).ok()); + reader::LogicalIndexReader reused; + + ASSERT_TRUE(good_segment.open_index(7, "body", &reused).ok()); + ASSERT_NE(nullptr, reused.reader()); + EXPECT_FALSE(good_segment.open_index(999, "missing", &reused).ok()); + ExpectClosedLogicalReader(reused); + + ASSERT_TRUE(good_segment.open_index(7, "body", &reused).ok()); + FailNextReadReader failing_reader(&good_local); + reader::SniiSegmentReader failing_segment; + ASSERT_TRUE(reader::SniiSegmentReader::open(&failing_reader, &failing_segment).ok()); + failing_reader.fail_next_read(); + EXPECT_FALSE(failing_segment.open_index(7, "body", &reused).ok()); + ExpectClosedLogicalReader(reused); + + for (const auto& [stage, bad_file] : corrupt_files) { + ASSERT_TRUE(good_segment.open_index(7, "body", &reused).ok()) << stage; + const std::string bad_path = WriteTemp(bad_file); + io::LocalFileReader bad_local; + ASSERT_TRUE(bad_local.open(bad_path).ok()) << stage; + reader::SniiSegmentReader bad_segment; + ASSERT_TRUE(reader::SniiSegmentReader::open(&bad_local, &bad_segment).ok()) << stage; + EXPECT_FALSE(bad_segment.open_index(7, "body", &reused).ok()) << stage; + ExpectClosedLogicalReader(reused); + std::remove(bad_path.c_str()); + } + std::remove(good_path.c_str()); +} + +TEST(SniiCompoundWriter, AppendFailurePoisonsWriterBeforeAnyValidFooter) { + // Append 1 is bootstrap; 2/3/4 are Core/STI/DBD, 5 is the raw directory, + // and 6 is the footer. The fake fails before writing any bytes. + for (size_t fail_on_append = 2; fail_on_append <= 6; ++fail_on_append) { + VerifyAppendFailurePoisonsWriter(fail_on_append); + } +} + +TEST(SniiCompoundWriter, ReopeningLogicalReaderClearsPreviousCommonGramsState) { + auto with_common_grams = EmptyIndex(7, "with"); + with_common_grams.common_grams_metadata = CompleteCommonGramsMetadata(); + const auto file = BuildIndexes({with_common_grams, EmptyIndex(8, std::string("without"))}); + const std::string path = WriteTemp(file); + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + reader::SniiSegmentReader segment; + ASSERT_TRUE(reader::SniiSegmentReader::open(&local, &segment).ok()); + + reader::LogicalIndexReader reused; + ASSERT_TRUE(segment.open_index(7, "with", &reused).ok()); + ASSERT_NE(nullptr, reused.common_grams_metadata()); + ASSERT_TRUE(segment.open_index(8, "without", &reused).ok()); + EXPECT_EQ(nullptr, reused.common_grams_metadata()); + EXPECT_EQ(CommonGramsPostingPolicy::kNone, reused.common_grams_posting_policy()); + std::remove(path.c_str()); +} + +namespace { + +// Oracle for the multi-super-block read-back test. "hot" is a very high-df term; +// at df=70000 it uses adaptive 1024-doc windows (df >= kAdaptiveWindowDfThreshold) +// -> ~69 windows -> >1 super-block at group_size=64. "spark" appears at position +// 1 immediately after "hot" in docs where d % 9 == 0, so the phrase "hot spark" +// holds exactly in those docs. "rare" is a tiny low-df term. +struct BigCorpus { + uint32_t doc_count = 70000; + std::vector hot_docs; // every doc + std::vector spark_docs; // d % 9 == 0 + std::vector rare_docs = {3, 17, 99, 1000, 19999}; + std::vector phrase_oracle; // docs where "hot spark" is consecutive +}; + +BigCorpus MakeBigCorpus() { + BigCorpus c; + for (uint32_t d = 0; d < c.doc_count; ++d) { + c.hot_docs.push_back(d); + if (d % 9 == 0) { + c.spark_docs.push_back(d); + c.phrase_oracle.push_back(d); // "hot"@0 then "spark"@1 -> phrase hit + } + } + return c; +} + +// Builds a TermPostings with all freq=1 and a single position per doc. +TermPostings MakePosTerm(const std::string& term, const std::vector& docs, uint32_t pos) { + TermPostings tp; + tp.term = term; + tp.docids = docs; + tp.freqs.assign(docs.size(), 1); + tp.positions_flat.assign(docs.size(), pos); // one position per doc, flat + return tp; +} + +SniiIndexInput MakeBigIndex(const BigCorpus& c) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = c.doc_count; + // One DICT block per term so every term is a block first-term (covered by the + // sampled-term index regardless of order), mirroring the other read tests. + in.target_dict_block_bytes = 1; + // Terms must be lexicographically sorted. + in.terms.push_back(MakePosTerm("hot", c.hot_docs, /*pos=*/0)); // huge -> windowed + in.terms.push_back(MakePosTerm("rare", c.rare_docs, /*pos=*/0)); // tiny -> inline + in.terms.push_back(MakePosTerm("spark", c.spark_docs, /*pos=*/1)); // -> windowed + return in; +} + +} // namespace + +// PHASE A read-back self-validation: a high-df term spanning MANY windows across +// MULTIPLE super-blocks. Asserts (a) sum of window doc_counts == df, (b) the +// windows tile the posting in order, (c) locate_window resolves the covering +// window, and (d) term_query / phrase_query agree with the oracle. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiCompoundWriter, MultiSuperBlockReadBack) { + const BigCorpus c = MakeBigCorpus(); + const std::string path = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(MakeBigIndex(c)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + + io::LocalFileReader local; + ASSERT_TRUE(local.open(path).ok()); + io::MeteredFileReader metered(&local); + reader::SniiSegmentReader seg; + ASSERT_TRUE(reader::SniiSegmentReader::open(&metered, &seg).ok()); + reader::LogicalIndexReader idx; + ASSERT_TRUE(seg.open_index(1, "body", &idx).ok()); + + // Resolve "hot" to its DictEntry + the block's frq_base. + bool found = false; + DictEntry hot; + uint64_t frq_base = 0, prx_base = 0; + ASSERT_TRUE(idx.lookup("hot", &found, &hot, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + EXPECT_EQ(hot.df, c.doc_count); + EXPECT_EQ(hot.enc, DictEntryEnc::kWindowed); + EXPECT_TRUE(hot.has_sb); + + // L0 tiering: this index's bsbf filter is tiny (<= kBsbfResidentMaxBytes), so it is + // loaded resident at open and an absent-term lookup is rejected IN MEMORY with zero + // reads (no per-lookup round). Loop a few absents to skip the rare false positive. + bool saw_resident_reject = false; + for (int i = 0; i < 8 && !saw_resident_reject; ++i) { + metered.reset_metrics(); + bool af = true; + DictEntry ad; + uint64_t afb = 0, apb = 0; + ASSERT_TRUE(idx.lookup("absent-zzz-" + std::to_string(i), &af, &ad, &afb, &apb).ok()); + if (!af) { + EXPECT_EQ(metered.metrics().read_at_calls, 0U); + EXPECT_EQ(metered.metrics().serial_rounds, 0U); + saw_resident_reject = true; + } + } + EXPECT_TRUE(saw_resident_reject); + metered.reset_metrics(); + + // L1 tiering (P1 cold-read fix): force the bloom NON-resident by lowering the + // resident threshold to 0, then re-open the SAME index through a reader that + // counts every physical read touching the bsbf section. With the fix the bloom + // is skipped ENTIRELY -- open does NOT read the 28B header and lookup does NOT + // issue a 32B probe -- so ZERO reads touch the bsbf window, yet a present term + // is still found and absent terms are still not-found, all via sti -> dict. + const RegionRef bsbf_ref = idx.section_refs().bsbf; + ASSERT_GT(bsbf_ref.length, kBsbfHeaderSize); // a real (small) filter exists on disk + ::setenv("SNII_BSBF_RESIDENT_MAX", "0", /*overwrite=*/1); + { + io::LocalFileReader l1_local; + ASSERT_TRUE(l1_local.open(path).ok()); + WindowTouchCountingReader counting(&l1_local, bsbf_ref.offset, bsbf_ref.length); + reader::SniiSegmentReader seg_l1; + ASSERT_TRUE(reader::SniiSegmentReader::open(&counting, &seg_l1).ok()); + reader::LogicalIndexReader idx_l1; + ASSERT_TRUE(seg_l1.open_index(1, "body", &idx_l1).ok()); + // open() must not have read the bsbf header (28B) on the non-resident path. + EXPECT_EQ(counting.window_touches(), 0U) << "non-resident open must skip the bsbf header"; + + // Present term: still found via sti -> dict, with no bloom involved. + bool pf = false; + DictEntry pe; + uint64_t pfb = 0, ppb = 0; + ASSERT_TRUE(idx_l1.lookup("hot", &pf, &pe, &pfb, &ppb).ok()); + EXPECT_TRUE(pf); + + // Absent terms: not found via dict. "absent-zzz-*" sorts before every + // sample (out-of-range sti reject); "rzz" sorts inside the term range so + // sti routes it to a real dict block that then misses -- both return absent + // and NEITHER probes the bloom. + for (int i = 0; i < 8; ++i) { + bool af = true; + DictEntry ad; + uint64_t afb = 0, apb = 0; + ASSERT_TRUE( + idx_l1.lookup("absent-zzz-" + std::to_string(i), &af, &ad, &afb, &apb).ok()); + EXPECT_FALSE(af) << "out-of-range absent i=" << i; + } + bool rf = true; + DictEntry rd; + uint64_t rfb = 0, rpb = 0; + ASSERT_TRUE(idx_l1.lookup("rzz", &rf, &rd, &rfb, &rpb).ok()); + EXPECT_FALSE(rf) << "in-range absent term must miss in the dict"; + + // No lookup may have probed the bsbf section: the bloom is skipped, not + // read on demand. This is the core of the P1 cold-read fix. + EXPECT_EQ(counting.window_touches(), 0U) + << "non-resident lookups must not probe the bsbf section"; + } + ::unsetenv("SNII_BSBF_RESIDENT_MAX"); + metered.reset_metrics(); + + // Fetch + parse the two-level prelude. + const uint64_t prelude_abs = + idx.section_refs().posting_region.offset + frq_base + hot.frq_off_delta; + std::vector prelude_bytes; + ASSERT_TRUE(local.read_at(prelude_abs, hot.prelude_len, &prelude_bytes).ok()); + FrqPreludeReader prelude; + ASSERT_TRUE(FrqPreludeReader::open(Slice(prelude_bytes), &prelude).ok()); + EXPECT_GT(prelude.n_super_blocks(), 1U) << "expected >1 super-block"; + + // (a)+(b): decode every window's dd region from the dd-block; doc_counts sum to + // df and the concatenated docids equal the full ascending posting [0, doc_count). + const uint64_t dd_block_start = prelude_abs + hot.prelude_len; + std::vector tiled; + uint64_t summed = 0; + uint64_t expect_win_base = 0; + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta m; + ASSERT_TRUE(prelude.window(w, &m).ok()); + EXPECT_EQ(m.win_base, expect_win_base) << "win_base w=" << w; + std::vector ddbytes; + ASSERT_TRUE(local.read_at(dd_block_start + m.dd_off, m.dd_disk_len, &ddbytes).ok()); + std::vector wdocs; + ASSERT_TRUE(decode_dd_region(Slice(ddbytes), DdMetaOf(m), m.win_base, &wdocs).ok()); + ASSERT_EQ(wdocs.size(), m.doc_count) << "w=" << w; + summed += m.doc_count; + expect_win_base = m.last_docid; + tiled.insert(tiled.end(), wdocs.begin(), wdocs.end()); + } + EXPECT_EQ(summed, c.doc_count); + ASSERT_EQ(tiled.size(), c.doc_count); + EXPECT_EQ(tiled, c.hot_docs); + + // (c): locate_window returns the window actually containing the docid. + const uint32_t probes[] = {0, 255, 256, 257, 5000, 16383, 16384, 19999}; + for (uint32_t docid : probes) { + bool lfound = false; + uint32_t w = 0; + ASSERT_TRUE(prelude.locate_window(docid, &lfound, &w).ok()); + ASSERT_TRUE(lfound) << "docid=" << docid; + WindowMeta m; + ASSERT_TRUE(prelude.window(w, &m).ok()); + EXPECT_GE(static_cast(docid), m.win_base + (w == 0 ? 0 : 1)) << "docid=" << docid; + EXPECT_LE(docid, m.last_docid) << "docid=" << docid; + } + // past-end -> not found. + { + bool lfound = true; + uint32_t w = 0; + ASSERT_TRUE(prelude.locate_window(c.doc_count, &lfound, &w).ok()); + EXPECT_FALSE(lfound); + } + + // (d): term_query / phrase_query agree with the oracle (full-read path). + std::vector hot_docs; + ASSERT_TRUE(query::term_query(idx, "hot", &hot_docs).ok()); + EXPECT_EQ(hot_docs, c.hot_docs); + + std::vector spark_docs; + ASSERT_TRUE(query::term_query(idx, "spark", &spark_docs).ok()); + EXPECT_EQ(spark_docs, c.spark_docs); + + std::vector rare_docs; + ASSERT_TRUE(query::term_query(idx, "rare", &rare_docs).ok()); + EXPECT_EQ(rare_docs, c.rare_docs); + + std::vector phrase_docs; + ASSERT_TRUE(query::phrase_query(idx, {"hot", "spark"}, &phrase_docs).ok()); + EXPECT_EQ(phrase_docs, c.phrase_oracle); + + // reversed phrase must be absent ("spark"@1 never precedes "hot"@0). + std::vector reversed; + ASSERT_TRUE(query::phrase_query(idx, {"spark", "hot"}, &reversed).ok()); + EXPECT_TRUE(reversed.empty()); + + std::remove(path.c_str()); +} + +// BYTE-IDENTICAL OUTPUT GUARANTEE: the section bytes (DICT region, .frq POD, +// .prx POD) now stream through scratch temp files instead of in-RAM vectors, and +// the xfilter is built from per-term hashes instead of retained strings. The +// produced container must be deterministic and byte-for-byte stable: building the +// same input twice yields identical container bytes (proves the temp-file path +// reproduces the exact same layout/offsets/content as itself, with no in-RAM +// residue affecting the bytes). The big multi-super-block index exercises every +// section type (windowed pod_ref + inline + multi-block dict + prx). +TEST(SniiCompoundWriter, StreamedOutputIsDeterministicByteForByte) { + const BigCorpus c = MakeBigCorpus(); + + auto build_one = [&](const std::string& path) { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(MakeBigIndex(c)).ok()); + ASSERT_TRUE(cw.finish().ok()); + }; + + const std::string path_a = TempPath(); + const std::string path_b = TempPath(); + build_one(path_a); + build_one(path_b); + + const std::vector file_a = ReadAll(path_a); + const std::vector file_b = ReadAll(path_b); + ASSERT_FALSE(file_a.empty()); + EXPECT_EQ(file_a.size(), file_b.size()); + EXPECT_EQ(file_a, file_b) << "streamed container must be byte-identical run-to-run"; + + std::remove(path_a.c_str()); + std::remove(path_b.c_str()); +} + +// Determinism for a multi-index docs-positions container with multiple DICT +// blocks per index (target_dict_block_bytes forces real block cuts, exercising +// the dict scratch-file concat + the BlockRecord rel_offset directory math). +TEST(SniiCompoundWriter, StreamedMultiIndexDeterministic) { + auto build_one = [&](const std::string& path) { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(path).ok()); + SniiCompoundWriter cw(&w); + SniiIndexInput a = MakeIndex(10, "title", 30); + SniiIndexInput b = MakeIndex(11, "body", 30); + a.target_dict_block_bytes = 64; // force multiple dict blocks + b.target_dict_block_bytes = 64; + ASSERT_TRUE(cw.add_logical_index(a).ok()); + ASSERT_TRUE(cw.add_logical_index(b).ok()); + ASSERT_TRUE(cw.finish().ok()); + }; + const std::string p1 = TempPath(); + const std::string p2 = TempPath(); + build_one(p1); + build_one(p2); + EXPECT_EQ(ReadAll(p1), ReadAll(p2)); + std::remove(p1.c_str()); + std::remove(p2.c_str()); +} + +// A flipped on-disk byte in the bsbf filter (header or bitset) or in the meta region +// must be detected on the read path, not silently propagated as a wrong probe result. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiCompoundWriter, ChecksumCorruptionDetectedOnRead) { + const std::string good = TempPath(); + { + io::LocalFileWriter w; + ASSERT_TRUE(w.open(good).ok()); + SniiCompoundWriter cw(&w); + ASSERT_TRUE(cw.add_logical_index(MakeIndex(7, "body", 30)).ok()); + ASSERT_TRUE(cw.finish().ok()); + } + const std::vector file = ReadAll(good); + std::remove(good.c_str()); + + TailFields tail; + MetadataDirectory directory; + ASSERT_TRUE(ParseDirectory(file, &tail, &directory).ok()); + const auto* metadata_ref = directory.find(7, "body"); + ASSERT_NE(metadata_ref, nullptr); + CoreMetadata core; + SampledTermIndexReader sti; + DictBlockDirectoryReader dbd; + ASSERT_TRUE(ParseMetadataGroup(file, *metadata_ref, &core, &sti, &dbd).ok()); + const auto bsbf = core.section_refs.bsbf; + ASSERT_GT(bsbf.length, doris::snii::format::kBsbfHeaderSize); // small filter -> L0 + + // The clean file opens fine. + auto opens_ok = [](const std::vector& bytes) { + const std::string p = WriteTemp(bytes); + io::LocalFileReader r; + EXPECT_TRUE(r.open(p).ok()); + reader::SniiSegmentReader seg; + Status sopen = reader::SniiSegmentReader::open(&r, &seg); + bool ok = sopen.ok(); + if (ok) { + reader::LogicalIndexReader idx; + ok = seg.open_index(7, "body", &idx).ok(); + } + std::remove(p.c_str()); + return ok; + }; + EXPECT_TRUE(opens_ok(file)); + + // (1) bsbf BITSET byte flip -> L0 open verifies the bitset crc -> rejected. + { + std::vector bad = file; + bad[bsbf.offset + doris::snii::format::kBsbfHeaderSize] ^= 0xFF; + EXPECT_FALSE(opens_ok(bad)); + } + // (2) bsbf HEADER byte flip (num_bytes field, covered by header crc) -> rejected. + { + std::vector bad = file; + bad[bsbf.offset + 8] ^= 0xFF; + EXPECT_FALSE(opens_ok(bad)); + } + // (3) Core byte flip -> logical-index open verifies the Core frame checksum. + { + std::vector bad = file; + bad[metadata_ref->core_metadata.offset] ^= 0xFF; + EXPECT_FALSE(opens_ok(bad)); + } +} diff --git a/be/test/storage/index/snii/writer/snii_no_bigram_writer_test.cpp b/be/test/storage/index/snii/writer/snii_no_bigram_writer_test.cpp new file mode 100644 index 00000000000000..dee6152ea9c56e --- /dev/null +++ b/be/test/storage/index/snii/writer/snii_no_bigram_writer_test.cpp @@ -0,0 +1,379 @@ +// 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 + +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "core/block/block.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" +#include "io/fs/file_writer.h" +#include "io/fs/local_file_system.h" +#include "runtime/exec_env.h" +#include "storage/index/index_file_writer.h" +#include "storage/index/index_writer.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/rowset/rowset_writer_context.h" +#include "storage/segment/segment_writer.h" +#include "storage/segment/vertical_segment_writer.h" +#include "storage/tablet/tablet_schema.h" +#include "storage/types.h" +#include "util/slice.h" + +namespace doris::segment_v2 { +namespace { + +constexpr int64_t kIndexId = 17; +constexpr const char* kTestDir = "./ut_dir/snii_no_bigram_writer_test"; +constexpr const char* kTmpRoot = "./ut_dir/snii_no_bigram_writer_tmp"; + +const std::vector kRows = { + "alpha beta gamma", + "alpha gamma beta", + "zeta alpha beta", +}; +const std::vector kExpectedPhraseDocs = {0, 2}; + +class SniiNoBigramWriter : public testing::Test { +protected: + static void SetUpTestSuite() { + auto fs = io::global_local_filesystem(); + ASSERT_TRUE(fs->delete_directory(kTmpRoot).ok()); + ASSERT_TRUE(fs->create_directory(kTmpRoot).ok()); + + std::vector paths; + paths.emplace_back(kTmpRoot, 1024 * 1024); + auto tmp_file_dirs = std::make_unique(paths); + ASSERT_TRUE(tmp_file_dirs->init().ok()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + // ExecEnv retains this owner process-wide, so its root outlives the suite. + } + + void SetUp() override { + auto fs = io::global_local_filesystem(); + ASSERT_TRUE(fs->delete_directory(kTestDir).ok()); + ASSERT_TRUE(fs->create_directory(kTestDir).ok()); + } + + void TearDown() override { + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); + } + + TabletSchemaSPtr create_phrase_schema() const { + auto schema = std::make_shared(); + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(DUP_KEYS); + schema_pb.set_num_short_key_columns(1); + schema_pb.set_num_rows_per_row_block(1024); + schema_pb.set_compress_kind(COMPRESS_NONE); + schema_pb.set_next_column_unique_id(2); + schema_pb.set_inverted_index_storage_format(InvertedIndexStorageFormatPB::SNII); + schema->init_from_pb(schema_pb); + + TabletColumn key_column; + key_column.set_name("c1"); + key_column.set_unique_id(0); + key_column.set_type(FieldType::OLAP_FIELD_TYPE_INT); + key_column.set_length(4); + key_column.set_index_length(4); + key_column.set_is_key(true); + key_column.set_is_nullable(false); + schema->append_column(key_column); + + TabletColumn text_column; + text_column.set_name("c2"); + text_column.set_unique_id(1); + text_column.set_type(FieldType::OLAP_FIELD_TYPE_VARCHAR); + text_column.set_length(65535); + text_column.set_is_key(false); + text_column.set_is_nullable(false); + schema->append_column(text_column); + + schema->append_index(create_phrase_index()); + return schema; + } + + TabletIndex create_phrase_index() const { + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(kIndexId); + pb.set_index_name("idx_c2"); + pb.add_col_unique_id(1); + pb.mutable_properties()->insert({"parser", "english"}); + pb.mutable_properties()->insert({"lower_case", "true"}); + pb.mutable_properties()->insert({"support_phrase", "true"}); + TabletIndex index; + index.init_from_pb(pb); + return index; + } + + Block create_block() const { + auto keys = ColumnInt32::create(); + auto values = ColumnString::create(); + for (size_t i = 0; i < kRows.size(); ++i) { + keys->insert_value(static_cast(i + 1)); + values->insert_data(kRows[i].data(), kRows[i].size()); + } + + Block block; + block.insert( + ColumnWithTypeAndName(std::move(keys), std::make_shared(), "c1")); + block.insert( + ColumnWithTypeAndName(std::move(values), std::make_shared(), "c2")); + return block; + } + + std::unique_ptr create_index_file_writer(const std::string& name, + std::string* path) const { + *path = std::string(kTestDir) + "/" + name + ".idx"; + io::FileWriterPtr file_writer; + const Status status = io::global_local_filesystem()->create_file(*path, &file_writer); + EXPECT_TRUE(status.ok()) << status; + if (!status.ok()) { + return nullptr; + } + return std::make_unique( + io::global_local_filesystem(), *path, "test_rowset", /*seg_id=*/0, + InvertedIndexStorageFormatPB::SNII, std::move(file_writer), + /*can_use_ram_dir=*/true, /*tablet_id=*/300); + } + + void write_horizontal(const std::string& name, DataWriteType write_type, + std::string* index_path) const { + auto schema = create_phrase_schema(); + auto index_file_writer = create_index_file_writer(name, index_path); + ASSERT_NE(index_file_writer, nullptr); + + const std::string data_path = std::string(kTestDir) + "/" + name + ".dat"; + io::FileWriterPtr data_file_writer; + ASSERT_TRUE(io::global_local_filesystem()->create_file(data_path, &data_file_writer).ok()); + + RowsetWriterContext rowset_context; + rowset_context.write_type = write_type; + SegmentWriterOptions options; + options.write_type = write_type; + options.rowset_ctx = &rowset_context; + SegmentWriter writer(data_file_writer.get(), /*segment_id=*/0, schema, nullptr, nullptr, + options, index_file_writer.get()); + ASSERT_TRUE(writer.init().ok()); + Block block = create_block(); + ASSERT_TRUE(writer.append_block(&block, 0, block.rows()).ok()); + + uint64_t segment_size = 0; + uint64_t index_size = 0; + ASSERT_TRUE(writer.finalize(&segment_size, &index_size).ok()); + ASSERT_TRUE(index_file_writer->begin_close().ok()); + ASSERT_TRUE(index_file_writer->finish_close().ok()); + } + + void write_vertical(const std::string& name, DataWriteType write_type, + std::string* index_path) const { + auto schema = create_phrase_schema(); + auto index_file_writer = create_index_file_writer(name, index_path); + ASSERT_NE(index_file_writer, nullptr); + + const std::string data_path = std::string(kTestDir) + "/" + name + ".dat"; + io::FileWriterPtr data_file_writer; + ASSERT_TRUE(io::global_local_filesystem()->create_file(data_path, &data_file_writer).ok()); + + RowsetWriterContext rowset_context; + rowset_context.write_type = write_type; + VerticalSegmentWriterOptions options; + options.write_type = write_type; + options.rowset_ctx = &rowset_context; + VerticalSegmentWriter writer(data_file_writer.get(), /*segment_id=*/0, schema, nullptr, + nullptr, options, index_file_writer.get()); + ASSERT_TRUE(writer.init().ok()); + Block block = create_block(); + ASSERT_TRUE(writer.batch_block(&block, 0, block.rows()).ok()); + ASSERT_TRUE(writer.write_batch().ok()); + + uint64_t segment_size = 0; + uint64_t index_size = 0; + ASSERT_TRUE(writer.finalize(&segment_size, &index_size).ok()); + ASSERT_TRUE(index_file_writer->begin_close().ok()); + ASSERT_TRUE(index_file_writer->finish_close().ok()); + } + + void write_scalar_index(const std::string& name, bool set_direct_load, + std::string* index_path) const { + auto schema = create_phrase_schema(); + TabletIndex index = create_phrase_index(); + auto index_file_writer = create_index_file_writer(name, index_path); + ASSERT_NE(index_file_writer, nullptr); + + std::unique_ptr writer; + ASSERT_TRUE(IndexColumnWriter::create(&schema->column(1), &writer, index_file_writer.get(), + &index) + .ok()); + if (set_direct_load) { + writer->set_direct_load(true); + } + std::vector values; + values.reserve(kRows.size()); + for (const auto& row : kRows) { + values.emplace_back(row); + } + ASSERT_TRUE(writer->add_values("c2", values.data(), values.size()).ok()); + ASSERT_TRUE(writer->finish().ok()); + ASSERT_TRUE(index_file_writer->begin_close().ok()); + ASSERT_TRUE(index_file_writer->finish_close().ok()); + } + + void write_array_index(const std::string& name, std::string* index_path) const { + TabletColumn array_column; + array_column.set_name("arr"); + array_column.set_unique_id(1); + array_column.set_type(FieldType::OLAP_FIELD_TYPE_ARRAY); + array_column.set_is_nullable(false); + TabletColumn item_column; + item_column.set_name("item"); + item_column.set_type(FieldType::OLAP_FIELD_TYPE_VARCHAR); + item_column.set_length(65535); + item_column.set_is_nullable(false); + array_column.add_sub_column(item_column); + + TabletIndex index = create_phrase_index(); + auto index_file_writer = create_index_file_writer(name, index_path); + ASSERT_NE(index_file_writer, nullptr); + std::unique_ptr writer; + ASSERT_TRUE( + IndexColumnWriter::create(&array_column, &writer, index_file_writer.get(), &index) + .ok()); + writer->set_direct_load(true); + + std::vector values; + values.reserve(kRows.size()); + for (const auto& row : kRows) { + values.emplace_back(row); + } + const std::vector offsets = {0, 1, 2, 3}; + ASSERT_TRUE(writer->add_array_values(field_type_size(item_column.type()), values.data(), + /*nested_null_map=*/nullptr, + reinterpret_cast(offsets.data()), + kRows.size()) + .ok()); + ASSERT_TRUE(writer->finish().ok()); + ASSERT_TRUE(index_file_writer->begin_close().ok()); + ASSERT_TRUE(index_file_writer->finish_close().ok()); + } + + void assert_only_unigrams(const std::string& path) const { + ::doris::snii::io::LocalFileReader file; + ASSERT_TRUE(file.open(path).ok()); + ::doris::snii::reader::SniiSegmentReader segment; + ASSERT_TRUE(::doris::snii::reader::SniiSegmentReader::open(&file, &segment).ok()); + ::doris::snii::reader::LogicalIndexReader index; + ASSERT_TRUE(segment.open_index(kIndexId, /*index_suffix=*/"", &index).ok()); + + std::vector<::doris::snii::reader::LogicalIndexReader::PrefixHit> hits; + ASSERT_TRUE(index.prefix_terms("", &hits).ok()); + ASSERT_FALSE(hits.empty()); + for (const auto& hit : hits) { + EXPECT_FALSE(::doris::snii::format::is_phrase_bigram_term(hit.term)) << hit.term; + } + std::vector docids; + ASSERT_TRUE(::doris::snii::query::phrase_query(index, {"alpha", "beta"}, &docids).ok()); + EXPECT_EQ(docids, kExpectedPhraseDocs); + } +}; + +TEST_F(SniiNoBigramWriter, HorizontalDirectLoadWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE( + write_horizontal("horizontal_direct", DataWriteType::TYPE_DIRECT, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, HorizontalCompactionWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE( + write_horizontal("horizontal_compaction", DataWriteType::TYPE_COMPACTION, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, HorizontalSchemaChangeWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE( + write_horizontal("horizontal_schema_change", DataWriteType::TYPE_SCHEMA_CHANGE, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, VerticalDirectLoadWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE(write_vertical("vertical_direct", DataWriteType::TYPE_DIRECT, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, VerticalCompactionWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE( + write_vertical("vertical_compaction", DataWriteType::TYPE_COMPACTION, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, VerticalSchemaChangeWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE( + write_vertical("vertical_schema_change", DataWriteType::TYPE_SCHEMA_CHANGE, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, AddIndexNoHintWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE(write_scalar_index("add_index_no_hint", false, &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST_F(SniiNoBigramWriter, DirectArrayWritesOnlyUnigrams) { + std::string path; + ASSERT_NO_FATAL_FAILURE(write_array_index("direct_array", &path)); + ASSERT_NO_FATAL_FAILURE(assert_only_unigrams(path)); +} + +TEST(SniiNoBigramWriterLifecycle, ExecEnvTmpDirectoryRemainsUsableAfterWriterSuite) { + auto* tmp_file_dirs = ExecEnv::GetInstance()->get_tmp_file_dirs(); + ASSERT_NE(tmp_file_dirs, nullptr); + const io::Path tmp_dir = tmp_file_dirs->get_tmp_file_dir(); + + auto fs = io::global_local_filesystem(); + bool tmp_dir_exists = false; + ASSERT_TRUE(fs->exists(tmp_dir, &tmp_dir_exists).ok()); + ASSERT_TRUE(tmp_dir_exists) << tmp_dir; + + const io::Path probe_path = tmp_dir / "snii_no_bigram_writer_lifecycle_probe"; + io::FileWriterPtr probe_writer; + ASSERT_TRUE(fs->create_file(probe_path, &probe_writer).ok()); + ASSERT_TRUE(probe_writer->close().ok()); + ASSERT_TRUE(fs->delete_file(probe_path).ok()); +} + +} // namespace +} // namespace doris::segment_v2 diff --git a/be/test/storage/index/snii/writer/snii_prx_tier_writer_test.cpp b/be/test/storage/index/snii/writer/snii_prx_tier_writer_test.cpp new file mode 100644 index 00000000000000..262ec252bd62e4 --- /dev/null +++ b/be/test/storage/index/snii/writer/snii_prx_tier_writer_test.cpp @@ -0,0 +1,290 @@ +// 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. + +// SNII prx-tier writer tests (patch C, see +// config::snii_prx_zstd_level_direct_load): a DIRECT load compresses the prx +// region at the cheaper load-tier zstd level, everything else keeps +// snii_prx_zstd_level. Contract pinned here: +// 1. The tier only changes prx BYTES, never SEMANTICS: a direct segment and +// its full-level twin answer position-dependent queries identically. +// 2. Non-direct paths (no hint / explicit not-direct) ignore the load-tier +// config completely -- byte-identical outputs whatever its value, so +// compaction / schema change / ADD INDEX segments are untouched. +// 3. The level is read at flush (same semantics as snii_prx_zstd_level): a +// mid-load change lands on the in-flight segment; the direct-load BIT +// itself stays captured-once. +// 4. The load-tier level is clamped to [3, 19] like the base level. + +#include + +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "io/fs/local_file_system.h" +#include "storage/index/index_file_writer.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/snii_index_writer.h" +#include "storage/tablet/tablet_schema.h" +#include "util/slice.h" + +namespace doris::segment_v2 { +namespace { + +constexpr int64_t kIndexId = 7; +constexpr const char* kTestDir = "./ut_dir/snii_prx_tier_writer_test"; + +void assert_ok(const Status& status) { + ASSERT_TRUE(status.ok()) << status.to_string(); +} + +// Positions-capable fulltext index: the only shape with a prx region to tier. +void init_phrase_index_meta(TabletIndex* meta) { + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(kIndexId); + pb.set_index_name("prx_tier_idx"); + pb.add_col_unique_id(0); + pb.mutable_properties()->insert({"parser", "english"}); + pb.mutable_properties()->insert({"lower_case", "true"}); + pb.mutable_properties()->insert({"support_phrase", "true"}); + meta->init_from_pb(pb); +} + +enum class DirectLoadHint { kNone, kDirect, kNotDirect }; + +// Same production call order as the segment writer: init() before +// set_direct_load, which precedes every row. `flip_after_hint` +// (optional) hot-changes snii_prx_zstd_level_direct_load right after the hint +// -- the earliest instant a live mInt32 change can land -- and keeps it through +// finish(), pinning the level's flush-read semantics. +void write_segment(const std::string& path, const TabletIndex& meta, + const std::vector& rows, DirectLoadHint hint, + const int32_t* flip_after_hint = nullptr) { + io::FileWriterPtr file_writer; + assert_ok(io::global_local_filesystem()->create_file(path, &file_writer)); + IndexFileWriter index_file_writer(io::global_local_filesystem(), path, "test_rowset", + /*seg_id=*/0, InvertedIndexStorageFormatPB::SNII, + std::move(file_writer), /*can_use_ram_dir=*/true, + /*tablet_id=*/300); + + SniiIndexColumnWriter writer(&index_file_writer, &meta, /*single_field=*/true); + assert_ok(writer.init()); + if (hint == DirectLoadHint::kDirect) { + writer.set_direct_load(true); + } else if (hint == DirectLoadHint::kNotDirect) { + writer.set_direct_load(false); + } + if (flip_after_hint != nullptr) { + config::snii_prx_zstd_level_direct_load = *flip_after_hint; + } + + std::vector slices; + slices.reserve(rows.size()); + for (const std::string& row : rows) { + slices.emplace_back(row); + } + assert_ok(writer.add_values("c1", slices.data(), slices.size())); + assert_ok(writer.finish()); + assert_ok(index_file_writer.begin_close()); + assert_ok(index_file_writer.finish_close()); +} + +// 400 rows x 30 tokens over a 40-word vocabulary with a deterministic pattern: +// enough position payload for the zstd level to visibly change the prx bytes, +// repetitive enough that higher levels find more to squeeze. +std::vector positional_rows() { + static const char* kVocab[] = { + "alpha", "beta", "gamma", "delta", "epsil", "zeta", "eta", "theta", "iota", "kappa", + "lam", "mu", "nu", "xi", "omic", "pi", "rho", "sigma", "tau", "upsil", + "phi", "chi", "psi", "omega", "one", "two", "three", "four", "five", "six", + "seven", "eight", "nine", "ten", "red", "green", "blue", "cyan", "lime", "teal"}; + constexpr size_t kVocabSize = sizeof(kVocab) / sizeof(kVocab[0]); + std::vector rows; + rows.reserve(400); + for (uint32_t i = 0; i < 400; ++i) { + std::string row; + for (uint32_t j = 0; j < 30; ++j) { + if (j > 0) { + row += ' '; + } + row += kVocab[(i * 31 + j * 7 + (i * j) % 5) % kVocabSize]; + } + rows.push_back(std::move(row)); + } + return rows; +} + +void open_index(const std::string& path, ::doris::snii::io::LocalFileReader* file, + ::doris::snii::reader::SniiSegmentReader* segment, + ::doris::snii::reader::LogicalIndexReader* idx) { + assert_ok(file->open(path)); + assert_ok(::doris::snii::reader::SniiSegmentReader::open(file, segment)); + assert_ok(segment->open_index(static_cast(kIndexId), /*index_suffix=*/"", idx)); +} + +std::vector run_phrase(const ::doris::snii::reader::LogicalIndexReader& idx, + const std::vector& terms) { + std::vector docids; + EXPECT_TRUE(::doris::snii::query::phrase_query(idx, terms, &docids).ok()); + return docids; +} + +std::string read_file_bytes(const std::string& path) { + std::ifstream in(path, std::ios::binary); + EXPECT_TRUE(in.good()) << path; + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +class SniiPrxTierWriterTest : public testing::Test { +protected: + void SetUp() override { + assert_ok(io::global_local_filesystem()->delete_directory(kTestDir)); + assert_ok(io::global_local_filesystem()->create_directory(kTestDir)); + _saved_prx_level = config::snii_prx_zstd_level; + _saved_load_level = config::snii_prx_zstd_level_direct_load; + init_phrase_index_meta(&_meta); + } + + void TearDown() override { + config::snii_prx_zstd_level = _saved_prx_level; + config::snii_prx_zstd_level_direct_load = _saved_load_level; + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); + } + + std::string test_path(const std::string& name) const { + return std::string(kTestDir) + "/" + name + ".idx"; + } + + TabletIndex _meta; + +private: + int32_t _saved_prx_level = 9; + int32_t _saved_load_level = 3; +}; + +// Contract 1: the tier changes prx bytes, never query semantics. Widest level +// split (base 19 vs load 3) so the size delta cannot vanish into zstd noise. +TEST_F(SniiPrxTierWriterTest, DirectLoadUsesLoadTierPrxLevelWithIdenticalAnswers) { + const std::vector rows = positional_rows(); + config::snii_prx_zstd_level = 19; + config::snii_prx_zstd_level_direct_load = 3; + + const std::string full_path = test_path("full_level"); + const std::string direct_path = test_path("direct_level"); + write_segment(full_path, _meta, rows, DirectLoadHint::kNone); + write_segment(direct_path, _meta, rows, DirectLoadHint::kDirect); + + const std::string full_bytes = read_file_bytes(full_path); + const std::string direct_bytes = read_file_bytes(direct_path); + // Level 3 compresses the prx region less than level 19: the direct segment + // must be strictly larger (this is ALSO the observable proving the tier + // actually took effect -- there is no level field in the format). + EXPECT_GT(direct_bytes.size(), full_bytes.size()); + + ::doris::snii::io::LocalFileReader full_file, direct_file; + ::doris::snii::reader::SniiSegmentReader full_segment, direct_segment; + ::doris::snii::reader::LogicalIndexReader full_idx, direct_idx; + open_index(full_path, &full_file, &full_segment, &full_idx); + open_index(direct_path, &direct_file, &direct_segment, &direct_idx); + + // Position-dependent answers must be identical and non-trivial across + // several adjacent pairs of the generated pattern. + bool any_nonempty = false; + for (const auto& phrase : std::vector> { + {"alpha", "theta"}, {"kappa", "sigma"}, {"one", "eight"}, {"red", "cyan"}}) { + const std::vector expect = run_phrase(full_idx, phrase); + EXPECT_EQ(run_phrase(direct_idx, phrase), expect) << phrase[0] << ' ' << phrase[1]; + any_nonempty |= !expect.empty(); + } + EXPECT_TRUE(any_nonempty) << "test corpus produced no phrase matches -- assertions vacuous"; +} + +// Contract 2: non-direct paths ignore the load-tier config completely. +TEST_F(SniiPrxTierWriterTest, NonDirectPathsIgnoreLoadTierLevel) { + const std::vector rows = positional_rows(); + config::snii_prx_zstd_level = 9; + + config::snii_prx_zstd_level_direct_load = 3; + const std::string none_lo = test_path("none_lo"); + const std::string notdirect_lo = test_path("notdirect_lo"); + write_segment(none_lo, _meta, rows, DirectLoadHint::kNone); + write_segment(notdirect_lo, _meta, rows, DirectLoadHint::kNotDirect); + + config::snii_prx_zstd_level_direct_load = 19; + const std::string none_hi = test_path("none_hi"); + write_segment(none_hi, _meta, rows, DirectLoadHint::kNone); + + const std::string baseline = read_file_bytes(none_lo); + ASSERT_FALSE(baseline.empty()); + EXPECT_EQ(read_file_bytes(notdirect_lo), baseline); // explicit not-direct == no hint + EXPECT_EQ(read_file_bytes(none_hi), baseline); // load-tier value is inert here +} + +// Contract 3: the LEVEL is read at flush (same semantics as +// snii_prx_zstd_level): a change landing after the captured hint but before +// finish() takes effect -- equal to the constant-config twin of the NEW value. +TEST_F(SniiPrxTierWriterTest, MidLoadLevelChangeLandsAtFlush) { + const std::vector rows = positional_rows(); + config::snii_prx_zstd_level = 9; + + config::snii_prx_zstd_level_direct_load = 9; + const std::string ref9 = test_path("flip_ref9"); + write_segment(ref9, _meta, rows, DirectLoadHint::kDirect); + + config::snii_prx_zstd_level_direct_load = 3; // captured hint under level 3 ... + const int32_t flip_to = 9; // ... flipped to 9 before any row + const std::string flipped = test_path("flip_to9"); + write_segment(flipped, _meta, rows, DirectLoadHint::kDirect, &flip_to); + + EXPECT_EQ(read_file_bytes(flipped), read_file_bytes(ref9)); +} + +// Contract 4: the load-tier level is clamped to [3, 19]. +TEST_F(SniiPrxTierWriterTest, LoadTierLevelClamped) { + const std::vector rows = positional_rows(); + config::snii_prx_zstd_level = 9; + + config::snii_prx_zstd_level_direct_load = 1; // below floor -> 3 + const std::string lvl_under = test_path("clamp_under"); + write_segment(lvl_under, _meta, rows, DirectLoadHint::kDirect); + config::snii_prx_zstd_level_direct_load = 3; + const std::string lvl_floor = test_path("clamp_floor"); + write_segment(lvl_floor, _meta, rows, DirectLoadHint::kDirect); + EXPECT_EQ(read_file_bytes(lvl_under), read_file_bytes(lvl_floor)); + + config::snii_prx_zstd_level_direct_load = 25; // above ceiling -> 19 + const std::string lvl_over = test_path("clamp_over"); + write_segment(lvl_over, _meta, rows, DirectLoadHint::kDirect); + config::snii_prx_zstd_level_direct_load = 19; + const std::string lvl_ceil = test_path("clamp_ceil"); + write_segment(lvl_ceil, _meta, rows, DirectLoadHint::kDirect); + EXPECT_EQ(read_file_bytes(lvl_over), read_file_bytes(lvl_ceil)); +} + +} // namespace +} // namespace doris::segment_v2 diff --git a/be/test/storage/index/snii/writer/snii_writer_golden_bytes_test.cpp b/be/test/storage/index/snii/writer/snii_writer_golden_bytes_test.cpp new file mode 100644 index 00000000000000..4b37cce59e5bba --- /dev/null +++ b/be/test/storage/index/snii/writer/snii_writer_golden_bytes_test.cpp @@ -0,0 +1,406 @@ +// 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. + +// GOLDEN BYTE pins for the ordinary SNII writer path. Each test +// writes ONE segment from a FIXED corpus through the production writer stack +// and asserts the output file's FNV-1a-64 digest against a recorded constant +// harvested after B2 stopped feeding hidden bigrams. Any change to tokenization +// semantics, position accounting, ignore_above / empty-value handling, or +// ordinary on-disk encoding flips the digest. +// +// The corpus deliberately hits the edge lanes: empty value (analyzed: skipped; +// keyword: a VALID empty token), punctuation-only row (zero analyzed tokens), +// >ignore_above value (keyword: skipped row), long token, repeated terms +// (position increments), unicode/mixed text, multiple add_values +// batches, and interleaved add_nulls runs. +// +// If a digest changes INTENTIONALLY (format or analyzer change), re-harvest by +// running the test and copying the "actual=" value from the failure message -- +// and say so loudly in the commit message. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "io/fs/local_file_system.h" +#include "storage/index/index_file_writer.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/snii_index_writer.h" +#include "storage/index/snii/writer/posting_window_emitter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" +#include "storage/tablet/tablet_schema.h" +#include "util/slice.h" + +namespace doris::segment_v2 { +namespace { + +constexpr int64_t kIndexId = 9; +constexpr const char* kTestDir = "./ut_dir/snii_writer_golden_bytes_test"; + +void assert_ok(const Status& status) { + ASSERT_TRUE(status.ok()) << status.to_string(); +} + +TabletIndex make_meta(const std::map& properties) { + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(kIndexId); + pb.set_index_name("golden_idx"); + pb.add_col_unique_id(0); + for (const auto& [k, v] : properties) { + pb.mutable_properties()->insert({k, v}); + } + TabletIndex meta; + meta.init_from_pb(pb); + return meta; +} + +uint64_t fnv1a64(const std::string& bytes) { + uint64_t h = 1469598103934665603ULL; + for (unsigned char c : bytes) { + h ^= c; + h *= 1099511628211ULL; + } + return h; +} + +uint64_t fnv1a64(std::span bytes) { + uint64_t h = 1469598103934665603ULL; + for (const uint8_t c : bytes) { + h ^= c; + h *= 1099511628211ULL; + } + return h; +} + +struct ShapeImage { + std::vector bytes; + doris::snii::format::DictEntryKind kind = doris::snii::format::DictEntryKind::kPodRef; + doris::snii::format::DictEntryEnc encoding = doris::snii::format::DictEntryEnc::kSlim; + std::vector window_docs; +}; + +doris::snii::writer::TermPostings make_shape_term(std::string term, uint32_t doc_count, + bool irregular_docids = false) { + doris::snii::writer::TermPostings postings; + postings.term = std::move(term); + postings.docids.reserve(doc_count); + postings.freqs.assign(doc_count, 1); + postings.positions_flat.assign(doc_count, 0); + uint32_t docid = 0; + for (uint32_t i = 0; i < doc_count; ++i) { + postings.docids.push_back(docid); + docid += irregular_docids ? 1U + ((i * 7919U) & 0xFFFFU) : 1U; + } + return postings; +} + +ShapeImage build_shape_image(doris::snii::writer::SniiIndexInput input, + std::string_view lookup_term) { + using doris::snii::Slice; + using doris::snii::format::DictEntry; + using doris::snii::format::FrqPreludeReader; + using doris::snii::reader::LogicalIndexReader; + using doris::snii::reader::SniiSegmentReader; + using doris::snii::snii_test::MemoryFile; + using doris::snii::writer::SniiCompoundWriter; + + const uint64_t index_id = input.index_id; + const std::string index_suffix = input.index_suffix; + MemoryFile file; + SniiCompoundWriter compound(&file); + assert_ok(compound.add_logical_index(input)); + assert_ok(compound.finish()); + + SniiSegmentReader segment; + LogicalIndexReader index; + assert_ok(SniiSegmentReader::open(&file, &segment)); + assert_ok(segment.open_index(index_id, index_suffix, &index)); + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index.lookup(lookup_term, &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + + ShapeImage image { + .bytes = file.data(), .kind = entry.kind, .encoding = entry.enc, .window_docs = {}}; + if (entry.enc == doris::snii::format::DictEntryEnc::kWindowed) { + const auto& posting = index.section_refs().posting_region; + const uint64_t prelude_offset = posting.offset + frq_base + entry.frq_off_delta; + EXPECT_LE(prelude_offset + entry.prelude_len, image.bytes.size()); + if (prelude_offset + entry.prelude_len > image.bytes.size()) { + return image; + } + FrqPreludeReader prelude; + assert_ok(FrqPreludeReader::open( + Slice(image.bytes.data() + prelude_offset, entry.prelude_len), &prelude)); + image.window_docs.reserve(prelude.n_windows()); + for (uint32_t i = 0; i < prelude.n_windows(); ++i) { + doris::snii::format::WindowMeta window; + assert_ok(prelude.window(i, &window)); + image.window_docs.push_back(window.doc_count); + } + } + return image; +} + +doris::snii::writer::SniiIndexInput make_shape_input( + uint64_t index_id, std::string suffix, doris::snii::writer::TermPostings term, + doris::snii::format::IndexConfig config = + doris::snii::format::IndexConfig::kDocsPositions) { + doris::snii::writer::SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = std::move(suffix); + input.config = config; + input.doc_count = term.docids.empty() ? 0 : term.docids.back() + 1; + input.write_freq = false; + input.terms.push_back(std::move(term)); + return input; +} + +doris::snii::writer::SniiIndexInput make_recut_input(uint64_t index_id, std::string suffix, + uint32_t doc_count) { + constexpr uint32_t kPrefixDocs = 8192; + constexpr uint32_t kFarPosition = 1U << 28; + auto term = make_shape_term("recut", doc_count); + term.positions_flat.clear(); + term.positions_flat.reserve(kPrefixDocs + 2 * (doc_count - kPrefixDocs)); + for (uint32_t i = 0; i < doc_count; ++i) { + if (i < kPrefixDocs) { + term.positions_flat.push_back(0); + continue; + } + term.freqs[i] = 2; + term.positions_flat.push_back(0); + term.positions_flat.push_back(kFarPosition); + } + auto input = make_shape_input(index_id, std::move(suffix), std::move(term)); + input.prx_window_limits = { + .max_docs = 1024, + .max_positions = 2048, + .max_uncomp_bytes = 2048, + }; + return input; +} + +std::string read_file_bytes(const std::string& path) { + std::ifstream in(path, std::ios::binary); + EXPECT_TRUE(in.good()) << path; + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +void add_batch(SniiIndexColumnWriter* writer, const std::vector& rows) { + std::vector slices; + slices.reserve(rows.size()); + for (const std::string& row : rows) { + slices.emplace_back(row); + } + assert_ok(writer->add_values("c1", slices.data(), slices.size())); +} + +// The fixed corpus: two add_values batches with add_nulls runs interleaved. +void feed_corpus(SniiIndexColumnWriter* writer) { + add_batch(writer, { + "hello world hello doris", + "", // analyzed: no tokens; keyword: valid EMPTY token + "The QUICK brown-fox; jumped!! over_the lazy dog 42 times", + "重复 重复 重复 词元 Doris 数据库 全文检索 mixed 中英 tokens", + }); + assert_ok(writer->add_nulls(3)); + add_batch(writer, { + std::string(300, 'x'), // keyword: > ignore_above(256) -> skipped + "single", + "!!! ??? ,,,", // analyzed: zero tokens survive + "hello world again and again and again", + }); + assert_ok(writer->add_nulls(1)); +} + +// Writes the corpus through the production stack and returns the segment +// file's digest. +uint64_t golden_digest(const std::string& name, const TabletIndex& meta) { + const std::string path = std::string(kTestDir) + "/" + name + ".idx"; + io::FileWriterPtr file_writer; + EXPECT_TRUE(io::global_local_filesystem()->create_file(path, &file_writer).ok()); + IndexFileWriter index_file_writer(io::global_local_filesystem(), path, "golden_rowset", + /*seg_id=*/0, InvertedIndexStorageFormatPB::SNII, + std::move(file_writer), /*can_use_ram_dir=*/true, + /*tablet_id=*/900); + SniiIndexColumnWriter writer(&index_file_writer, &meta, /*single_field=*/true); + EXPECT_TRUE(writer.init().ok()); + feed_corpus(&writer); + EXPECT_TRUE(writer.finish().ok()); + EXPECT_TRUE(index_file_writer.begin_close().ok()); + EXPECT_TRUE(index_file_writer.finish_close().ok()); + const std::string bytes = read_file_bytes(path); + EXPECT_FALSE(bytes.empty()); + return fnv1a64(bytes); +} + +class SniiWriterGoldenBytes : public testing::Test { +protected: + void SetUp() override { + // Pin every live config the write path reads, so the digests do not move + // under future default changes. + _saved_dict_lvl = config::snii_dict_block_zstd_level; + _saved_prx_lvl = config::snii_prx_zstd_level; + _saved_prx_load_lvl = config::snii_prx_zstd_level_direct_load; + assert_ok(io::global_local_filesystem()->delete_directory(kTestDir)); + assert_ok(io::global_local_filesystem()->create_directory(kTestDir)); + config::snii_dict_block_zstd_level = 3; + config::snii_prx_zstd_level = 3; + config::snii_prx_zstd_level_direct_load = 3; + _saved_write_freq = config::snii_positions_index_write_freq; + config::snii_positions_index_write_freq = false; + } + + void TearDown() override { + config::snii_dict_block_zstd_level = _saved_dict_lvl; + config::snii_prx_zstd_level = _saved_prx_lvl; + config::snii_prx_zstd_level_direct_load = _saved_prx_load_lvl; + config::snii_positions_index_write_freq = _saved_write_freq; + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); + } + +private: + int32_t _saved_dict_lvl = 3; + int32_t _saved_prx_lvl = 3; + int32_t _saved_prx_load_lvl = 3; + bool _saved_write_freq = false; +}; + +// Whole-image digests re-harvested for the protobuf v1 metadata layout. They +// still pin posting bytes together with every framing, directory, and metadata +// byte, so future format changes remain explicit. +constexpr uint64_t kGoldenEnglishPhrase = 0x0adb7bf49bed5dc2ULL; +constexpr uint64_t kGoldenUnicodePhrase = 0x2ea85ae3a736665cULL; +constexpr uint64_t kGoldenKeywordDocsOnly = 0xcdfd89278e7a7979ULL; + +TEST_F(SniiWriterGoldenBytes, EnglishPhrase) { + const TabletIndex meta = + make_meta({{"parser", "english"}, {"lower_case", "true"}, {"support_phrase", "true"}}); + const uint64_t digest = golden_digest("english_phrase", meta); + EXPECT_EQ(digest, kGoldenEnglishPhrase) + << "actual=0x" << std::hex << digest + << " -- token-path output changed; see file header before re-harvesting"; +} + +TEST_F(SniiWriterGoldenBytes, UnicodePhrase) { + const TabletIndex meta = + make_meta({{"parser", "unicode"}, {"lower_case", "true"}, {"support_phrase", "true"}}); + const uint64_t digest = golden_digest("unicode_phrase", meta); + EXPECT_EQ(digest, kGoldenUnicodePhrase) + << "actual=0x" << std::hex << digest + << " -- token-path output changed; see file header before re-harvesting"; +} + +TEST_F(SniiWriterGoldenBytes, KeywordDocsOnly) { + const TabletIndex meta = make_meta({{"ignore_above", "256"}}); + const uint64_t digest = golden_digest("keyword_docs_only", meta); + EXPECT_EQ(digest, kGoldenKeywordDocsOnly) + << "actual=0x" << std::hex << digest + << " -- token-path output changed; see file header before re-harvesting"; +} + +TEST_F(SniiWriterGoldenBytes, PostingShapeMatrixCompleteImageDigest) { + doris::snii::writer::testing::reset_window_emitter_counters(); + auto docs_only_term = make_shape_term("docs-only", 512); + docs_only_term.positions_flat.clear(); + docs_only_term.retain_positions = false; + const std::array images = { + build_shape_image(make_shape_input(101, "inline", make_shape_term("inline", 1)), + "inline"), + build_shape_image(make_shape_input(102, "slim", make_shape_term("slim", 511, true)), + "slim"), + build_shape_image(make_shape_input(103, "df-511", make_shape_term("df-511", 511)), + "df-511"), + build_shape_image(make_shape_input(104, "df-512", make_shape_term("df-512", 512)), + "df-512"), + build_shape_image(make_shape_input(105, "df-8191", make_shape_term("df-8191", 8191)), + "df-8191"), + build_shape_image(make_shape_input(106, "df-8192", make_shape_term("df-8192", 8192)), + "df-8192"), + build_shape_image(make_recut_input(107, "recut-full-tail", 8192 + 1024), "recut"), + build_shape_image(make_recut_input(108, "recut-partial-tail", 8192 + 512), "recut"), + }; + const ShapeImage docs_only = + build_shape_image(make_shape_input(109, "docs-only", std::move(docs_only_term), + doris::snii::format::IndexConfig::kDocsOnly), + "docs-only"); + + EXPECT_EQ(images[0].kind, doris::snii::format::DictEntryKind::kInline); + EXPECT_EQ(images[0].encoding, doris::snii::format::DictEntryEnc::kSlim); + EXPECT_EQ(images[1].kind, doris::snii::format::DictEntryKind::kPodRef); + EXPECT_EQ(images[1].encoding, doris::snii::format::DictEntryEnc::kSlim); + EXPECT_EQ(images[2].encoding, doris::snii::format::DictEntryEnc::kSlim); + EXPECT_EQ(images[3].window_docs, (std::vector {256, 256})); + EXPECT_EQ(images[4].window_docs.size(), 32U); + EXPECT_EQ(images[4].window_docs.back(), 255U); + EXPECT_EQ(images[5].window_docs, (std::vector(8, 1024))); + ASSERT_GT(images[6].window_docs.size(), 9U); + ASSERT_GT(images[7].window_docs.size(), 9U); + EXPECT_TRUE(std::ranges::all_of(images[6].window_docs.begin(), + images[6].window_docs.begin() + 8, + [](uint32_t docs) { return docs == 1024; })); + EXPECT_TRUE(std::ranges::all_of(images[7].window_docs.begin(), + images[7].window_docs.begin() + 8, + [](uint32_t docs) { return docs == 1024; })); + EXPECT_LT(images[6].window_docs[8], 1024U); + EXPECT_LT(images[7].window_docs[8], 512U); + EXPECT_EQ(docs_only.encoding, doris::snii::format::DictEntryEnc::kWindowed); + + constexpr std::array names = { + "inline", "slim", "df-511", "df-512", "df-8191", + "df-8192", "recut-full", "recut-tail", "docs-only", + }; + constexpr std::array expected = { + 0x3d00d59799c7d0adULL, 0x1ae78d4f5bcfe8b9ULL, 0x2b6f0cf4ba73bfb0ULL, + 0xc3e82019faf77965ULL, 0x8152796902268ea2ULL, 0xe9219cd6881137a9ULL, + 0xc7da487463843f7aULL, 0xc65096369e752f3eULL, 0x70e35f7c3b9c42a1ULL, + }; + for (size_t i = 0; i < images.size(); ++i) { + const uint64_t actual = fnv1a64(images[i].bytes); + EXPECT_EQ(actual, expected[i]) + << names[i] << " actual=0x" << std::hex << actual + << "; the complete SNII image changed; never refresh only a local region"; + } + const uint64_t docs_only_actual = fnv1a64(docs_only.bytes); + EXPECT_EQ(docs_only_actual, expected.back()) + << names.back() << " actual=0x" << std::hex << docs_only_actual + << "; the complete SNII image changed; never refresh only a local region"; + EXPECT_EQ(doris::snii::writer::testing::window_emitter_finished_terms(), 6U); +} + +} // namespace +} // namespace doris::segment_v2 diff --git a/be/test/storage/index/snii/writer/spill_run_codec_test.cpp b/be/test/storage/index/snii/writer/spill_run_codec_test.cpp new file mode 100644 index 00000000000000..7e68202539d398 --- /dev/null +++ b/be/test/storage/index/snii/writer/spill_run_codec_test.cpp @@ -0,0 +1,1254 @@ +// 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 "storage/index/snii/writer/spill_run_codec.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +// doris::snii::Status was deleted in the Doris integration (R01); the codec now returns +// doris::Status. Corruption is surfaced via the INVERTED_INDEX_FILE_CORRUPTED +// error code (verified against the integrated spill_run_codec.cpp), not a generic +// CORRUPTION code, so the corruption assertions below check that code explicitly. +using doris::Status; +using doris::snii::writer::compact_runs; +using doris::snii::writer::merge_run_sources; +using doris::snii::writer::MemoryReporter; +using doris::snii::writer::RunReader; +using doris::snii::writer::RunWriter; +using doris::snii::writer::TermPostings; + +namespace { + +std::string RunPath() { + static int counter = 0; + return "/tmp/snii_runcodec_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".run"; +} + +// RAII temp file: removed on scope exit so the suite leaves no debris. +struct TempRun { + std::string path = RunPath(); + ~TempRun() { std::remove(path.c_str()); } +}; + +uint64_t FileSize(const std::string& path) { + struct stat st {}; + EXPECT_EQ(::stat(path.c_str(), &st), 0); + return st.st_size < 0 ? 0 : static_cast(st.st_size); +} + +// A run record is keyed by term-id; this pairs the id with the postings so the +// test can both write (by id) and assert (the resolved string round-trips). +struct IdTerm { + uint32_t id; + TermPostings tp; +}; + +TermPostings MakeTerm(std::vector docids, std::vector freqs, + std::vector> positions = {}) { + TermPostings tp; + tp.docids = std::move(docids); + tp.freqs = std::move(freqs); + tp.set_positions_per_doc(positions); // flatten per-doc lists into positions_flat + // The codec derives the run record shape from retain_positions (the + // authoritative flag), not from whether positions happen to be empty; a + // no-positions term must carry retain_positions=false or the reader's + // has_positions=false open rejects the kPositioned record. + tp.retain_positions = !tp.positions_flat.empty(); + return tp; +} + +// Computes the term-id -> lexicographic rank array over a dense vocab, mirroring +// SpimiTermBuffer::ensure_string_rank(). MergeRuns now takes this dense integer rank +// as its heap/gather key (instead of comparing vocab strings inline), so the tests +// hand it the same lexicographic rank the production caller derives from the vocab. +std::vector LexRank(const std::vector& vocab) { + std::vector order(vocab.size()); + std::iota(order.begin(), order.end(), 0U); + std::ranges::sort(order, [&](uint32_t a, uint32_t b) { return vocab[a] < vocab[b]; }); + std::vector rank(vocab.size(), 0U); + for (uint32_t r = 0; r < order.size(); ++r) { + rank[order[r]] = r; + } + return rank; +} + +Status MergeRuns(const std::vector& run_paths, const std::vector& vocab, + const std::vector& string_rank, bool has_positions, + const std::function& fn) { + return merge_run_sources(run_paths, vocab, string_rank, has_positions, + [&](doris::snii::writer::StreamedTermPostings&& streamed) { + TermPostings materialized; + RETURN_IF_ERROR(doris::snii::writer::materialize_streamed_term( + std::move(streamed), &materialized)); + fn(std::move(materialized)); + return Status::OK(); + }); +} + +// Writes a single run from `terms` (by id) and reads it back, asserting an exact +// round-trip of every field. The reader leaves current().term empty (runs store +// only the id), so the term-id is checked via current_id(). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +void RoundTrip(const std::vector& terms, bool has_positions) { + TempRun run; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + for (const auto& t : terms) { + ASSERT_TRUE(w.write_term(t.id, t.tp).ok()); + } + ASSERT_TRUE(w.close().ok()); + } + RunReader r; + ASSERT_TRUE(r.open(run.path, has_positions).ok()); + for (const auto& expect : terms) { + ASSERT_FALSE(r.exhausted()); + EXPECT_EQ(r.current_id(), expect.id); + // Positions are LAZY: the count is known after advance(), the bytes only after + // materialize_positions(). + EXPECT_EQ(r.current_pos_count(), expect.tp.positions_flat.size()); + ASSERT_TRUE(r.materialize_positions().ok()); + const TermPostings& got = r.current(); + EXPECT_EQ(got.docids, expect.tp.docids); + EXPECT_EQ(got.freqs, expect.tp.freqs); + if (has_positions) { + EXPECT_EQ(got.positions_flat, expect.tp.positions_flat); + } + ASSERT_TRUE(r.advance().ok()); + } + EXPECT_TRUE(r.exhausted()); +} + +} // namespace + +// DoS prevention: a corrupt/truncated run whose n_docs length varint decodes to an +// absurd value must yield Corruption (bounded by the run's file size), NOT an +// uncaught std::bad_alloc from read_raw_u32's resize(). No docid data follows the +// huge count, so without the file-size bound this would resize() to ~4e9 u32s. +TEST(SniiSpillRunCodec, CorruptDocCountIsCorruptionNotBadAlloc) { + TempRun run; + { + // NOLINTBEGIN(clang-analyzer-unix.Stream): closed on the success path; only an + // ASSERT failure would skip fclose, which aborts the test anyway. + std::FILE* f = std::fopen(run.path.c_str(), "wb"); + ASSERT_NE(f, nullptr); + uint8_t buf[16]; + size_t n = 0; + n += doris::snii::encode_varint64(0, buf + n); // term_id = 0 + n += doris::snii::encode_varint64(0xFFFFFFFFULL, buf + n); // n_docs ~= 4e9, no data follows + ASSERT_EQ(std::fwrite(buf, 1, n, f), n); + std::fclose(f); + // NOLINTEND(clang-analyzer-unix.Stream) + } + RunReader r; + const Status s = r.open(run.path, /*has_positions=*/false); // open() -> advance() + EXPECT_TRUE(s.is()) << s; +} + +// Empty run: open succeeds, immediately exhausted, merge yields nothing. +TEST(SniiSpillRunCodec, EmptyRun) { + TempRun run; + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.close().ok()); + RunReader r; + ASSERT_TRUE(r.open(run.path, /*has_positions=*/true).ok()); + EXPECT_TRUE(r.exhausted()); +} + +// Single doc, with positions: smallest non-trivial record round-trips. +TEST(SniiSpillRunCodec, SingleDocWithPositions) { + RoundTrip({{.id = 7, .tp = MakeTerm({7}, {3}, {{0, 4, 9}})}}, /*has_positions=*/true); +} + +// Docs-only run (no positions): positions field is zero and decode skips it. +TEST(SniiSpillRunCodec, NoPositions) { + RoundTrip( + {{.id = 0, .tp = MakeTerm({0, 5, 99}, {1, 2, 1})}, {.id = 1, .tp = MakeTerm({3}, {4})}}, + /*has_positions=*/false); +} + +TEST(SniiSpillRunCodec, PositionedRunRejectsPositionCountMismatch) { + TempRun run; + { + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm({7}, {2}, {{3}})).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + RunReader reader; + const Status status = reader.open(run.path, /*has_positions=*/true); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiSpillRunCodec, DocsOnlyRunRejectsPositionPayload) { + TempRun run; + { + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm({7}, {1}, {{3}})).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + RunReader reader; + const Status status = reader.open(run.path, /*has_positions=*/false); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiSpillRunCodec, DuplicateDocidWithinOneRunIsCorruption) { + TempRun run; + { + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm({7, 7}, {1, 1})).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + RunReader reader; + const Status status = reader.open(run.path, /*has_positions=*/false); + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiSpillRunCodec, MergeRunSourcesAccountsReadersAndReleasesOnSuccess) { + constexpr uint32_t kDocsPerRun = 4096; + const std::vector vocab = {"wide"}; + const std::vector rank = {0}; + TempRun first; + TempRun second; + for (size_t run = 0; run < 2; ++run) { + std::vector docids(kDocsPerRun); + std::iota(docids.begin(), docids.end(), static_cast(run) * kDocsPerRun); + std::vector freqs(kDocsPerRun, 1); + RunWriter writer; + ASSERT_TRUE(writer.open(run == 0 ? first.path : second.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm(std::move(docids), std::move(freqs))).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + const uint64_t encoded_input_bytes = FileSize(first.path) + FileSize(second.path); + MemoryReporter reporter; + int64_t during_callback = 0; + ASSERT_TRUE(merge_run_sources( + {first.path, second.path}, vocab, rank, /*has_positions=*/false, + [&](doris::snii::writer::StreamedTermPostings&& streamed) { + during_callback = reporter.current_bytes(); + return doris::snii::writer::consume_streamed_term(std::move(streamed)); + }, + {}, &reporter) + .ok()); + EXPECT_GT(during_callback, static_cast(encoded_input_bytes)); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpillRunCodec, RunReaderDocidReservationFailureReleasesAllCharges) { + constexpr uint32_t kDocs = 20000; + const std::vector vocab = {"term"}; + const std::vector rank = {0}; + TempRun run; + { + std::vector docids(kDocs); + std::iota(docids.begin(), docids.end(), 0U); + std::vector freqs(kDocs, 1); + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm(std::move(docids), std::move(freqs))).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/100U << 10, + MemoryReporter::CapPolicy::kHardLimit); + const Status status = merge_run_sources( + {run.path}, vocab, rank, /*has_positions=*/false, + [](doris::snii::writer::StreamedTermPostings&& streamed) { + return doris::snii::writer::consume_streamed_term(std::move(streamed)); + }, + {}, &reporter); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpillRunCodec, RunReaderFrequencyReservationFailureReleasesAllCharges) { + constexpr uint32_t kDocs = 20000; + const std::vector vocab = {"term"}; + TempRun run; + { + std::vector docids(kDocs); + std::iota(docids.begin(), docids.end(), 0U); + std::vector freqs(kDocs, 1); + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm(std::move(docids), std::move(freqs))).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/160U << 10, + MemoryReporter::CapPolicy::kHardLimit); + { + RunReader reader(&reporter); + const Status status = reader.open(run.path, /*has_positions=*/false); + EXPECT_TRUE(status.is()) << status; + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpillRunCodec, RunReaderMaterializedPositionsAreAccountedAndReleased) { + constexpr uint32_t kDocs = 4096; + TempRun run; + { + TermPostings postings; + postings.docids.resize(kDocs); + std::iota(postings.docids.begin(), postings.docids.end(), 0U); + postings.freqs.assign(kDocs, 2); + postings.positions_flat.resize(2 * kDocs); + std::iota(postings.positions_flat.begin(), postings.positions_flat.end(), 0U); + postings.retain_positions = true; + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, postings).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + MemoryReporter reporter; + { + RunReader reader(&reporter); + ASSERT_TRUE(reader.open(run.path, /*has_positions=*/true).ok()); + const int64_t before_positions = reporter.current_bytes(); + ASSERT_TRUE(reader.materialize_positions().ok()); + EXPECT_GE(reporter.current_bytes() - before_positions, + static_cast(2 * kDocs * sizeof(uint32_t))); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpillRunCodec, RunWriterStreamsWideTermWithinAccountedBound) { + constexpr uint32_t kDocs = (1U << 20) + 1; + TempRun run; + std::vector docids(kDocs); + std::iota(docids.begin(), docids.end(), 0U); + TermPostings postings = MakeTerm(std::move(docids), {}); + + int64_t observed = 0; + int64_t peak = 0; + MemoryReporter reporter([&](int64_t delta) { + observed += delta; + peak = std::max(peak, observed); + }); + RunWriter writer(&reporter); + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, postings).ok()); + ASSERT_TRUE(writer.close().ok()); + EXPECT_GT(FileSize(run.path), 4U << 20); + EXPECT_LE(peak, static_cast((4U << 20) + 64)); + EXPECT_EQ(observed, 0); + EXPECT_EQ(reporter.current_bytes(), 0); + + RunReader reader; + ASSERT_TRUE(reader.open(run.path, /*has_positions=*/true).ok()); + EXPECT_EQ(reader.current().docids, postings.docids); + EXPECT_TRUE(reader.current().freqs.empty()); + EXPECT_TRUE(reader.current().positions_flat.empty()); + ASSERT_TRUE(reader.advance().ok()); + EXPECT_TRUE(reader.exhausted()); +} + +TEST(SniiSpillRunCodec, RunWriterGrowsStagingBufferGeometrically) { + constexpr uint32_t kTerms = 4096; + TempRun run; + const TermPostings postings = MakeTerm({7}, {}); + + int64_t observed = 0; + uint32_t positive_reservations = 0; + MemoryReporter reporter([&](int64_t delta) { + observed += delta; + if (delta > 0) { + ++positive_reservations; + } + }); + + RunWriter writer(&reporter); + ASSERT_TRUE(writer.open(run.path).ok()); + for (uint32_t term_id = 0; term_id < kTerms; ++term_id) { + ASSERT_TRUE(writer.write_term(term_id, postings).ok()); + } + ASSERT_TRUE(writer.close().ok()); + + EXPECT_GT(FileSize(run.path), 0); + EXPECT_LE(positive_reservations, 32U); + EXPECT_EQ(observed, 0); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpillRunCodec, CompactRunsMergedPostingReservationHonorsHardLimitAndReleases) { + constexpr uint32_t kDocsPerRun = 4096; + TempRun first; + TempRun second; + TempRun output; + for (size_t run = 0; run < 2; ++run) { + std::vector docids(kDocsPerRun); + std::iota(docids.begin(), docids.end(), static_cast(run) * kDocsPerRun); + RunWriter writer; + ASSERT_TRUE(writer.open(run == 0 ? first.path : second.path).ok()); + ASSERT_TRUE(writer.write_term(0, MakeTerm(std::move(docids), {})).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/80U << 10, + MemoryReporter::CapPolicy::kHardLimit); + const Status status = compact_runs({first.path, second.path}, {0}, /*has_positions=*/true, + output.path, &reporter); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(reporter.current_bytes(), 0); +} + +// Several terms with varied widths round-trip in ascending id order. +TEST(SniiSpillRunCodec, MultiTermRoundTrip) { + RoundTrip( + { + {.id = 0, .tp = MakeTerm({0, 1, 2}, {1, 1, 1}, {{0}, {1}, {2}})}, + {.id = 1, .tp = MakeTerm({10}, {2}, {{3, 8}})}, + {.id = 2, .tp = MakeTerm({4, 100}, {2, 1}, {{0, 1}, {7}})}, + }, + /*has_positions=*/true); +} + +// K-way merge: a term-id present in EVERY run is concatenated in ascending run +// order; an id present in only ONE run passes through unchanged. The merged +// stream is ordered by each id's VOCAB STRING and the string is resolved onto +// the emitted TermPostings. +TEST(SniiSpillRunCodec, MergeConcatenatesAcrossRuns) { + // Vocab: id 0 -> "common", 1 -> "only0", 2 -> "zzz". Ordered by string: + // "common" < "only0" < "zzz", which happens to match id order here. + const std::vector vocab = {"common", "only0", "zzz"}; + TempRun r0, r1, r2; + // Each run covers a strictly later docid range for the shared id 0. + { + RunWriter w; + ASSERT_TRUE(w.open(r0.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({0, 1}, {1, 2}, {{0}, {1, 2}})).ok()); + ASSERT_TRUE(w.write_term(1, MakeTerm({3}, {1}, {{5}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + { + RunWriter w; + ASSERT_TRUE(w.open(r1.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({5}, {1}, {{0}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + { + RunWriter w; + ASSERT_TRUE(w.open(r2.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({8, 9}, {1, 1}, {{0}, {0}})).ok()); + ASSERT_TRUE(w.write_term(2, MakeTerm({2}, {1}, {{4}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + + std::vector merged; + ASSERT_TRUE(MergeRuns({r0.path, r1.path, r2.path}, vocab, LexRank(vocab), + /*has_positions=*/true, + [&](TermPostings&& tp) { merged.push_back(std::move(tp)); }) + .ok()); + + ASSERT_EQ(merged.size(), 3U); + EXPECT_EQ(merged[0].term, "common"); + EXPECT_EQ(merged[0].docids, (std::vector {0, 1, 5, 8, 9})); + EXPECT_EQ(merged[0].freqs, (std::vector {1, 2, 1, 1, 1})); + // Flat positions: doc0{0} doc1{1,2} doc5{0} doc8{0} doc9{0}. + EXPECT_EQ(merged[0].positions_flat, (std::vector {0, 1, 2, 0, 0, 0})); + EXPECT_EQ(std::vector(merged[0].doc_positions(1).begin(), + merged[0].doc_positions(1).end()), + (std::vector {1, 2})); + EXPECT_EQ(merged[1].term, "only0"); + EXPECT_EQ(merged[1].docids, (std::vector {3})); + EXPECT_EQ(merged[2].term, "zzz"); + EXPECT_EQ(merged[2].docids, (std::vector {2})); +} + +// BOUNDARY COALESCE with FLAT positions: a spill that falls BETWEEN two tokens of +// the SAME doc leaves that doc ending one run and beginning the next with the same +// docid. The merge must fold them into ONE doc whose positions concatenate (run +// order) into the correct flat layout -- the trickiest flat-positions merge path. +TEST(SniiSpillRunCodec, MergeCoalescesBoundaryDocPositionsFlat) { + const std::vector vocab = {"alpha"}; + TempRun r0, r1; + { + RunWriter w; + ASSERT_TRUE(w.open(r0.path).ok()); + // doc 0 (pos 0,7), doc 1 first half (pos 1) -- doc 1 continues in r1. + ASSERT_TRUE(w.write_term(0, MakeTerm({0, 1}, {2, 1}, {{0, 7}, {1}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + { + RunWriter w; + ASSERT_TRUE(w.open(r1.path).ok()); + // doc 1 second half (pos 4,9), then doc 2 (pos 3). + ASSERT_TRUE(w.write_term(0, MakeTerm({1, 2}, {2, 1}, {{4, 9}, {3}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + std::vector merged; + ASSERT_TRUE(MergeRuns({r0.path, r1.path}, vocab, LexRank(vocab), /*has_positions=*/true, + [&](TermPostings&& tp) { merged.push_back(std::move(tp)); }) + .ok()); + ASSERT_EQ(merged.size(), 1U); + EXPECT_EQ(merged[0].docids, (std::vector {0, 1, 2})); + // doc 1 coalesced: freq 1 + 2 = 3, positions 1,4,9 (run order). + EXPECT_EQ(merged[0].freqs, (std::vector {2, 3, 1})); + // Flat: doc0{0,7} doc1{1,4,9} doc2{3}. + EXPECT_EQ(merged[0].positions_flat, (std::vector {0, 7, 1, 4, 9, 3})); + EXPECT_EQ(std::vector(merged[0].doc_positions(1).begin(), + merged[0].doc_positions(1).end()), + (std::vector {1, 4, 9})); +} + +// The merge order follows the VOCAB STRING, not the numeric id: ids whose +// strings sort in the opposite order are emitted lexicographically. +TEST(SniiSpillRunCodec, MergeOrdersByVocabStringNotId) { + // id 0 -> "zebra", id 1 -> "apple": string order is apple(1) < zebra(0). + const std::vector vocab = {"zebra", "apple"}; + TempRun r0; + { + RunWriter w; + ASSERT_TRUE(w.open(r0.path).ok()); + // Written in run order by string: apple(1) before zebra(0). + ASSERT_TRUE(w.write_term(1, MakeTerm({2}, {1})).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({5}, {1})).ok()); + ASSERT_TRUE(w.close().ok()); + } + std::vector order; + ASSERT_TRUE(MergeRuns({r0.path}, vocab, LexRank(vocab), /*has_positions=*/false, + [&](TermPostings&& tp) { order.push_back(tp.term); }) + .ok()); + EXPECT_EQ(order, (std::vector {"apple", "zebra"})); +} + +// Lazy positions: stream_positions yields the SAME bytes as the materialized +// block, even when pulled in awkward (non-block-aligned) chunk sizes that straddle +// the reader's internal 64 KiB window boundaries. +TEST(SniiSpillRunCodec, StreamPositionsMatchesMaterialized) { + TempRun run; + // One wide term: 5000 docs, freq 3 each -> 15000 flat positions spanning several + // internal read windows. + std::vector docids, freqs, flat; + for (uint32_t d = 0; d < 5000; ++d) { + docids.push_back(d); + freqs.push_back(3); + flat.push_back(d * 7 + 0); + flat.push_back(d * 7 + 1); + flat.push_back(d * 7 + 2); + } + TermPostings tp; + tp.docids = docids; + tp.freqs = freqs; + tp.positions_flat = flat; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(0, tp).ok()); + ASSERT_TRUE(w.close().ok()); + } + RunReader r; + ASSERT_TRUE(r.open(run.path, /*has_positions=*/true).ok()); + ASSERT_EQ(r.current_pos_count(), flat.size()); + ASSERT_EQ(r.positions_remaining(), flat.size()); + // Pull in odd chunks (7, 1000, 7, 1000, ...) until drained. + std::vector got; + std::vector chunks = {7, 1000, 7, 1000}; + size_t ci = 0; + while (r.positions_remaining() > 0) { + size_t want = std::min(chunks[ci % chunks.size()], + static_cast(r.positions_remaining())); + ++ci; + std::vector buf(want); + ASSERT_TRUE(r.stream_positions(buf.data(), want).ok()); + got.insert(got.end(), buf.begin(), buf.end()); + } + EXPECT_EQ(got, flat); + EXPECT_TRUE(r.positions_drained()); + ASSERT_TRUE(r.advance().ok()); + EXPECT_TRUE(r.exhausted()); +} + +// advance() after a PARTIALLY-streamed term skips the unread positions and lands +// on the next record correctly. +TEST(SniiSpillRunCodec, PartialStreamThenAdvanceSkipsRemainder) { + TempRun run; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({0, 1, 2}, {2, 2, 2}, {{10, 11}, {20, 21}, {30, 31}})) + .ok()); + ASSERT_TRUE(w.write_term(1, MakeTerm({9}, {1}, {{99}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + RunReader r; + ASSERT_TRUE(r.open(run.path, /*has_positions=*/true).ok()); + ASSERT_EQ(r.current_id(), 0U); + // Pull only the first two positions, then advance -- the remaining 4 are skipped. + std::vector buf(2); + ASSERT_TRUE(r.stream_positions(buf.data(), 2).ok()); + EXPECT_EQ(buf, (std::vector {10, 11})); + ASSERT_TRUE(r.advance().ok()); + ASSERT_FALSE(r.exhausted()); + EXPECT_EQ(r.current_id(), 1U); + ASSERT_TRUE(r.materialize_positions().ok()); + EXPECT_EQ(r.current().positions_flat, (std::vector {99})); +} + +namespace { + +Status DrainStreamed(doris::snii::writer::StreamedTermPostings&& source, TermPostings* output) { + return doris::snii::writer::materialize_streamed_term(std::move(source), output); +} + +} // namespace + +// WIDE-TERM STREAMING == MATERIALIZED (byte-identity proof at the postings level): +// a term with df >= kSlimDfThreshold split across several runs (with a boundary doc +// straddling a spill) must yield IDENTICAL docids/freqs/positions whether the merge +// fills the writer-owned window source or materializes the retained helper path. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiSpillRunCodec, MergeWideTermStreamsIdenticalToMaterialized) { + const std::vector vocab = {"wide"}; + // Build a wide term (df ~ 2000) sharded across 3 runs, with the LAST doc of each + // run continuing as the FIRST doc of the next (boundary-doc coalesce). + TempRun r0, r1, r2; + auto shard = [&](TempRun& run, uint32_t lo, uint32_t hi, uint32_t carry_first) { + TermPostings tp; + for (uint32_t d = lo; d < hi; ++d) { + tp.docids.push_back(d); + // Boundary docs (lo when it's a carry) get freq 1 here; otherwise freq 2. + const uint32_t fc = 2; + tp.freqs.push_back(fc); + for (uint32_t k = 0; k < fc; ++k) { + tp.positions_flat.push_back(d * 13 + k); + } + } + (void)carry_first; + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(0, tp).ok()); + ASSERT_TRUE(w.close().ok()); + }; + // Ranges chosen so doc 700 ends r0 AND begins r1 (boundary), doc 1400 likewise. + // Encode the boundary by repeating that docid at the seam with extra positions. + { + TermPostings a; + for (uint32_t d = 0; d <= 700; ++d) { + a.docids.push_back(d); + a.freqs.push_back(2); + a.positions_flat.push_back(d * 13); + a.positions_flat.push_back(d * 13 + 1); + } + RunWriter w; + ASSERT_TRUE(w.open(r0.path).ok()); + ASSERT_TRUE(w.write_term(0, a).ok()); + ASSERT_TRUE(w.close().ok()); + } + { + TermPostings b; + // doc 700 continues here (boundary): extra positions for it, then 701..1400. + b.docids.push_back(700); + b.freqs.push_back(1); + b.positions_flat.push_back(700 * 13 + 2); + for (uint32_t d = 701; d <= 1400; ++d) { + b.docids.push_back(d); + b.freqs.push_back(2); + b.positions_flat.push_back(d * 13); + b.positions_flat.push_back(d * 13 + 1); + } + RunWriter w; + ASSERT_TRUE(w.open(r1.path).ok()); + ASSERT_TRUE(w.write_term(0, b).ok()); + ASSERT_TRUE(w.close().ok()); + } + { + TermPostings c; + c.docids.push_back(1400); + c.freqs.push_back(1); + c.positions_flat.push_back(1400 * 13 + 2); + for (uint32_t d = 1401; d <= 2100; ++d) { + c.docids.push_back(d); + c.freqs.push_back(2); + c.positions_flat.push_back(d * 13); + c.positions_flat.push_back(d * 13 + 1); + } + RunWriter w; + ASSERT_TRUE(w.open(r2.path).ok()); + ASSERT_TRUE(w.write_term(0, c).ok()); + ASSERT_TRUE(w.close().ok()); + } + (void)shard; + + const std::vector paths = {r0.path, r1.path, r2.path}; + TermPostings materialized, streamed; + ASSERT_TRUE( + MergeRuns(paths, vocab, LexRank(vocab), /*has_positions=*/true, [&](TermPostings&& tp) { + materialized = std::move(tp); + }).ok()); + ASSERT_TRUE(merge_run_sources(paths, vocab, LexRank(vocab), /*has_positions=*/true, + [&](doris::snii::writer::StreamedTermPostings&& source) { + return DrainStreamed(std::move(source), &streamed); + }) + .ok()); + + // Both paths must produce identical docids, freqs, and positions. + EXPECT_GE(materialized.docids.size(), 512U); // wide enough to take the stream path + EXPECT_EQ(materialized.docids, streamed.docids); + EXPECT_EQ(materialized.freqs, streamed.freqs); + EXPECT_EQ(materialized.positions_flat, streamed.positions_flat); + // Boundary doc 700 coalesced: freq 2 (r0) + 1 (r1) = 3, positions in run order. + const auto it = std::ranges::find(materialized.docids, 700U); + ASSERT_NE(it, materialized.docids.end()); + const size_t bi = static_cast(it - materialized.docids.begin()); + EXPECT_EQ(materialized.freqs[bi], 3U); +} + +// A run record whose term-id is >= vocab.size() must make MergeRuns return +// Corruption (NOT index a vocab[id] out of bounds, which is UB / a crash). The +// id is decoded as a perfectly valid varint, so it is the in-merge vocab-range +// check -- not varint decode -- that must fire. This guards both the heap-seed +// range check and the post-advance one by placing the bad id as the SECOND term +// (the first term seeds the heap fine; the bad id is reached after advance()). +TEST(SniiSpillRunCodec, MergeTermIdOutOfVocabIsCorruption) { + const std::vector vocab = {"only"}; // valid ids: {0} + TempRun run; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({0}, {1}, {{0}})).ok()); // id 0: OK + ASSERT_TRUE(w.write_term(5, MakeTerm({9}, {1}, {{0}})).ok()); // id 5: out of range + ASSERT_TRUE(w.close().ok()); + } + std::vector merged; + const Status s = MergeRuns({run.path}, vocab, LexRank(vocab), /*has_positions=*/true, + [&](TermPostings&& tp) { merged.push_back(std::move(tp)); }); + EXPECT_TRUE(s.is()) << s; +} + +// And when the BAD id is the FIRST record of a run, the heap-seed range check (in +// MergeRuns, before any term is emitted) must fire -- still Corruption, no UB. +TEST(SniiSpillRunCodec, MergeFirstTermIdOutOfVocabIsCorruption) { + const std::vector vocab = {"a", "b"}; // valid ids: {0,1} + TempRun run; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(9, MakeTerm({0}, {1}, {{0}})).ok()); // id 9: out of range + ASSERT_TRUE(w.close().ok()); + } + std::vector merged; + const Status s = MergeRuns({run.path}, vocab, LexRank(vocab), /*has_positions=*/true, + [&](TermPostings&& tp) { merged.push_back(std::move(tp)); }); + EXPECT_TRUE(s.is()) << s; +} + +// A positioned record whose declared n_pos differs from sum(freqs) is rejected +// while open() loads the first record, before any allocation or position read. +TEST(SniiSpillRunCodec, NPosExceedsFileIsCorruption) { + TempRun run; + { + // NOLINTBEGIN(clang-analyzer-unix.Stream): closed on the success path; only an + // ASSERT failure would skip fclose, which aborts the test anyway. + std::FILE* f = std::fopen(run.path.c_str(), "wb"); + ASSERT_NE(f, nullptr); + uint8_t buf[40]; + size_t n = 0; + n += doris::snii::encode_varint64(0, buf + n); // term_id = 0 + n += doris::snii::encode_varint64(2, buf + n); // shape = kPositioned + n += doris::snii::encode_varint64(1, buf + n); // n_docs = 1 + // docid[0] = 0 and freq[0] = 1 as RAW LE u32 blocks (matching the writer). + const uint32_t one_docid = 0, one_freq = 1; + std::memcpy(buf + n, &one_docid, sizeof(uint32_t)); + n += sizeof(uint32_t); + std::memcpy(buf + n, &one_freq, sizeof(uint32_t)); + n += sizeof(uint32_t); + n += doris::snii::encode_varint64(0xFFFFFFFFULL, buf + n); // n_pos ~= 4e9, no data follows + ASSERT_EQ(std::fwrite(buf, 1, n, f), n); + std::fclose(f); + // NOLINTEND(clang-analyzer-unix.Stream) + } + RunReader r; + const Status s = r.open(run.path, /*has_positions=*/true); + EXPECT_TRUE(s.is()) << s; +} + +// A truncated position block must fail through the source contract while every +// unwritten slot in the value-initialized writer buffer remains deterministic. +TEST(SniiSpillRunCodec, TruncatedPositionsFailWithoutUninitializedTail) { + const std::vector vocab = {"wide"}; + constexpr uint32_t kDocs = 600; + TempRun run; + { + TermPostings postings; + for (uint32_t docid = 0; docid < kDocs; ++docid) { + postings.docids.push_back(docid); + postings.freqs.push_back(1); + postings.positions_flat.push_back(docid * 3 + 1); + } + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, postings).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + struct stat file_stat {}; + ASSERT_EQ(::stat(run.path.c_str(), &file_stat), 0); + ASSERT_EQ(::truncate(run.path.c_str(), + file_stat.st_size - static_cast(100 * sizeof(uint32_t))), + 0); + + bool source_called = false; + std::vector positions; + const Status status = merge_run_sources( + {run.path}, vocab, LexRank(vocab), /*has_positions=*/true, + [&](doris::snii::writer::StreamedTermPostings&& streamed) { + source_called = true; + doris::snii::writer::TermPostingBuffer buffer(nullptr); + bool exhausted = false; + Status fill_status = streamed.source->fill(kDocs, &buffer, &exhausted); + positions.assign(buffer.positions_flat().begin(), buffer.positions_flat().end()); + return fill_status; + }); + + EXPECT_TRUE(status.is()) << status; + ASSERT_TRUE(source_called); + ASSERT_EQ(positions.size(), kDocs); + EXPECT_TRUE(std::all_of(positions.end() - 100, positions.end(), + [](uint32_t value) { return value == 0; })); +} + +// A truncated run file is rejected by decode (anti-corruption on bytes we read). +TEST(SniiSpillRunCodec, TruncatedRunIsCorruption) { + TempRun run; + { + RunWriter w; + ASSERT_TRUE(w.open(run.path).ok()); + ASSERT_TRUE(w.write_term(0, MakeTerm({0, 1, 2}, {1, 1, 1}, {{0}, {0}, {0}})).ok()); + ASSERT_TRUE(w.write_term(1, MakeTerm({4}, {1}, {{0}})).ok()); + ASSERT_TRUE(w.close().ok()); + } + // Chop the file so the second record promises more bytes than remain. + ASSERT_EQ(::truncate(run.path.c_str(), 4), 0); + RunReader r; + Status s = r.open(run.path, /*has_positions=*/true); + while (s.ok() && !r.exhausted()) { + s = r.advance(); + } + EXPECT_FALSE(s.ok()); +} + +// Returning from the callback without draining the borrowed source is rejected. +TEST(SniiSpillRunCodec, MergeRunSourcesRejectsUnconsumedSource) { + TempRun run; + TermPostings postings; + for (uint32_t docid = 0; docid < 1000; ++docid) { + postings.docids.push_back(docid); + postings.freqs.push_back(1); + postings.positions_flat.push_back(docid); + } + { + RunWriter writer; + ASSERT_TRUE(writer.open(run.path).ok()); + ASSERT_TRUE(writer.write_term(0, postings).ok()); + ASSERT_TRUE(writer.close().ok()); + } + + const std::vector vocab = {"wide"}; + const Status status = merge_run_sources( + {run.path}, vocab, LexRank(vocab), /*has_positions=*/true, + [&](doris::snii::writer::StreamedTermPostings&&) { return Status::OK(); }); + EXPECT_TRUE(status.is()) << status; +} + +// =========================================================================== +// SniiSpillMergeTest -- T15: MergeRuns keyed on the integer string_rank array. +// +// MergeRuns now takes a precomputed `string_rank` (term-id -> lexicographic rank) +// and keys its heap/gather on that dense 4 B integer array instead of comparing +// vocab strings inline. These cases prove (a) the key is the integer rank array +// (FM-04: a deliberately NON-lexicographic rank permutation drives the emit +// order), (b) output stays byte-identical when the rank is the lexicographic one +// (FM-01..FM-03, FM-09), (c) the wide-term streamed path is unaffected (FM-05), +// (d) the error/boundary paths (FM-06..FM-08), and (e) end-to-end spill == +// in-memory through SpimiTermBuffer's production wiring (FM-10). +// =========================================================================== + +namespace { + +// Writes one run file from (term-id, postings) pairs in the given order. The caller +// supplies them sorted by the MERGE KEY (the spill writer's contract: ascending by +// the same rank MergeRuns will use). Asserts on any I/O failure. +void WriteRun(const std::string& path, const std::vector& terms) { + RunWriter w; + ASSERT_TRUE(w.open(path).ok()); + for (const auto& t : terms) { + ASSERT_TRUE(w.write_term(t.id, t.tp).ok()); + } + ASSERT_TRUE(w.close().ok()); +} + +// K-way merges `paths` under the integer `rank` key, collecting every emitted term +// into `out` with positions materialized so callers can compare flat arrays. +Status CollectMerge(const std::vector& paths, const std::vector& vocab, + const std::vector& rank, bool has_positions, + std::vector* out) { + return MergeRuns(paths, vocab, rank, has_positions, + [&](TermPostings&& tp) { out->push_back(std::move(tp)); }); +} + +} // namespace + +// FM-04 (KEY PROOF): the heap/gather key is the integer string_rank ARRAY, not the +// vocab strings. The vocab sorts lexicographically as a(id1) < b(id0) < c(id2), but +// we pass a DIFFERENT permutation (id0->0, id2->1, id1->2), so the rank order is +// b(id0), c(id2), a(id1) -- matching neither the vocab string order (a,b,c) nor the +// numeric id order. The two runs hold DISJOINT-but-overlapping term sets so the heap +// must actually interleave them; the emitted order must follow the rank array. The +// OLD vocab-string comparator, fed these rank-sorted runs, would instead emit +// b,a,c, so this sequence equality FAILS on the un-optimized code and PASSES once +// the comparator keys on string_rank. +TEST(SniiSpillMergeTest, MergeRunsOrdersByStringRankInteger) { + const std::vector vocab = {"b", "a", "c"}; + const std::vector rank = {0, 2, 1}; // id0->0, id1->2, id2->1 + TempRun r0, r1; + // Each run sorted ascending by the merge key (rank): run0 = id0(0), id1(2); + // run1 = id0(0), id2(1). id0 ("b") appears in BOTH runs (integer-id gather). + WriteRun(r0.path, {{.id = 0, .tp = MakeTerm({0}, {1})}, {.id = 1, .tp = MakeTerm({2}, {1})}}); + WriteRun(r1.path, {{.id = 0, .tp = MakeTerm({5}, {1})}, {.id = 2, .tp = MakeTerm({3}, {1})}}); + std::vector merged; + ASSERT_TRUE( + CollectMerge({r0.path, r1.path}, vocab, rank, /*has_positions=*/false, &merged).ok()); + ASSERT_EQ(merged.size(), 3U); + std::vector order; + for (const auto& m : merged) { + order.push_back(m.term); + } + EXPECT_EQ(order, (std::vector {"b", "c", "a"})); // strictly the rank order + // id0 ("b") gathered across both runs in run order -> ascending docids. + EXPECT_EQ(merged[0].term, "b"); + EXPECT_EQ(merged[0].docids, (std::vector {0, 5})); + EXPECT_EQ(merged[1].term, "c"); + EXPECT_EQ(merged[1].docids, (std::vector {3})); + EXPECT_EQ(merged[2].term, "a"); + EXPECT_EQ(merged[2].docids, (std::vector {2})); +} + +// FM-01: lexicographic rank reproduces dictionary order; an id present in several +// runs concatenates in run order (docids stay ascending). No positions. +TEST(SniiSpillMergeTest, MergeByLexRankConcatenatesNoPositions) { + const std::vector vocab = {"banana", "apple", "cherry"}; // ids 0,1,2 + const std::vector rank = LexRank(vocab); // apple(1) < banana(0) < cherry(2) + TempRun r0, r1; + // Each run sorted by lex rank: apple(1), then banana(0) / cherry(2). + WriteRun(r0.path, + {{.id = 1, .tp = MakeTerm({0, 2}, {1, 1})}, {.id = 0, .tp = MakeTerm({1}, {1})}}); + WriteRun(r1.path, + {{.id = 1, .tp = MakeTerm({5}, {1})}, {.id = 2, .tp = MakeTerm({3, 9}, {1, 1})}}); + std::vector merged; + ASSERT_TRUE( + CollectMerge({r0.path, r1.path}, vocab, rank, /*has_positions=*/false, &merged).ok()); + ASSERT_EQ(merged.size(), 3U); + EXPECT_EQ(merged[0].term, "apple"); + EXPECT_EQ(merged[0].docids, (std::vector {0, 2, 5})); // r0{0,2} ++ r1{5} + EXPECT_EQ(merged[0].freqs, (std::vector {1, 1, 1})); + EXPECT_EQ(merged[1].term, "banana"); + EXPECT_EQ(merged[1].docids, (std::vector {1})); + EXPECT_EQ(merged[2].term, "cherry"); + EXPECT_EQ(merged[2].docids, (std::vector {3, 9})); +} + +// FM-02: same shape with positions -- positions_flat materializes correctly per term +// (document order, partitioned by freqs). +TEST(SniiSpillMergeTest, MergeByLexRankWithPositions) { + const std::vector vocab = {"banana", "apple", "cherry"}; + const std::vector rank = LexRank(vocab); + TempRun r0, r1; + WriteRun(r0.path, {{.id = 1, .tp = MakeTerm({0, 2}, {2, 1}, {{3, 4}, {7}})}, + {.id = 0, .tp = MakeTerm({1}, {1}, {{5}})}}); + WriteRun(r1.path, {{.id = 1, .tp = MakeTerm({5}, {2}, {{0, 9}})}, + {.id = 2, .tp = MakeTerm({3}, {1}, {{6}})}}); + std::vector merged; + ASSERT_TRUE( + CollectMerge({r0.path, r1.path}, vocab, rank, /*has_positions=*/true, &merged).ok()); + ASSERT_EQ(merged.size(), 3U); + EXPECT_EQ(merged[0].term, "apple"); + EXPECT_EQ(merged[0].docids, (std::vector {0, 2, 5})); + EXPECT_EQ(merged[0].freqs, (std::vector {2, 1, 2})); + // doc0{3,4} doc2{7} doc5{0,9} + EXPECT_EQ(merged[0].positions_flat, (std::vector {3, 4, 7, 0, 9})); + EXPECT_EQ(merged[1].positions_flat, (std::vector {5})); + EXPECT_EQ(merged[2].positions_flat, (std::vector {6})); +} + +// FM-03: a doc split across a spill boundary (last doc of run0 == first doc of run1) +// coalesces into one entry (freqs summed, positions spliced in run order). The merge +// key is the integer rank, but the concat boundary path is unchanged. +TEST(SniiSpillMergeTest, MergeCoalescesBoundaryDoc) { + const std::vector vocab = {"x"}; + const std::vector rank = LexRank(vocab); // {0} + TempRun r0, r1; + // doc 0, then doc 4 (first half). doc 4 continues in r1. + WriteRun(r0.path, {{.id = 0, .tp = MakeTerm({0, 4}, {1, 2}, {{1}, {2, 3}})}}); + // doc 4 (second half pos 8), then doc 7. + WriteRun(r1.path, {{.id = 0, .tp = MakeTerm({4, 7}, {1, 1}, {{8}, {9}})}}); + std::vector merged; + ASSERT_TRUE( + CollectMerge({r0.path, r1.path}, vocab, rank, /*has_positions=*/true, &merged).ok()); + ASSERT_EQ(merged.size(), 1U); + EXPECT_EQ(merged[0].docids, (std::vector {0, 4, 7})); // one entry per docid + EXPECT_EQ(merged[0].freqs, (std::vector {1, 3, 1})); // doc4: 2 + 1 = 3 + EXPECT_EQ(merged[0].positions_flat, (std::vector {1, 2, 3, 8, 9})); // doc4: 2,3,8 +} + +// FM-05: a wide term split across runs fills writer-owned source windows with +// byte-identical docids, freqs, and positions. +TEST(SniiSpillMergeTest, MergeWideTermStreamsMatchesMaterialized) { + const std::vector vocab = {"wide"}; + const std::vector rank = LexRank(vocab); + TempRun r0, r1; + { + TermPostings a; + for (uint32_t d = 0; d <= 450; ++d) { // docs 0..450, freq 2 each + a.docids.push_back(d); + a.freqs.push_back(2); + a.positions_flat.push_back(d * 5); + a.positions_flat.push_back(d * 5 + 1); + } + WriteRun(r0.path, {{.id = 0, .tp = a}}); + } + { + TermPostings b; + b.docids.push_back(450); // boundary doc continues from r0 + b.freqs.push_back(1); + b.positions_flat.push_back(450 * 5 + 2); + for (uint32_t d = 451; d <= 900; ++d) { + b.docids.push_back(d); + b.freqs.push_back(2); + b.positions_flat.push_back(d * 5); + b.positions_flat.push_back(d * 5 + 1); + } + WriteRun(r1.path, {{.id = 0, .tp = b}}); + } + const std::vector paths = {r0.path, r1.path}; + TermPostings materialized, streamed; + ASSERT_TRUE(MergeRuns(paths, vocab, rank, /*has_positions=*/true, [&](TermPostings&& tp) { + materialized = std::move(tp); + }).ok()); + ASSERT_TRUE(merge_run_sources(paths, vocab, rank, /*has_positions=*/true, + [&](doris::snii::writer::StreamedTermPostings&& source) { + return DrainStreamed(std::move(source), &streamed); + }) + .ok()); + EXPECT_GE(materialized.docids.size(), + static_cast(doris::snii::format::kSlimDfThreshold)); + EXPECT_EQ(materialized.docids, streamed.docids); + EXPECT_EQ(materialized.freqs, streamed.freqs); + EXPECT_EQ(materialized.positions_flat, streamed.positions_flat); + // Boundary doc 450 coalesced: freq 2 (r0) + 1 (r1) = 3. + const auto it = std::ranges::find(materialized.docids, 450U); + ASSERT_NE(it, materialized.docids.end()); + EXPECT_EQ(materialized.freqs[static_cast(it - materialized.docids.begin())], 3U); +} + +// FM-06: a single run passes through unchanged; an empty run and an empty run-set +// both emit nothing and return OK (degenerate inputs). +TEST(SniiSpillMergeTest, MergeSingleRunAndEmptyInputs) { + const std::vector vocab = {"a", "b"}; + const std::vector rank = LexRank(vocab); + TempRun r0; + WriteRun(r0.path, + {{.id = 0, .tp = MakeTerm({1, 2}, {1, 1})}, {.id = 1, .tp = MakeTerm({3}, {1})}}); + std::vector merged; + ASSERT_TRUE(CollectMerge({r0.path}, vocab, rank, /*has_positions=*/false, &merged).ok()); + ASSERT_EQ(merged.size(), 2U); + EXPECT_EQ(merged[0].term, "a"); + EXPECT_EQ(merged[0].docids, (std::vector {1, 2})); + EXPECT_EQ(merged[1].term, "b"); + EXPECT_EQ(merged[1].docids, (std::vector {3})); + + // Empty run (no terms) -> fn never invoked. + TempRun empty; + WriteRun(empty.path, {}); + int calls = 0; + ASSERT_TRUE(MergeRuns({empty.path}, vocab, rank, /*has_positions=*/false, [&](TermPostings&&) { + ++calls; + }).ok()); + EXPECT_EQ(calls, 0); + + // No run paths at all -> also OK, zero calls. + calls = 0; + ASSERT_TRUE(MergeRuns({}, vocab, rank, /*has_positions=*/false, [&](TermPostings&&) { + ++calls; + }).ok()); + EXPECT_EQ(calls, 0); +} + +// FM-07: a run term-id >= vocab.size() is rejected as Corruption -- the +// current_id() < vocab.size() guards remain, so string_rank[term_id] is never +// indexed out of range. +TEST(SniiSpillMergeTest, MergeOutOfRangeTermIdIsCorruption) { + const std::vector vocab = {"only"}; // valid id 0 + const std::vector rank = LexRank(vocab); // size 1 + TempRun run; + WriteRun(run.path, {{.id = 0, .tp = MakeTerm({0}, {1})}, + {.id = 3, .tp = MakeTerm({9}, {1})}}); // id 3 out of range + std::vector merged; + const Status s = MergeRuns({run.path}, vocab, rank, /*has_positions=*/false, + [&](TermPostings&& tp) { merged.push_back(std::move(tp)); }); + EXPECT_TRUE(s.is()) << s; +} + +// FM-08: string_rank sized differently from vocab is an InternalError, rejected at the +// entry guard before any run is opened or any term emitted (the T15 guard). +TEST(SniiSpillMergeTest, MergeRankVocabSizeMismatchIsInternal) { + const std::vector vocab = {"a", "b", "c"}; + const std::vector rank = {0, 1}; // size 2 != vocab size 3 + TempRun run; + WriteRun(run.path, {{.id = 0, .tp = MakeTerm({0}, {1})}}); + int calls = 0; + const Status s = MergeRuns({run.path}, vocab, rank, /*has_positions=*/false, + [&](TermPostings&&) { ++calls; }); + EXPECT_TRUE(s.is()) << s; + EXPECT_EQ(calls, 0); // rejected before emitting anything +} + +// FM-09 (equivalence baseline): a richer scenario -- multiple terms across multiple +// runs, positions, and a boundary-doc overlap -- compared field-by-field against the +// hand-computed expected merged stream. With the lexicographic rank this pins the +// byte-identical output (== the old vocab-string-keyed semantics). +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiSpillMergeTest, MergeProducesByteIdenticalOutput) { + const std::vector vocab = {"delta", "alpha", "charlie"}; // ids 0,1,2 + const std::vector rank = LexRank(vocab); // alpha(1) merged; + ASSERT_TRUE( + CollectMerge({r0.path, r1.path, r2.path}, vocab, rank, /*has_positions=*/true, &merged) + .ok()); + ASSERT_EQ(merged.size(), 3U); + // alpha (id1): r0 docs{0,3} ++ r1 doc{3} -> doc3 coalesces (freq 2 + 1 = 3). + EXPECT_EQ(merged[0].term, "alpha"); + EXPECT_EQ(merged[0].docids, (std::vector {0, 3})); + EXPECT_EQ(merged[0].freqs, (std::vector {1, 3})); + EXPECT_EQ(merged[0].positions_flat, + (std::vector {2, 0, 5, 8})); // doc0{2} doc3{0,5,8} + // charlie (id2): r1 only. + EXPECT_EQ(merged[1].term, "charlie"); + EXPECT_EQ(merged[1].docids, (std::vector {2})); + EXPECT_EQ(merged[1].freqs, (std::vector {2})); + EXPECT_EQ(merged[1].positions_flat, (std::vector {1, 4})); + // delta (id0): r0 doc{1} ++ r2 docs{6,7}. + EXPECT_EQ(merged[2].term, "delta"); + EXPECT_EQ(merged[2].docids, (std::vector {1, 6, 7})); + EXPECT_EQ(merged[2].freqs, (std::vector {1, 1, 1})); + EXPECT_EQ(merged[2].positions_flat, (std::vector {9, 0, 0})); +} + +// FM-10 (end-to-end): a borrowed-vocab SpimiTermBuffer fed the SAME tokens produces +// byte-identical merged postings whether it stays in memory (threshold 0) or spills +// to many runs (tiny threshold) and goes through the rank-keyed k-way merge. This +// drives the production wiring (SpimiTermBuffer::merge_runs -> ensure_string_rank -> +// MergeRuns(string_rank_)) and proves spill == in-memory under the new integer key. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiSpillMergeTest, SpillMergeEqualsInMemory) { + using doris::snii::writer::SpimiTermBuffer; + // 6-id vocab in a NON-lexicographic id order, so the derived rank permutation is + // non-trivial (id order != string order). + const std::vector vocab = {"m", "g", "t", "a", "p", "c"}; + auto feed = [&](SpimiTermBuffer& buf) { + // Globally ascending docids; per term ascending docids; per (term,doc) 1..3 + // consecutive tokens (freq) with ascending positions. A sparse mask leaves some + // (term,doc) cells empty so terms get varied df and the merge must interleave. + for (uint32_t d = 0; d < 9; ++d) { + for (uint32_t id = 0; id < static_cast(vocab.size()); ++id) { + if (((d * 5 + id * 3) % 4) == 1) { + continue; // sparse: skip some (term,doc) + } + const uint32_t freq = 1 + ((d + id) % 3); // 1..3 tokens in this doc + for (uint32_t k = 0; k < freq; ++k) { + buf.add_token(id, d, /*pos=*/d * 50 + id * 7 + k); + } + } + } + }; + + std::vector in_memory; + { + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill_threshold_bytes=*/0); + feed(buf); + ASSERT_TRUE( + buf.for_each_term_sorted([&](doris::snii::writer::StreamedTermPostings&& source) { + TermPostings postings; + RETURN_IF_ERROR(DrainStreamed(std::move(source), &postings)); + in_memory.push_back(std::move(postings)); + return Status::OK(); + }).ok()); + EXPECT_EQ(buf.run_count_for_test(), 0U); // pure in-memory: no spill + } + + std::vector spilled; + size_t runs = 0; + { + // Tiny threshold: the first 32 KiB arena block immediately exceeds it, so a + // spill fires repeatedly -> many small runs (each id lands in several runs and a + // multi-token doc straddles run seams -> exercises boundary-doc coalesce). + SpimiTermBuffer buf(&vocab, /*has_positions=*/true, /*spill_threshold_bytes=*/1); + feed(buf); + ASSERT_TRUE( + buf.for_each_term_sorted([&](doris::snii::writer::StreamedTermPostings&& source) { + TermPostings postings; + RETURN_IF_ERROR(DrainStreamed(std::move(source), &postings)); + spilled.push_back(std::move(postings)); + return Status::OK(); + }).ok()); + runs = buf.run_count_for_test(); + } + EXPECT_GT(runs, 1U); // the spill path actually fired multiple runs + + ASSERT_EQ(in_memory.size(), spilled.size()); + for (size_t i = 0; i < in_memory.size(); ++i) { + EXPECT_EQ(in_memory[i].term, spilled[i].term) << "term index " << i; + EXPECT_EQ(in_memory[i].docids, spilled[i].docids) << "docids of " << in_memory[i].term; + EXPECT_EQ(in_memory[i].freqs, spilled[i].freqs) << "freqs of " << in_memory[i].term; + EXPECT_EQ(in_memory[i].positions_flat, spilled[i].positions_flat) + << "positions of " << in_memory[i].term; + } +} diff --git a/be/test/storage/index/snii/writer/spillable_byte_buffer_test.cpp b/be/test/storage/index/snii/writer/spillable_byte_buffer_test.cpp new file mode 100644 index 00000000000000..d2be29d640d97f --- /dev/null +++ b/be/test/storage/index/snii/writer/spillable_byte_buffer_test.cpp @@ -0,0 +1,128 @@ +// 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 "storage/index/snii/writer/spillable_byte_buffer.h" + +#include + +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_writer.h" + +using doris::snii::writer::SpillableByteBuffer; +using doris::snii::Slice; +using doris::Status; + +namespace { + +// In-RAM sink: collects everything appended so a test can compare stream_into's +// output against the exact bytes that were fed in. +class CollectWriter : public doris::snii::io::FileWriter { +public: + Status append(Slice s) override { + bytes_.insert(bytes_.end(), s.data(), s.data() + s.size()); + written_ += s.size(); + return Status::OK(); + } + Status finalize() override { return Status::OK(); } + uint64_t bytes_written() const override { return written_; } + const std::vector& bytes() const { return bytes_; } + +private: + std::vector bytes_; + uint64_t written_ = 0; +}; + +std::vector Pattern(size_t n, uint8_t seed) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = static_cast((i * 31 + seed) & 0xFF); + } + return v; +} + +// Feeds `blocks` chunks of `block` bytes through a buffer with the given cap, then +// asserts: size() == total, spilled() matches expectation, and stream_into() +// reproduces the exact concatenation (RAM-resident or read back from the temp). +void RoundTrip(uint64_t cap, size_t block, int blocks, bool expect_spill) { + SpillableByteBuffer buf(cap, "test"); + std::vector expect; + for (int b = 0; b < blocks; ++b) { + const auto chunk = Pattern(block, static_cast(b)); + ASSERT_TRUE(buf.append(Slice(chunk)).ok()); + expect.insert(expect.end(), chunk.begin(), chunk.end()); + } + EXPECT_EQ(buf.size(), expect.size()); + EXPECT_EQ(buf.spilled(), expect_spill); + ASSERT_TRUE(buf.seal().ok()); + CollectWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + EXPECT_EQ(out.bytes(), expect); +} + +// append_move adopts the caller's vector; verify identical bytes for RAM + spill. +void RoundTripMove(uint64_t cap, size_t block, int blocks, bool expect_spill) { + SpillableByteBuffer buf(cap, "test"); + std::vector expect; + for (int b = 0; b < blocks; ++b) { + auto chunk = Pattern(block, static_cast(b + 7)); + expect.insert(expect.end(), chunk.begin(), chunk.end()); + ASSERT_TRUE(buf.append_move(std::move(chunk)).ok()); + } + EXPECT_EQ(buf.size(), expect.size()); + EXPECT_EQ(buf.spilled(), expect_spill); + ASSERT_TRUE(buf.seal().ok()); + CollectWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + EXPECT_EQ(out.bytes(), expect); +} + +} // namespace + +TEST(SniiSpillableByteBuffer, StaysInRamUnderCap) { + RoundTrip(/*cap=*/1U << 20, /*block=*/4096, /*blocks=*/4, /*expect_spill=*/false); +} + +TEST(SniiSpillableByteBuffer, SpillsOverCapAndRoundTripsByteForByte) { + // 8 x 4 KiB = 32 KiB through an 8 KiB cap -> spills after the 2nd block. + RoundTrip(/*cap=*/8192, /*block=*/4096, /*blocks=*/8, /*expect_spill=*/true); +} + +TEST(SniiSpillableByteBuffer, MaxCapNeverSpills) { + RoundTrip(/*cap=*/UINT64_MAX, /*block=*/65536, /*blocks=*/8, /*expect_spill=*/false); +} + +TEST(SniiSpillableByteBuffer, MoveAppendStaysInRam) { + RoundTripMove(/*cap=*/1U << 20, /*block=*/4096, /*blocks=*/4, /*expect_spill=*/false); +} + +TEST(SniiSpillableByteBuffer, MoveAppendSpillsAndRoundTrips) { + RoundTripMove(/*cap=*/8192, /*block=*/4096, /*blocks=*/8, /*expect_spill=*/true); +} + +TEST(SniiSpillableByteBuffer, EmptyBufferStreamsNothing) { + SpillableByteBuffer buf(1U << 20, "test"); + EXPECT_EQ(buf.size(), 0U); + EXPECT_FALSE(buf.spilled()); + ASSERT_TRUE(buf.seal().ok()); + CollectWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + EXPECT_TRUE(out.bytes().empty()); +} diff --git a/be/test/storage/index/snii/writer/spimi_global_memory_limiter_test.cpp b/be/test/storage/index/snii/writer/spimi_global_memory_limiter_test.cpp new file mode 100644 index 00000000000000..eb2c783d14bde2 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_global_memory_limiter_test.cpp @@ -0,0 +1,542 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/compact_posting_pool.h" +#include "storage/index/snii/writer/global_memory_limiter.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +// G09: process-wide build-RAM limiter. The per-writer gate-2 cap bounds ONE +// SPIMI accumulator; a concurrent load keeps (tablets x concurrency) writers +// alive at once, none of which may ever reach its own cap (wikipedia at +// concurrency 16: 100+ writers x 300-500 MB, ~41 GiB, zero per-writer spills). +// These tests pin the limiter's contract: +// (1) REGISTRY: register / absolute-report / unregister maintain the exact +// total and entry count; reports for unregistered flags are ignored; +// (2) SELECTION: over budget flags the largest-ARENA eligible buffers +// (arena >= the victim floor -- only the arena is reclaimable by a +// forced spill) until the flagged arena covers the overage -- counting +// already-pending flags -- and under budget (or budget 0 = off) never +// flags anything; +// (3) HONOR: the owner's next add_token observes a pending request, spills +// (run_count increments; global_forced_spills seam bumps) BYPASSING the +// G08 anti-churn floor but respecting the FORCED-SPILL FLOOR, and +// clears the flag; requests are advisory; +// (4) LIFETIME: attach registers the current resident bytes, the debounced +// report path keeps the registry equal to resident_bytes(), and the +// destructor un-registers; +// (5) THREADS: concurrent register / report / unregister (with flags dying +// right after unregister) is race-free -- the TSAN canary; +// (6) ANTI-STORM (the conc=16 wikipedia field failure): the floor makes a +// below-floor request a pending NO-OP, the cooldown exempts a +// just-spilled buffer until its arena regrows past the floor, and the +// budget-sanity check degrades to log-once-and-stop when the per-writer +// budget share cannot be met by spilling persistent-dominated buffers. +using doris::snii::writer::CompactPostingPool; +using doris::snii::writer::GlobalMemoryLimiter; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::StreamedTermPostings; +using doris::Status; +using doris::snii::writer::TermPostings; + +namespace snii_testing = doris::snii::writer::testing; + +namespace { + +constexpr int64_t kMiB = 1LL << 20; + +// Distinct short ordinary terms ("uaa", "uab", ...). +std::string unigram(uint32_t i) { + std::string s = "u"; + s += static_cast('a' + i % 26); + s += static_cast('a' + (i / 26) % 26); + s += static_cast('a' + (i / 676) % 26); + return s; +} + +// ---- (1) registry bookkeeping --------------------------------------------- + +TEST(SniiGlobalMemoryLimiter, RegistryAddRemoveTotal) { + GlobalMemoryLimiter lim; // budget defaults to 0 (off): pure tracking + std::atomic a {false}; + std::atomic b {false}; + EXPECT_EQ(lim.total_bytes(), 0); + EXPECT_EQ(lim.registered_count(), 0U); + + lim.register_buffer(&a, 100, 40); + lim.register_buffer(&b, 60, 10); + EXPECT_EQ(lim.total_bytes(), 160); + EXPECT_EQ(lim.registered_count(), 2U); + + lim.report(&a, 150, 90); // ABSOLUTE totals, not deltas + EXPECT_EQ(lim.total_bytes(), 210); + + lim.register_buffer(&a, 40, 5); // re-register updates in place, no duplicate + EXPECT_EQ(lim.total_bytes(), 100); + EXPECT_EQ(lim.registered_count(), 2U); + + lim.unregister_buffer(&a); + EXPECT_EQ(lim.total_bytes(), 60); + EXPECT_EQ(lim.registered_count(), 1U); + + lim.unregister_buffer(&a); // double-unregister is harmless + lim.report(&a, 999, 999); // a report for an unregistered flag is ignored + EXPECT_EQ(lim.total_bytes(), 60); + EXPECT_EQ(lim.registered_count(), 1U); + + lim.unregister_buffer(&b); + EXPECT_EQ(lim.total_bytes(), 0); + EXPECT_EQ(lim.registered_count(), 0U); + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); +} + +// ---- (2) victim selection --------------------------------------------------- + +// MiB-scale entries with arena == resident keep the pre-floor selection shape +// visible while every entry clears the default 64 MiB victim floor AND every +// budget keeps budget/count above the 96 MiB/writer sanity minimum. +TEST(SniiGlobalMemoryLimiter, OverBudgetFlagsLargestUntilOverageCovered) { + GlobalMemoryLimiter lim; + lim.set_budget_bytes(1000 * kMiB); + std::atomic a {false}; // the largest + std::atomic b {false}; + std::atomic c {false}; + lim.register_buffer(&a, 900 * kMiB, 900 * kMiB); + lim.register_buffer(&b, 500 * kMiB, 500 * kMiB); + lim.register_buffer(&c, 100 * kMiB, 100 * kMiB); // total 1500 MiB > 1000 MiB + // Registration alone never flags: reacting belongs to the report path. + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); + EXPECT_FALSE(c.load()); + + // Overage 500 MiB: the largest (a, 900 MiB of arena) alone covers it; b + // and c spared. + lim.report(&c, 100 * kMiB, 100 * kMiB); + EXPECT_TRUE(a.load()); + EXPECT_FALSE(b.load()); + EXPECT_FALSE(c.load()); + + // a's request is still pending -- its arena counts toward coverage, so a + // re-report while still over budget must NOT widen the victim set. + lim.report(&c, 100 * kMiB, 100 * kMiB); + EXPECT_FALSE(b.load()); + EXPECT_FALSE(c.load()); + + // Deeper overage (budget 300 MiB -> overage 1200 MiB): a (900) is no + // longer enough, b joins (900 + 500 >= 1200), c is still spared. + lim.set_budget_bytes(300 * kMiB); + lim.report(&c, 100 * kMiB, 100 * kMiB); + EXPECT_TRUE(a.load()); + EXPECT_TRUE(b.load()); + EXPECT_FALSE(c.load()); +} + +TEST(SniiGlobalMemoryLimiter, UnderBudgetOrDisabledNeverFlags) { + GlobalMemoryLimiter lim; + lim.set_budget_bytes(1000 * kMiB); + std::atomic a {false}; + std::atomic b {false}; + lim.register_buffer(&a, 600 * kMiB, 600 * kMiB); + lim.register_buffer(&b, 300 * kMiB, 300 * kMiB); + lim.report(&a, 700 * kMiB, 700 * kMiB); // total == budget: at (not over) budget + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); + + lim.set_budget_bytes(0); // 0 = limiter off, no matter how large the total + lim.report(&a, 100000 * kMiB, 100000 * kMiB); + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); +} + +// The field storm's root selection bug: victims were ranked by RESIDENT total, +// which is dominated by PERSISTENT (non-spillable) vocabulary and slot structures. +// Victims must be ranked by their reclaimable ARENA, and a buffer below the +// victim floor is not a victim at all. +TEST(SniiGlobalMemoryLimiter, VictimsSelectedByArenaNotResident) { + GlobalMemoryLimiter lim; + lim.set_budget_bytes(600 * kMiB); // 2 writers -> 300 MiB/writer: sane budget + std::atomic persistent_heavy {false}; + std::atomic arena_heavy {false}; + // Larger RESIDENT, tiny arena (8 MiB < the 64 MiB default floor). + lim.register_buffer(&persistent_heavy, 1000 * kMiB, 8 * kMiB); + // Smaller resident, large reclaimable arena. + lim.register_buffer(&arena_heavy, 500 * kMiB, 300 * kMiB); + + lim.report(&persistent_heavy, 1000 * kMiB, 8 * kMiB); // total 1500 > 600 + EXPECT_TRUE(arena_heavy.load()) << "the reclaimable-arena holder is the victim"; + EXPECT_FALSE(persistent_heavy.load()) + << "a below-floor arena must never be flagged, however large its resident total"; +} + +// PER-BUFFER COOLDOWN: after a victim's forced spill its arena is ~0; further +// over-budget reports must NOT re-flag it until the arena regrows past the +// floor. (No timer: eligibility IS the cooldown.) +TEST(SniiGlobalMemoryLimiter, CooldownSkipsJustSpilledBufferUntilArenaRegrows) { + GlobalMemoryLimiter lim; + lim.set_budget_bytes(200 * kMiB); // 1 writer -> 200 MiB/writer: sane budget + std::atomic flag {false}; + lim.register_buffer(&flag, 400 * kMiB, 200 * kMiB); + + lim.report(&flag, 400 * kMiB, 200 * kMiB); // over budget, arena >= floor + EXPECT_TRUE(flag.load()); + flag.store(false); // the owner honors: spill, clear the flag... + // ...and its next report shows the arena reclaimed but the PERSISTENT + // remainder still over budget (the field shape). + lim.report(&flag, 210 * kMiB, 0); + EXPECT_FALSE(flag.load()) << "cooldown: a spilled (empty-arena) buffer is exempt"; + lim.report(&flag, 240 * kMiB, 32 * kMiB); // regrown, still below the floor + EXPECT_FALSE(flag.load()) << "still exempt below the floor"; + lim.report(&flag, 280 * kMiB, 72 * kMiB); // regrown PAST the 64 MiB floor + EXPECT_TRUE(flag.load()) << "eligible again once a floor's worth is reclaimable"; +} + +// BUDGET SANITY: when budget / registered_count < the per-writer useful +// minimum, spilling cannot meet the budget (persistent structures dominate) -- +// the limiter must degrade (log once, flag nothing) and recover when the +// ratio does. +TEST(SniiGlobalMemoryLimiter, BudgetDegradationStopsFlaggingAndRecovers) { + GlobalMemoryLimiter lim; + // 3 writers on a 240 MiB budget: 80 MiB/writer < the 96 MiB minimum. + lim.set_budget_bytes(240 * kMiB); + std::atomic a {false}; + std::atomic b {false}; + std::atomic c {false}; + lim.register_buffer(&a, 200 * kMiB, 150 * kMiB); + lim.register_buffer(&b, 200 * kMiB, 150 * kMiB); + lim.register_buffer(&c, 200 * kMiB, 150 * kMiB); + EXPECT_FALSE(lim.budget_degraded()); + + lim.report(&a, 200 * kMiB, 150 * kMiB); // total 600 > 240: over budget... + EXPECT_TRUE(lim.budget_degraded()) << "80 MiB/writer < 96 MiB: degrade"; + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); + EXPECT_FALSE(c.load()); + lim.report(&b, 200 * kMiB, 150 * kMiB); // sustained episode: still no flags + EXPECT_FALSE(a.load()); + EXPECT_FALSE(b.load()); + EXPECT_FALSE(c.load()); + + // One writer drains: 240 MiB / 2 = 120 MiB/writer >= 96 MiB -- recovered. + lim.unregister_buffer(&c); + lim.report(&a, 200 * kMiB, 150 * kMiB); // total 400 > 240, sane ratio + EXPECT_FALSE(lim.budget_degraded()); + EXPECT_TRUE(a.load() || b.load()) << "flagging must resume after recovery"; +} + +// ---- (3) the owner honors a pending request --------------------------------- + +TEST(SniiSpimiGlobalSpill, OwnerHonorsRequestAtNextTokenAndClears) { + snii_testing::reset_global_forced_spills(); + // Unlimited local threshold: the per-writer gate can never fire, so any + // spill below is attributable to the global request alone. The G08 + // anti-churn floor (arena >= cap/4) is bypassed by construction here -- + // the arena holds a single 32 KiB block, far below any production cap/4. + // The forced-spill floor is dropped to its one-block minimum: THIS test + // pins the honor mechanics, not the floor (see the floor tests below). + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + buf.set_forced_spill_min_arena_bytes(1); + for (uint32_t d = 0; d < 200; ++d) { + buf.add_token("hot", /*docid=*/d, /*pos=*/0); + } + ASSERT_TRUE(buf.status().ok()); + ASSERT_EQ(buf.run_count_for_test(), 0U); + + buf.request_global_spill_for_test(); // what the limiter does cross-thread + EXPECT_TRUE(buf.global_spill_requested_for_test()); + EXPECT_EQ(snii_testing::global_forced_spills(), 0U); + + buf.add_token("hot", /*docid=*/200, /*pos=*/0); // next token observes it + EXPECT_EQ(buf.run_count_for_test(), 1U) << "forced spill must cut a run"; + EXPECT_FALSE(buf.global_spill_requested_for_test()) << "honor must clear"; + EXPECT_EQ(snii_testing::global_forced_spills(), 1U); + + // Once cleared, later tokens spill nothing further (advisory, one-shot). + buf.add_token("hot", /*docid=*/201, /*pos=*/0); + EXPECT_EQ(buf.run_count_for_test(), 1U); + + // The forced spill changes WHEN a run was cut, never WHAT is emitted. + std::vector out = buf.finalize_sorted(); + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_EQ(out.size(), 1U); + EXPECT_EQ(out[0].term, "hot"); + ASSERT_EQ(out[0].docids.size(), 202U); + EXPECT_EQ(out[0].docids.front(), 0U); + EXPECT_EQ(out[0].docids.back(), 201U); +} + +TEST(SniiSpimiGlobalSpill, RequestOnEmptyArenaIsPendingUntilARunIsWritable) { + snii_testing::reset_global_forced_spills(); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + buf.set_forced_spill_min_arena_bytes(1); // one-block minimum still applies + // A request on an empty buffer has nothing to write: purely advisory, no + // run, flag stays pending (a drained owner would simply never observe it). + buf.request_global_spill_for_test(); + EXPECT_EQ(buf.run_count_for_test(), 0U); + EXPECT_TRUE(buf.global_spill_requested_for_test()); + // The first token claims the first 32 KiB arena block -- the "arena >= one + // block" floor is now met, so the pending request is honored right there. + buf.add_token("t", /*docid=*/0, /*pos=*/0); + EXPECT_EQ(buf.run_count_for_test(), 1U); + EXPECT_FALSE(buf.global_spill_requested_for_test()); + EXPECT_EQ(snii_testing::global_forced_spills(), 1U); + ASSERT_TRUE(buf.status().ok()); +} + +// ---- (6) forced-spill floor: below-floor requests are pending no-ops -------- + +// The field storm's honor-side bug: a request was honored with a single 32 KiB +// arena block, cutting a near-empty run. With the (default 64 MiB) floor, a +// request planted while the arena is small must spill NOTHING -- not per +// token, not once -- and stay pending. +TEST(SniiSpimiGlobalSpill, RequestBelowForcedSpillFloorIsNotHonored) { + snii_testing::reset_global_forced_spills(); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + // Production-default floor (64 MiB) -- far above this feed's arena. + ASSERT_EQ(buf.forced_spill_min_arena_bytes(), + SpimiTermBuffer::kDefaultForcedSpillMinArenaBytes); + buf.request_global_spill_for_test(); + for (uint32_t d = 0; d < 5000; ++d) { // a few arena blocks, << 64 MiB + buf.add_token(unigram(d % 300), /*docid=*/d, /*pos=*/0); + } + ASSERT_TRUE(buf.status().ok()); + EXPECT_EQ(buf.run_count_for_test(), 0U) << "below-floor request must be a no-op"; + EXPECT_EQ(snii_testing::global_forced_spills(), 0U); + EXPECT_TRUE(buf.global_spill_requested_for_test()) << "...that stays pending"; +} + +// Deferred honor: with a small floor (3 arena blocks), a pending request is a +// no-op until the arena regrows past the floor, then fires exactly once. +TEST(SniiSpimiGlobalSpill, RequestHonoredOnceArenaRegrowsPastFloor) { + snii_testing::reset_global_forced_spills(); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + const uint64_t kFloor = 3ULL * CompactPostingPool::kBlockSize; + buf.set_forced_spill_min_arena_bytes(kFloor); + buf.request_global_spill_for_test(); + uint32_t docid = 0; + // One block is not enough: feed until the arena holds one block and check + // the request is still pending, then keep feeding until the floor is met. + while (buf.status().ok() && buf.run_count_for_test() == 0 && docid < 200000) { + buf.add_token("hot", docid, /*pos=*/docid % 7); + ++docid; + } + ASSERT_TRUE(buf.status().ok()); + ASSERT_EQ(buf.run_count_for_test(), 1U) << "must fire once the floor is met"; + EXPECT_FALSE(buf.global_spill_requested_for_test()); + EXPECT_EQ(snii_testing::global_forced_spills(), 1U); + // The run was cut with at least a floor's worth of arena accumulated: the + // feed needed more than 2 blocks' worth of tokens (each token appends + // >= 2 bytes, so 2 blocks < 64 KiB of payload < the token count here). + EXPECT_GT(docid, (2ULL * CompactPostingPool::kBlockSize) / 4) + << "honor must not fire before the floor's worth of arena existed"; +} + +// ---- (6) the storm scenario end-to-end -------------------------------------- + +// The conc=16 wikipedia field failure in miniature, with PRODUCTION-DEFAULT +// anti-storm settings: a budget far below the writers' (persistent-dominated) +// resident sum and far below (96 MiB x writers). The old limiter flagged every +// buffer on every report and each honored with one 32 KiB block -- a storm of +// tiny runs. Now: the budget-sanity check degrades (log once, stop flagging), +// the victim floor exempts every small-arena buffer anyway, so ZERO forced +// spills and ZERO runs result. +TEST(SniiSpimiGlobalSpill, StormScenarioTinyBudgetSmallArenasProducesZeroForcedSpills) { + snii_testing::reset_global_forced_spills(); + GlobalMemoryLimiter lim; // production defaults: 64 MiB floor, 96 MiB/writer sanity + lim.set_budget_bytes(256 * 1024); + + constexpr size_t kBuffers = 6; + std::vector> buffers; + buffers.reserve(kBuffers); + for (size_t i = 0; i < kBuffers; ++i) { + auto buf = std::make_unique(/*has_positions=*/true, + /*spill_threshold_bytes=*/0); + buf->attach_global_limiter(&lim); + buffers.push_back(std::move(buf)); + } + // Interleaved feed: every buffer's resident grows well past its budget + // share (the registry total exceeds 1 MiB almost immediately), reports + // fire constantly, yet nothing may be flagged or spilled. + for (uint32_t round = 0; round < 400; ++round) { + for (auto& buf : buffers) { + for (uint32_t k = 0; k < 8; ++k) { + buf->add_token(unigram((round * 8 + k) % 2000), /*docid=*/round * 8 + k, + /*pos=*/0); + } + } + } + ASSERT_GT(lim.total_bytes(), lim.budget_bytes()) << "the scenario must be over budget"; + EXPECT_TRUE(lim.budget_degraded()) << "256 KiB / 6 writers is an unmeetable budget"; + EXPECT_EQ(snii_testing::global_forced_spills(), 0U) << "no forced-spill storm"; + for (auto& buf : buffers) { + ASSERT_TRUE(buf->status().ok()); + EXPECT_EQ(buf->run_count_for_test(), 0U) << "no runs were cut"; + EXPECT_FALSE(buf->global_spill_requested_for_test()) << "no flags were planted"; + } +} + +// ---- (4) attach / report / detach lifetime ---------------------------------- + +TEST(SniiSpimiGlobalSpill, AttachRegistersReportsTrackResidentAndDtorUnregisters) { + GlobalMemoryLimiter lim; // budget 0: tracking only, no flags + { + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + buf.attach_global_limiter(&lim); + EXPECT_EQ(lim.registered_count(), 1U); + EXPECT_EQ(lim.total_bytes(), static_cast(buf.resident_bytes_for_test())); + + for (uint32_t i = 0; i < 300; ++i) { + buf.add_token(unigram(i), /*docid=*/i, /*pos=*/0); + } + ASSERT_TRUE(buf.status().ok()); + // The debounced report path forwards ABSOLUTE totals: at rest the + // registry equals the buffer's real resident bytes exactly. + EXPECT_EQ(lim.total_bytes(), static_cast(buf.resident_bytes_for_test())); + + buf.attach_global_limiter(&lim); // at-most-once: ignored + EXPECT_EQ(lim.registered_count(), 1U); + } + // Destruction un-registers: nothing leaks into the process-wide total. + EXPECT_EQ(lim.registered_count(), 0U); + EXPECT_EQ(lim.total_bytes(), 0); +} + +// End-to-end: two attached buffers, one small and one that grows past the +// budget. The limiter must flag the larger-ARENA grower once its arena clears +// the victim floor (the small buffer's single arena block never does); the +// grower's own next token honors the request (its local threshold is +// unlimited, the G08 floor bypassed); the small buffer is never flagged and +// never spills. Degradation is disabled: the KiB-scale budget is intentional. +TEST(SniiSpimiGlobalSpill, LimiterFlagsLargestOwnerSpillsSmallBufferSpared) { + snii_testing::reset_global_forced_spills(); + GlobalMemoryLimiter lim; // declared BEFORE the buffers: outlives them + lim.set_budget_bytes(256 * 1024); + const int64_t kFloor = 2LL * CompactPostingPool::kBlockSize; // 64 KiB + lim.set_min_victim_arena_bytes(kFloor); + lim.set_min_useful_budget_per_writer_bytes(0); // KiB-scale test budget + + SpimiTermBuffer small(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + small.set_forced_spill_min_arena_bytes(static_cast(kFloor)); + small.attach_global_limiter(&lim); + SpimiTermBuffer big(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + big.set_forced_spill_min_arena_bytes(static_cast(kFloor)); + big.attach_global_limiter(&lim); + + for (uint32_t i = 0; i < 50; ++i) { // ~1 arena block: forever below the floor + small.add_token(unigram(i), /*docid=*/i, /*pos=*/0); + } + ASSERT_TRUE(small.status().ok()); + + // Distinct terms grow big's resident (vocab + intern + slots + arena) past + // the budget and its ARENA past the floor; the over-budget report then + // flags big (the largest eligible arena), and big's own add path honors + // the request. Bounded feed with an early exit. The bound stays below + // unigram()'s 17576 distinct strings so every fed term is DISTINCT (the + // drain-count assertion relies on that); three tokens per term grow the + // arena ~3x faster than the vocab, so the floor is met well within it. + uint32_t fed = 0; + uint32_t docid = 0; + for (uint32_t k = 0; k < 16000 && big.run_count_for_test() == 0; ++k, ++fed) { + big.add_token(unigram(k), docid++, /*pos=*/0); + big.add_token(unigram(k), docid++, /*pos=*/1); + big.add_token(unigram(k), docid++, /*pos=*/2); + } + ASSERT_TRUE(big.status().ok()); + ASSERT_GT(big.run_count_for_test(), 0U) << "grower must be forced to spill"; + EXPECT_FALSE(big.global_spill_requested_for_test()) << "honor must clear"; + EXPECT_GE(snii_testing::global_forced_spills(), 1U); + EXPECT_EQ(small.run_count_for_test(), 0U) << "small buffer must be spared"; + EXPECT_FALSE(small.global_spill_requested_for_test()); + + // The forced spill released the grower's arena and reported the drop: the + // registry total is back to the exact resident sum of both buffers. + EXPECT_EQ(lim.total_bytes(), static_cast(small.resident_bytes_for_test()) + + static_cast(big.resident_bytes_for_test())); + + // Both buffers still drain cleanly (the small one in memory, the big one + // through its forced run + k-way merge). + size_t small_terms = 0; + ASSERT_TRUE(small.for_each_term_sorted([&small_terms](StreamedTermPostings&& source) { + RETURN_IF_ERROR( + doris::snii::writer::consume_streamed_term(std::move(source))); + ++small_terms; + return Status::OK(); + }).ok()); + EXPECT_EQ(small_terms, 50U); + size_t big_terms = 0; + ASSERT_TRUE(big.for_each_term_sorted([&big_terms](StreamedTermPostings&& source) { + RETURN_IF_ERROR( + doris::snii::writer::consume_streamed_term(std::move(source))); + ++big_terms; + return Status::OK(); + }).ok()); + EXPECT_EQ(big_terms, fed); +} + +// ---- (5) thread-safety canary ------------------------------------------------ + +// Concurrent register / report / unregister with a tiny budget so cross-thread +// flagging constantly targets flags that die right after their unregister -- +// the exact lifetime the mutex must protect. Run under TSAN this is the G09 +// race canary; under ASAN it still catches any touch-after-unregister. The +// floor and the sanity minimum are dropped so the byte-scale entries keep +// producing flags (the point is flag traffic, not selection policy). +TEST(SniiGlobalMemoryLimiter, ConcurrentRegisterReportUnregisterIsClean) { + GlobalMemoryLimiter lim; + lim.set_budget_bytes(1); // everything is over budget: flags fly + lim.set_min_victim_arena_bytes(1); + lim.set_min_useful_budget_per_writer_bytes(0); + constexpr int kThreads = 8; + constexpr int kIters = 400; + std::vector threads; + threads.reserve(kThreads); + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&lim, t] { + for (int i = 0; i < kIters; ++i) { + std::atomic flag {false}; + lim.register_buffer(&flag, 1000 + t, 1000 + t); + for (int r = 0; r < 4; ++r) { + lim.report(&flag, 1000 + t + r, 1000 + t + r); + if (flag.load(std::memory_order_relaxed)) { + // The owner honoring: observe and clear on its thread. + flag.store(false, std::memory_order_relaxed); + } + } + lim.unregister_buffer(&flag); + // `flag` is destroyed here -- after unregister_buffer returned, + // the limiter must never touch it again. + } + }); + } + for (auto& th : threads) { + th.join(); + } + EXPECT_EQ(lim.total_bytes(), 0); + EXPECT_EQ(lim.registered_count(), 0U); +} + +} // namespace diff --git a/be/test/storage/index/snii/writer/spimi_locality_bench_test.cpp b/be/test/storage/index/snii/writer/spimi_locality_bench_test.cpp new file mode 100644 index 00000000000000..4ef4a4ff411a52 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_locality_bench_test.cpp @@ -0,0 +1,280 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +namespace doris::snii::writer { +namespace { + +uint64_t env_u64(const char* name, uint64_t default_value) { + const char* value = std::getenv(name); + if (value == nullptr || *value == '\0') { + return default_value; + } + return static_cast(std::strtoull(value, nullptr, 10)); +} + +constexpr uint64_t kFnvOffset = 1469598103934665603ULL; +constexpr uint64_t kFnvPrime = 1099511628211ULL; + +uint64_t fnv_bytes(uint64_t hash, const void* data, size_t size) { + const auto* bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + hash ^= bytes[i]; + hash *= kFnvPrime; + } + return hash; +} + +uint64_t fnv_u64(uint64_t hash, uint64_t value) { + return fnv_bytes(hash, &value, sizeof(value)); +} + +struct BenchCorpus { + std::vector vocab; + std::vector tokens; + uint32_t tokens_per_doc = 200; +}; + +const BenchCorpus& corpus() { + static const BenchCorpus corpus = [] { + BenchCorpus result; + const uint64_t token_count = env_u64("SNII_BENCH_TOKENS", 300000); + constexpr uint32_t kVocabSize = 262144; + std::mt19937_64 rng(0x5a11d5eedULL); + + result.vocab.reserve(kVocabSize); + std::uniform_int_distribution extra_len(0, 6); + std::uniform_int_distribution letter(0, 24); + for (uint32_t id = 0; id < kVocabSize; ++id) { + std::string term; + uint32_t value = id; + do { + term.push_back(static_cast('b' + (value % 25))); + value /= 25; + } while (value != 0); + term.push_back('a'); + for (int i = 0, n = extra_len(rng); i < n; ++i) { + term.push_back(static_cast('b' + letter(rng))); + } + result.vocab.push_back(std::move(term)); + } + + std::vector cdf(kVocabSize); + double sum = 0.0; + for (uint32_t rank = 1; rank <= kVocabSize; ++rank) { + sum += 1.0 / std::pow(static_cast(rank), 1.07); + cdf[rank - 1] = sum; + } + std::uniform_real_distribution sample(0.0, sum); + result.tokens.reserve(token_count); + for (uint64_t i = 0; i < token_count; ++i) { + const auto it = std::lower_bound(cdf.begin(), cdf.end(), sample(rng)); + result.tokens.push_back(static_cast(it - cdf.begin())); + } + return result; + }(); + return corpus; +} + +struct FeedResult { + uint64_t total_ns = 0; + uint64_t unique_terms = 0; + uint64_t total_tokens = 0; + uint64_t digest = 0; +}; + +FeedResult run_feed(bool drain_digest) { + const BenchCorpus& input = corpus(); + FeedResult result; + MemoryReporter reporter(nullptr, /*cap_bytes=*/0); + SpimiTermBuffer buffer(/*has_positions=*/true, /*spill_threshold_bytes=*/0, &reporter); + + const auto begin = std::chrono::steady_clock::now(); + uint32_t docid = 0; + for (size_t token = 0; token < input.tokens.size();) { + const size_t doc_end = + std::min(input.tokens.size(), token + static_cast(input.tokens_per_doc)); + uint32_t position = 0; + for (; token < doc_end; ++token, ++position) { + buffer.add_token(input.vocab[input.tokens[token]], docid, position); + } + ++docid; + } + result.total_ns = static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - begin) + .count()); + result.unique_terms = buffer.unique_terms(); + result.total_tokens = buffer.total_tokens(); + EXPECT_TRUE(buffer.status().ok()) << buffer.status().to_string(); + + if (drain_digest) { + uint64_t hash = kFnvOffset; + const auto status = buffer.for_each_term_sorted([&](StreamedTermPostings&& source) { + TermPostings postings; + RETURN_IF_ERROR(materialize_streamed_term(std::move(source), &postings)); + hash = fnv_bytes(hash, postings.term.data(), postings.term.size()); + hash = fnv_u64(hash, postings.docids.size()); + hash = fnv_bytes(hash, postings.docids.data(), + postings.docids.size() * sizeof(uint32_t)); + hash = fnv_bytes(hash, postings.freqs.data(), postings.freqs.size() * sizeof(uint32_t)); + hash = fnv_bytes(hash, postings.positions_flat.data(), + postings.positions_flat.size() * sizeof(uint32_t)); + return Status::OK(); + }); + EXPECT_TRUE(status.ok()) << status.to_string(); + result.digest = hash; + } + return result; +} + +uint64_t median_ns(std::vector values) { + std::ranges::sort(values); + return values[values.size() / 2]; +} + +void print_result(const char* label, size_t repetition, const FeedResult& result) { + const double tokens = static_cast(result.total_tokens); + printf("[bench] %-16s rep=%zu total=%9.3f ms (%.1f ns/add)\n", label, repetition, + result.total_ns / 1e6, result.total_ns / tokens); +} + +} // namespace + +TEST(SniiSpimiLocalityBenchTest, FeedBaseline) { + const size_t repetitions = static_cast(env_u64("SNII_BENCH_REPS", 3)); + std::vector total; + uint64_t digest = 0; + uint64_t unique_terms = 0; + for (size_t repetition = 0; repetition < repetitions; ++repetition) { + const FeedResult result = run_feed(/*drain_digest=*/repetition == 0); + if (repetition == 0) { + digest = result.digest; + unique_terms = result.unique_terms; + } else { + ASSERT_EQ(unique_terms, result.unique_terms); + } + total.push_back(result.total_ns); + print_result("baseline", repetition, result); + } + printf("[bench] baseline median: total=%9.3f ms unique_terms=%llu digest=%016llx\n", + median_ns(total) / 1e6, static_cast(unique_terms), + static_cast(digest)); + ASSERT_NE(digest, 0U); +} + +TEST(SniiSpimiLocalityBenchTest, PrefetchToggleAB) { +#ifndef SNII_G11_PREFETCH + GTEST_SKIP() << "prefetch candidate not compiled in"; +#else + const size_t repetitions = static_cast(env_u64("SNII_BENCH_REPS", 3)); + std::vector disabled_total; + std::vector enabled_total; + uint64_t disabled_digest = 0; + uint64_t enabled_digest = 0; + for (size_t repetition = 0; repetition < repetitions; ++repetition) { + testing::set_bench_disable_g11_prefetch(true); + const FeedResult disabled = run_feed(/*drain_digest=*/repetition == 0); + testing::set_bench_disable_g11_prefetch(false); + const FeedResult enabled = run_feed(/*drain_digest=*/repetition == 0); + if (repetition == 0) { + disabled_digest = disabled.digest; + enabled_digest = enabled.digest; + } + disabled_total.push_back(disabled.total_ns); + enabled_total.push_back(enabled.total_ns); + print_result("prefetch-OFF", repetition, disabled); + print_result("prefetch-ON", repetition, enabled); + } + testing::set_bench_disable_g11_prefetch(false); + printf("[bench] prefetch A/B: OFF=%9.3f ms ON=%9.3f ms\n", median_ns(disabled_total) / 1e6, + median_ns(enabled_total) / 1e6); + ASSERT_EQ(disabled_digest, enabled_digest); +#endif +} + +TEST(SniiSpimiLocalityBenchTest, InternSetReserveUpperBound) { + const size_t repetitions = static_cast(env_u64("SNII_BENCH_REPS", 3)); + const BenchCorpus& input = corpus(); + + struct Hash { + using is_transparent = void; + const std::vector* vocab; + size_t operator()(std::string_view value) const noexcept { + return fnv_bytes(kFnvOffset, value.data(), value.size()); + } + size_t operator()(uint32_t id) const noexcept { + return operator()(std::string_view((*vocab)[id])); + } + }; + struct Equal { + using is_transparent = void; + const std::vector* vocab; + bool operator()(uint32_t lhs, uint32_t rhs) const noexcept { return lhs == rhs; } + bool operator()(std::string_view lhs, uint32_t rhs) const noexcept { + return lhs == std::string_view((*vocab)[rhs]); + } + }; + + std::unordered_set distinct_ids(input.tokens.begin(), input.tokens.end()); + const size_t distinct = distinct_ids.size(); + const auto run = [&](bool reserve) { + const auto begin = std::chrono::steady_clock::now(); + std::unordered_set intern(0, Hash {&input.vocab}, + Equal {&input.vocab}); + if (reserve) { + intern.reserve(distinct); + } + for (uint32_t id : input.tokens) { + const std::string& term = input.vocab[id]; + if (intern.find(std::string_view(term)) == intern.end()) { + intern.insert(id); + } + } + EXPECT_EQ(intern.size(), distinct); + return static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - begin) + .count()); + }; + + std::vector no_reserve; + std::vector with_reserve; + for (size_t repetition = 0; repetition < repetitions; ++repetition) { + no_reserve.push_back(run(false)); + with_reserve.push_back(run(true)); + } + printf("[bench] intern-set: distinct=%zu no-reserve=%9.3f ms reserve=%9.3f ms\n", distinct, + median_ns(no_reserve) / 1e6, median_ns(with_reserve) / 1e6); +} + +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/spimi_resident_accounting_test.cpp b/be/test/storage/index/snii/writer/spimi_resident_accounting_test.cpp new file mode 100644 index 00000000000000..eafe75dfb9dd13 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_resident_accounting_test.cpp @@ -0,0 +1,161 @@ +// 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 + +#include +#include +#include + +#include "storage/index/snii/writer/memory_reporter.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +using doris::snii::writer::MemoryReporter; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::StreamedTermPostings; +using doris::snii::writer::TermPostings; +using doris::Status; + +namespace { + +std::string ordinary_term(uint32_t id) { + return "ordinary-prefix-sharing-resident-term-" + std::to_string(id) + "-payload"; +} + +void expect_same_stream(const std::vector& got, const std::vector& want, + const char* label) { + ASSERT_EQ(got.size(), want.size()) << label; + for (size_t i = 0; i < got.size(); ++i) { + EXPECT_EQ(got[i].term, want[i].term) << label << " term order diverged at " << i; + EXPECT_EQ(got[i].docids, want[i].docids) << label << " " << got[i].term; + EXPECT_EQ(got[i].freqs, want[i].freqs) << label << " " << got[i].term; + EXPECT_EQ(got[i].positions_flat, want[i].positions_flat) << label << " " << got[i].term; + } +} + +TEST(SniiSpimiResidentAccounting, ResidentCoversOwnedVocabularyAndIsMonotone) { + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + + constexpr uint32_t kTerms = 4000; + uint64_t previous = buf.resident_bytes_for_test(); + uint64_t owned_string_payload_floor = 0; + for (uint32_t id = 0; id < kTerms; ++id) { + const std::string term = ordinary_term(id); + owned_string_payload_floor += term.size() + 1; + buf.add_token(term, /*docid=*/0, /*pos=*/id); + const uint64_t current = buf.resident_bytes_for_test(); + ASSERT_GE(current, previous) << "resident_bytes decreased at ordinary term " << id; + previous = current; + } + + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + EXPECT_EQ(buf.unique_terms(), kTerms); + EXPECT_GE(buf.resident_bytes_for_test(), owned_string_payload_floor); +} + +TEST(SniiSpimiResidentAccounting, TinyThresholdSpillsOnOwnedVocabularyGrowth) { + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/64 * 1024); + + uint32_t terms_fed = 0; + for (; terms_fed < 4000 && buf.run_count_for_test() == 0; ++terms_fed) { + buf.add_token(ordinary_term(terms_fed), /*docid=*/0, /*pos=*/terms_fed); + } + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_GT(buf.run_count_for_test(), 0U); + + size_t terms_seen = 0; + ASSERT_TRUE(buf.for_each_term_sorted([&terms_seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR( + doris::snii::writer::consume_streamed_term(std::move(source))); + ++terms_seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(terms_seen, terms_fed); + EXPECT_TRUE(buf.status().ok()) << buf.status().to_string(); +} + +TEST(SniiSpimiResidentAccounting, OwnedVocabularySpillEqualsNoSpillControl) { + SpimiTermBuffer spilled(/*has_positions=*/true, /*spill_threshold_bytes=*/64 * 1024); + SpimiTermBuffer control(/*has_positions=*/true, /*spill_threshold_bytes=*/0); + + constexpr uint32_t kTerms = 800; + for (SpimiTermBuffer* buf : {&spilled, &control}) { + for (uint32_t id = 0; id < kTerms; ++id) { + buf->add_token(ordinary_term(id), /*docid=*/1, /*pos=*/id); + } + for (uint32_t id = 0; id < kTerms; ++id) { + buf->add_token(ordinary_term(id), /*docid=*/2, /*pos=*/id + 1); + } + ASSERT_TRUE(buf->status().ok()) << buf->status().to_string(); + } + ASSERT_GE(spilled.run_count_for_test(), 1U); + ASSERT_EQ(control.run_count_for_test(), 0U); + + const std::vector got = spilled.finalize_sorted(); + const std::vector want = control.finalize_sorted(); + ASSERT_TRUE(spilled.status().ok()) << spilled.status().to_string(); + ASSERT_TRUE(control.status().ok()) << control.status().to_string(); + ASSERT_EQ(want.size(), kTerms); + expect_same_stream(got, want, "owned-vocabulary spill"); +} + +TEST(SniiSpimiResidentAccounting, OwnedModeReporterNetsToZeroAfterInMemoryDrain) { + MemoryReporter reporter; + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0, &reporter); + for (uint32_t id = 0; id < 600; ++id) { + buf.add_token(ordinary_term(id), /*docid=*/0, /*pos=*/id); + buf.add_token(ordinary_term(id), /*docid=*/1, /*pos=*/id + 1); + } + + EXPECT_GT(reporter.current_bytes(), 0); + size_t terms_seen = 0; + ASSERT_TRUE(buf.for_each_term_sorted([&terms_seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR( + doris::snii::writer::consume_streamed_term(std::move(source))); + ++terms_seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(terms_seen, 600U); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiSpimiResidentAccounting, OwnedModeReporterNetsToZeroAfterSpilledDrain) { + MemoryReporter reporter(/*consume_release=*/nullptr, /*cap_bytes=*/64 * 1024, + MemoryReporter::CapPolicy::kSpillThreshold); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/0, &reporter); + + uint32_t terms_fed = 0; + for (; terms_fed < 4000 && buf.run_count_for_test() < 2; ++terms_fed) { + buf.add_token(ordinary_term(terms_fed), /*docid=*/0, /*pos=*/terms_fed); + } + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_GE(buf.run_count_for_test(), 1U); + EXPECT_GT(reporter.current_bytes(), 0); + + size_t terms_seen = 0; + ASSERT_TRUE(buf.for_each_term_sorted([&terms_seen](StreamedTermPostings&& source) { + RETURN_IF_ERROR( + doris::snii::writer::consume_streamed_term(std::move(source))); + ++terms_seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(terms_seen, terms_fed); + EXPECT_EQ(reporter.current_bytes(), 0); +} + +} // namespace diff --git a/be/test/storage/index/snii/writer/spimi_run_cap_test.cpp b/be/test/storage/index/snii/writer/spimi_run_cap_test.cpp new file mode 100644 index 00000000000000..588d8adf10d728 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_run_cap_test.cpp @@ -0,0 +1,198 @@ +// 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 + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +// G09 run-file FD hygiene: a SPIMI buffer's spill runs are ALL (re)opened +// simultaneously -- one fd each, held for the whole k-way merge -- so an +// unbounded run count across ~100 concurrent writers exhausted the BE nofile +// rlimit in the conc=16 wikipedia field failure ('Too many open files' at run +// reopen). The cap (set_max_run_files / config +// snii_spill_max_run_files_per_buffer) merge-compacts the accumulated runs +// into ONE before a new spill would exceed it. These tests pin: +// * the cap is HONORED: the tracked run count never exceeds it, and the +// compaction seam records each collapse; +// * compaction is INVISIBLE in the output: a capped buffer drains the +// byte-identical term stream (terms, docids, freqs, positions) of an +// uncapped control fed the same tokens; +// * cap 0 disables compaction entirely (pre-cap behavior). +using doris::Status; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::StreamedTermPostings; +using doris::snii::writer::TermPostings; +using doris::snii::writer::materialize_streamed_term; + +namespace snii_testing = doris::snii::writer::testing; + +namespace { + +// Distinct short ordinary terms ("uaa", "uab", ...). +std::string unigram(uint32_t i) { + std::string s = "u"; + s += static_cast('a' + i % 26); + s += static_cast('a' + (i / 26) % 26); + s += static_cast('a' + (i / 676) % 26); + return s; +} + +// One feed step at docid k: the shared term "hot" gets TWO tokens in the SAME +// doc (freq 2 -- a spill lands between them in the per-token spill regime +// below, exercising boundary-doc coalescing across run seams and through a +// compaction) and unigram(k) gets one token (a distinct term per step). With +// the tests' tiny 1 KiB local threshold, resident (>= one fresh 32 KiB arena +// block) exceeds the cap on EVERY token and the arena floor (one block) is +// met as soon as a chain claims its block -- so every token cuts a run: +// N fed tokens == N run files, densely exercising the cap machinery with tiny +// step counts (an uncapped 30-step feed already holds 90 runs). +void feed_step(SpimiTermBuffer* buf, uint32_t k) { + buf->add_token("hot", /*docid=*/k, /*pos=*/0); + buf->add_token("hot", /*docid=*/k, /*pos=*/1); + buf->add_token(unigram(k), /*docid=*/k, /*pos=*/2); +} + +// Bounds the honored-cap search loop; far below unigram()'s 17576 distinct +// strings so every step's unigram is DISTINCT (the drained-term-count +// assertions rely on that). +constexpr uint32_t kMaxSteps = 16000; + +struct DrainedTerm { + std::vector docids; + std::vector freqs; + std::vector positions; +}; + +std::map drain(SpimiTermBuffer* buf) { + std::map out; + Status s = buf->for_each_term_sorted([&out](StreamedTermPostings&& source) { + TermPostings tp; + RETURN_IF_ERROR(materialize_streamed_term(std::move(source), &tp)); + DrainedTerm& t = out[tp.term]; + t.docids = tp.docids; + t.freqs = tp.freqs; + t.positions = tp.positions_flat; + return Status::OK(); + }); + EXPECT_TRUE(s.ok()) << s.to_string(); + return out; +} + +TEST(SniiSpimiRunCap, CapIsHonoredAndCompactionSeamFires) { + snii_testing::reset_run_compactions(); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/1024); + constexpr size_t kCap = 3; + buf.set_max_run_files(kCap); + + // Feed until the cap has forced a compaction (bounded); the cap must hold + // at EVERY step, so track the running maximum of the run count. + size_t max_runs_seen = 0; + uint32_t steps = 0; + while (steps < kMaxSteps && snii_testing::run_compactions() < 1) { + feed_step(&buf, steps); + ++steps; + max_runs_seen = std::max(max_runs_seen, buf.run_count_for_test()); + } + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_GE(snii_testing::run_compactions(), 1U) + << "the cap was never hit within " << steps << " steps"; + EXPECT_LE(max_runs_seen, kCap) << "run count exceeded the cap"; + + // The capped buffer still drains every term exactly once with the full + // posting content. + std::map got = drain(&buf); + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + ASSERT_EQ(got.size(), static_cast(steps) + 1U); // distinct unigrams + "hot" + const DrainedTerm& hot = got.at("hot"); + ASSERT_EQ(hot.docids.size(), static_cast(steps)); + EXPECT_EQ(hot.docids.front(), 0U); + EXPECT_EQ(hot.docids.back(), steps - 1); + for (uint32_t f : hot.freqs) { + ASSERT_EQ(f, 2U) << "boundary-doc coalescing must survive compaction"; + } + ASSERT_EQ(hot.positions.size(), 2ULL * steps); + EXPECT_EQ(hot.positions[0], 0U); + EXPECT_EQ(hot.positions[1], 1U); + const DrainedTerm& first = got.at(unigram(0)); + ASSERT_EQ(first.docids.size(), 1U); + EXPECT_EQ(first.docids[0], 0U); + ASSERT_EQ(first.positions.size(), 1U); + EXPECT_EQ(first.positions[0], 2U); +} + +TEST(SniiSpimiRunCap, CompactedDrainMatchesUncappedControl) { + snii_testing::reset_run_compactions(); + // 30 steps == 90 tokens == 90 runs uncapped (per-token spills, see + // feed_step): plenty of cap-2 compactions on the capped side while the + // control's 90-way merge stays trivial. + constexpr uint32_t kSteps = 30; + + SpimiTermBuffer capped(/*has_positions=*/true, /*spill_threshold_bytes=*/1024); + capped.set_max_run_files(2); // aggressive: compact on nearly every spill + for (uint32_t k = 0; k < kSteps; ++k) { + feed_step(&capped, k); + } + ASSERT_TRUE(capped.status().ok()) << capped.status().to_string(); + ASSERT_GE(snii_testing::run_compactions(), 1U); + + SpimiTermBuffer control(/*has_positions=*/true, /*spill_threshold_bytes=*/1024); + control.set_max_run_files(0); // uncapped: the pre-cap spill pipeline + for (uint32_t k = 0; k < kSteps; ++k) { + feed_step(&control, k); + } + ASSERT_TRUE(control.status().ok()) << control.status().to_string(); + ASSERT_GT(control.run_count_for_test(), 2U) << "the control must hold many runs"; + + std::map got = drain(&capped); + std::map want = drain(&control); + ASSERT_TRUE(capped.status().ok()) << capped.status().to_string(); + ASSERT_TRUE(control.status().ok()) << control.status().to_string(); + ASSERT_EQ(got.size(), want.size()); + for (const auto& [term, w] : want) { + auto it = got.find(term); + ASSERT_NE(it, got.end()) << "missing term: " << term; + EXPECT_EQ(it->second.docids, w.docids) << term; + EXPECT_EQ(it->second.freqs, w.freqs) << term; + EXPECT_EQ(it->second.positions, w.positions) << term; + } +} + +TEST(SniiSpimiRunCap, ZeroCapDisablesCompaction) { + snii_testing::reset_run_compactions(); + SpimiTermBuffer buf(/*has_positions=*/true, /*spill_threshold_bytes=*/1024); + buf.set_max_run_files(0); + constexpr uint32_t kSteps = 30; // 90 per-token spill runs, uncapped + for (uint32_t k = 0; k < kSteps; ++k) { + feed_step(&buf, k); + } + ASSERT_TRUE(buf.status().ok()) << buf.status().to_string(); + EXPECT_GT(buf.run_count_for_test(), 2U); + EXPECT_EQ(snii_testing::run_compactions(), 0U); + // Still drains cleanly through the plain multi-run merge. + std::map got = drain(&buf); + EXPECT_EQ(got.size(), static_cast(kSteps) + 1U); +} + +} // namespace diff --git a/be/test/storage/index/snii/writer/spimi_spill_writer_test.cpp b/be/test/storage/index/snii/writer/spimi_spill_writer_test.cpp new file mode 100644 index 00000000000000..9555962c7c39d7 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_spill_writer_test.cpp @@ -0,0 +1,324 @@ +// 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 +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::writer; +using doris::Status; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_spill_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +std::vector ReadAll(const std::string& path) { + io::LocalFileReader r; + EXPECT_TRUE(r.open(path).ok()); + std::vector out; + EXPECT_TRUE(r.read_at(0, r.size(), &out).ok()); + return out; +} + +// Deterministic (term, doc, pos) stream with globally ascending docids. Mixes +// high-df ("alpha", every doc), mid-df, multi-token docs, and a term whose docs +// straddle arbitrary spill boundaries -- so the spill path exercises a term in +// one run, a term in every run, and a spill boundary mid-term. +void Feed(SpimiTermBuffer* buf, uint32_t doc_count) { + for (uint32_t d = 0; d < doc_count; ++d) { + buf->add_token("alpha", d, 0); // every doc (spans all runs) + buf->add_token("alpha", d, 7); // freq 2 in every doc + if (d % 2 == 0) { + buf->add_token("beta", d, 1); + } + if (d % 3 == 0) { + buf->add_token("gamma", d, 2); + } + if (d % 5 == 0) { + buf->add_token("delta", d, 3); + buf->add_token("delta", d, 9); + } + if (d == 1) { + buf->add_token("singleton", d, 4); // only one doc, one run + } + } +} + +SniiIndexInput BaseInput(uint32_t doc_count) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = doc_count; + in.target_dict_block_bytes = 512; // force several DICT blocks + return in; +} + +std::vector WriteContainer(const SniiIndexInput& in) { + const std::string path = TempPath(); + io::LocalFileWriter writer; + EXPECT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + EXPECT_TRUE(compound.add_logical_index(in).ok()); + EXPECT_TRUE(compound.finish().ok()); + std::vector bytes = ReadAll(path); + std::remove(path.c_str()); + return bytes; +} + +// Builds a container by STREAMING a freshly-fed buffer at the given spill +// threshold (0 == unlimited), returning the container bytes. +std::vector BuildStreamed(uint32_t docs, size_t spill_bytes) { + SpimiTermBuffer buf(/*has_positions=*/true, spill_bytes); + Feed(&buf, docs); + SniiIndexInput in = BaseInput(docs); + in.term_source = &buf; + return WriteContainer(in); +} + +// Opens a container from bytes and runs a couple of queries for a sanity match. +struct OpenedIndex { + std::string path; + io::LocalFileReader local; + std::unique_ptr metered; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader index; + + ~OpenedIndex() { + if (!path.empty()) { + std::remove(path.c_str()); + } + } +}; + +void OpenFromBytes(const std::vector& bytes, OpenedIndex* out) { + out->path = TempPath(); + io::LocalFileWriter w; + ASSERT_TRUE(w.open(out->path).ok()); + ASSERT_TRUE(w.append(Slice(bytes)).ok()); + ASSERT_TRUE(w.finalize().ok()); + ASSERT_TRUE(out->local.open(out->path).ok()); + out->metered = std::make_unique(&out->local); + ASSERT_TRUE(reader::SniiSegmentReader::open(out->metered.get(), &out->segment).ok()); + ASSERT_TRUE(out->segment.open_index(1, "body", &out->index).ok()); +} + +std::vector TermDocs(OpenedIndex* idx, const std::string& term) { + std::vector docs; + const Status s = query::term_query(idx->index, term, &docs); + EXPECT_TRUE(s.ok()) << "term=" << term << " err=" << s.to_string(); + return docs; +} + +} // namespace + +// CORE GUARANTEE: building the SAME corpus with a tiny spill threshold (forcing +// many spills + a final k-way merge) yields a container that is BYTE-FOR-BYTE +// identical to the unlimited in-memory build. +TEST(SniiSpimiSpillWriter, SpilledMatchesUnlimitedBytes) { + constexpr uint32_t kDocs = 400; + const std::vector unlimited = BuildStreamed(kDocs, /*spill=*/0); + // ~4 KiB threshold forces dozens of spills across this corpus. + const std::vector spilled = BuildStreamed(kDocs, /*spill=*/4096); + + ASSERT_EQ(unlimited.size(), spilled.size()); + EXPECT_EQ(unlimited, spilled); +} + +// Regression: the wide term "alpha" (every doc, freq 2) can have a doc split across +// a spill seam, so its PRE-coalesce total_docs reaches >= kSlimDfThreshold (512) +// while its POST-coalesce df stays below it. The merge gate keyed on total_docs +// would then STREAM positions (leaving positions_flat empty) into the slim writer +// path, which reads positions_flat directly -> segfault. Sweeping docs around the +// 512 boundary x several spill thresholds exercises the band; byte-identity vs the +// unlimited build proves both no-crash and identical output. +TEST(SniiSpimiSpillWriter, WideTermDfNearThresholdAcrossSeamMatchesUnlimited) { + for (uint32_t docs = 505; docs <= 515; ++docs) { + const std::vector unlimited = BuildStreamed(docs, /*spill=*/0); + for (size_t spill : {size_t {1024}, size_t {2048}, size_t {4096}}) { + const std::vector spilled = BuildStreamed(docs, spill); + ASSERT_EQ(unlimited.size(), spilled.size()) << "docs=" << docs << " spill=" << spill; + EXPECT_EQ(unlimited, spilled) << "docs=" << docs << " spill=" << spill; + } + } +} + +// Identical bytes regardless of threshold size: a mid threshold (one or two +// spills) must also match the unlimited build exactly. +TEST(SniiSpimiSpillWriter, MidThresholdAlsoIdentical) { + constexpr uint32_t kDocs = 400; + const std::vector unlimited = BuildStreamed(kDocs, /*spill=*/0); + const std::vector mid = BuildStreamed(kDocs, /*spill=*/32 * 1024); + EXPECT_EQ(unlimited, mid); +} + +// An extremely small threshold (spill almost every token) still produces the +// identical container -- stresses many single-term runs and the merge. +TEST(SniiSpimiSpillWriter, ExtremeSpillIdentical) { + constexpr uint32_t kDocs = 200; + const std::vector unlimited = BuildStreamed(kDocs, /*spill=*/0); + const std::vector tiny = BuildStreamed(kDocs, /*spill=*/1); + EXPECT_EQ(unlimited, tiny); +} + +// No-positions config with spilling still matches the in-memory build. +TEST(SniiSpimiSpillWriter, NoPositionsSpillIdentical) { + constexpr uint32_t kDocs = 300; + auto build = [&](size_t spill) { + SpimiTermBuffer buf(/*has_positions=*/false, spill); + for (uint32_t d = 0; d < kDocs; ++d) { + buf.add_token("alpha", d, 0); + if (d % 2 == 0) { + buf.add_token("beta", d, 0); + } + if (d % 4 == 0) { + buf.add_token("gamma", d, 0); + } + } + SniiIndexInput in = BaseInput(kDocs); + in.config = IndexConfig::kDocsOnly; // docids only (no positions section) + in.term_source = &buf; + return WriteContainer(in); + }; + EXPECT_EQ(build(0), build(2048)); +} + +// Queries against the spilled and unlimited containers return identical results. +// Uses >= kSlimDfThreshold docs so the high-df term is windowed (populating the +// .prx POD) -- the segment reader infers has_positions from a non-empty prx POD. +TEST(SniiSpimiSpillWriter, QueriesMatchAcrossSpill) { + constexpr uint32_t kDocs = 700; + OpenedIndex un, sp; + OpenFromBytes(BuildStreamed(kDocs, 0), &un); + OpenFromBytes(BuildStreamed(kDocs, 4096), &sp); + + for (const char* term : {"alpha", "beta", "gamma", "delta", "singleton", "absent"}) { + EXPECT_EQ(TermDocs(&un, term), TermDocs(&sp, term)) << "term=" << term; + } + std::vector un_phrase, sp_phrase; + ASSERT_TRUE(query::phrase_query(un.index, {"alpha", "alpha"}, &un_phrase).ok()); + ASSERT_TRUE(query::phrase_query(sp.index, {"alpha", "alpha"}, &sp_phrase).ok()); + EXPECT_EQ(un_phrase, sp_phrase); +} + +// WIDE-TERM STREAMED MERGE byte-identity: with enough docs that the high-df term +// "alpha" exceeds kSlimDfThreshold, the spilled k-way merge takes the WIDE-term +// direct run-source path instead of materializing +// the term's full positions_flat. The produced container MUST stay byte-for-byte +// identical to the unlimited in-memory build (which materializes positions). This +// is the direct regression guard for the merge-phase peak-RSS streaming change. +TEST(SniiSpimiSpillWriter, WideTermStreamedMergeMatchesUnlimitedBytes) { + // 1200 docs -> alpha df=1200 (>= kSlimDfThreshold 512) -> windowed + streamed. + constexpr uint32_t kDocs = 1200; + static_assert(kDocs >= doris::snii::format::kSlimDfThreshold, "alpha must be a wide term"); + const std::vector unlimited = BuildStreamed(kDocs, /*spill=*/0); + // Small threshold forces many spills, so alpha lands in EVERY run and the merge + // coalesces it across runs via the streamed pump. + const std::vector spilled = BuildStreamed(kDocs, /*spill=*/8192); + ASSERT_EQ(unlimited.size(), spilled.size()); + EXPECT_EQ(unlimited, spilled); +} + +// Single-doc corpus with spilling: an edge corpus must not crash or diverge. +TEST(SniiSpimiSpillWriter, SingleDocCorpus) { + SpimiTermBuffer un(/*has_positions=*/true, 0); + SpimiTermBuffer sp(/*has_positions=*/true, 1); + for (auto* b : {&un, &sp}) { + b->add_token("only", 0, 0); + b->add_token("only", 0, 1); + b->add_token("word", 0, 2); + } + SniiIndexInput un_in = BaseInput(1); + un_in.term_source = &un; + SniiIndexInput sp_in = BaseInput(1); + sp_in.term_source = &sp; + EXPECT_EQ(WriteContainer(un_in), WriteContainer(sp_in)); +} + +// finalize_sorted (the materialized accessor) also reflects spilled runs and is +// byte-identical to the in-memory result at the postings level. +TEST(SniiSpimiSpillWriter, FinalizeSortedMatchesAcrossSpill) { + constexpr uint32_t kDocs = 150; + SpimiTermBuffer un(/*has_positions=*/true, 0); + SpimiTermBuffer sp(/*has_positions=*/true, 256); + Feed(&un, kDocs); + Feed(&sp, kDocs); + const std::vector a = un.finalize_sorted(); + const std::vector b = sp.finalize_sorted(); + ASSERT_EQ(a.size(), b.size()); + for (size_t i = 0; i < a.size(); ++i) { + EXPECT_EQ(a[i].term, b[i].term); + EXPECT_EQ(a[i].docids, b[i].docids); + EXPECT_EQ(a[i].freqs, b[i].freqs); + EXPECT_EQ(a[i].positions_flat, b[i].positions_flat); + } + EXPECT_TRUE(un.status().ok()); + EXPECT_TRUE(sp.status().ok()); +} + +// Gate-2 trigger now compares REAL resident bytes (pool_.arena_bytes() + +// slot_of_.capacity()*4) against the configured cap, NOT the old per-token +// estimate. With a small cap, the first 32 KiB arena block immediately exceeds it, +// so at least one spill fires; the unlimited (cap=0) build never spills. +TEST(SniiSpimiSpillWriter, ArenaByteCapTriggersSpill) { + constexpr uint32_t kDocs = 400; + + SpimiTermBuffer capped(/*has_positions=*/true, /*spill=*/4096); + Feed(&capped, kDocs); + SniiIndexInput capped_in = BaseInput(kDocs); + capped_in.term_source = &capped; + const std::vector capped_bytes = WriteContainer(capped_in); + EXPECT_TRUE(capped.status().ok()); + // A spill ran: real resident (>= one 32 KiB block) crossed the 4 KiB cap. + EXPECT_GE(capped.run_count_for_test(), 1U); + + SpimiTermBuffer unlimited(/*has_positions=*/true, /*spill=*/0); + Feed(&unlimited, kDocs); + SniiIndexInput unlimited_in = BaseInput(kDocs); + unlimited_in.term_source = &unlimited; + const std::vector unlimited_bytes = WriteContainer(unlimited_in); + EXPECT_TRUE(unlimited.status().ok()); + // Unlimited never spills (corpus fits well under the arena hard-stop). + EXPECT_EQ(unlimited.run_count_for_test(), 0U); + + // And the metric switch did not change the output: byte-for-byte identical. + ASSERT_EQ(capped_bytes.size(), unlimited_bytes.size()); + EXPECT_EQ(capped_bytes, unlimited_bytes); +} diff --git a/be/test/storage/index/snii/writer/spimi_streaming_writer_test.cpp b/be/test/storage/index/snii/writer/spimi_streaming_writer_test.cpp new file mode 100644 index 00000000000000..42567006249bf8 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_streaming_writer_test.cpp @@ -0,0 +1,249 @@ +// 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 +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/io/local_file.h" +#include "storage/index/snii/writer/logical_index_writer.h" +#include "storage/index/snii/writer/posting_window_emitter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +using namespace doris::snii; +using namespace doris::snii::format; +using namespace doris::snii::writer; + +namespace { + +std::string TempPath() { + static int counter = 0; + return "/tmp/snii_stream_test_" + std::to_string(getpid()) + "_" + std::to_string(counter++) + + ".idx"; +} + +std::vector ReadAll(const std::string& path) { + io::LocalFileReader r; + EXPECT_TRUE(r.open(path).ok()); + std::vector out; + EXPECT_TRUE(r.read_at(0, r.size(), &out).ok()); + return out; +} + +// Feeds a deterministic (term, doc, pos) stream into a SPIMI buffer. Docids +// arrive in ascending order per term (the normal tokenizer contract); some +// terms span many docs so both slim and (with enough docs) windowed paths and +// the DICT block splitter are exercised. +void Feed(SpimiTermBuffer* buf, uint32_t doc_count) { + for (uint32_t d = 0; d < doc_count; ++d) { + buf->add_token("alpha", d, 0); // every doc: high df + if (d % 2 == 0) { + buf->add_token("beta", d, 1); // half the docs + } + if (d % 7 == 0) { + buf->add_token("gamma", d, 2); + buf->add_token("gamma", d, 5); // freq 2 in this doc + } + if (d == 3 || d == 4) { + buf->add_token("delta", d, d); // tiny df + } + } +} + +// Writes a single-index container from a SniiIndexInput and returns the bytes. +std::vector WriteContainer(const SniiIndexInput& in) { + const std::string path = TempPath(); + io::LocalFileWriter writer; + EXPECT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + EXPECT_TRUE(compound.add_logical_index(in).ok()); + EXPECT_TRUE(compound.finish().ok()); + std::vector bytes = ReadAll(path); + std::remove(path.c_str()); + return bytes; +} + +std::vector WriteStreamedContainer(SniiIndexInput in, TermPostings postings) { + const std::string path = TempPath(); + io::LocalFileWriter writer; + EXPECT_TRUE(writer.open(path).ok()); + SniiCompoundWriter compound(&writer); + SniiStreamedIndexSession* session = nullptr; + EXPECT_TRUE(compound.begin_streamed_index(std::move(in), &session).ok()); + SpanTermPostingSource source(postings.docids, postings.freqs, postings.positions_flat); + EXPECT_TRUE(session->push_term(StreamedTermPostings { + .term = std::move(postings.term), + .retain_positions = postings.retain_positions, + .source = &source, + }) + .ok()); + EXPECT_TRUE(session->finish().ok()); + EXPECT_TRUE(compound.finish().ok()); + std::vector bytes = ReadAll(path); + std::remove(path.c_str()); + return bytes; +} + +SniiIndexInput BaseInput(uint32_t doc_count) { + SniiIndexInput in; + in.index_id = 1; + in.index_suffix = "body"; + in.config = IndexConfig::kDocsPositions; + in.doc_count = doc_count; + in.target_dict_block_bytes = 512; // force several DICT blocks + return in; +} + +} // namespace + +// The streaming term_source path must produce a BYTE-IDENTICAL container to the +// materialized terms vector path: the flat-array accumulator + stream-finalize +// must not change a single output byte. +TEST(SniiSpimiStreamingWriter, StreamingMatchesMaterializedBytes) { + constexpr uint32_t kDocs = 300; + + SpimiTermBuffer mat_buf(/*has_positions=*/true); + Feed(&mat_buf, kDocs); + SniiIndexInput mat_in = BaseInput(kDocs); + mat_in.terms = mat_buf.finalize_sorted(); + const std::vector mat_bytes = WriteContainer(mat_in); + + SpimiTermBuffer stream_buf(/*has_positions=*/true); + Feed(&stream_buf, kDocs); + SniiIndexInput stream_in = BaseInput(kDocs); + stream_in.term_source = &stream_buf; + const std::vector stream_bytes = WriteContainer(stream_in); + + ASSERT_EQ(mat_bytes.size(), stream_bytes.size()); + EXPECT_EQ(mat_bytes, stream_bytes); +} + +// A high-df term is consumed through bounded posting windows instead of first +// materializing a complete positions_flat vector. The streamed build must remain +// byte-identical to the explicit materialized test path, including every PRX byte. +TEST(SniiSpimiStreamingWriter, StreamedPositionsMatchMaterializedBytesHighDf) { + // The wide term spans many posting windows. The small term exercises the slim + // encoding shape in the same container. + constexpr uint32_t kDocs = 80000; + auto feed = [](SpimiTermBuffer* buf) { + for (uint32_t d = 0; d < kDocs; ++d) { + buf->add_token("hot", d, 0); + if (d % 1000 == 0) { + buf->add_token("cold", d, 1); + } + } + }; + + SpimiTermBuffer mat_buf(/*has_positions=*/true); + feed(&mat_buf); + SniiIndexInput mat_in = BaseInput(kDocs); + mat_in.terms = mat_buf.finalize_sorted(); // materialized: positions_flat + const std::vector mat_bytes = WriteContainer(mat_in); + + SpimiTermBuffer stream_buf(/*has_positions=*/true); + feed(&stream_buf); + SniiIndexInput stream_in = BaseInput(kDocs); + stream_in.term_source = &stream_buf; + const std::vector stream_bytes = WriteContainer(stream_in); + + ASSERT_EQ(mat_bytes.size(), stream_bytes.size()); + EXPECT_EQ(mat_bytes, stream_bytes); +} + +// A low-df term with many positions stays on the same bounded source contract while +// selecting the slim encoding shape. Byte identity covers that cross-product. +TEST(SniiSpimiStreamingWriter, StreamedLowDfHighNtokMatchesMaterialized) { + constexpr uint32_t kDocs = 200; + constexpr uint32_t kReps = 400; + auto feed = [](SpimiTermBuffer* buf) { + for (uint32_t d = 0; d < kDocs; ++d) { + for (uint32_t p = 0; p < kReps; ++p) { + buf->add_token("rep", d, p); + } + } + }; + + SpimiTermBuffer mat_buf(/*has_positions=*/true); + feed(&mat_buf); + SniiIndexInput mat_in = BaseInput(kDocs); + mat_in.terms = mat_buf.finalize_sorted(); + const std::vector mat_bytes = WriteContainer(mat_in); + + SpimiTermBuffer stream_buf(/*has_positions=*/true); + feed(&stream_buf); + SniiIndexInput stream_in = BaseInput(kDocs); + stream_in.term_source = &stream_buf; + const std::vector stream_bytes = WriteContainer(stream_in); + + ASSERT_EQ(mat_bytes.size(), stream_bytes.size()); + EXPECT_EQ(mat_bytes, stream_bytes); +} + +// The streaming path drains its source: after build the buffer is empty. +TEST(SniiSpimiStreamingWriter, StreamingConsumesSource) { + SpimiTermBuffer buf(/*has_positions=*/true); + Feed(&buf, 50); + EXPECT_GT(buf.unique_terms(), 0U); + + SniiIndexInput in = BaseInput(50); + in.term_source = &buf; + LogicalIndexWriter writer(in); + // build() streams the posting region straight into a FileWriter sink; this test + // only asserts the source is drained, so a throwaway temp sink suffices. + const std::string post_path = TempPath(); + io::LocalFileWriter post; + ASSERT_TRUE(post.open(post_path).ok()); + ASSERT_TRUE(writer.build(&post).ok()); + EXPECT_EQ(buf.unique_terms(), 0U); + std::remove(post_path.c_str()); +} + +// The import/SPIMI source path and the compaction-style pushed session must +// converge on the one window emitter after their producer-specific plumbing. +TEST(SniiSpimiStreamingWriter, SpimiAndStreamedSessionUseSameWindowEmitter) { + constexpr uint32_t kDocs = 1024; + + doris::snii::writer::testing::reset_window_emitter_counters(); + SpimiTermBuffer spimi(/*has_positions=*/true); + for (uint32_t doc = 0; doc < kDocs; ++doc) { + spimi.add_token("hot", doc, 0); + } + SniiIndexInput spimi_input = BaseInput(kDocs); + spimi_input.term_source = &spimi; + const std::vector spimi_bytes = WriteContainer(spimi_input); + EXPECT_EQ(doris::snii::writer::testing::window_emitter_finished_terms(), 1U); + + doris::snii::writer::testing::reset_window_emitter_counters(); + TermPostings postings; + postings.term = "hot"; + postings.docids.resize(kDocs); + std::iota(postings.docids.begin(), postings.docids.end(), 0); + postings.freqs.assign(kDocs, 1); + postings.positions_flat.assign(kDocs, 0); + const std::vector streamed_bytes = + WriteStreamedContainer(BaseInput(kDocs), std::move(postings)); + EXPECT_EQ(doris::snii::writer::testing::window_emitter_finished_terms(), 1U); + EXPECT_EQ(streamed_bytes, spimi_bytes); +} diff --git a/be/test/storage/index/snii/writer/spimi_term_buffer_test.cpp b/be/test/storage/index/snii/writer/spimi_term_buffer_test.cpp new file mode 100644 index 00000000000000..d845f5495578e1 --- /dev/null +++ b/be/test/storage/index/snii/writer/spimi_term_buffer_test.cpp @@ -0,0 +1,860 @@ +// 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 "storage/index/snii/writer/spimi_term_buffer.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/term_posting_test_utils.h" + +using doris::snii::writer::MemoryReporter; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::StreamedTermPostings; +using doris::snii::writer::TermPostingBuffer; +using doris::snii::writer::TermPostings; +using doris::Status; +namespace ErrorCode = doris::ErrorCode; + +namespace { + +void expect_docids_for_term(const std::vector& postings, std::string_view term, + const std::vector& expected_docids) { + const auto found = std::ranges::find(postings, term, &TermPostings::term); + ASSERT_NE(found, postings.end()); + EXPECT_EQ(found->docids, expected_docids); +} + +} // namespace + +// Tokens accumulate into sorted terms with ascending docids and per-doc positions. +TEST(SniiSpimiTermBuffer, AccumulateAndSort) { + SpimiTermBuffer buf(/*has_positions=*/true); + // doc 0: "banana apple apple" + buf.add_token("banana", 0, 0); + buf.add_token("apple", 0, 1); + buf.add_token("apple", 0, 2); + // doc 1: "apple cherry" + buf.add_token("apple", 1, 0); + buf.add_token("cherry", 1, 1); + + EXPECT_EQ(buf.unique_terms(), 3U); + EXPECT_EQ(buf.total_tokens(), 5U); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 3U); + // Sorted lexicographically: apple, banana, cherry. + EXPECT_EQ(terms[0].term, "apple"); + EXPECT_EQ(terms[1].term, "banana"); + EXPECT_EQ(terms[2].term, "cherry"); + + // apple: docs 0 (freq 2, pos {1,2}) and 1 (freq 1, pos {0}). + const TermPostings& apple = terms[0]; + ASSERT_EQ(apple.docids.size(), 2U); + EXPECT_EQ(apple.docids[0], 0U); + EXPECT_EQ(apple.freqs[0], 2U); + ASSERT_EQ(apple.doc_positions(0).size(), 2U); + EXPECT_EQ(apple.doc_positions(0)[0], 1U); + EXPECT_EQ(apple.doc_positions(0)[1], 2U); + EXPECT_EQ(apple.docids[1], 1U); + EXPECT_EQ(apple.freqs[1], 1U); +} + +// Without positions, docs-only postings are SETS: a repeated occurrence inside +// the same document is discarded at accumulate time (it would only cost arena / +// spill bytes -- nothing consumes a docs-only tf), so per-doc freq stays 1. +TEST(SniiSpimiTermBuffer, DocsOnlyNoPositions) { + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token("x", 0, 0); + buf.add_token("x", 0, 1); + buf.add_token("x", 2, 0); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, "x"); + ASSERT_EQ(terms[0].docids.size(), 2U); + EXPECT_EQ(terms[0].docids[0], 0U); + EXPECT_EQ(terms[0].freqs[0], 1U); + EXPECT_EQ(terms[0].docids[1], 2U); + EXPECT_EQ(terms[0].freqs[1], 1U); + EXPECT_TRUE(terms[0].positions_flat.empty()); +} + +TEST(SniiSpimiTermBuffer, Empty) { + SpimiTermBuffer buf(true); + EXPECT_EQ(buf.unique_terms(), 0U); + EXPECT_TRUE(buf.finalize_sorted().empty()); +} + +TEST(SniiSpimiTermBufferTest, OwnedTermInternerMaterializesDistinctTermsOnce) { + namespace snii_testing = doris::snii::writer::testing; + + MemoryReporter reporter; + { + SpimiTermBuffer empty_buffer(/*has_positions=*/false, /*spill_threshold_bytes=*/0, + &reporter); + EXPECT_EQ(empty_buffer.resident_bytes_for_test(), 0U); + EXPECT_EQ(reporter.current_bytes(), 0); + } + EXPECT_EQ(reporter.current_bytes(), 0); + + snii_testing::reset_vocab_string_materialization_count(); + SpimiTermBuffer buffer(/*has_positions=*/false, /*spill_threshold_bytes=*/1); + buffer.add_token("alpha", 0, 0); + ASSERT_EQ(buffer.run_count_for_test(), 1U); + buffer.add_token("beta", 1, 0); + buffer.add_token("alpha", 2, 0); + buffer.add_token("long-materialized-term", 3, 0); + buffer.add_token("beta", 4, 0); + EXPECT_EQ(snii_testing::vocab_string_materialization_count(), 3U); + + std::vector postings = buffer.finalize_sorted(); + ASSERT_EQ(postings.size(), 3U); + expect_docids_for_term(postings, "alpha", {0U, 2U}); + expect_docids_for_term(postings, "beta", {1U, 4U}); + expect_docids_for_term(postings, "long-materialized-term", {3U}); +} + +TEST(SniiSpimiTermBufferTest, OwnedTermInternerUsesTermIdSlotsAndPreservesHashCollisions) { + namespace snii_testing = doris::snii::writer::testing; + + EXPECT_EQ(SpimiTermBuffer::owned_term_key_size_for_test(), 4U); + snii_testing::reset_vocab_string_materialization_count(); + snii_testing::reset_owned_term_full_byte_comparison_count(); + SpimiTermBuffer buffer(/*has_positions=*/false); + buffer.set_owned_term_hash_mask_for_test(0); + buffer.add_token("", 0, 0); + buffer.add_token("short", 1, 0); + buffer.add_token("abcdefgh-left", 2, 0); + buffer.add_token("abcdefgh-rght", 3, 0); + buffer.add_token("", 4, 0); + buffer.add_token("short", 5, 0); + buffer.add_token("abcdefgh-left", 6, 0); + buffer.add_token("abcdefgh-rght", 7, 0); + EXPECT_GT(snii_testing::owned_term_full_byte_comparison_count(), 0U); + EXPECT_EQ(snii_testing::vocab_string_materialization_count(), 4U); + + const std::vector postings = buffer.finalize_sorted(); + ASSERT_EQ(postings.size(), 4U); + expect_docids_for_term(postings, "", {0U, 4U}); + expect_docids_for_term(postings, "short", {1U, 5U}); + expect_docids_for_term(postings, "abcdefgh-left", {2U, 6U}); + expect_docids_for_term(postings, "abcdefgh-rght", {3U, 7U}); +} + +TEST(SniiSpimiTermBufferTest, OwnedTermInternerFailureLeavesTableReusable) { + namespace snii_testing = doris::snii::writer::testing; + + snii_testing::reset_vocab_string_materialization_count(); + MemoryReporter reporter; + SpimiTermBuffer buffer(/*has_positions=*/false, /*spill_threshold_bytes=*/0, &reporter); + + snii_testing::fail_next_owned_term_reserve(); + EXPECT_THROW(buffer.add_token("reserve-failure-term", 0, 0), std::bad_alloc); + EXPECT_EQ(snii_testing::vocab_string_materialization_count(), 0U); + EXPECT_GT(buffer.resident_bytes_for_test(), 0U); + EXPECT_EQ(reporter.current_bytes(), buffer.resident_bytes_for_test()); + + snii_testing::fail_next_owned_term_emplace(); + EXPECT_THROW(buffer.add_token("emplace-failure-term", 1, 0), std::bad_alloc); + EXPECT_EQ(snii_testing::vocab_string_materialization_count(), 0U); + EXPECT_EQ(reporter.current_bytes(), buffer.resident_bytes_for_test()); + + buffer.add_token("emplace-failure-term", 2, 0); + buffer.add_token("", 3, 0); + buffer.add_token("emplace-failure-term", 4, 0); + EXPECT_EQ(snii_testing::vocab_string_materialization_count(), 2U); + + const std::vector postings = buffer.finalize_sorted(); + ASSERT_EQ(postings.size(), 2U); + expect_docids_for_term(postings, "", {3U}); + expect_docids_for_term(postings, "emplace-failure-term", {2U, 4U}); +} + +// Feeds the same token stream into two buffers and asserts the streaming +// for_each_term_sorted produces the byte-identical postings finalize_sorted +// returns (same flat-array refactor must not change observable output). +TEST(SniiSpimiTermBuffer, StreamingMatchesMaterialized) { + auto feed = [](SpimiTermBuffer& b) { + b.add_token("banana", 0, 0); + b.add_token("apple", 0, 1); + b.add_token("apple", 0, 2); + b.add_token("apple", 1, 0); + b.add_token("cherry", 1, 1); + b.add_token("apple", 5, 3); + b.add_token("apple", 5, 7); + b.add_token("banana", 9, 0); + }; + SpimiTermBuffer mat(/*has_positions=*/true); + SpimiTermBuffer strm(/*has_positions=*/true); + feed(mat); + feed(strm); + + std::vector material = mat.finalize_sorted(); + std::vector streamed; + Status st = strm.for_each_term_sorted([&](StreamedTermPostings&& source) { + TermPostings postings; + RETURN_IF_ERROR(materialize_streamed_term(std::move(source), &postings)); + streamed.push_back(std::move(postings)); + return Status::OK(); + }); + EXPECT_TRUE(st.ok()); + + ASSERT_EQ(material.size(), streamed.size()); + for (size_t i = 0; i < material.size(); ++i) { + EXPECT_EQ(material[i].term, streamed[i].term); + EXPECT_EQ(material[i].docids, streamed[i].docids); + EXPECT_EQ(material[i].freqs, streamed[i].freqs); + EXPECT_EQ(material[i].positions_flat, streamed[i].positions_flat); + } + // apple: docs {0(pos 1,2), 1(pos 0), 5(pos 3,7)} -> positions re-sliced by freq. + ASSERT_EQ(streamed[0].term, "apple"); + ASSERT_EQ(streamed[0].docids.size(), 3U); + EXPECT_EQ(streamed[0].freqs, (std::vector {2U, 1U, 2U})); + EXPECT_EQ(std::vector(streamed[0].doc_positions(2).begin(), + streamed[0].doc_positions(2).end()), + (std::vector {3U, 7U})); +} + +TEST(SniiSpimiTermBufferTest, PositionedArenaSourceDecodesEachTokenOnce) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.add_token("term", 2, 3); + buffer.add_token("term", 2, 7); + buffer.add_token("term", 5, 1); + buffer.add_token("term", 9, 2); + buffer.add_token("term", 9, 4); + buffer.add_token("term", 9, 8); + + doris::snii::writer::testing::reset_compact_chain_varint_decode_count(); + std::vector docids; + std::vector freqs; + std::vector positions; + const Status status = buffer.for_each_term_sorted([&](StreamedTermPostings&& term) { + TermPostingBuffer transfer(/*memory_reporter=*/nullptr); + bool exhausted = false; + while (!exhausted) { + transfer.clear_reuse(); + RETURN_IF_ERROR(term.source->fill(/*target_docs=*/2, &transfer, &exhausted)); + docids.insert(docids.end(), transfer.docids().begin(), transfer.docids().end()); + freqs.insert(freqs.end(), transfer.freqs().begin(), transfer.freqs().end()); + positions.insert(positions.end(), transfer.positions_flat().begin(), + transfer.positions_flat().end()); + } + return Status::OK(); + }); + + ASSERT_TRUE(status.ok()) << status; + EXPECT_EQ(docids, (std::vector {2, 5, 9})); + EXPECT_EQ(freqs, (std::vector {2, 1, 3})); + EXPECT_EQ(positions, (std::vector {3, 7, 1, 2, 4, 8})); + EXPECT_EQ(doris::snii::writer::testing::compact_chain_varint_decode_count(), 9U); +} + +// Out-of-order docid GROUPS (each doc's tokens stay contiguous, but the docids +// are not non-decreasing) are tolerated and reordered once at finalize, with +// each doc carrying its own positions (defensive fallback path, e.g. a merge of +// pre-sorted runs). Tokens for a single docid are always contiguous. +TEST(SniiSpimiTermBuffer, OutOfOrderDocidsSortedAtFinalize) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token("t", 5, 50); // doc 5 group (contiguous) + buf.add_token("t", 5, 51); + buf.add_token("t", 1, 10); // doc 1 group, arrives after doc 5 + buf.add_token("t", 1, 11); + buf.add_token("t", 3, 30); // doc 3 group + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + const TermPostings& t = terms[0]; + EXPECT_EQ(t.docids, (std::vector {1U, 3U, 5U})); + EXPECT_EQ(t.freqs, (std::vector {2U, 1U, 2U})); + EXPECT_EQ(t.positions_flat, (std::vector {10U, 11U, 30U, 50U, 51U})); +} + +// BORROWED-vocab id path: feeding raw term-ids (no per-token string work) +// produces the SAME lexicographically sorted postings as the string path. The +// vocab order (apple < banana < cherry) drives the emitted order, NOT the id +// order (banana=0, apple=1, cherry=2). +TEST(SniiSpimiTermBuffer, TermIdPathMatchesStringPath) { + const std::vector vocab = {"banana", "apple", "cherry"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true); + // doc 0: "banana apple apple", doc 1: "apple cherry" -- by id. + buf.add_token(0, 0, 0); // banana + buf.add_token(1, 0, 1); // apple + buf.add_token(1, 0, 2); // apple + buf.add_token(1, 1, 0); // apple + buf.add_token(2, 1, 1); // cherry + + EXPECT_EQ(buf.unique_terms(), 3U); + EXPECT_EQ(buf.total_tokens(), 5U); + EXPECT_TRUE(buf.status().ok()); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 3U); + EXPECT_EQ(terms[0].term, "apple"); + EXPECT_EQ(terms[1].term, "banana"); + EXPECT_EQ(terms[2].term, "cherry"); + const TermPostings& apple = terms[0]; + ASSERT_EQ(apple.docids.size(), 2U); + EXPECT_EQ(apple.freqs[0], 2U); + EXPECT_EQ(std::vector(apple.doc_positions(0).begin(), apple.doc_positions(0).end()), + (std::vector {1U, 2U})); + EXPECT_EQ(apple.docids[1], 1U); + EXPECT_EQ(apple.freqs[1], 1U); +} + +// A term-id never touched is simply skipped (no empty term emitted); an empty +// vocab yields no terms and stays valid. +TEST(SniiSpimiTermBuffer, UntouchedIdSkippedAndEmptyVocab) { + const std::vector vocab = {"a", "b", "c", "d"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false); + buf.add_token(0, 0, 0); // a + buf.add_token(2, 1, 0); // c -- ids 1 (b) and 3 (d) never touched + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[1].term, "c"); + + const std::vector empty; + SpimiTermBuffer empty_buf(&empty, /*has_positions=*/false); + EXPECT_EQ(empty_buf.unique_terms(), 0U); + EXPECT_TRUE(empty_buf.finalize_sorted().empty()); + EXPECT_TRUE(empty_buf.status().ok()); +} + +// An out-of-range term-id is rejected: the token is ignored and an +// InvalidArgument is latched into status(). +TEST(SniiSpimiTermBuffer, OutOfRangeTermIdRejected) { + const std::vector vocab = {"x", "y"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/true); + buf.add_token(0, 0, 0); // valid + buf.add_token(5, 0, 1); // out of range -> ignored + latched + EXPECT_FALSE(buf.status().ok()); + EXPECT_EQ(buf.unique_terms(), 1U); + EXPECT_EQ(buf.total_tokens(), 1U); // the rejected token was not counted +} + +// The borrowed-vocab id path is byte-identical across a spill: a tiny threshold +// (many spills + k-way merge over term-id runs) must match the unlimited build. +TEST(SniiSpimiTermBuffer, TermIdSpillMatchesUnlimited) { + const std::vector vocab = {"alpha", "beta", "gamma", "delta"}; + auto feed = [&](SpimiTermBuffer& b) { + for (uint32_t d = 0; d < 300; ++d) { + b.add_token(0, d, 0); // alpha: every doc + b.add_token(0, d, 9); // freq 2 + if (d % 2 == 0) { + b.add_token(1, d, 1); // beta + } + if (d % 3 == 0) { + b.add_token(2, d, 2); // gamma + } + if (d % 5 == 0) { + b.add_token(3, d, 3); // delta + } + } + }; + SpimiTermBuffer un(&vocab, /*has_positions=*/true, /*spill=*/0); + SpimiTermBuffer sp(&vocab, /*has_positions=*/true, /*spill=*/256); + feed(un); + feed(sp); + const std::vector a = un.finalize_sorted(); + const std::vector b = sp.finalize_sorted(); + ASSERT_EQ(a.size(), b.size()); + for (size_t i = 0; i < a.size(); ++i) { + EXPECT_EQ(a[i].term, b[i].term); + EXPECT_EQ(a[i].docids, b[i].docids); + EXPECT_EQ(a[i].freqs, b[i].freqs); + EXPECT_EQ(a[i].positions_flat, b[i].positions_flat); + } + EXPECT_TRUE(un.status().ok()); + EXPECT_TRUE(sp.status().ok()); +} + +// for_each_term_sorted drains the buffer term-by-term: after each callback the +// consumed term is gone, so at most one term's arrays remain materialized. +TEST(SniiSpimiTermBuffer, StreamingDrainsAndShrinks) { + SpimiTermBuffer buf(/*has_positions=*/false); + for (uint32_t d = 0; d < 100; ++d) { + buf.add_token("a", d, 0); + buf.add_token("b", d, 0); + buf.add_token("c", d, 0); + } + EXPECT_EQ(buf.unique_terms(), 3U); + std::vector remaining_after_each; + size_t seen = 0; + EXPECT_TRUE(buf.for_each_term_sorted([&](StreamedTermPostings&& source) { + TermPostings tp; + RETURN_IF_ERROR(materialize_streamed_term(std::move(source), &tp)); + ++seen; + EXPECT_EQ(tp.docids.size(), 100U); + remaining_after_each.push_back(buf.unique_terms()); + return Status::OK(); + }).ok()); + EXPECT_EQ(seen, 3U); + // After consuming each of the 3 terms, the live count drops 2,1,0. + EXPECT_EQ(remaining_after_each, (std::vector {2U, 1U, 0U})); + EXPECT_EQ(buf.unique_terms(), 0U); +} + +// A REVISITED docid (the out-of-order defensive path actually re-touches a doc: +// feed 5,1,5) MUST coalesce into ONE entry per docid -- summed freq, positions +// concatenated in document order -- matching the k-way merge path and the writer's +// strictly-ascending precondition. Without coalescing this yielded docids +// {1,5,5} (duplicate, unsorted) that the writer later rejects. +TEST(SniiSpimiTermBuffer, RevisitedDocidCoalescesWithPositions) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token("t", 5, 50); // doc 5, first visit + buf.add_token("t", 5, 51); + buf.add_token("t", 1, 10); // doc 1 + buf.add_token("t", 5, 52); // doc 5 REVISITED (a fresh doc-group, same docid) + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + const TermPostings& t = terms[0]; + // Exactly one entry per docid, strictly ascending. + EXPECT_EQ(t.docids, (std::vector {1U, 5U})); + // doc 5 freq = 2 (first visit) + 1 (revisit) = 3; doc 1 freq = 1. + EXPECT_EQ(t.freqs, (std::vector {1U, 3U})); + // Positions in document order: doc 1 {10}, then doc 5's two visits in arrival + // order {50,51} then {52}. + EXPECT_EQ(t.positions_flat, (std::vector {10U, 50U, 51U, 52U})); + // doc_positions slices stay consistent with the merged freqs. + EXPECT_EQ(std::vector(t.doc_positions(1).begin(), t.doc_positions(1).end()), + (std::vector {50U, 51U, 52U})); + EXPECT_TRUE(buf.status().ok()); +} + +// Same revisit, positions disabled: docids coalesce to a unique ascending set. +// Docs-only postings carry set semantics, so the revisited doc keeps freq 1 +// instead of summing occurrences. +TEST(SniiSpimiTermBuffer, RevisitedDocidCoalescesNoPositions) { + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token("t", 5, 0); + buf.add_token("t", 1, 0); + buf.add_token("t", 5, 0); // revisit doc 5 + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {1U, 5U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U, 1U})); + EXPECT_TRUE(terms[0].positions_flat.empty()); +} + +// The coalesced out-of-order output satisfies the writer's strictly-ascending +// docid precondition: docids are unique AND strictly increasing (the exact check +// LogicalIndexWriter::validate_term enforces). This is the contract the fix +// restores -- previously a revisited docid produced a non-ascending list. +TEST(SniiSpimiTermBuffer, RevisitedDocidProducesStrictlyAscending) { + SpimiTermBuffer buf(/*has_positions=*/true); + // A messy revisit pattern: 9,3,9,3,1. + buf.add_token("w", 9, 0); + buf.add_token("w", 3, 0); + buf.add_token("w", 9, 1); + buf.add_token("w", 3, 1); + buf.add_token("w", 1, 0); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + const TermPostings& t = terms[0]; + ASSERT_FALSE(t.docids.empty()); + for (size_t i = 1; i < t.docids.size(); ++i) { + EXPECT_LT(t.docids[i - 1], t.docids[i]) << "docids must be strictly ascending"; + } + EXPECT_EQ(t.docids, (std::vector {1U, 3U, 9U})); + EXPECT_EQ(t.freqs, (std::vector {1U, 2U, 2U})); + // Total positions equals total tokens (sum of freqs) -- nothing dropped. + uint64_t total_freq = 0; + for (uint32_t f : t.freqs) { + total_freq += f; + } + EXPECT_EQ(t.positions_flat.size(), total_freq); +} + +// Hardening: add_token(string_view) on a BORROWED-vocab buffer is rejected (it +// would otherwise grow the owned vocab out of step with the borrowed one and +// corrupt the build). The token is ignored and an error is latched. +TEST(SniiSpimiTermBuffer, AddTokenStringViewRejectedInBorrowedMode) { + const std::vector vocab = {"a", "b"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false); + buf.add_token(0, 0, 0); // valid id-path token + buf.add_token(std::string_view("a"), 1, 0); // illegal on a borrowed-vocab buffer + EXPECT_FALSE(buf.status().ok()); + // The string-view token was ignored: only the one id-path token counts. + EXPECT_EQ(buf.total_tokens(), 1U); + EXPECT_EQ(buf.unique_terms(), 1U); +} + +// A spill's open() I/O failure surfaces as a latched error in status(). We force +// open() to fail deterministically by exhausting every free +// file descriptor below the soft limit, so the spill's ::open returns EMFILE. +TEST(SniiSpimiTermBuffer, SpillOpenIoFailureLatched) { + // Tiny threshold so the very first token triggers a spill_to_run(). + SpimiTermBuffer buf(/*has_positions=*/false, /*spill_threshold_bytes=*/1); + + // Cap the soft limit low so we can exhaust the fd table cheaply, then hold it. + struct rlimit saved {}; + ASSERT_EQ(getrlimit(RLIMIT_NOFILE, &saved), 0); + struct rlimit tight = saved; + tight.rlim_cur = 64; // small, but >= the few gtest/std fds already open + tight.rlim_cur = std::min(tight.rlim_cur, saved.rlim_max); + ASSERT_EQ(setrlimit(RLIMIT_NOFILE, &tight), 0); + + // Open /dev/null until the table is full: every free fd below the limit is now + // taken, so the next ::open (the spill's) cannot get one -> EMFILE. + std::vector hogs; + for (;;) { + int fd = ::open("/dev/null", O_RDONLY); + if (fd < 0) { + break; // table exhausted + } + hogs.push_back(fd); + } + ASSERT_FALSE(hogs.empty()); + + buf.add_token("z", 0, 0); // triggers a spill whose RunWriter::open must fail + + // Release the hog fds and restore the limit before asserting (so gtest I/O works). + for (int fd : hogs) { + ::close(fd); + } + ASSERT_EQ(setrlimit(RLIMIT_NOFILE, &saved), 0); + + EXPECT_FALSE(buf.status().ok()) << "spill open() failure must latch an error"; +} + +// Double-drain safety: a SECOND drain (finalize_sorted after for_each_term_sorted, +// or a second finalize_sorted) must NOT silently re-emit or emit a wrong stream; +// it returns empty / no callbacks AND latches an error. +TEST(SniiSpimiTermBuffer, DoubleDrainIsRejected) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token("a", 0, 0); + buf.add_token("b", 0, 0); + + std::vector first = buf.finalize_sorted(); + ASSERT_EQ(first.size(), 2U); + EXPECT_TRUE(buf.status().ok()); + + // Second finalize_sorted: empty result + latched error. + std::vector second = buf.finalize_sorted(); + EXPECT_TRUE(second.empty()); + EXPECT_FALSE(buf.status().ok()); + + // for_each_term_sorted after a drain also emits nothing; now returns an error. + size_t seen = 0; + EXPECT_FALSE(buf.for_each_term_sorted([&](StreamedTermPostings&&) { + ++seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(seen, 0U); +} + +// Double-drain via for_each_term_sorted first, then finalize_sorted: same guard. +TEST(SniiSpimiTermBuffer, DoubleDrainStreamingThenMaterialized) { + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token("a", 0, 0); + + size_t seen = 0; + EXPECT_TRUE(buf.for_each_term_sorted([&](StreamedTermPostings&& source) { + RETURN_IF_ERROR(consume_streamed_term(std::move(source))); + ++seen; + return Status::OK(); + }).ok()); + EXPECT_EQ(seen, 1U); + EXPECT_TRUE(buf.status().ok()); + + std::vector again = buf.finalize_sorted(); + EXPECT_TRUE(again.empty()); + EXPECT_FALSE(buf.status().ok()); +} + +TEST(SniiSpimiTermBuffer, StreamingConsumerFailurePropagatesForMemoryAndSpill) { + for (size_t spill_threshold : {size_t {0}, size_t {1}}) { + SpimiTermBuffer buffer(/*has_positions=*/true, spill_threshold); + buffer.add_token("alpha", 0, 0); + buffer.add_token("beta", 1, 0); + size_t calls = 0; + const Status status = buffer.for_each_term_sorted([&](StreamedTermPostings&&) { + ++calls; + return Status::Error("injected consumer failure"); + }); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(calls, 1U); + EXPECT_EQ(buffer.unique_terms(), 0U); + } +} + +// BYTE-IDENTICAL guard for normal ascending input: the streaming and materialized +// drains over an ASCENDING feed (the common, valid path -- NOT the out-of-order +// path the coalescing fix touches) must produce identical docids/freqs/positions. +// This asserts the fix did not perturb the normal path's produced postings. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiSpimiTermBuffer, AscendingInputByteIdenticalAcrossDrains) { + auto feed = [](SpimiTermBuffer& b) { + for (uint32_t d = 0; d < 50; ++d) { + b.add_token("apple", d, d * 2); + b.add_token("apple", d, d * 2 + 1); // freq 2 per doc + if (d % 2 == 0) { + b.add_token("banana", d, d); + } + if (d % 3 == 0) { + b.add_token("cherry", d, d + 100); + } + } + }; + SpimiTermBuffer mat(/*has_positions=*/true); + SpimiTermBuffer strm(/*has_positions=*/true); + feed(mat); + feed(strm); + + std::vector material = mat.finalize_sorted(); + std::vector streamed; + Status st = strm.for_each_term_sorted([&](StreamedTermPostings&& source) { + TermPostings postings; + RETURN_IF_ERROR(materialize_streamed_term(std::move(source), &postings)); + streamed.push_back(std::move(postings)); + return Status::OK(); + }); + EXPECT_TRUE(st.ok()); + + ASSERT_EQ(material.size(), streamed.size()); + for (size_t i = 0; i < material.size(); ++i) { + EXPECT_EQ(material[i].term, streamed[i].term); + EXPECT_EQ(material[i].docids, streamed[i].docids); + EXPECT_EQ(material[i].freqs, streamed[i].freqs); + EXPECT_EQ(material[i].positions_flat, streamed[i].positions_flat); + } + // Spot-check apple stayed exactly one entry per ascending docid (no coalescing + // path was taken for this valid feed). + ASSERT_EQ(material[0].term, "apple"); + EXPECT_EQ(material[0].docids.size(), 50U); + for (uint32_t f : material[0].freqs) { + EXPECT_EQ(f, 2U); + } + EXPECT_TRUE(mat.status().ok()); + EXPECT_TRUE(strm.status().ok()); +} + +// --------------------------------------------------------------------------- +// T17: MemoryReporter per-token zero-delta debounce. +// +// accumulate() reports its REAL resident-byte delta (posting arena + the +// vocab-sized slot index) to the writer-level MemoryReporter once per token. The +// arena grows only ~every 32 KiB block and the borrowed-vocab slot index is +// fixed-capacity, so the vast majority of tokens see delta==0. report_arena_delta() +// now SKIPS the locked fetch_add for those (debounce). These tests pin the +// deterministic op-count (report() calls == arena-growth events, never per token), +// the byte-level equivalence (current_bytes() unchanged), and the REDLINE: over_cap() +// is still evaluated UNCONDITIONALLY every token (never gated on the local delta), so +// a dict-side push over the unified cap still triggers a spill when this buffer's own +// delta is 0. A MemoryReporter built with a counting consume_release lambda exposes +// the exact per-token report() count / delta values as a deterministic seam. +// --------------------------------------------------------------------------- + +// FV-1 (deterministic op-count + functional): feeding 100 same-doc tokens issues +// exactly TWO report() calls -- one ctor delta (the resident slot index) and one for +// the first token (its 32 KiB arena block plus the G08-charged first-touch Term-slot / +// touched-list growth, a few dozen bytes) -- and NEVER a zero-delta report. Before the +// debounce, tokens 2..100 each issued report(0): 101 calls, 99 of them zero. +TEST(SniiSpimiTermBufferTest, AccumulateIssuesNoZeroDeltaReport) { + std::vector deltas; + MemoryReporter rep([&deltas](int64_t d) { deltas.push_back(d); }, /*cap_bytes=*/0); + const std::vector vocab = {"a"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &rep); + + for (int i = 0; i < 100; ++i) { + buf.add_token(0, /*docid=*/1, /*pos=*/0); // same term, same doc + } + + // 1 ctor report (slot index) + 1 first-token report (first 32 KiB arena block + + // the one-term slot-pool/touched-list first-touch growth). The other 99 same-doc + // tokens leave resident unchanged -> debounced away. + ASSERT_EQ(deltas.size(), 2U); + EXPECT_GT(deltas[0], 0); // slot index resident bytes (vocab-sized) + // One CompactPostingPool block (1 << 15) dominates; the G08 slot-pool charge for + // the single touched term adds well under 128 B on top. + EXPECT_GE(deltas[1], 32768); + EXPECT_LT(deltas[1], 32768 + 128); + for (int64_t d : deltas) { + EXPECT_NE(d, 0) << "no zero-delta report() may be issued on the hot path"; + } + EXPECT_TRUE(buf.status().ok()); +} + +// FV-2 (equivalence + count stability): the report() COUNT is independent of the +// token count (100 vs 500 same-doc tokens both issue exactly 2 reports), and the +// resulting unified total is byte-identical (resident = first arena block + slot +// index + first-touch slot-pool growth, not a function of token count). The sum +// of issued deltas equals +// current_bytes() -- the MemoryReporter self-balancing invariant the debounce +// preserves. Snapshots are taken WHILE each buffer is live (before its dtor reports +// the final balancing negative). +TEST(SniiSpimiTermBufferTest, ReportedTotalMatchesResidentRegardlessOfTokenCount) { + const std::vector vocab = {"a"}; + + std::vector d100; + std::vector d500; + int64_t cur100 = 0; + int64_t cur500 = 0; + size_t count100 = 0; + size_t count500 = 0; + + { + MemoryReporter rep([&d100](int64_t d) { d100.push_back(d); }, /*cap_bytes=*/0); + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &rep); + for (int i = 0; i < 100; ++i) { + buf.add_token(0, /*docid=*/1, /*pos=*/0); + } + count100 = d100.size(); // snapshot before the dtor's balancing report + cur100 = rep.current_bytes(); + } + { + MemoryReporter rep([&d500](int64_t d) { d500.push_back(d); }, /*cap_bytes=*/0); + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &rep); + for (int i = 0; i < 500; ++i) { + buf.add_token(0, /*docid=*/1, /*pos=*/0); + } + count500 = d500.size(); + cur500 = rep.current_bytes(); + } + + // Count stability: neither buffer's report count scales with token count. + EXPECT_EQ(count100, 2U); + EXPECT_EQ(count500, 2U); + // Byte-level equivalence: identical unified total despite 5x the tokens. + EXPECT_EQ(cur100, cur500); + ASSERT_GE(d100.size(), 2U); + ASSERT_GE(d500.size(), 2U); + // Both: identical first-token delta -- one 32 KiB arena block plus the same + // G08 first-touch slot-pool/touched-list growth (the slot index cancels in + // this delta; the growth is token-count-independent, hence the equality). + EXPECT_EQ(d100[1], d500[1]); + EXPECT_GE(d100[1], 32768); + EXPECT_LT(d100[1], 32768 + 128); + // Self-balancing: live current_bytes() == sum of every delta issued so far. + int64_t sum100 = 0; + for (size_t i = 0; i < count100; ++i) { + sum100 += d100[i]; + } + EXPECT_EQ(cur100, sum100); + EXPECT_EQ(cur100, d100[0] + d100[1]); // slot index + first-token resident +} + +// FV-3 (REDLINE guard): the debounce skips report() ONLY -- it must NOT gate +// over_cap(). over_cap() reads the writer-level UNIFIED total (shared with the dict +// buffer), so a dict-side allocation can push the total over the cap while THIS +// buffer's local arena delta is 0. accumulate() must still evaluate over_cap() every +// token and spill. The dict-side growth is simulated with an external rep.report(). +TEST(SniiSpimiTermBufferTest, OverCapStillFiresWhenLocalArenaDeltaIsZero) { + // Cap just above one arena block + slot index, so token1 alone does NOT spill. + MemoryReporter rep(/*consume_release=*/nullptr, /*cap_bytes=*/32768 + 1000, + MemoryReporter::CapPolicy::kSpillThreshold); + const std::vector vocab = {"a"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &rep); + + // token1: arena grows to one 32 KiB block; unified total < cap -> no spill. + buf.add_token(0, /*docid=*/0, /*pos=*/0); + ASSERT_EQ(buf.run_count_for_test(), 0U); + + // Simulate dict-side growth pushing the UNIFIED total over the cap. The buffer's + // own reported_resident_ is unchanged by this external report. + rep.report(2000); + ASSERT_TRUE(rep.over_cap()); + + // token2: same doc -> this buffer's local arena delta is 0 (report() debounced), + // but over_cap() is still evaluated unconditionally and is now true -> spill. + buf.add_token(0, /*docid=*/0, /*pos=*/0); + EXPECT_EQ(buf.run_count_for_test(), 1U) + << "over_cap() must NOT be gated on the local arena delta"; + EXPECT_TRUE(buf.status().ok()); +} + +// FV-4 (boundary): an EMPTY borrowed vocab has a zero-capacity slot index, so the +// ctor's resident delta is 0 and is debounced away -- no report() is issued and +// construction does not crash. Before the debounce the ctor issued report(0). +TEST(SniiSpimiTermBufferTest, EmptyVocabReportsNoDelta) { + std::vector deltas; + MemoryReporter rep([&deltas](int64_t d) { deltas.push_back(d); }, /*cap_bytes=*/0); + const std::vector empty_vocab; + SpimiTermBuffer buf(&empty_vocab, /*has_positions=*/false, /*spill=*/0, &rep); + + // Zero-capacity slot index + empty arena -> resident 0 -> ctor delta 0 -> skipped. + EXPECT_TRUE(deltas.empty()); + EXPECT_EQ(buf.unique_terms(), 0U); + EXPECT_TRUE(buf.finalize_sorted().empty()); + EXPECT_TRUE(buf.status().ok()); +} + +// FV-5 (no-reporter path unaffected): with a null reporter, report_arena_delta() is a +// no-op and finalize_sorted() still produces the correct postings. Confirms the +// debounce change did not perturb the off-Doris path. +TEST(SniiSpimiTermBufferTest, NullReporterFinalizeIsCorrect) { + const std::vector vocab = {"a"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, /*reporter=*/nullptr); + + for (int i = 0; i < 100; ++i) { + buf.add_token(0, /*docid=*/1, /*pos=*/0); // same term, same doc + } + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[0].docids, (std::vector {1U})); + // Docs-only set semantics: the 99 same-doc repeats are discarded at + // accumulate time, so the doc's freq is 1 (not the occurrence count). + EXPECT_EQ(terms[0].freqs, (std::vector {1U})); + EXPECT_TRUE(terms[0].positions_flat.empty()); + EXPECT_TRUE(buf.status().ok()); +} + +// FV-6 (spill/drain negative path + self-balance): a forced spill emits NONZERO +// negative deltas (arena reset, then slot-index free), never a zero. After a full +// drain the reporter's unified total returns to 0 -- the debounce preserves the +// self-balancing invariant (no leaked positive). The merged postings are correct. +TEST(SniiSpimiTermBufferTest, SpillNegativeDeltasHaveNoZeroAndBalance) { + std::vector deltas; + // Cap below one arena block, so the first token's block forces a spill. + MemoryReporter rep([&deltas](int64_t d) { deltas.push_back(d); }, /*cap_bytes=*/1000, + MemoryReporter::CapPolicy::kSpillThreshold); + const std::vector vocab = {"a"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false, /*spill=*/0, &rep); + + buf.add_token(0, /*docid=*/0, /*pos=*/0); // grows one block, then spills + ASSERT_EQ(buf.run_count_for_test(), 1U) << "the over-cap token must spill"; + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[0].docids, (std::vector {0U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U})); + EXPECT_TRUE(buf.status().ok()); + + // No zero-delta report on any path (grow, spill-negative, or merge-free). + for (int64_t d : deltas) { + EXPECT_NE(d, 0) << "spill/drain deltas must all be nonzero"; + } + // Self-balancing: after a full drain every reported byte has been returned. + EXPECT_EQ(rep.current_bytes(), 0); +} diff --git a/be/test/storage/index/snii/writer/temp_dir_test.cpp b/be/test/storage/index/snii/writer/temp_dir_test.cpp new file mode 100644 index 00000000000000..5c0d42e1f332e2 --- /dev/null +++ b/be/test/storage/index/snii/writer/temp_dir_test.cpp @@ -0,0 +1,46 @@ +// 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 "storage/index/snii/writer/temp_dir.h" + +#include + +#include +#include + +namespace doris::snii::writer { +namespace { + +// ExecEnv's tmp_file_dirs are initialized by the global SniiTmpDirEnvironment +// (be/test/storage/index/snii/snii_test_env.cpp). +TEST(SniiTempDir, ResolvesIntoSniiSubdirOfConfiguredTmpDir) { + const std::string dir = resolve_temp_dir(); + ASSERT_FALSE(dir.empty()); + // SNII gets its own "snii" subdirectory so it does not crowd the tmp root. + ASSERT_GE(dir.size(), 5U); + EXPECT_EQ(0, dir.compare(dir.size() - 5, 5, "/snii")) << dir; + // resolve_temp_dir() creates the directory, so it stats successfully. + EXPECT_NE(temp_dir_available_bytes(dir), UINT64_MAX) << dir; +} + +TEST(SniiTempDir, AvailableBytesStatsRealDirAndSentinelsOnFailure) { + EXPECT_NE(temp_dir_available_bytes("/tmp"), UINT64_MAX); // real dir -> some free + EXPECT_EQ(temp_dir_available_bytes("/snii_no_such_dir_xyzzy"), UINT64_MAX); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/term_posting_source_test.cpp b/be/test/storage/index/snii/writer/term_posting_source_test.cpp new file mode 100644 index 00000000000000..49af0d3642728a --- /dev/null +++ b/be/test/storage/index/snii/writer/term_posting_source_test.cpp @@ -0,0 +1,245 @@ +// 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 "storage/index/snii/writer/term_posting_source.h" + +#include + +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/writer/memory_reporter.h" + +namespace doris::snii::writer { +namespace { + +template +void expect_values(std::span actual, const std::array& expected) { + ASSERT_EQ(actual.size(), expected.size()); + EXPECT_TRUE(std::equal(actual.begin(), actual.end(), expected.begin())); +} + +TEST(SniiTermPostingBufferTest, AppendsPositionedChunksAndReusesReservedCapacity) { + MemoryReporter reporter; + { + TermPostingBuffer buffer(&reporter); + const std::array first_docids {1, 4}; + const std::array first_freqs {2, 1}; + const std::array first_positions {0, 3, 9}; + ASSERT_TRUE(buffer.append(first_docids, first_freqs, first_positions).ok()); + EXPECT_EQ(buffer.document_count(), 2U); + expect_values(buffer.docids(), first_docids); + expect_values(buffer.freqs(), first_freqs); + expect_values(buffer.positions_flat(), first_positions); + + const int64_t charged_bytes = reporter.current_bytes(); + EXPECT_GT(charged_bytes, 0); + buffer.clear_reuse(); + EXPECT_TRUE(buffer.empty()); + EXPECT_EQ(reporter.current_bytes(), charged_bytes); + + const std::array second_docids {7}; + const std::array second_freqs {1}; + const std::array second_positions {2}; + ASSERT_TRUE(buffer.append(second_docids, second_freqs, second_positions).ok()); + expect_values(buffer.docids(), second_docids); + EXPECT_EQ(reporter.current_bytes(), charged_bytes); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiTermPostingBufferTest, AcceptsDocsOnlyAndRejectsInvalidShapesWithoutMutation) { + TermPostingBuffer buffer(nullptr); + const std::array docids {2, 9}; + ASSERT_TRUE(buffer.append(docids, {}, {}).ok()); + expect_values(buffer.docids(), docids); + EXPECT_TRUE(buffer.freqs().empty()); + EXPECT_TRUE(buffer.positions_flat().empty()); + + const std::array wrong_freqs {1}; + EXPECT_TRUE(buffer.append(docids, wrong_freqs, {}).is()); + expect_values(buffer.docids(), docids); + + const std::array freqs {1, 2}; + const std::array too_few_positions {0, 1}; + EXPECT_TRUE(buffer.append(docids, freqs, too_few_positions).is()); + expect_values(buffer.docids(), docids); + + buffer.clear_reuse(); + ASSERT_TRUE(buffer.append(docids, freqs, {}).ok()); + expect_values(buffer.docids(), docids); + expect_values(buffer.freqs(), freqs); + EXPECT_TRUE(buffer.positions_flat().empty()); +} + +TEST(SniiTermPostingBufferTest, RejectsReplacementBeforeGrowingPastHardLimit) { + constexpr uint64_t kInitialBytes = 6 * sizeof(uint32_t); + MemoryReporter reporter(nullptr, kInitialBytes, MemoryReporter::CapPolicy::kHardLimit); + { + TermPostingBuffer buffer(&reporter); + const std::array docids {1, 2}; + const std::array freqs {1, 1}; + const std::array positions {0, 0}; + ASSERT_TRUE(buffer.append(docids, freqs, positions).ok()); + ASSERT_EQ(reporter.current_bytes(), static_cast(kInitialBytes)); + + EXPECT_TRUE(buffer.append(docids, freqs, positions).is()); + expect_values(buffer.docids(), docids); + expect_values(buffer.freqs(), freqs); + expect_values(buffer.positions_flat(), positions); + EXPECT_EQ(reporter.current_bytes(), static_cast(kInitialBytes)); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiTermPostingBufferTest, ProvidesDirectWritableFillStorage) { + MemoryReporter reporter; + TermPostingBuffer buffer(&reporter); + MutableTermPostingSpan span; + + ASSERT_TRUE(buffer.grow_uninitialized(/*document_count=*/2, /*has_freqs=*/true, + /*position_count=*/3, &span) + .ok()); + ASSERT_EQ(span.docids.size(), 2U); + ASSERT_EQ(span.freqs.size(), 2U); + ASSERT_EQ(span.positions_flat.size(), 3U); + span.docids[0] = 7; + span.docids[1] = 11; + span.freqs[0] = 1; + span.freqs[1] = 2; + span.positions_flat[0] = 3; + span.positions_flat[1] = 5; + span.positions_flat[2] = 8; + + expect_values(buffer.docids(), std::array {7, 11}); + expect_values(buffer.freqs(), std::array {1, 2}); + expect_values(buffer.positions_flat(), std::array {3, 5, 8}); + EXPECT_GT(reporter.current_bytes(), 0); +} + +TEST(SniiTermPostingBufferTest, ReleasesOversizedLaneBeforeLoweringCharge) { + constexpr size_t kRetainedCapacity = 8192; + MemoryReporter reporter; + TermPostingBuffer buffer(&reporter); + MutableTermPostingSpan span; + ASSERT_TRUE(buffer.grow_uninitialized(/*document_count=*/2, /*has_freqs=*/true, + /*position_count=*/kRetainedCapacity + 1, &span) + .ok()); + ASSERT_EQ(reporter.current_bytes(), + static_cast((kRetainedCapacity + 5) * sizeof(uint32_t))); + + buffer.clear_reuse_and_release_excess(kRetainedCapacity); + EXPECT_TRUE(buffer.empty()); + EXPECT_EQ(reporter.current_bytes(), 4 * static_cast(sizeof(uint32_t))); + + ASSERT_TRUE(buffer.grow_uninitialized(/*document_count=*/1, /*has_freqs=*/true, + /*position_count=*/1, &span) + .ok()); + EXPECT_EQ(reporter.current_bytes(), 5 * static_cast(sizeof(uint32_t))); +} + +TEST(SniiTermPostingBufferTest, AppendsPositionsIncrementallyToWritableFill) { + TermPostingBuffer buffer(nullptr); + MutableTermPostingSpan documents; + ASSERT_TRUE(buffer.grow_uninitialized(/*document_count=*/2, /*has_freqs=*/true, + /*position_count=*/0, &documents) + .ok()); + documents.docids[0] = 3; + documents.docids[1] = 8; + documents.freqs[0] = 1; + documents.freqs[1] = 2; + + ASSERT_TRUE(buffer.append_position(4).ok()); + ASSERT_TRUE(buffer.append_position(2).ok()); + ASSERT_TRUE(buffer.append_position(9).ok()); + + expect_values(buffer.docids(), std::array {3, 8}); + expect_values(buffer.freqs(), std::array {1, 2}); + expect_values(buffer.positions_flat(), std::array {4, 2, 9}); +} + +TEST(SniiTermPostingBufferTest, RejectsPositionGrowthBeforeExceedingHardLimit) { + constexpr uint64_t kHardLimitBytes = 5 * sizeof(uint32_t); + MemoryReporter reporter(nullptr, kHardLimitBytes, MemoryReporter::CapPolicy::kHardLimit); + { + TermPostingBuffer buffer(&reporter); + MutableTermPostingSpan document; + ASSERT_TRUE(buffer.grow_uninitialized(/*document_count=*/1, /*has_freqs=*/true, + /*position_count=*/0, &document) + .ok()); + document.docids[0] = 7; + document.freqs[0] = 2; + ASSERT_TRUE(buffer.append_position(1).ok()); + ASSERT_TRUE(buffer.append_position(5).ok()); + ASSERT_EQ(reporter.current_bytes(), 4 * static_cast(sizeof(uint32_t))); + + EXPECT_TRUE(buffer.append_position(9).is()); + expect_values(buffer.docids(), std::array {7}); + expect_values(buffer.freqs(), std::array {2}); + expect_values(buffer.positions_flat(), std::array {1, 5}); + EXPECT_EQ(reporter.current_bytes(), 4 * static_cast(sizeof(uint32_t))); + } + EXPECT_EQ(reporter.current_bytes(), 0); +} + +TEST(SniiTermPostingBufferTest, GrowsIncrementalRunsGeometrically) { + size_t positive_reservations = 0; + MemoryReporter reporter([&](int64_t delta) { + if (delta > 0) { + ++positive_reservations; + } + }); + TermPostingBuffer buffer(&reporter); + for (uint32_t docid = 0; docid < 8; ++docid) { + const std::array docids {docid}; + const std::array freqs {1}; + const std::array positions {docid}; + ASSERT_TRUE(buffer.append(docids, freqs, positions).ok()); + } + + EXPECT_EQ(buffer.document_count(), 8U); + EXPECT_EQ(positive_reservations, 7U); +} + +TEST(SniiTermPostingSourceTest, SpanSourceFillsExactWindowsWithoutOwningPostings) { + const std::vector docids {1, 4, 9}; + const std::vector freqs {2, 1, 3}; + const std::vector positions {0, 2, 1, 3, 7, 8}; + SpanTermPostingSource source(docids, freqs, positions); + TermPostingBuffer buffer(nullptr); + bool exhausted = false; + + ASSERT_TRUE(source.fill(2, &buffer, &exhausted).ok()); + EXPECT_FALSE(exhausted); + expect_values(buffer.docids(), std::array {1, 4}); + expect_values(buffer.freqs(), std::array {2, 1}); + expect_values(buffer.positions_flat(), std::array {0, 2, 1}); + + buffer.clear_reuse(); + ASSERT_TRUE(source.fill(2, &buffer, &exhausted).ok()); + EXPECT_TRUE(exhausted); + expect_values(buffer.docids(), std::array {9}); + expect_values(buffer.freqs(), std::array {3}); + expect_values(buffer.positions_flat(), std::array {3, 7, 8}); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii/writer/term_posting_test_utils.h b/be/test/storage/index/snii/writer/term_posting_test_utils.h new file mode 100644 index 00000000000000..a45b8f959665a3 --- /dev/null +++ b/be/test/storage/index/snii/writer/term_posting_test_utils.h @@ -0,0 +1,63 @@ +// 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. + +#pragma once + +#include + +#include "common/status.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::writer { + +inline Status materialize_streamed_term(StreamedTermPostings&& streamed, TermPostings* output, + uint32_t window_docs = 1024) { + if (output == nullptr || streamed.source == nullptr || window_docs == 0) { + return Status::Error( + "test posting materializer: invalid arguments"); + } + *output = TermPostings(); + output->term = std::move(streamed.term); + output->retain_positions = streamed.retain_positions; + TermPostingBuffer buffer(nullptr); + bool exhausted = false; + while (!exhausted) { + buffer.clear_reuse(); + RETURN_IF_ERROR(streamed.source->fill(window_docs, &buffer, &exhausted)); + output->docids.insert(output->docids.end(), buffer.docids().begin(), buffer.docids().end()); + output->freqs.insert(output->freqs.end(), buffer.freqs().begin(), buffer.freqs().end()); + output->positions_flat.insert(output->positions_flat.end(), buffer.positions_flat().begin(), + buffer.positions_flat().end()); + } + return Status::OK(); +} + +inline Status consume_streamed_term(StreamedTermPostings&& streamed, uint32_t window_docs = 1024) { + if (streamed.source == nullptr || window_docs == 0) { + return Status::Error( + "test posting consumer: invalid arguments"); + } + TermPostingBuffer buffer(nullptr); + bool exhausted = false; + while (!exhausted) { + buffer.clear_reuse(); + RETURN_IF_ERROR(streamed.source->fill(window_docs, &buffer, &exhausted)); + } + return Status::OK(); +} + +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii_b2_or_read_test.cpp b/be/test/storage/index/snii_b2_or_read_test.cpp new file mode 100644 index 00000000000000..6fedd4900f7ece --- /dev/null +++ b/be/test/storage/index/snii_b2_or_read_test.cpp @@ -0,0 +1,346 @@ +// 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. + +// SNII Batch 2 -- multi-term OR read-amplification + streaming docid-union (T09). +// +// Covers three reader-only, byte-identical changes: +// (1) emit_docid_union streams each posting into a dedup-capable sink (Roaring) +// over ONE shared fetch round -- dense-full windows stay runs via +// append_range -- instead of materializing a per-term vector + K-way merge. +// (2) union_sorted_many reserves by summed input size (single allocation), +// capped by reserve_cap so heavily-overlapping inputs do not over-reserve. +// (3) the OR resolve path threads one request-scoped DictBlockCache through its +// per-term lookups, so terms sharing a DICT block read+decode it once. +// +// All assertions are deterministic (op-counts, capacities, set equality, I/O round +// counts through MeteredFileReader / MemoryFile). No wall-clock gates. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/io/metered_file_reader.h" +#include "storage/index/snii/query/boolean_query.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/docid_set_ops.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +using namespace doris::snii; +using namespace doris::snii::snii_test; +using doris::Status; + +namespace qi = doris::snii::query::internal; + +namespace { + +// A dedup-capable sink (dedups()==true) that records how each posting was handed +// off: append_range for run-preserving dense windows, append_sorted otherwise. The +// collected docids land in a std::set so equality checks are order-independent +// (a real Roaring sink also dedups/orders internally). +class CountingDedupSink final : public query::DocIdSink { +public: + Status append_sorted(std::span docids) override { + ++sorted_calls; + for (uint32_t docid : docids) { + ids.insert(docid); + } + return Status::OK(); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + ++range_calls; + for (uint64_t docid = first; docid < last_exclusive; ++docid) { + ids.insert(static_cast(docid)); + } + return Status::OK(); + } + + bool dedups() const override { return true; } + + std::set ids; + size_t sorted_calls = 0; + size_t range_calls = 0; +}; + +std::vector closed_range(uint32_t begin, uint32_t end_exclusive) { + std::vector out; + out.reserve(end_exclusive - begin); + for (uint32_t docid = begin; docid < end_exclusive; ++docid) { + out.push_back(docid); + } + return out; +} + +// Opens the standard 9000-doc fixture over a MeteredFileReader so I/O rounds and +// remote GETs can be measured. build_reader() writes the index bytes into `file`; +// a fresh segment/index is then re-opened over a metered wrapper of those bytes. +struct MeteredIndex { + MemoryFile file; + std::unique_ptr metered; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; +}; + +void OpenMeteredFixture(MeteredIndex* fx, size_t block_size) { + reader::SniiSegmentReader seg0; + reader::LogicalIndexReader idx0; + assert_ok(build_reader(&fx->file, &seg0, &idx0)); + fx->metered = std::make_unique(&fx->file, block_size); + assert_ok(reader::SniiSegmentReader::open(fx->metered.get(), &fx->segment)); + assert_ok(fx->segment.open_index(7, "Body", &fx->idx)); +} + +} // namespace + +// --------------------------------------------------------------------------- +// T09 F16 -- union_sorted_many reserves by total, capped by reserve_cap. +// --------------------------------------------------------------------------- + +TEST(SniiB2OrRead, UnionReservesByTotalForDisjointLists) { + // Three disjoint sorted lists (sum == 12); union == sum. + const std::vector> lists = {{0, 3, 6, 9}, {1, 4, 7, 10}, {2, 5, 8, 11}}; + const size_t total = 12; + + const std::vector got = qi::union_sorted_many(lists); + + EXPECT_EQ(got.size(), total); + // reserve(total) + exactly `total` pushes -> a single allocation, so capacity is + // exactly total (was geometric growth off reserve(largest)). + EXPECT_EQ(got.capacity(), total); + EXPECT_TRUE(std::ranges::is_sorted(got)); +} + +TEST(SniiB2OrRead, UnionRespectsReserveCapOnHeavyOverlap) { + // >8 identical lists -> heap path; total == 40 but the union is only 4 elements. + const std::vector> lists(10, std::vector {1, 2, 3, 4}); + const size_t union_size = 4; + + const std::vector got = qi::union_sorted_many(lists, /*reserve_cap=*/union_size); + + EXPECT_EQ(got.size(), union_size); + // Capped at reserve_cap rather than over-reserving 10x (= 40) on overlap. + EXPECT_EQ(got.capacity(), union_size); + EXPECT_EQ(got, (std::vector {1, 2, 3, 4})); +} + +TEST(SniiB2OrRead, UnionContentUnchangedByReserveFix) { + const std::vector> lists = {{0, 2, 4, 6, 8}, {1, 3, 5}, {4, 5, 6, 7}}; + std::set expected_set; + for (const std::vector& list : lists) { + expected_set.insert(list.begin(), list.end()); + } + const std::vector want(expected_set.begin(), expected_set.end()); + + EXPECT_EQ(qi::union_sorted_many(lists), want); +} + +// --------------------------------------------------------------------------- +// T09 F15 -- dedups() capability gate + streaming OR. +// --------------------------------------------------------------------------- + +TEST(SniiB2OrRead, DedupCapabilityGate) { + std::vector backing; + query::VectorDocIdSink vector_sink(backing); + EXPECT_FALSE(vector_sink.dedups()) << "plain vector sink needs materialize+merge"; + + CountingDedupSink dedup_sink; + EXPECT_TRUE(dedup_sink.dedups()) << "Roaring-style sink dedups/orders natively"; +} + +TEST(SniiB2OrRead, MultiTermOrPreservesDenseRangeToDedupSink) { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &segment, &idx)); + + CountingDedupSink sink; + assert_ok(query::boolean_or(idx, {"failed", "sparse_left"}, &sink)); + + // "failed" is a dense full posting (docids 0..8999): its dense-full windows must + // stream in via append_range (run-preserving), not be expanded element-by-element + // through a merge accumulator -- the old path issued only append_sorted(acc). + EXPECT_GE(sink.range_calls, 1u); + + // failed covers every doc, so the union is the full doc range. + const std::vector got(sink.ids.begin(), sink.ids.end()); + EXPECT_EQ(got, closed_range(0, 9000)); +} + +TEST(SniiB2OrRead, MultiTermOrStreamingMatchesMergePath) { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &segment, &idx)); + + const std::vector terms = {"needle", "sparse_left", "driver"}; + + // Streaming path: dedup-capable sink -> emit_docid_postings_streamed. + CountingDedupSink sink; + assert_ok(query::boolean_or(idx, terms, &sink)); + const std::vector streamed(sink.ids.begin(), sink.ids.end()); + + // Merge path: vector out -> build_docid_union + union_sorted_many. + std::vector merged; + assert_ok(query::boolean_or(idx, terms, &merged)); + + EXPECT_TRUE(std::ranges::is_sorted(merged)); + EXPECT_FALSE(merged.empty()); + EXPECT_EQ(streamed, merged) << "streaming OR must yield the identical docid set"; +} + +TEST(SniiB2OrRead, NonDedupSinkFallsBackToMergeContract) { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &segment, &idx)); + + const std::vector terms = {"needle", "sparse_left"}; + + // A non-dedup sink (VectorDocIdSink) routed through the sink overload must keep + // the merge path so its single globally-sorted-deduplicated span contract holds. + std::vector via_sink; + query::VectorDocIdSink vector_sink(via_sink); + assert_ok(query::boolean_or(idx, terms, &vector_sink)); + + std::vector via_vector; + assert_ok(query::boolean_or(idx, terms, &via_vector)); + + EXPECT_TRUE(std::ranges::is_sorted(via_sink)); + EXPECT_EQ(via_sink, via_vector); +} + +TEST(SniiB2OrRead, SingleTermAndBoundaryInputs) { + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &segment, &idx)); + + // Single resolved term -> exact known docids. + CountingDedupSink single; + assert_ok(query::boolean_or(idx, {"needle"}, &single)); + EXPECT_EQ((std::vector(single.ids.begin(), single.ids.end())), + (std::vector {100, 101, 102, 6000})); + + // Empty terms -> OK, sink untouched. + CountingDedupSink empty; + assert_ok(query::boolean_or(idx, std::vector {}, &empty)); + EXPECT_TRUE(empty.ids.empty()); + + // All-miss term -> OK, sink untouched (resolve skips absent terms). + CountingDedupSink miss; + assert_ok(query::boolean_or(idx, {"zzz_absent_term"}, &miss)); + EXPECT_TRUE(miss.ids.empty()); + + // Null sink -> InvalidArgument. + query::DocIdSink* null_sink = nullptr; + const Status st = query::boolean_or(idx, {"failed"}, null_sink); + EXPECT_FALSE(st.ok()); + EXPECT_TRUE(st.is()) << st.to_string(); +} + +// --------------------------------------------------------------------------- +// Read amplification -- one shared fetch round for the multi-term OR postings. +// --------------------------------------------------------------------------- + +TEST(SniiB2OrRead, MultiTermOrIssuesSingleSerialRound) { + MeteredIndex fx; + OpenMeteredFixture(&fx, /*block_size=*/256); + const std::vector terms = {"failed", "driver", "sparse_left"}; + + // Per-term baseline: each term resolved + read on its own -> its own fetch round. + fx.metered->reset_metrics(); + for (const std::string& term : terms) { + std::vector docs; + assert_ok(query::term_query(fx.idx, term, &docs)); + } + const io::IoMetrics per_term = fx.metered->metrics(); + + // Batched OR: one shared fetch round for all postings' docid reads. + fx.metered->reset_metrics(); + std::vector got; + assert_ok(query::boolean_or(fx.idx, terms, &got)); + const io::IoMetrics batched = fx.metered->metrics(); + + EXPECT_EQ(batched.serial_rounds, 1u) << "multi-term OR must read all postings in one round"; + EXPECT_GE(per_term.serial_rounds, 2u); + EXPECT_GT(per_term.serial_rounds, batched.serial_rounds); + // Coalescing never increases physical GETs or bytes vs the per-term path. + EXPECT_LE(batched.range_gets, per_term.range_gets); + EXPECT_LE(batched.remote_bytes, per_term.remote_bytes); + + // Result set is unchanged: failed(0..8999) covers driver(0..7999) and sparse_left. + EXPECT_EQ(got, closed_range(0, 9000)); +} + +TEST(SniiB2OrRead, MultiTermOrDedupsSharedDictBlockReads) { + // Force on-demand DICT blocks so each lookup reads its block from the file + // (resident DICT blocks would serve lookups from memory and hide the effect). + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + reader::SniiSegmentReader segment; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &segment, &idx)); + + const format::RegionRef dict = idx.section_refs().dict_region; + const auto dict_reads = [&]() { + size_t count = 0; + for (const MemoryFile::Read& r : file.reads()) { + if (r.offset >= dict.offset && r.offset < dict.offset + dict.length) { + ++count; + } + } + return count; + }; + + const std::vector terms = {"123", "almost", "driver", "failed", + "needle", "order", "ordinal", "repeat", + "sparse_left", "sparse_right", "trace"}; + + // Per-term baseline: term_query resolves each DICT block on its own (cache=null) + // -> one DICT-region read per term. + file.clear_reads(); + for (const std::string& term : terms) { + std::vector docs; + assert_ok(query::term_query(idx, term, &docs)); + } + const size_t per_term_dict_reads = dict_reads(); + + // Multi-term OR threads one request-scoped DictBlockCache through resolve, so each + // unique DICT block is read + decoded once for the whole OR. + file.clear_reads(); + std::vector got; + assert_ok(query::boolean_or(idx, terms, &got)); + const size_t or_dict_reads = dict_reads(); + + EXPECT_EQ(per_term_dict_reads, terms.size()); + EXPECT_GE(or_dict_reads, 1u); + EXPECT_LT(or_dict_reads, per_term_dict_reads) + << "OR must not re-read a DICT block already decoded for another term"; + EXPECT_FALSE(got.empty()); +} diff --git a/be/test/storage/index/snii_b2_t07_dict_test.cpp b/be/test/storage/index/snii_b2_t07_dict_test.cpp new file mode 100644 index 00000000000000..03cbd88de59db5 --- /dev/null +++ b/be/test/storage/index/snii_b2_t07_dict_test.cpp @@ -0,0 +1,634 @@ +// 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. + +// T07 -- DICT entry key-first decode (exact find_term + prefix streaming +// early-stop). These tests assert three things the optimization promises: +// 1. byte-identical decode: decode_dict_entry == decode_dict_entry_key + +// decode_dict_entry_rest, and visit_prefix_range streams the same entries +// decode_all materializes; +// 2. body-decode op-count drops: find_term materializes 1 (hit) / 0 (miss) +// bodies instead of every scanned entry; visit_prefix_range materializes +// only the accepted/emitted bodies (anchor-jump + early-stop); +// 3. results unchanged at the reader/query layer (prefix_terms / prefix_query +// route through the rewired streaming path). + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii { +namespace { + +using namespace doris::snii::format; // NOLINT(google-build-using-namespace) +using namespace doris::snii::snii_test; // NOLINT(google-build-using-namespace) +using doris::Status; + +// ---- entry builders (valid, self-consistent locators so they round-trip) ---- + +DictEntry MakePodRef(std::string term, uint32_t df, uint64_t frq_off, uint64_t prx_off = 0) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kPodRef; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.ttf_delta = df * 2; + e.max_freq = 9; + e.frq_off_delta = frq_off; + e.frq_len = 128; + e.frq_docs_len = 64; // dd region on-disk length (<= frq_len) + e.dd_meta.uncomp_len = 70; + e.dd_meta.crc = 0xABCD1234U; // pod_ref regions keep their per-region crc + e.freq_meta.uncomp_len = 40; + e.freq_meta.crc = 0x55AA00FFU; + e.prx_off_delta = prx_off; + e.prx_len = 64; + return e; +} + +DictEntry MakePodRefWindowed(std::string term, uint32_t df, uint64_t frq_off) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kPodRef; + e.enc = DictEntryEnc::kWindowed; + e.df = df; + e.ttf_delta = df * 2; + e.max_freq = 5; + e.frq_off_delta = frq_off; + e.frq_len = 200; + e.prelude_len = 10; // 0 < prelude_len <= frq_docs_len + e.frq_docs_len = 120; // <= frq_len + e.prx_off_delta = 0; + e.prx_len = 50; + return e; +} + +DictEntry MakeInline(std::string term, uint32_t df) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kInline; + e.enc = DictEntryEnc::kSlim; + e.df = df; + e.ttf_delta = df * 3; + e.max_freq = 7; + e.frq_bytes = {0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80}; + e.inline_dd_disk_len = 5; // <= frq_bytes.size() + e.dd_meta.uncomp_len = 12; + e.freq_meta.uncomp_len = 6; + e.prx_bytes = {0xAA, 0xBB, 0xCC}; + return e; +} + +// Full field comparison of two DECODED entries (used to prove the key-first path +// produces byte-identical output to the full / decode_all path). +void ExpectEntryEq(const DictEntry& a, const DictEntry& b) { + EXPECT_EQ(a.term, b.term); + EXPECT_EQ(a.kind, b.kind); + EXPECT_EQ(a.enc, b.enc); + EXPECT_EQ(a.has_sb, b.has_sb); + EXPECT_EQ(a.df, b.df); + EXPECT_EQ(a.ttf_delta, b.ttf_delta); + EXPECT_EQ(a.max_freq, b.max_freq); + EXPECT_EQ(a.frq_off_delta, b.frq_off_delta); + EXPECT_EQ(a.frq_len, b.frq_len); + EXPECT_EQ(a.prelude_len, b.prelude_len); + EXPECT_EQ(a.frq_docs_len, b.frq_docs_len); + EXPECT_EQ(a.prx_off_delta, b.prx_off_delta); + EXPECT_EQ(a.prx_len, b.prx_len); + EXPECT_EQ(a.inline_dd_disk_len, b.inline_dd_disk_len); + EXPECT_EQ(a.frq_bytes, b.frq_bytes); + EXPECT_EQ(a.prx_bytes, b.prx_bytes); + EXPECT_EQ(a.dd_meta.zstd, b.dd_meta.zstd); + EXPECT_EQ(a.dd_meta.uncomp_len, b.dd_meta.uncomp_len); + EXPECT_EQ(a.dd_meta.disk_len, b.dd_meta.disk_len); + EXPECT_EQ(a.dd_meta.crc, b.dd_meta.crc); + EXPECT_EQ(a.dd_meta.verify_crc, b.dd_meta.verify_crc); + EXPECT_EQ(a.freq_meta.zstd, b.freq_meta.zstd); + EXPECT_EQ(a.freq_meta.uncomp_len, b.freq_meta.uncomp_len); + EXPECT_EQ(a.freq_meta.disk_len, b.freq_meta.disk_len); + EXPECT_EQ(a.freq_meta.crc, b.freq_meta.crc); + EXPECT_EQ(a.freq_meta.verify_crc, b.freq_meta.verify_crc); +} + +std::vector BuildBlock(const std::vector& entries, IndexTier tier, + bool has_positions, uint64_t frq_base, uint64_t prx_base, + uint32_t anchor_interval) { + DictBlockBuilder builder(tier, has_positions, frq_base, prx_base, anchor_interval); + for (const auto& e : entries) { + builder.add_entry(e); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +// A no-op key predicate (accept all); passing an empty std::function exercises +// the accept-all short-circuit in visit_prefix_range. +const std::function kAcceptAll {}; + +// Collects the streamed entries' terms; never stops early. +std::function collect_terms(std::vector* out) { + return [out](DictEntry&& e, bool* stop) -> Status { + out->push_back(e.term); + *stop = false; + return Status::OK(); + }; +} + +// ---- F-1: key/rest split and decode_all are byte-identical ---- + +// Entry level: decode_dict_entry (full) vs decode_dict_entry_key + _rest, over a +// front-coded sequence mixing pod_ref(slim/windowed) and inline entries. +TEST(SniiB2T07Dict, KeyRestSplitDecodesByteIdenticalToFull) { + const IndexTier tier = IndexTier::kT2; + const std::vector entries = { + MakePodRef("alpha", 3, 0, 0), + MakeInline("alphabet", 5), + MakePodRefWindowed("beta", 7, 100), + MakePodRef("betas", 9, 300, 80), + MakeInline("gamma", 2), + }; + + ByteSink sink; + std::string prev; + for (const auto& e : entries) { + assert_ok(encode_dict_entry(e, prev, tier, &sink)); + prev = e.term; + } + const std::vector buf = sink.buffer(); + + // Baseline: full decode. + std::vector full; + { + ByteSource src {Slice(buf)}; + std::string p; + for (size_t i = 0; i < entries.size(); ++i) { + DictEntry e; + assert_ok(decode_dict_entry(&src, p, tier, &e)); + p = e.term; + full.push_back(std::move(e)); + } + EXPECT_TRUE(src.eof()); + } + + // key + rest decode: byte-identical fields, exactly one body decode per entry. + reset_dict_entry_counters(); + std::vector split; + { + ByteSource src {Slice(buf)}; + std::string p; + for (size_t i = 0; i < entries.size(); ++i) { + DictEntry e; + size_t body_start = 0; + uint64_t total = 0; + assert_ok(decode_dict_entry_key(&src, p, &e, &body_start, &total)); + assert_ok(decode_dict_entry_rest(&src, tier, body_start, total, &e)); + p = e.term; + split.push_back(std::move(e)); + } + EXPECT_TRUE(src.eof()); + } + EXPECT_EQ(dict_entry_body_decode_count(), entries.size()); + + ASSERT_EQ(full.size(), split.size()); + for (size_t i = 0; i < full.size(); ++i) { + ExpectEntryEq(full[i], split[i]); + } +} + +// skip_dict_entry_body advances across entries (front coding still rebuilds each +// term) without materializing any body. +TEST(SniiB2T07Dict, SkipDictEntryBodyAdvancesAndCountsNoBody) { + const IndexTier tier = IndexTier::kT2; + const std::vector entries = {MakePodRef("aa", 1, 0, 0), MakeInline("ab", 2), + MakePodRefWindowed("ac", 3, 50)}; + ByteSink sink; + std::string prev; + for (const auto& e : entries) { + assert_ok(encode_dict_entry(e, prev, tier, &sink)); + prev = e.term; + } + const std::vector buf = sink.buffer(); + + reset_dict_entry_counters(); + ByteSource src {Slice(buf)}; + std::string p; + std::vector terms; + for (size_t i = 0; i < entries.size(); ++i) { + DictEntry e; + size_t body_start = 0; + uint64_t total = 0; + assert_ok(decode_dict_entry_key(&src, p, &e, &body_start, &total)); + terms.push_back(e.term); + assert_ok(skip_dict_entry_body(&src, body_start, total)); + p = e.term; + } + EXPECT_TRUE(src.eof()); + EXPECT_EQ(dict_entry_body_decode_count(), 0U); // keys only + ASSERT_EQ(terms.size(), 3U); + EXPECT_EQ(terms[0], "aa"); + EXPECT_EQ(terms[1], "ab"); + EXPECT_EQ(terms[2], "ac"); +} + +// Block level: visit_prefix_range("") streams every entry, byte-identical to +// decode_all, across multiple anchor segments and mixed entry kinds. +TEST(SniiB2T07Dict, VisitPrefixRangeEmptyPrefixMatchesDecodeAll) { + std::vector entries; + for (int i = 0; i < 40; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "term_%02d", i); + if (i % 3 == 0) { + entries.push_back(MakeInline(buf, static_cast(i + 1))); + } else if (i % 3 == 1) { + entries.push_back(MakePodRef(buf, static_cast(i + 1), + static_cast(i) * 10, + static_cast(i) * 4)); + } else { + entries.push_back(MakePodRefWindowed(buf, static_cast(i + 1), + static_cast(i) * 10)); + } + } + const std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 8192, 16384, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader)); + + std::vector golden; + assert_ok(reader.decode_all(&golden)); + + std::vector streamed; + bool exhausted = false; + assert_ok(reader.visit_prefix_range( + "", kAcceptAll, + [&](DictEntry&& e, bool* stop) -> Status { + streamed.push_back(std::move(e)); + *stop = false; + return Status::OK(); + }, + &exhausted)); + EXPECT_FALSE(exhausted); // the empty prefix never leaves the range + ASSERT_EQ(streamed.size(), golden.size()); + for (size_t i = 0; i < golden.size(); ++i) { + ExpectEntryEq(golden[i], streamed[i]); + } +} + +// ---- F-2 / F-3: find_term materializes only the matched body ---- + +TEST(SniiB2T07Dict, FindTermDecodesOnlyMatchedBody) { + std::vector entries; + for (int i = 0; i < 16; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "k%02d", i); + entries.push_back( + MakePodRef(buf, static_cast(i + 1), static_cast(i) * 8)); + } + const std::vector bytes = + BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); // 1 anchor + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + // First / middle / last entry: each materializes exactly one body even though + // the scan walks up to 16 keys. + for (const char* t : {"k00", "k07", "k15"}) { + reset_dict_entry_counters(); + bool found = false; + DictEntry out; + assert_ok(reader.find_term(t, &found, &out)); + EXPECT_TRUE(found) << t; + EXPECT_EQ(out.term, t) << t; + EXPECT_EQ(dict_entry_body_decode_count(), 1U) << t; + } +} + +TEST(SniiB2T07Dict, FindTermMissDecodesNoBody) { + std::vector entries; // k00, k02, ... k30 -> gaps between keys + for (int i = 0; i < 32; i += 2) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "k%02d", i); + entries.push_back(MakePodRef(buf, 1, static_cast(i) * 4)); + } + ASSERT_EQ(entries.size(), 16U); + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + bool found = true; + DictEntry out; + + // Gap inside the range: stop at the first key past target, no body decode. + reset_dict_entry_counters(); + assert_ok(reader.find_term("k15", &found, &out)); + EXPECT_FALSE(found); + EXPECT_EQ(dict_entry_body_decode_count(), 0U); + + // Below the first anchor term: rejected by the anchor search, no scan at all. + reset_dict_entry_counters(); + found = true; + assert_ok(reader.find_term("a", &found, &out)); + EXPECT_FALSE(found); + EXPECT_EQ(dict_entry_body_decode_count(), 0U); + + // Above the last term: scan keys to eof, still no body decode. + reset_dict_entry_counters(); + found = true; + assert_ok(reader.find_term("z", &found, &out)); + EXPECT_FALSE(found); + EXPECT_EQ(dict_entry_body_decode_count(), 0U); +} + +// find_term's matched body is byte-identical to the decode_all entry, across +// anchor boundaries (mixed kinds). +TEST(SniiB2T07Dict, FindTermReturnsEntryMatchingDecodeAll) { + std::vector entries; + for (int i = 0; i < 40; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "term_%02d", i); + if (i % 2 == 0) { + entries.push_back(MakeInline(buf, static_cast(i + 1))); + } else { + entries.push_back(MakePodRef(buf, static_cast(i + 1), + static_cast(i) * 10, + static_cast(i) * 4)); + } + } + const std::vector bytes = BuildBlock(entries, IndexTier::kT2, true, 4096, 8192, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT2, true, &reader)); + + std::vector golden; + assert_ok(reader.decode_all(&golden)); + for (const auto& g : golden) { + bool found = false; + DictEntry out; + assert_ok(reader.find_term(g.term, &found, &out)); + ASSERT_TRUE(found) << g.term; + ExpectEntryEq(g, out); + } +} + +// ---- F-6 / anchor-jump / accept_key / early-stop on visit_prefix_range ---- + +TEST(SniiB2T07Dict, VisitPrefixRangeStopsAtBoundaryDecodingOnlyMatches) { + const std::vector entries = { + MakePodRef("aa1", 1, 0), MakePodRef("aa2", 1, 10), MakePodRef("ab1", 1, 20), + MakeInline("ab2", 2), MakePodRef("ab3", 1, 40), MakePodRef("ac1", 1, 50), + MakePodRef("ac2", 1, 60), + }; + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + reset_dict_entry_counters(); + std::vector got; + bool exhausted = false; + assert_ok(reader.visit_prefix_range("ab", kAcceptAll, collect_terms(&got), &exhausted)); + EXPECT_TRUE(exhausted); // saw "ac1" past the range + ASSERT_EQ(got.size(), 3U); + EXPECT_EQ(got[0], "ab1"); + EXPECT_EQ(got[1], "ab2"); + EXPECT_EQ(got[2], "ab3"); + // aa* skipped (key-only), ac* never reached -> exactly the 3 ab* bodies. + EXPECT_EQ(dict_entry_body_decode_count(), 3U); +} + +TEST(SniiB2T07Dict, VisitPrefixRangeAnchorJumpSkipsPrePrefixSegments) { + std::vector entries; + for (int i = 0; i < 32; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "aa%03d", i); + entries.push_back(MakePodRef(buf, 1, static_cast(i))); + } + for (int i = 0; i < 4; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "zz%03d", i); + entries.push_back(MakePodRef(buf, 1, 1000 + static_cast(i))); + } + // anchor_interval=8 -> several whole "aa*" anchor segments are jumped, never + // decoded; the segment straddling the prefix is scanned key-only. + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 8); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + reset_dict_entry_counters(); + std::vector got; + bool exhausted = false; + assert_ok(reader.visit_prefix_range("zz", kAcceptAll, collect_terms(&got), &exhausted)); + ASSERT_EQ(got.size(), 4U); + EXPECT_EQ(got[0], "zz000"); + EXPECT_EQ(got[3], "zz003"); + EXPECT_FALSE(exhausted); // zz003 is the last entry -> block ended in range + EXPECT_EQ(dict_entry_body_decode_count(), 4U); // only the zz* bodies +} + +TEST(SniiB2T07Dict, VisitPrefixRangeAcceptKeySkipsRejectedBodies) { + std::vector entries; + for (int i = 0; i < 10; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "ab%d", i); + entries.push_back(MakePodRef(buf, 1, static_cast(i) * 5)); + } + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + reset_dict_entry_counters(); + std::vector got; + bool exhausted = false; + const std::function accept_even = [](std::string_view t) { + return ((t.back() - '0') % 2) == 0; + }; + assert_ok(reader.visit_prefix_range("ab", accept_even, collect_terms(&got), &exhausted)); + ASSERT_EQ(got.size(), 5U); // ab0, ab2, ab4, ab6, ab8 + EXPECT_EQ(got[0], "ab0"); + EXPECT_EQ(got[4], "ab8"); + // Rejected keys skip their body: 5 bodies, not 10. + EXPECT_EQ(dict_entry_body_decode_count(), 5U); +} + +TEST(SniiB2T07Dict, VisitPrefixRangeEarlyStopAfterK) { + std::vector entries; + for (int i = 0; i < 10; ++i) { + char buf[16]; + std::snprintf(buf, sizeof(buf), "ab%d", i); + entries.push_back(MakePodRef(buf, 1, static_cast(i) * 5)); + } + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 0, 0, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + reset_dict_entry_counters(); + std::vector got; + bool exhausted = false; + int n = 0; + assert_ok(reader.visit_prefix_range( + "ab", kAcceptAll, + [&](DictEntry&& e, bool* stop) -> Status { + got.push_back(e.term); + ++n; + *stop = (n >= 3); + return Status::OK(); + }, + &exhausted)); + ASSERT_EQ(got.size(), 3U); + EXPECT_FALSE(exhausted); // stopped by the visitor, not a boundary + EXPECT_EQ(dict_entry_body_decode_count(), 3U); // no bodies past the stop +} + +// ---- F-5 / F-7: single-entry block + corrupt entry_len ---- + +TEST(SniiB2T07Dict, SingleEntryBlockVisitAndFind) { + const std::vector entries = {MakePodRef("solo", 7, 0)}; + const std::vector bytes = BuildBlock(entries, IndexTier::kT1, false, 4096, 0, 16); + DictBlockReader reader; + assert_ok(DictBlockReader::open(Slice(bytes), IndexTier::kT1, false, &reader)); + + reset_dict_entry_counters(); + bool found = false; + DictEntry out; + assert_ok(reader.find_term("solo", &found, &out)); + EXPECT_TRUE(found); + EXPECT_EQ(out.term, "solo"); + EXPECT_EQ(dict_entry_body_decode_count(), 1U); + + reset_dict_entry_counters(); + std::vector got; + bool exhausted = false; + assert_ok(reader.visit_prefix_range("so", kAcceptAll, collect_terms(&got), &exhausted)); + ASSERT_EQ(got.size(), 1U); + EXPECT_EQ(got[0], "solo"); + EXPECT_FALSE(exhausted); + EXPECT_EQ(dict_entry_body_decode_count(), 1U); +} + +TEST(SniiB2T07Dict, DecodeDictEntryKeyRejectsCorruptEntryLen) { + // entry_len varint decodes to 127 but only 3 bytes remain after it. + const std::vector buf = {0x7F, 0x00, 0x00, 0x00}; + ByteSource src {Slice(buf)}; + DictEntry e; + size_t body_start = 0; + uint64_t total = 0; + const Status s = decode_dict_entry_key(&src, "", &e, &body_start, &total); + EXPECT_FALSE(s.ok()); +} + +// ---- reader / query level: results unchanged + body-decode bounded ---- + +// prefix_terms routes through the rewired visit_prefix_range; it returns the same +// ordered hits and materializes only the matched bodies. +TEST(SniiB2T07DictReader, PrefixTermsResultsUnchangedAndStreamBounded) { + ScopedEnv resident_bsbf("SNII_BSBF_RESIDENT_MAX", "1048576"); + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &seg, &idx)); + + std::vector hits; + reset_dict_entry_counters(); + assert_ok(idx.prefix_terms("ord", &hits, 10)); + ASSERT_EQ(hits.size(), 2U); + EXPECT_EQ(hits[0].term, "order"); + EXPECT_EQ(hits[1].term, "ordinal"); + // Resident block: only the two prefix matches' bodies are materialized. + EXPECT_EQ(dict_entry_body_decode_count(), 2U); +} + +TEST(SniiB2T07DictReader, PrefixTermsEmptyPrefixEnumeratesAllSorted) { + ScopedEnv resident_bsbf("SNII_BSBF_RESIDENT_MAX", "1048576"); + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &seg, &idx)); + + std::vector hits; + assert_ok(idx.prefix_terms("", &hits, 0)); + std::vector terms; + terms.reserve(hits.size()); + for (const auto& h : hits) { + terms.push_back(h.term); + } + const std::vector expected = {"123", "almost", "driver", "failed", + "needle", "order", "ordinal", "repeat", + "sparse_left", "sparse_right", "trace"}; + EXPECT_EQ(terms, expected); +} + +TEST(SniiB2T07DictReader, LookupDecodesOnlyMatchedBody) { + ScopedEnv resident_bsbf("SNII_BSBF_RESIDENT_MAX", "1048576"); + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &seg, &idx)); + + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + reset_dict_entry_counters(); + assert_ok(idx.lookup("trace", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + EXPECT_EQ(entry.term, "trace"); + EXPECT_EQ(dict_entry_body_decode_count(), 1U); // only "trace" body materialized +} + +// End-to-end: a prefix query over the rewired streaming path equals the union of +// the expanded terms' postings. +TEST(SniiB2T07DictReader, PrefixQueryMatchesTermUnion) { + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(build_reader(&file, &seg, &idx)); + + std::vector prefix_docids; + assert_ok(query::prefix_query(idx, "ord", &prefix_docids)); + + std::vector order_docids; + std::vector ordinal_docids; + assert_ok(query::term_query(idx, "order", &order_docids)); + assert_ok(query::term_query(idx, "ordinal", &ordinal_docids)); + + std::vector expected; + std::set_union(order_docids.begin(), order_docids.end(), ordinal_docids.begin(), + ordinal_docids.end(), std::back_inserter(expected)); + EXPECT_EQ(prefix_docids, expected); +} + +} // namespace +} // namespace doris::snii diff --git a/be/test/storage/index/snii_b2_t08_wildcard_test.cpp b/be/test/storage/index/snii_b2_t08_wildcard_test.cpp new file mode 100644 index 00000000000000..87ce313c99f71b --- /dev/null +++ b/be/test/storage/index/snii_b2_t08_wildcard_test.cpp @@ -0,0 +1,359 @@ +// 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. + +// T08 -- wildcard matcher scratch reuse. +// +// Proves the request-scoped internal::WildcardMatcher (a) matches bit-for-bit +// identically to the former per-call DP in wildcard_query.cpp, and (b) reuses its +// two DP scratch rows across every visited term so a whole-dictionary scan +// performs O(1) heap allocations (<= 2) instead of O(2N). A header-only +// CountingAllocator gives the deterministic allocation counts; a byte-for-byte +// copy of the original DP serves as the equivalence oracle. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/wildcard_matcher.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii::query { +using doris::Status; // RETURN_IF_ERROR / Status::OK() expand to a bare Status. +namespace { + +// Shared reader/writer fixtures live in snii_query_test_util.h; pull the ones +// this suite needs into scope unqualified. +using snii_test::assert_ok; +using snii_test::build_reader; +using snii_test::MemoryFile; + +// Minimal counting allocator: every allocate() bumps a per-type static counter so +// tests can assert exact heap-allocation counts without overriding global new. +// Stateless (all instances compare equal), so std::vector::swap stays a pointer +// swap that performs no allocation -- exactly what the matcher relies on. +template +struct CountingAllocator { + using value_type = T; + + CountingAllocator() noexcept = default; + // Converting (rebind) constructor: intentionally non-explicit, as required by + // the Allocator named requirement / std::allocator_traits. + template + CountingAllocator(const CountingAllocator& /*other*/) noexcept {} // NOLINT(*-explicit-*) + + T* allocate(std::size_t n) { + ++s_total_allocs; + ++s_live; + return static_cast(::operator new(n * sizeof(T))); + } + void deallocate(T* p, std::size_t /*n*/) noexcept { + --s_live; + ::operator delete(p); + } + + template + bool operator==(const CountingAllocator& /*other*/) const noexcept { + return true; + } + template + bool operator!=(const CountingAllocator& /*other*/) const noexcept { + return false; + } + + static void reset() { + s_total_allocs = 0; + s_live = 0; + } + static std::size_t total_allocs() { return s_total_allocs; } + static std::size_t live() { return s_live; } + + static inline std::size_t s_total_allocs = 0; + static inline std::size_t s_live = 0; +}; + +// Byte-for-byte copy of the former wildcard_query.cpp DP (templated only so the +// baseline-characterization test can count its per-call allocations). This is the +// equivalence oracle the optimized matcher must reproduce exactly. +template > +bool wildcard_match_dp_reference(std::string_view pattern, std::string_view text) { + std::vector prev(text.size() + 1, 0); + std::vector curr(text.size() + 1, 0); + prev[0] = 1; + + for (char p : pattern) { + std::fill(curr.begin(), curr.end(), 0); + if (p == '*') { + curr[0] = prev[0]; + for (size_t i = 1; i <= text.size(); ++i) { + curr[i] = prev[i] || curr[i - 1]; + } + } else { + for (size_t i = 1; i <= text.size(); ++i) { + curr[i] = prev[i - 1] && (p == '?' || p == text[i - 1]); + } + } + prev.swap(curr); + } + return prev[text.size()] != 0; +} + +// All strings of length 0..max_len over `alphabet`, for exhaustive equivalence. +std::vector all_strings_up_to(std::string_view alphabet, size_t max_len) { + std::vector out; + out.emplace_back(); + size_t level_begin = 0; + for (size_t len = 1; len <= max_len; ++len) { + const size_t level_end = out.size(); + for (size_t i = level_begin; i < level_end; ++i) { + for (char c : alphabet) { + out.push_back(out[i] + c); + } + } + level_begin = level_end; + } + return out; +} + +// `count` terms whose first (warmup) entry is the longest (`max_len`); every later +// term is <= max_len, so a matcher that reuses scratch reallocates only on the +// first call. The lengths still vary across terms (cycling 0..max_len). +std::vector make_varied_length_terms(size_t count, size_t max_len) { + std::vector terms; + terms.reserve(count); + terms.push_back(std::string(max_len, 'a')); + for (size_t i = 1; i < count; ++i) { + const size_t len = i % (max_len + 1); + terms.push_back(std::string(len, static_cast('a' + (i % 26)))); + } + return terms; +} + +// W-EQ-DP: the optimized matcher reproduces the reference DP bit-for-bit over an +// exhaustive small-alphabet battery (covers "", leading/trailing '*'/'?', +// consecutive "**", '?' interplay) plus realistic dictionary patterns/terms. One +// matcher is reused across all terms of a pattern, so this also proves scratch +// reuse never corrupts a result. +TEST(SniiWildcardQueryTest, MatcherEquivalentToReferenceDp) { + const std::vector exhaustive_patterns = all_strings_up_to("ab*?", 4); + const std::vector exhaustive_texts = all_strings_up_to("ab", 4); + for (const std::string& pattern : exhaustive_patterns) { + internal::WildcardMatcher<> matcher(pattern); + for (const std::string& text : exhaustive_texts) { + EXPECT_EQ(matcher(text), wildcard_match_dp_reference(pattern, text)) + << "pattern=\"" << pattern << "\" text=\"" << text << "\""; + } + } + + const std::vector patterns = { + "", "*", "?", "ord*", "?rder", "*failed*order*", "ordinal", + "order", "**a**", "a?b", "*a*", "?ailed", "sparse_*", "*_left"}; + const std::vector terms = {"", + "order", + "ordinal", + "failed", + "needle", + "driver", + "almost", + "123", + "repeat", + "sparse_left", + "sparse_right", + "trace", + "ordering", + std::string(40, 'a')}; + for (const std::string& pattern : patterns) { + internal::WildcardMatcher<> matcher(pattern); + for (const std::string& text : terms) { + EXPECT_EQ(matcher(text), wildcard_match_dp_reference(pattern, text)) + << "pattern=\"" << pattern << "\" text=\"" << text << "\""; + } + } +} + +// W-EMPTY-PAT: an empty pattern matches only the empty string. +TEST(SniiWildcardQueryTest, EmptyPatternMatchesOnlyEmptyText) { + internal::WildcardMatcher<> matcher(""); + EXPECT_TRUE(matcher("")); + EXPECT_FALSE(matcher("a")); +} + +// W-STAR-ONLY: "*" matches the empty string and any non-empty string. +TEST(SniiWildcardQueryTest, StarMatchesEverything) { + internal::WildcardMatcher<> matcher("*"); + EXPECT_TRUE(matcher("")); + EXPECT_TRUE(matcher("x")); + EXPECT_TRUE(matcher("xyz")); +} + +// W-QMARK: "?" matches exactly one byte. +TEST(SniiWildcardQueryTest, QuestionMarkMatchesExactlyOneByte) { + internal::WildcardMatcher<> matcher("?"); + EXPECT_FALSE(matcher("")); + EXPECT_TRUE(matcher("a")); + EXPECT_FALSE(matcher("ab")); +} + +// W-CONSEC-STAR: consecutive '*' degrade gracefully. +TEST(SniiWildcardQueryTest, ConsecutiveStars) { + internal::WildcardMatcher<> matcher("**a**"); + EXPECT_TRUE(matcher("a")); + EXPECT_TRUE(matcher("xax")); + EXPECT_FALSE(matcher("b")); +} + +// W-ANCHOR: a literal pattern is anchored at both ends (full match only). +TEST(SniiWildcardQueryTest, LiteralIsFullyAnchored) { + internal::WildcardMatcher<> matcher("ab"); + EXPECT_TRUE(matcher("ab")); + EXPECT_FALSE(matcher("abc")); + EXPECT_FALSE(matcher("xab")); +} + +// Perf (deterministic): the matcher allocates its two scratch rows once and reuses +// them across every term -- total heap allocations stay <= 2 and are independent +// of the term count N. +TEST(SniiWildcardQueryTest, MatcherReusesScratchAcrossTerms) { + using Alloc = CountingAllocator; + + auto allocs_for_n = [](size_t n) { + Alloc::reset(); + internal::WildcardMatcher matcher("*a*"); + for (const std::string& term : make_varied_length_terms(n, /*max_len=*/64)) { + matcher(term); + } + return Alloc::total_allocs(); + }; + + EXPECT_LE(allocs_for_n(1000), 2U); + // N-independent: exactly the two scratch rows, whether N=10 or N=1000. + EXPECT_EQ(allocs_for_n(10), 2U); + EXPECT_EQ(allocs_for_n(1000), 2U); + EXPECT_EQ(Alloc::live(), 0U); // matcher destroyed each lambda call: no leak. +} + +// Perf (deterministic): once warmed up with the longest term, the scratch capacity +// never changes across subsequent shorter/equal-length terms (no realloc). +TEST(SniiWildcardQueryTest, MatcherScratchCapacityStable) { + internal::WildcardMatcher<> matcher("*x*"); + matcher(std::string(64, 'a')); // warmup with the longest term + const size_t cap_after_warmup = matcher.scratch_capacity(); + EXPECT_GE(cap_after_warmup, 65U); + for (size_t len : {size_t {0}, size_t {1}, size_t {7}, size_t {63}, size_t {64}}) { + matcher(std::string(len, 'b')); + EXPECT_EQ(matcher.scratch_capacity(), cap_after_warmup); + } +} + +// Perf baseline (deterministic, contrast): the former per-call DP constructs two +// std::vectors every call, so N terms cost exactly 2N allocations -- the cost the +// reused matcher above eliminates. +TEST(SniiWildcardQueryTest, PerCallDpReferenceAllocatesTwicePerTerm) { + using Alloc = CountingAllocator; + Alloc::reset(); + const std::vector terms = make_varied_length_terms(/*count=*/100, /*max_len=*/64); + for (const std::string& term : terms) { + wildcard_match_dp_reference("*a*", term); + } + EXPECT_EQ(Alloc::total_allocs(), 2U * terms.size()); + EXPECT_EQ(Alloc::live(), 0U); +} + +// W-RESULT: end-to-end, "ord*" returns the sorted deduplicated union of the +// "order" and "ordinal" docid sets (independently computed expected set). +TEST(SniiWildcardQueryTest, WildcardResultIsTermUnion) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector order_docids; + std::vector ordinal_docids; + assert_ok(term_query(index_reader, "order", &order_docids)); + assert_ok(term_query(index_reader, "ordinal", &ordinal_docids)); + std::vector expected; + std::set_union(order_docids.begin(), order_docids.end(), ordinal_docids.begin(), + ordinal_docids.end(), std::back_inserter(expected)); + + std::vector docids; + assert_ok(wildcard_query(index_reader, "ord*", &docids)); + EXPECT_EQ(docids, expected); +} + +// W-QMARK-FULL: a leading '?' forces a full-dictionary scan; "?rder" matches only +// the 5-byte "order" term (not 7-byte "ordinal"). +TEST(SniiWildcardQueryTest, LeadingQuestionMarkResult) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector order_docids; + assert_ok(term_query(index_reader, "order", &order_docids)); + + std::vector docids; + assert_ok(wildcard_query(index_reader, "?rder", &docids)); + EXPECT_EQ(docids, order_docids); +} + +// W-MAXEXP: max_expansions caps the number of expanded terms; terms enumerate in +// sorted order, so "*" with max_expansions=1 yields only the first term "123". +TEST(SniiWildcardQueryTest, MaxExpansionsCapsExpansion) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector first_term_docids; + assert_ok(term_query(index_reader, "123", &first_term_docids)); + + std::vector docids; + assert_ok(wildcard_query(index_reader, "*", &docids, /*max_expansions=*/1)); + EXPECT_EQ(docids, first_term_docids); +} + +// W-NULL-OUT / W-NULL-SINK: null output and null sink return InvalidArgument +// (no crash, no throw). +TEST(SniiWildcardQueryTest, NullArgumentsReturnInvalidArgument) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector* const null_docids = nullptr; + EXPECT_TRUE(wildcard_query(index_reader, "a*", null_docids) + .is()); + + DocIdSink* const null_sink = nullptr; + EXPECT_TRUE( + wildcard_query(index_reader, "a*", null_sink).is()); +} + +} // namespace +} // namespace doris::snii::query diff --git a/be/test/storage/index/snii_b2_t10_cover_test.cpp b/be/test/storage/index/snii_b2_t10_cover_test.cpp new file mode 100644 index 00000000000000..94f1ecf2a8fdb8 --- /dev/null +++ b/be/test/storage/index/snii_b2_t10_cover_test.cpp @@ -0,0 +1,376 @@ +// 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. + +// T10 -- select_covering_windows as a monotonic two-pointer cursor (O(C + N)). +// Covers: (a) the isolated cursor core select_covering_windows_cursor and the +// FrqPreludeReader::select_covering_windows member produce results element-for-element +// identical to the legacy per-candidate locate_window + run-collapse oracle; (b) the +// window_probe_count() seam shows the cursor is O(C + N) (not O(C * N)) and bounded by +// C + N regardless of group_size, while the legacy in-block rescan grows with G; (c) the +// packed win_last_docid_ catalogue is byte-identical to WindowMeta.last_docid; and (d) the +// wired phrase query path is unchanged. + +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii { +namespace { + +using namespace snii_test; // assert_ok, build_reader, MemoryFile, ... + +using format::build_frq_prelude; +using format::FrqPreludeColumns; +using format::FrqPreludeReader; +using format::select_covering_windows_cursor; +using format::WindowMeta; + +// Trusted oracle: the legacy locate_window-per-candidate + run-collapse semantics, +// computed with a dumb per-candidate linear scan (first window whose last_docid >= d). +std::vector oracle_select(const std::vector& win_last_docid, + const std::vector& candidates) { + std::vector out; + uint32_t last = UINT32_MAX; + for (uint32_t d : candidates) { + bool found = false; + uint32_t hit = 0; + for (uint32_t w = 0; w < win_last_docid.size(); ++w) { + if (d <= win_last_docid[w]) { + found = true; + hit = w; + break; + } + } + if (!found) continue; + if (hit != last) { + out.push_back(hit); + last = hit; + } + } + return out; +} + +// Derives sb_last_docid exactly as FrqPreludeReader::open does: the absolute last docid +// of each super-block's last window. +std::vector derive_sb_last_docid(const std::vector& win_last_docid, + uint32_t group_size) { + std::vector sb; + const uint32_t n = static_cast(win_last_docid.size()); + if (n == 0) return sb; + const uint32_t n_super = (n + group_size - 1) / group_size; + sb.reserve(n_super); + for (uint32_t s = 0; s < n_super; ++s) { + const uint32_t last_win = std::min((s + 1) * group_size, n) - 1; + sb.push_back(win_last_docid[last_win]); + } + return sb; +} + +// Runs the pure cursor core over directory arrays derived from win_last_docid. +std::vector run_cursor(const std::vector& win_last_docid, uint32_t group_size, + const std::vector& candidates) { + const std::vector sb = derive_sb_last_docid(win_last_docid, group_size); + std::vector out; + select_covering_windows_cursor(win_last_docid.data(), + static_cast(win_last_docid.size()), sb.data(), + static_cast(sb.size()), group_size, candidates, &out); + return out; +} + +// Builds a real prelude from strictly-increasing absolute last docids (doc_count=1 and +// contiguous 1-byte dd/freq regions satisfy the prelude width / layout validators). +void make_test_prelude(const std::vector& last_docids, uint32_t group_size, + FrqPreludeReader* reader) { + FrqPreludeColumns cols; + cols.has_freq = true; + cols.has_prx = false; + cols.group_size = group_size; + uint64_t dd_running = 0; + uint64_t freq_running = 0; + for (uint32_t v : last_docids) { + WindowMeta m; + m.last_docid = v; + m.doc_count = 1; + m.dd_off = dd_running; + m.dd_disk_len = 1; + m.dd_uncomp_len = 1; + m.crc_dd = v; + dd_running += 1; + m.freq_off = freq_running; + m.freq_disk_len = 1; + m.freq_uncomp_len = 1; + m.crc_freq = v + 1; + freq_running += 1; + cols.windows.push_back(m); + } + ByteSink sink; + assert_ok(build_frq_prelude(cols, &sink)); + assert_ok(FrqPreludeReader::open(sink.view(), reader)); +} + +// Production oracle: per-candidate locate_window + run-collapse over a real prelude. +std::vector locate_window_select(const FrqPreludeReader& prelude, + const std::vector& candidates) { + std::vector out; + uint32_t last = UINT32_MAX; + for (uint32_t d : candidates) { + bool found = false; + uint32_t w = 0; + assert_ok(prelude.locate_window(d, &found, &w)); + if (!found) continue; + if (w != last) { + out.push_back(w); + last = w; + } + } + return out; +} + +// FV-01: randomized equivalence of the cursor with the oracle across N and group_size. +TEST(SniiB2T10CoverTest, CursorMatchesOracleOnRandomAscendingSets) { + std::mt19937 rng(0x7710u); + const uint32_t kNs[] = {1, 2, 8, 64, 200, 4000}; + const uint32_t kGs[] = {1, 8, 64}; + std::uniform_int_distribution step(1, 5); // strictly increasing last_docid + std::uniform_int_distribution keep(0, 3); // ~1/4 of docids become candidates + for (uint32_t n : kNs) { + std::vector win_last(n); + uint32_t acc = 0; + for (uint32_t w = 0; w < n; ++w) { + acc += step(rng); + win_last[w] = acc; + } + const uint32_t max_docid = acc + 3; // a few candidates land past the last window + std::vector cands; + for (uint32_t d = 0; d <= max_docid; ++d) { + if (keep(rng) == 0) cands.push_back(d); // ascending, distinct, some same-window runs + } + const std::vector expected = oracle_select(win_last, cands); + for (uint32_t g : kGs) { + EXPECT_EQ(run_cursor(win_last, g, cands), expected) << "n=" << n << " g=" << g; + } + } +} + +// FV-02: the reader member (real prelude, production path) agrees with per-candidate +// locate_window over the same prelude. +TEST(SniiB2T10CoverTest, MemberMatchesRealLocateWindow) { + std::vector last_docids; + uint32_t acc = 0; + for (uint32_t w = 0; w < 300; ++w) { + acc += 1 + (w % 7); + last_docids.push_back(acc); + } + FrqPreludeReader prelude; + make_test_prelude(last_docids, /*group_size=*/64, &prelude); + + std::vector cands; + for (uint32_t d = 0; d <= acc + 5; d += 3) cands.push_back(d); + + std::vector got; + prelude.select_covering_windows(cands, &got); + EXPECT_EQ(got, locate_window_select(prelude, cands)); + EXPECT_EQ(got, oracle_select(last_docids, cands)); +} + +// FV-03: single window -- candidates inside emit {0}, candidates past it emit {}. +TEST(SniiB2T10CoverTest, SingleWindow) { + const std::vector win_last = {100}; + EXPECT_EQ(run_cursor(win_last, 64, {0, 50, 100}), (std::vector {0})); + EXPECT_EQ(run_cursor(win_last, 64, {0, 50, 100}), oracle_select(win_last, {0, 50, 100})); + EXPECT_TRUE(run_cursor(win_last, 64, {101, 200}).empty()); +} + +// FV-04: a candidate at docid 0 / at the first window's last docid still emits window 0. +TEST(SniiB2T10CoverTest, CandidateInFirstWindow) { + const std::vector win_last = {10, 20, 30}; + EXPECT_EQ(run_cursor(win_last, 2, {0, 10}), (std::vector {0})); + EXPECT_EQ(run_cursor(win_last, 2, {0, 10}), oracle_select(win_last, {0, 10})); +} + +// FV-05: candidates past the last window are dropped (early break == locate_window miss). +TEST(SniiB2T10CoverTest, CandidatesPastLastWindowAreDropped) { + const std::vector win_last = {10, 20, 30}; + const std::vector cands = {5, 25, 31, 40, 100}; // 5->w0, 25->w2, rest miss + EXPECT_EQ(run_cursor(win_last, 2, cands), (std::vector {0, 2})); + EXPECT_EQ(run_cursor(win_last, 2, cands), oracle_select(win_last, cands)); +} + +// FV-06: empty candidate list -> empty result (output is cleared first). +TEST(SniiB2T10CoverTest, EmptyCandidates) { + const std::vector win_last = {10, 20}; + const std::vector sb = derive_sb_last_docid(win_last, 64); + std::vector out = {99, 98, 97}; // pre-filled: must be cleared + select_covering_windows_cursor(win_last.data(), 2, sb.data(), static_cast(sb.size()), + 64, {}, &out); + EXPECT_TRUE(out.empty()); +} + +// FV-07: zero windows -> empty result, no crash (mirrors locate_window's empty early-out). +TEST(SniiB2T10CoverTest, EmptyWindows) { + const std::vector win_last; + EXPECT_TRUE(run_cursor(win_last, 64, {0, 5, 100}).empty()); +} + +// FV-08: sparse candidates crossing several super-block boundaries; cursor, member, and +// real locate_window all agree (boundary jump correctness). +TEST(SniiB2T10CoverTest, SuperBlockBoundaryCrossingSparse) { + std::vector last_docids; // 10 windows, stride 100 -> window w covers [w*100, ..+99] + for (uint32_t w = 0; w < 10; ++w) last_docids.push_back((w + 1) * 100 - 1); + FrqPreludeReader prelude; + make_test_prelude(last_docids, /*group_size=*/4, &prelude); // 3 super-blocks + const std::vector cands = {50, 450, 850, 999}; // windows 0, 4, 8, 9 + std::vector got; + prelude.select_covering_windows(cands, &got); + EXPECT_EQ(got, (std::vector {0, 4, 8, 9})); + EXPECT_EQ(got, locate_window_select(prelude, cands)); + EXPECT_EQ(run_cursor(last_docids, 4, cands), got); +} + +// FV-09: every window is covered -> the full {0..N-1} set, de-duplicated. +TEST(SniiB2T10CoverTest, DenseEveryWindowCovered) { + std::vector win_last; + std::vector cands; + std::vector expected; + for (uint32_t w = 0; w < 16; ++w) { + win_last.push_back(w); // window w covers exactly docid w + cands.push_back(w); + expected.push_back(w); + } + EXPECT_EQ(run_cursor(win_last, 4, cands), expected); + EXPECT_EQ(run_cursor(win_last, 4, cands), oracle_select(win_last, cands)); +} + +// FV-10: the packed win_last_docid_ catalogue is byte-identical to WindowMeta.last_docid. +TEST(SniiB2T10CoverTest, PackedLastDocidMatchesWindowMeta) { + std::vector last_docids; + uint32_t acc = 0; + for (uint32_t w = 0; w < 130; ++w) { // spans >1 super-block at G=64 + acc += 1 + (w % 3); + last_docids.push_back(acc); + } + FrqPreludeReader prelude; + make_test_prelude(last_docids, /*group_size=*/64, &prelude); + ASSERT_EQ(prelude.n_windows(), 130u); + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta m; + assert_ok(prelude.window(w, &m)); + EXPECT_EQ(prelude.window_last_docid(w), m.last_docid) << "w=" << w; + } +} + +// Deterministic complexity invariant: probe_count <= C + N for both dense and sparse +// candidate sets (i.e. O(C + N), far below O(C * N)); the sparse case additionally stays +// below N, proving the super-block jump avoids a full linear scan. +TEST(SniiB2T10CoverTest, CursorProbeCountStaysLinearInCandidatesPlusWindows) { + { + const uint32_t n = 256; // dense: C == N * stride + std::vector win_last(n); + for (uint32_t w = 0; w < n; ++w) win_last[w] = w * 4 + 3; + std::vector cands; + for (uint32_t d = 0; d < n * 4; ++d) cands.push_back(d); + format::testing::reset_window_probe_count(); + const std::vector got = run_cursor(win_last, 64, cands); + const uint64_t probes = format::testing::window_probe_count(); + EXPECT_LE(probes, static_cast(cands.size()) + n); + EXPECT_EQ(got, oracle_select(win_last, cands)); + } + { + const uint32_t n = 4000; // sparse: C << N + std::vector win_last(n); + for (uint32_t w = 0; w < n; ++w) win_last[w] = w; + const std::vector cands = {0, 1000, 2000, 3000, 3999}; + format::testing::reset_window_probe_count(); + const std::vector got = run_cursor(win_last, 64, cands); + const uint64_t probes = format::testing::window_probe_count(); + EXPECT_LE(probes, static_cast(cands.size()) + n); + EXPECT_LT(probes, n); // super-block jump prevents an O(N) linear walk + EXPECT_EQ(got, oracle_select(win_last, cands)); + } +} + +// Deterministic before/after contrast: the legacy per-candidate locate_window probe count +// grows with group_size and exceeds C + N, whereas the cursor stays <= C + N for every G +// and yields the same covering-window set regardless of G. +TEST(SniiB2T10CoverTest, LegacyProbeGrowsWithGroupSizeButCursorStaysLinear) { + const uint32_t n = 200; + std::vector last_docids(n); + std::vector cands(n); + for (uint32_t w = 0; w < n; ++w) { + last_docids[w] = w; // window w covers exactly docid w + cands[w] = w; // one candidate per window + } + const uint64_t cn = static_cast(cands.size()) + n; // C + N + + FrqPreludeReader prelude_g8; + FrqPreludeReader prelude_g64; + make_test_prelude(last_docids, 8, &prelude_g8); + make_test_prelude(last_docids, 64, &prelude_g64); + + format::testing::reset_window_probe_count(); + (void)locate_window_select(prelude_g8, cands); + const uint64_t legacy_g8 = format::testing::window_probe_count(); + + format::testing::reset_window_probe_count(); + (void)locate_window_select(prelude_g64, cands); + const uint64_t legacy_g64 = format::testing::window_probe_count(); + + EXPECT_GT(legacy_g64, legacy_g8); // deeper in-block rescans with larger G + EXPECT_GT(legacy_g8, cn); // even G=8 already exceeds C + N + EXPECT_GT(legacy_g64, cn); + + format::testing::reset_window_probe_count(); + const std::vector cur_g8 = run_cursor(last_docids, 8, cands); + const uint64_t cursor_g8 = format::testing::window_probe_count(); + + format::testing::reset_window_probe_count(); + const std::vector cur_g64 = run_cursor(last_docids, 64, cands); + const uint64_t cursor_g64 = format::testing::window_probe_count(); + + EXPECT_LE(cursor_g8, cn); // bounded by C + N regardless of G + EXPECT_LE(cursor_g64, cn); + EXPECT_LT(cursor_g64, legacy_g64); // cursor crushes the legacy probe count + EXPECT_EQ(cur_g8, cur_g64); // identical covering-window set across G + EXPECT_EQ(cur_g8, oracle_select(last_docids, cands)); +} + +// FV-11: the wired phrase query result is unchanged (no regression through the windowed +// docid-conjunction path that now uses the cursor member). +TEST(SniiB2T10CoverTest, WiredPhraseQueryResultUnchanged) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(query::phrase_query(index_reader, {"failed", "order"}, &docids)); + EXPECT_EQ(docids, (std::vector {5000, 7000, 8000})); +} + +} // namespace +} // namespace doris::snii diff --git a/be/test/storage/index/snii_clucene_guard_test.cpp b/be/test/storage/index/snii_clucene_guard_test.cpp new file mode 100644 index 00000000000000..da4daf7361ea23 --- /dev/null +++ b/be/test/storage/index/snii_clucene_guard_test.cpp @@ -0,0 +1,176 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +namespace fs = std::filesystem; + +const std::regex kCluceneRefPattern(R"(clucene|lucene::|CL_NS| locate_snii_core_dir() { + std::vector candidates; +#ifdef SNII_CORE_SOURCE_DIR + // Optional exact path injected by the build system (highest priority). If injected + // by CMake it must be a quoted string literal, e.g. + // add_definitions(-DSNII_CORE_SOURCE_DIR="${CMAKE_SOURCE_DIR}/src/storage/index/snii"). + candidates.emplace_back(SNII_CORE_SOURCE_DIR); +#endif + + std::error_code ec; + fs::path source_path(__FILE__); + if (source_path.is_relative()) { + source_path = fs::absolute(source_path, ec); + } + // Walk up from this test's source directory; matches either /be/... or + // /... so it is robust to layout changes without hardcoding the depth. + for (fs::path dir = source_path.parent_path(); !dir.empty();) { + candidates.push_back(dir / "be" / "src" / "storage" / "index" / "snii"); + candidates.push_back(dir / "src" / "storage" / "index" / "snii"); + if (dir == dir.root_path()) { + break; + } + dir = dir.parent_path(); + } + + // Fallback: the DORIS_HOME env var (consistent with existing test utilities). + if (const char* doris_home = std::getenv("DORIS_HOME"); doris_home != nullptr) { + candidates.push_back(fs::path(doris_home) / "be" / "src" / "storage" / "index" / "snii"); + } + + for (const auto& candidate : candidates) { + if (fs::is_directory(candidate, ec)) { + auto canonical = fs::weakly_canonical(candidate, ec); + return ec ? candidate : canonical; + } + } + return std::nullopt; +} + +struct CluceneHit { + std::string file; + int line = 0; + std::string text; +}; + +// Reads one file and collects matching lines; returns false if it cannot be opened +// (treated as a scan gap). +bool scan_file_for_clucene(const fs::path& file, std::vector* hits) { + std::ifstream input(file); + if (!input.is_open()) { + return false; + } + std::string line; + int line_no = 0; + while (std::getline(input, line)) { + ++line_no; + if (std::regex_search(line, kCluceneRefPattern)) { + hits->push_back({file.string(), line_no, line}); + } + } + return true; +} + +} // namespace + +// The SNII core engine (the be/src/storage/index/snii module subdirs) must contain +// zero CLucene references; CLucene is allowed only in the top-level Doris integration +// layer (snii_doris_adapter / snii_index_*). +TEST(SniiCluceneDecouplingGuardTest, CoreHasNoCluceneRef) { + const auto snii_core_dir = locate_snii_core_dir(); + ASSERT_TRUE(snii_core_dir.has_value()) + << "could not locate the SNII core source directory be/src/storage/index/snii. " + << "Run the test inside the source tree, or set DORIS_HOME / -DSNII_CORE_SOURCE_DIR. " + << "(__FILE__=" << __FILE__ << ")"; + + std::error_code ec; + fs::recursive_directory_iterator it(*snii_core_dir, ec); + ASSERT_TRUE(!ec) << "cannot iterate " << snii_core_dir->string() << ": " << ec.message(); + const fs::recursive_directory_iterator end; + + std::vector hits; + std::vector unreadable; + int scanned_files = 0; + for (; it != end; it.increment(ec)) { + ASSERT_TRUE(!ec) << "error iterating " << snii_core_dir->string() << ": " << ec.message(); + const fs::directory_entry& entry = *it; + std::error_code stat_ec; + if (!entry.is_regular_file(stat_ec) || stat_ec) { + continue; + } + // Guard only the SNII core engine (the module subdirs), which must stay + // CLucene-free. Skip the top-level Doris integration files (snii_doris_adapter + // / snii_index_*), which legitimately bridge to CLucene, and non-source files + // (docs/*.md). + const std::string ext = entry.path().extension().string(); + if (ext != ".h" && ext != ".cpp" && ext != ".cc") { + continue; + } + if (entry.path().parent_path() == *snii_core_dir) { + continue; + } + // Narrow exemption: compaction eligibility resolves the DESTINATION + // index's analyzer through Doris's analyzer stack, whose factory is + // CLucene-backed and throws CLuceneError. That resolution is + // integration work that currently lives one directory deep; until it + // moves behind an integration-layer factory, this single file may + // reference CLucene. Keep the exemption to exactly this file so any + // new core reference still fails the guard. + if (entry.path().parent_path().filename() == "compaction" && + entry.path().filename() == "eligibility.cpp") { + continue; + } + ++scanned_files; + if (!scan_file_for_clucene(entry.path(), &hits)) { + unreadable.push_back(entry.path().string()); + } + } + + // Sanity: we must actually scan some files; otherwise a broken locate would make + // the guard silently pass. + ASSERT_GT(scanned_files, 0) << "scanned no files under " << snii_core_dir->string() + << "; the guard is ineffective."; + + // Unreadable files create scan gaps (possible missed detections); surface them + // non-fatally. + EXPECT_TRUE(unreadable.empty()) + << "some files could not be read (possible missed detections): " << unreadable.size() + << " total, first is " << (unreadable.empty() ? std::string {} : unreadable.front()); + + if (!hits.empty()) { + std::ostringstream oss; + oss << "the SNII core " << snii_core_dir->string() << " has " << hits.size() + << " CLucene reference(s) (must be 0; violates the SNII/CLucene decoupling " + "constraint). CLucene is allowed only in the integration layer:\n"; + for (const auto& hit : hits) { + oss << " " << hit.file << ":" << hit.line << ": " << hit.text << "\n"; + } + FAIL() << oss.str(); + } +} diff --git a/be/test/storage/index/snii_crc32c_equiv_test.cpp b/be/test/storage/index/snii_crc32c_equiv_test.cpp new file mode 100644 index 00000000000000..b17c8fe6e96ff2 --- /dev/null +++ b/be/test/storage/index/snii_crc32c_equiv_test.cpp @@ -0,0 +1,145 @@ +// 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. + +// R05-crc32c: proves the snii crc32c header (now a thin inline wrapper over the +// Doris/Google crc32c thirdparty) yields byte-identical checksums to both (a) the +// well-known canonical CRC32C test vectors and (b) the Doris ::crc32c library it +// delegates to. (a) guards the on-disk format value; (b) guards the delegation +// contract for every call site (one-shot, seeded extend, and segmented chaining). +// These checks are deterministic and need no pre-captured golden fixture. + +#include +#include + +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/crc32c.h" + +namespace { + +using doris::snii::Slice; + +// Fills n deterministic bytes from a seeded engine. The exact bytes are +// irrelevant to the equivalence asserts (both sides hash the same buffer); the +// fixed seed only keeps runs reproducible. +std::vector make_bytes(std::mt19937_64& rng, size_t n) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = static_cast(rng() & 0xFFU); + } + return v; +} + +// Lengths exercising the slice-by-8 main loop plus every <8-byte tail remainder, +// the 4-byte step, sub-block sizes, and several >1KB buffers. +constexpr size_t kLengths[] = {0, 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, + 33, 63, 64, 65, 127, 128, 255, 256, 1023, 1024, 1025, 4096, 5000}; + +} // namespace + +// (a) Hardcoded canonical CRC32C vectors. A reuse that altered the value -- and +// therefore the on-disk format -- would break these. +TEST(SniiCrc32cEquiv, StandardVectors) { + // Classic CRC-32C / iSCSI check value for the ASCII string "123456789". + { + const char* s = "123456789"; + EXPECT_EQ(0xE3069283U, doris::snii::crc32c(Slice(reinterpret_cast(s), 9))); + } + // Empty input conditions to 0 (Extend(0, ., 0)). + EXPECT_EQ(0x00000000U, doris::snii::crc32c(Slice(nullptr, 0))); + + // RFC 3720 section B.4 vectors, matching be/test/util/crc32c_test.cpp. + uint8_t buf[32]; + std::memset(buf, 0x00, sizeof(buf)); + EXPECT_EQ(0x8A9136AAU, doris::snii::crc32c(Slice(buf, sizeof(buf)))); + + std::memset(buf, 0xFF, sizeof(buf)); + EXPECT_EQ(0x62A8AB43U, doris::snii::crc32c(Slice(buf, sizeof(buf)))); + + for (int i = 0; i < 32; ++i) { + buf[i] = static_cast(i); + } + EXPECT_EQ(0x46DD794EU, doris::snii::crc32c(Slice(buf, sizeof(buf)))); + + for (int i = 0; i < 32; ++i) { + buf[i] = static_cast(31 - i); + } + EXPECT_EQ(0x113FDB5CU, doris::snii::crc32c(Slice(buf, sizeof(buf)))); +} + +// (b1) One-shot equivalence: doris::snii::crc32c(Slice(buf, n)) == ::crc32c::Crc32c(buf, n). +TEST(SniiCrc32cEquiv, OneShotMatchesDorisLib) { + std::mt19937_64 rng(0xC0FFEEULL); + for (size_t n : kLengths) { + const std::vector buf = make_bytes(rng, n); + const uint8_t* p = buf.data(); + EXPECT_EQ(doris::snii::crc32c(Slice(p, n)), ::crc32c::Crc32c(p, n)) << "len=" << n; + } +} + +// (b2) Seeded extend equivalence: doris::snii::crc32c_extend(prev, Slice) == ::crc32c::Extend(prev, ...), +// covering prev == 0 and several non-zero priors. +TEST(SniiCrc32cEquiv, ExtendMatchesDorisLib) { + std::mt19937_64 rng(0x12345678ULL); + const uint32_t priors[] = {0U, 1U, 0xFFFFFFFFU, 0xDEADBEEFU, 0x80000000U}; + for (uint32_t prev : priors) { + for (size_t n : kLengths) { + const std::vector buf = make_bytes(rng, n); + const uint8_t* p = buf.data(); + EXPECT_EQ(doris::snii::crc32c_extend(prev, Slice(p, n)), ::crc32c::Extend(prev, p, n)) + << "prev=" << prev << " len=" << n; + } + } +} + +// (b3) Segmented chaining: extend over a split equals the one-shot over the whole, +// and matches the Doris lib's Extend(Crc32c(prefix), suffix) -- i.e. the +// Extend(Extend(0, a), b) == Crc32c(a ++ b) property every framed section relies on. +TEST(SniiCrc32cEquiv, ExtendSplitEqualsWhole) { + std::mt19937_64 rng(0xABCDEFULL); + const size_t totals[] = {1, 2, 8, 9, 16, 17, 100, 1024, 4096}; + for (size_t n : totals) { + const std::vector buf = make_bytes(rng, n); + const uint8_t* p = buf.data(); + const uint32_t whole = doris::snii::crc32c(Slice(p, n)); + for (size_t k = 0; k <= n; ++k) { + const uint32_t chained = doris::snii::crc32c_extend(doris::snii::crc32c(Slice(p, k)), + Slice(p + k, n - k)); + EXPECT_EQ(whole, chained) << "n=" << n << " k=" << k; + EXPECT_EQ(chained, ::crc32c::Extend(::crc32c::Crc32c(p, k), p + k, n - k)) + << "n=" << n << " k=" << k; + } + } +} + +// (b4) Cross-decode both directions: a crc stamped by the snii wrapper verifies +// under the Doris lib, and a crc stamped by the Doris lib verifies under the snii +// wrapper -- the bidirectional guarantee that old and new on-disk bytes interop. +TEST(SniiCrc32cEquiv, CrossDecodeBidirectional) { + std::mt19937_64 rng(0x5A5A5A5AULL); + for (size_t n : kLengths) { + const std::vector buf = make_bytes(rng, n); + const uint8_t* p = buf.data(); + EXPECT_EQ(doris::snii::crc32c(Slice(p, n)), ::crc32c::Crc32c(p, n)) << "len=" << n; + EXPECT_EQ(::crc32c::Crc32c(p, n), doris::snii::crc32c(Slice(p, n))) << "len=" << n; + } +} diff --git a/be/test/storage/index/snii_doris_adapter_test.cpp b/be/test/storage/index/snii_doris_adapter_test.cpp new file mode 100644 index 00000000000000..14c42f3b8f792e --- /dev/null +++ b/be/test/storage/index/snii_doris_adapter_test.cpp @@ -0,0 +1,897 @@ +// 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 "storage/index/snii/snii_doris_adapter.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "io/fs/file_reader.h" +#include "io/fs/local_file_system.h" +#include "io/fs/path.h" +#include "io/io_common.h" +#include "runtime/exec_env.h" +#include "runtime/runtime_state.h" +#include "storage/index/index_file_reader.h" +#include "storage/index/inverted/inverted_index_cache.h" +#include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/reader/windowed_posting.h" +#include "storage/index/snii/snii_index_reader.h" +#include "storage/index/snii_query_test_util.h" +#include "storage/olap_common.h" +#include "storage/tablet/tablet_schema.h" +#include "util/slice.h" +#include "util/threadpool.h" + +namespace doris::segment_v2::snii_doris { +namespace { + +struct CapturedIOContext { + bool has_ctx = false; + bool is_inverted_index = false; + bool is_index_data = false; + bool read_file_cache = true; + bool is_disposable = false; + int64_t expiration_time = 0; + io::FileCacheStatistics* file_cache_stats = nullptr; +}; + +struct CapturedRead { + size_t offset = 0; + size_t len = 0; + CapturedIOContext io_ctx; + // Destination buffer the read wrote into; used to prove single-segment reads + // land directly in the caller's output (no temp/double buffer). + const void* dst = nullptr; + // Worker thread that served the read; supports parallel-path assertions. + std::thread::id thread_id; +}; + +class RecordingFileReader final : public io::FileReader { +public: + explicit RecordingFileReader(std::string data) : _data(std::move(data)) {} + + Status close() override { + _closed = true; + return Status::OK(); + } + + const io::Path& path() const override { return _path; } + size_t size() const override { return _data.size(); } + bool closed() const override { return _closed; } + int64_t mtime() const override { return 0; } + + // Safe to read after read_batch returns: all worker reads happen-before the + // caller (latch join) and writes are guarded by _reads_mu. + const std::vector& reads() const { return _reads; } + + void xor_byte(size_t offset, uint8_t mask) { _data[offset] ^= static_cast(mask); } + + void block_once_on_range_containing(size_t offset) { + std::lock_guard guard(_control_mu); + _block_offset = offset; + _block_armed = true; + _blocked = false; + _released = false; + _matching_reads = 0; + } + + bool wait_for_blocked_read(std::chrono::seconds timeout) { + std::unique_lock lock(_control_mu); + return _control_cv.wait_for(lock, timeout, [&] { return _blocked; }); + } + + void release_blocked_read() { + std::lock_guard guard(_control_mu); + _released = true; + _control_cv.notify_all(); + } + + size_t matching_read_count() const { + std::lock_guard guard(_control_mu); + return _matching_reads; + } + + void add_file_cache_stats_sentinel(size_t offset, int64_t sentinel) { + _file_cache_stats_sentinels.emplace_back(offset, sentinel); + } + +protected: + Status read_at_impl(size_t offset, Slice result, size_t* bytes_read, + const io::IOContext* io_ctx) override { + CapturedRead read; + read.offset = offset; + read.len = result.size; + read.dst = result.data; + read.thread_id = std::this_thread::get_id(); + if (io_ctx != nullptr) { + read.io_ctx.has_ctx = true; + read.io_ctx.is_inverted_index = io_ctx->is_inverted_index; + read.io_ctx.is_index_data = io_ctx->is_index_data; + read.io_ctx.read_file_cache = io_ctx->read_file_cache; + read.io_ctx.is_disposable = io_ctx->is_disposable; + read.io_ctx.expiration_time = io_ctx->expiration_time; + read.io_ctx.file_cache_stats = io_ctx->file_cache_stats; + } + + { + std::unique_lock lock(_control_mu); + if (_block_offset >= offset && _block_offset < offset + result.size) { + ++_matching_reads; + if (_block_armed) { + _block_armed = false; + _blocked = true; + _control_cv.notify_all(); + _control_cv.wait(lock, [&] { return _released; }); + } + } + } + + if (result.size > 0) { + std::memcpy(result.data, _data.data() + offset, result.size); + } + *bytes_read = result.size; + + // Parallel batch reads may invoke this from several pool threads at once; + // guard the capture log so the test infrastructure is race-free. + { + std::lock_guard guard(_reads_mu); + _reads.push_back(read); + } + // Configured before batch workers start and read-only thereafter. Each + // physical segment also gets its own stats slot, so these writes do not race. + if (io_ctx != nullptr && io_ctx->file_cache_stats != nullptr) { + for (const auto& [stats_offset, sentinel] : _file_cache_stats_sentinels) { + if (stats_offset == offset) { + io_ctx->file_cache_stats->inverted_index_num_remote_io_total += sentinel; + } + } + } + if (_fail_offset.has_value() && *_fail_offset >= offset && + *_fail_offset < offset + result.size) { + return Status::IOError("injected read failure at offset {}", *_fail_offset); + } + return Status::OK(); + } + +public: + // Fail any read whose range covers `offset` (after capturing it), so batch + // tests can make exactly one physical segment error out. + void set_fail_offset(size_t offset) { _fail_offset = offset; } + +private: + std::string _data; + io::Path _path = "/tmp/snii_doris_adapter_test.idx"; + bool _closed = false; + std::optional _fail_offset; + std::vector> _file_cache_stats_sentinels; + mutable std::mutex _reads_mu; + std::vector _reads; + mutable std::mutex _control_mu; + std::condition_variable _control_cv; + size_t _block_offset = std::numeric_limits::max(); + bool _block_armed = false; + bool _blocked = false; + bool _released = false; + size_t _matching_reads = 0; +}; + +class FollowerJoinedLatch { +public: + static void notify(void* opaque) noexcept { + auto* latch = static_cast(opaque); + { + std::lock_guard guard(latch->_mutex); + latch->_joined = true; + } + latch->_cv.notify_all(); + } + + bool wait_for(std::chrono::seconds timeout) { + std::unique_lock lock(_mutex); + return _cv.wait_for(lock, timeout, [&] { return _joined; }); + } + +private: + std::mutex _mutex; + std::condition_variable _cv; + bool _joined = false; +}; + +// Deterministic, position-dependent byte pattern so every range's expected bytes +// are just `pattern.substr(offset, len)`. +std::string make_pattern(size_t n) { + std::string s(n, '\0'); + for (size_t i = 0; i < n; ++i) { + s[i] = static_cast('a' + static_cast(i % 26)); + } + return s; +} + +std::string as_string(const std::vector& v) { + return std::string(v.begin(), v.end()); +} + +using PrxStatsSnapshot = std::array; + +PrxStatsSnapshot prx_stats_snapshot(const OlapReaderStatistics& stats) { + return {stats.snii_stats.prx_raw_frames, stats.snii_stats.prx_zstd_frames, + stats.snii_stats.prx_pfor_frames, stats.snii_stats.prx_plaintext_bytes, + stats.snii_stats.prx_total_docs, stats.snii_stats.prx_selected_docs, + stats.snii_stats.prx_total_positions, stats.snii_stats.prx_selected_positions, + stats.snii_stats.prx_fetch_ns, stats.snii_stats.prx_decode_ns, + stats.snii_stats.prx_phrase_verify_ns}; +} + +void set_prx_stats_sentinel(OlapReaderStatistics* stats) { + stats->snii_stats.prx_raw_frames = 101; + stats->snii_stats.prx_zstd_frames = 102; + stats->snii_stats.prx_pfor_frames = 103; + stats->snii_stats.prx_plaintext_bytes = 104; + stats->snii_stats.prx_total_docs = 105; + stats->snii_stats.prx_selected_docs = 106; + stats->snii_stats.prx_total_positions = 107; + stats->snii_stats.prx_selected_positions = 108; + stats->snii_stats.prx_fetch_ns = 109; + stats->snii_stats.prx_decode_ns = 110; + stats->snii_stats.prx_phrase_verify_ns = 111; +} + +std::vector bitmap_docids(const roaring::Roaring& bitmap) { + return {bitmap.begin(), bitmap.end()}; +} + +struct ActualQueryExecutionContext { + ActualQueryExecutionContext() { + TQueryOptions options; + options.enable_inverted_index_query_cache = false; + options.enable_inverted_index_searcher_cache = false; + runtime_state.set_query_options(options); + context->io_ctx = &io_ctx; + context->stats = &stats; + context->runtime_state = &runtime_state; + } + + OlapReaderStatistics stats; + io::IOContext io_ctx; + RuntimeState runtime_state; + IndexQueryContextPtr context = std::make_shared(); +}; + +void init_actual_path_index_meta(TabletIndex* meta) { + TabletIndexPB pb; + pb.set_index_type(IndexType::INVERTED); + pb.set_index_id(31); + pb.set_index_name("snii_doris_adapter_actual_path_idx"); + pb.add_col_unique_id(0); + pb.mutable_properties()->insert({"parser", "english"}); + pb.mutable_properties()->insert({"lower_case", "true"}); + pb.mutable_properties()->insert({"support_phrase", "true"}); + meta->init_from_pb(pb); +} + +Status resolve_prx_range(const doris::snii::reader::LogicalIndexReader& reader, + std::string_view term, doris::snii::io::Range* range) { + bool found = false; + doris::snii::format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + RETURN_IF_ERROR(reader.lookup(term, &found, &entry, &frq_base, &prx_base)); + if (!found) { + return Status::NotFound("term {} missing from actual-path fixture", term); + } + doris::snii::format::FrqPreludeReader prelude; + RETURN_IF_ERROR(doris::snii::reader::fetch_windowed_prelude(reader, entry, frq_base, &prelude)); + doris::snii::reader::WindowAbsRange window_range; + RETURN_IF_ERROR(doris::snii::reader::windowed_window_range(reader, entry, frq_base, prx_base, + prelude, 0, /*want_positions=*/true, + /*want_freq=*/false, &window_range)); + range->offset = window_range.prx_off; + range->len = window_range.prx_len; + return Status::OK(); +} + +// Restores the batch-read executor seam on scope exit so an injected pool never +// leaks into sibling tests, even if an assertion returns early. +struct IoPoolSeamGuard { + explicit IoPoolSeamGuard(ThreadPool* pool) { + DorisSniiFileReader::set_io_thread_pool_for_test(pool); + } + ~IoPoolSeamGuard() { DorisSniiFileReader::set_io_thread_pool_for_test(nullptr); } + IoPoolSeamGuard(const IoPoolSeamGuard&) = delete; + IoPoolSeamGuard& operator=(const IoPoolSeamGuard&) = delete; +}; + +} // namespace + +class SniiIndexReaderActualPathTest : public testing::Test { +protected: + void SetUp() override { + init_actual_path_index_meta(&_meta); + constexpr uint32_t kDocCount = 9000; + doris::snii::snii_test::MemoryFile memory_file; + doris::snii::writer::SniiIndexInput input; + input.index_id = 31; + input.config = doris::snii::format::IndexConfig::kDocsPositions; + input.doc_count = kDocCount; + input.terms = { + doris::snii::snii_test::make_term( + "failed", doris::snii::snii_test::docs_with_one_position(0, kDocCount, 0)), + doris::snii::snii_test::make_term( + "order", doris::snii::snii_test::docs_with_one_position(0, kDocCount, 1))}; + doris::snii::writer::SniiCompoundWriter writer(&memory_file); + const Status add_status = writer.add_logical_index(input); + ASSERT_TRUE(add_status.ok()) << add_status.to_string(); + const Status finish_status = writer.finish(); + ASSERT_TRUE(finish_status.ok()) << finish_status.to_string(); + const auto& bytes = memory_file.data(); + std::string data(reinterpret_cast(bytes.data()), bytes.size()); + _recording_reader = std::make_shared(std::move(data)); + _adapter_reader = std::make_shared(_recording_reader); + + auto segment_reader = std::make_unique(); + const Status open_status = doris::snii::reader::SniiSegmentReader::open( + _adapter_reader.get(), segment_reader.get()); + ASSERT_TRUE(open_status.ok()) << open_status.to_string(); + + _file_reader = std::make_shared(io::global_local_filesystem(), + "/tmp/snii_doris_adapter_actual_path", + InvertedIndexStorageFormatPB::SNII); + _file_reader->_snii_file_reader = _adapter_reader; + _file_reader->_snii_segment_reader = std::move(segment_reader); + _file_reader->_inited = true; + _index_reader = SniiIndexReader::create_shared(&_meta, _file_reader, + InvertedIndexReaderType::FULLTEXT); + + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); + _query_cache.reset(InvertedIndexQueryCache::create_global_cache(1024 * 1024, 1)); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_query_cache.get()); + } + + void TearDown() override { + _index_reader.reset(); + _file_reader.reset(); + _adapter_reader.reset(); + _recording_reader.reset(); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); + _query_cache.reset(); + } + + doris::snii::io::Range prx_range(std::string_view term) { + auto logical_reader = _file_reader->open_snii_index(&_meta); + EXPECT_TRUE(logical_reader.has_value()) << logical_reader.error(); + doris::snii::io::Range range; + if (logical_reader.has_value()) { + const Status status = resolve_prx_range(*logical_reader.value(), term, &range); + EXPECT_TRUE(status.ok()) << status.to_string(); + } + return range; + } + + Status run_phrase(const IndexQueryContextPtr& context, std::string_view column, + std::shared_ptr* bitmap) { + Field query_value = Field::create_field(std::string("failed order")); + return _index_reader->query(context, std::string(column), query_value, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, *bitmap); + } + + TabletIndex _meta; + std::shared_ptr _recording_reader; + std::shared_ptr _adapter_reader; + std::shared_ptr _file_reader; + std::shared_ptr _index_reader; + InvertedIndexQueryCache* _previous_query_cache = nullptr; + std::unique_ptr _query_cache; +}; + +TEST_F(SniiIndexReaderActualPathTest, SingleFlightFollowerLeavesPrxStatsUnchanged) { + const doris::snii::io::Range first_prx = prx_range("failed"); + ASSERT_GT(first_prx.len, 0U); + _recording_reader->block_once_on_range_containing(first_prx.offset); + + ActualQueryExecutionContext leader; + ActualQueryExecutionContext follower; + set_prx_stats_sentinel(&follower.stats); + const PrxStatsSnapshot follower_before = prx_stats_snapshot(follower.stats); + std::shared_ptr leader_bitmap; + std::shared_ptr follower_bitmap; + FollowerJoinedLatch follower_joined; + _index_reader->set_single_flight_follower_joined_observer_for_test(&FollowerJoinedLatch::notify, + &follower_joined); + + auto leader_future = std::async(std::launch::async, [&] { + return run_phrase(leader.context, "single_flight_content", &leader_bitmap); + }); + const bool leader_blocked = _recording_reader->wait_for_blocked_read(std::chrono::seconds(5)); + if (!leader_blocked) { + _recording_reader->release_blocked_read(); + const Status leader_status = leader_future.get(); + FAIL() << "leader did not reach configured PRX range: " << leader_status.to_string(); + } + + auto follower_future = std::async(std::launch::async, [&] { + return run_phrase(follower.context, "single_flight_content", &follower_bitmap); + }); + const bool follower_did_join = follower_joined.wait_for(std::chrono::seconds(5)); + if (!follower_did_join) { + _recording_reader->release_blocked_read(); + const Status leader_status = leader_future.get(); + const Status follower_status = follower_future.get(); + FAIL() << "follower did not join single-flight: leader=" << leader_status.to_string() + << ", follower=" << follower_status.to_string(); + } + + _recording_reader->release_blocked_read(); + const Status leader_status = leader_future.get(); + const Status follower_status = follower_future.get(); + ASSERT_TRUE(leader_status.ok()) << leader_status.to_string(); + ASSERT_TRUE(follower_status.ok()) << follower_status.to_string(); + ASSERT_NE(leader_bitmap, nullptr); + ASSERT_NE(follower_bitmap, nullptr); + EXPECT_EQ(bitmap_docids(*leader_bitmap), bitmap_docids(*follower_bitmap)); + EXPECT_GT(leader.stats.snii_stats.prx_plaintext_bytes, 0); + EXPECT_EQ(prx_stats_snapshot(follower.stats), follower_before); + EXPECT_EQ(_recording_reader->matching_read_count(), 1U); +} + +TEST_F(SniiIndexReaderActualPathTest, LaterCorruptFrameFlushesEarlierSuccessfulFrame) { + const doris::snii::io::Range first_prx = prx_range("failed"); + const doris::snii::io::Range second_prx = prx_range("order"); + ASSERT_GT(first_prx.len, 0U); + ASSERT_GT(second_prx.len, 0U); + ASSERT_NE(first_prx.offset, second_prx.offset); + const uint64_t second_crc_byte = second_prx.offset + second_prx.len - 1; + ASSERT_LT(second_crc_byte, _recording_reader->size()); + _recording_reader->xor_byte(static_cast(second_crc_byte), 0x80); + + ActualQueryExecutionContext execution; + std::shared_ptr bitmap; + const Status status = run_phrase(execution.context, "late_error_content", &bitmap); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(execution.stats.snii_stats.prx_raw_frames + + execution.stats.snii_stats.prx_zstd_frames + + execution.stats.snii_stats.prx_pfor_frames, + 1); + EXPECT_GT(execution.stats.snii_stats.prx_plaintext_bytes, 0); + EXPECT_GT(execution.stats.snii_stats.prx_total_docs, 0); + EXPECT_EQ(execution.stats.snii_stats.prx_phrase_verify_ns, 0); +} + +TEST(DorisSniiFileReaderTest, ReadAtPropagatesIndexIOContextAndRecordsStats) { + auto recording_reader = std::make_shared("0123456789abcdef"); + DorisSniiFileReader reader(recording_reader); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.is_disposable = true; + io_ctx.is_index_data = true; + io_ctx.read_file_cache = false; + io_ctx.file_cache_stats = &stats; + + std::vector out; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + auto status = reader.read_at(2, 5, &out); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + ASSERT_EQ(out.size(), 5); + EXPECT_EQ(std::string(out.begin(), out.end()), "23456"); + ASSERT_EQ(recording_reader->reads().size(), 1); + const auto& captured = recording_reader->reads()[0].io_ctx; + EXPECT_TRUE(captured.has_ctx); + EXPECT_TRUE(captured.is_inverted_index); + EXPECT_TRUE(captured.is_index_data); + EXPECT_FALSE(captured.read_file_cache); + EXPECT_TRUE(captured.is_disposable); + EXPECT_EQ(captured.file_cache_stats, &stats); + + EXPECT_EQ(stats.inverted_index_request_bytes, 5); + EXPECT_EQ(stats.inverted_index_read_bytes, 5); + EXPECT_EQ(stats.inverted_index_range_read_count, 1); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); +} + +TEST(DorisSniiFileReaderTest, ReadBatchRecordsLogicalAndCoalescedPhysicalIO) { + auto recording_reader = + std::make_shared("0123456789abcdefghijklmnopqrstuvwxyz"); + DorisSniiFileReader reader(recording_reader); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.file_cache_stats = &stats; + + std::vector> outs; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + std::vector<::doris::snii::io::Range> ranges {{0, 4}, {6, 3}, {20, 2}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + ASSERT_EQ(outs.size(), 3); + EXPECT_EQ(std::string(outs[0].begin(), outs[0].end()), "0123"); + EXPECT_EQ(std::string(outs[1].begin(), outs[1].end()), "678"); + EXPECT_EQ(std::string(outs[2].begin(), outs[2].end()), "kl"); + + ASSERT_EQ(recording_reader->reads().size(), 1); + EXPECT_EQ(recording_reader->reads()[0].offset, 0); + EXPECT_EQ(recording_reader->reads()[0].len, 22); + + EXPECT_EQ(stats.inverted_index_request_bytes, 9); + EXPECT_EQ(stats.inverted_index_read_bytes, 22); + EXPECT_EQ(stats.inverted_index_range_read_count, 1); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); +} + +// A direct-remote (NO_CACHE bypass) reader has no CachedRemoteFileReader below +// it, so it must count its own reads as physical remote bytes; a default reader +// must not (its wrapper does, or the bytes are local). +TEST(DorisSniiFileReaderTest, DirectRemoteReaderCountsPhysicalRemoteBytes) { + auto recording_reader = std::make_shared("0123456789abcdef"); + DorisSniiFileReader direct_reader(recording_reader, /*io_ctx=*/nullptr, + /*direct_remote_io=*/true); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.file_cache_stats = &stats; + + std::vector out; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + auto status = direct_reader.read_at(2, 5, &out); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + EXPECT_EQ(stats.inverted_index_read_bytes, 5); + EXPECT_EQ(stats.inverted_index_remote_physical_read_bytes, 5); + + io::FileCacheStatistics cached_stats; + io::IOContext cached_io_ctx; + cached_io_ctx.file_cache_stats = &cached_stats; + DorisSniiFileReader cached_reader(recording_reader); + { + DorisSniiFileReader::ScopedIOContext scope(&cached_io_ctx); + auto status = cached_reader.read_at(2, 5, &out); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + EXPECT_EQ(cached_stats.inverted_index_read_bytes, 5); + EXPECT_EQ(cached_stats.inverted_index_remote_physical_read_bytes, 0); +} + +// When one physical segment of a batch fails, lower-layer stats from every +// attempted segment must reach the caller, while logical counters record only +// what completed (including direct-remote physical bytes). +TEST(DorisSniiFileReaderTest, ReadBatchKeepsCompletedStatsWhenOneSegmentFails) { + // Two ranges spaced far beyond the coalescing gap form two physical segments. + constexpr int64_t successful_segment_sentinel = 101; + constexpr int64_t failed_segment_sentinel = 1009; + std::string data(8192, 'x'); + auto recording_reader = std::make_shared(std::move(data)); + recording_reader->add_file_cache_stats_sentinel(0, successful_segment_sentinel); + recording_reader->add_file_cache_stats_sentinel(6000, failed_segment_sentinel); + recording_reader->set_fail_offset(6000); + DorisSniiFileReader reader(recording_reader, /*io_ctx=*/nullptr, + /*direct_remote_io=*/true); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.file_cache_stats = &stats; + + std::vector> outs; + Status status; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + std::vector<::doris::snii::io::Range> ranges {{.offset = 0, .len = 4}, + {.offset = 6000, .len = 8}}; + status = reader.read_batch(ranges, &outs); + } + EXPECT_FALSE(status.ok()); + const auto& reads = recording_reader->reads(); + ASSERT_EQ(reads.size(), 2); + const auto successful_read = + std::ranges::find_if(reads, [](const CapturedRead& read) { return read.offset == 0; }); + const auto failed_read = std::ranges::find_if( + reads, [](const CapturedRead& read) { return read.offset == 6000; }); + ASSERT_NE(successful_read, reads.end()); + ASSERT_NE(failed_read, reads.end()); + ASSERT_NE(successful_read->io_ctx.file_cache_stats, nullptr); + ASSERT_NE(failed_read->io_ctx.file_cache_stats, nullptr); + EXPECT_NE(successful_read->io_ctx.file_cache_stats, &stats); + EXPECT_NE(failed_read->io_ctx.file_cache_stats, &stats); + EXPECT_NE(successful_read->io_ctx.file_cache_stats, failed_read->io_ctx.file_cache_stats); + // Logical counters include only the completed segment. Lower-layer stats + // written before either segment returns, including the failed one, are all merged. + EXPECT_EQ(stats.inverted_index_request_bytes, 4); + EXPECT_EQ(stats.inverted_index_read_bytes, 4); + EXPECT_EQ(stats.inverted_index_remote_physical_read_bytes, 4); + EXPECT_EQ(stats.inverted_index_range_read_count, 2); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); + EXPECT_EQ(stats.inverted_index_num_remote_io_total, + successful_segment_sentinel + failed_segment_sentinel); +} + +// FB-02: three ranges spaced >4096 apart form three disjoint physical segments. +// Each is still a separate physical read, but the batch is one concurrent round +// (F19): serial_read_rounds drops from K(==3) to 1. +TEST(DorisSniiFileReaderTest, ReadBatchIssuesSingleSerialRoundForDisjointSegments) { + const std::string data = make_pattern(20000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.file_cache_stats = &stats; + + std::vector> outs; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + std::vector<::doris::snii::io::Range> ranges {{0, 4}, {8192, 4}, {16384, 4}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + ASSERT_EQ(outs.size(), 3); + EXPECT_EQ(as_string(outs[0]), data.substr(0, 4)); + EXPECT_EQ(as_string(outs[1]), data.substr(8192, 4)); + EXPECT_EQ(as_string(outs[2]), data.substr(16384, 4)); + + EXPECT_EQ(recording_reader->reads().size(), 3); + EXPECT_EQ(stats.inverted_index_request_bytes, 12); + EXPECT_EQ(stats.inverted_index_read_bytes, 12); + EXPECT_EQ(stats.inverted_index_range_read_count, 3); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); +} + +// F27: a single-range group reads straight into the caller's output slot, with no +// temporary buffer and no second memcpy. Proven by destination-pointer identity. +TEST(DorisSniiFileReaderTest, ReadBatchSingleSegmentReadsInPlace) { + const std::string data = make_pattern(20000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + std::vector> outs; + std::vector<::doris::snii::io::Range> ranges {{100, 8}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + + ASSERT_EQ(outs.size(), 1); + EXPECT_EQ(as_string(outs[0]), data.substr(100, 8)); + ASSERT_EQ(recording_reader->reads().size(), 1); + // The read wrote directly into outs[0]'s storage (no double buffer). + EXPECT_EQ(recording_reader->reads()[0].dst, outs[0].data()); +} + +// FB-03: a batch mixing one coalesced group (temp + scatter) and one single-range +// group (direct read). Both branches produce correct bytes. +TEST(DorisSniiFileReaderTest, ReadBatchMixedSingleAndCoalescedGroups) { + const std::string data = make_pattern(20000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + std::vector> outs; + std::vector<::doris::snii::io::Range> ranges {{0, 4}, {4, 4}, {9000, 4}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + + ASSERT_EQ(outs.size(), 3); + EXPECT_EQ(as_string(outs[0]), data.substr(0, 4)); + EXPECT_EQ(as_string(outs[1]), data.substr(4, 4)); + EXPECT_EQ(as_string(outs[2]), data.substr(9000, 4)); + + // Two physical segments: the coalesced [0,8) group and the single [9000,9004). + ASSERT_EQ(recording_reader->reads().size(), 2); + // The single-range group [9000,9004) read directly into outs[2]. + const auto& reads = recording_reader->reads(); + bool single_in_place = false; + for (const auto& r : reads) { + if (r.offset == 9000) { + single_in_place = (r.dst == outs[2].data()); + } + } + EXPECT_TRUE(single_in_place); +} + +// FB-04: empty batch and zero-length ranges. No physical reads; zero-length slots +// stay empty; outs aligns 1:1 with the input. +TEST(DorisSniiFileReaderTest, ReadBatchHandlesEmptyAndZeroLengthRanges) { + auto recording_reader = std::make_shared(make_pattern(64)); + DorisSniiFileReader reader(recording_reader); + + std::vector> empty_outs; + auto empty_status = reader.read_batch({}, &empty_outs); + ASSERT_TRUE(empty_status.ok()) << empty_status.to_string(); + EXPECT_TRUE(empty_outs.empty()); + EXPECT_EQ(recording_reader->reads().size(), 0); + + std::vector> outs; + std::vector<::doris::snii::io::Range> ranges {{5, 0}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_EQ(outs.size(), 1); + EXPECT_TRUE(outs[0].empty()); + EXPECT_EQ(recording_reader->reads().size(), 0); +} + +// FB-05: an out-of-range request surfaces a corruption error and does not crash. +TEST(DorisSniiFileReaderTest, ReadBatchReturnsErrorForOutOfRange) { + auto recording_reader = std::make_shared(make_pattern(64)); + DorisSniiFileReader reader(recording_reader); + + std::vector> outs; + std::vector<::doris::snii::io::Range> ranges {{63, 100}}; + auto status = reader.read_batch(ranges, &outs); + EXPECT_FALSE(status.ok()); +} + +// FB-06: unsorted input still produces outputs in the caller's original index +// order (the skip-sort guard only avoids the sort when already sorted). +TEST(DorisSniiFileReaderTest, ReadBatchPreservesOriginalOrderForUnsortedInput) { + const std::string data = make_pattern(20000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + std::vector> outs; + std::vector<::doris::snii::io::Range> ranges {{16384, 2}, {0, 4}, {8192, 3}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + + ASSERT_EQ(outs.size(), 3); + EXPECT_EQ(as_string(outs[0]), data.substr(16384, 2)); + EXPECT_EQ(as_string(outs[1]), data.substr(0, 4)); + EXPECT_EQ(as_string(outs[2]), data.substr(8192, 3)); +} + +// IOContext passthrough: every per-segment read sees the caller's flags. Each +// segment is routed through a private FileCacheStatistics slot (so disjoint reads +// never race), which is then merged back into the caller's sink. +TEST(DorisSniiFileReaderTest, ReadBatchPropagatesIOContextFlagsPerSegment) { + const std::string data = make_pattern(20000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.is_index_data = true; + io_ctx.read_file_cache = false; + io_ctx.is_disposable = true; + io_ctx.expiration_time = 123; + io_ctx.file_cache_stats = &stats; + + std::vector> outs; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + std::vector<::doris::snii::io::Range> ranges {{0, 4}, {8192, 4}, {16384, 4}}; + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + const auto& reads = recording_reader->reads(); + ASSERT_EQ(reads.size(), 3); + for (const auto& r : reads) { + EXPECT_TRUE(r.io_ctx.has_ctx); + EXPECT_TRUE(r.io_ctx.is_inverted_index); + EXPECT_TRUE(r.io_ctx.is_index_data); + EXPECT_FALSE(r.io_ctx.read_file_cache); + EXPECT_TRUE(r.io_ctx.is_disposable); + EXPECT_EQ(r.io_ctx.expiration_time, 123); + // Per-segment private stats slot, not the caller's sink directly. + EXPECT_NE(r.io_ctx.file_cache_stats, nullptr); + EXPECT_NE(r.io_ctx.file_cache_stats, &stats); + } + // Aggregate counters still land on the caller's sink after the merge. + EXPECT_EQ(stats.inverted_index_range_read_count, 3); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); +} + +// FB-08: with a real injected pool, eight disjoint segments are read in parallel. +// All bytes are correct, every segment is a distinct physical read, and the batch +// is one concurrent round. Run under TSAN to prove thread-safety. +TEST(DorisSniiFileReaderConcurrencyTest, ParallelSegmentReadsAreThreadSafe) { + const std::string data = make_pattern(60000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + std::unique_ptr pool; + auto pool_st = + ThreadPoolBuilder("snii_batch_test").set_min_threads(2).set_max_threads(4).build(&pool); + ASSERT_TRUE(pool_st.ok()) << pool_st.to_string(); + IoPoolSeamGuard seam(pool.get()); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.file_cache_stats = &stats; + + std::vector<::doris::snii::io::Range> ranges; + for (size_t i = 0; i < 8; ++i) { + ranges.push_back({static_cast(i) * 8192, 4}); + } + + std::vector> outs; + { + DorisSniiFileReader::ScopedIOContext scope(&io_ctx); + auto status = reader.read_batch(ranges, &outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + ASSERT_EQ(outs.size(), 8); + for (size_t i = 0; i < 8; ++i) { + EXPECT_EQ(as_string(outs[i]), data.substr(i * 8192, 4)); + } + EXPECT_EQ(recording_reader->reads().size(), 8); + EXPECT_EQ(stats.inverted_index_range_read_count, 8); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); +} + +// FB-07: the injected-pool (parallel) path and the seam-nullptr path both produce +// the exact same outputs (and match ground truth) -- parallelism is invisible. +TEST(DorisSniiFileReaderConcurrencyTest, ParallelPathMatchesSerialPath) { + const std::string data = make_pattern(60000); + auto recording_reader = std::make_shared(data); + DorisSniiFileReader reader(recording_reader); + + std::vector<::doris::snii::io::Range> ranges; + std::vector> expected; + for (size_t i = 0; i < 8; ++i) { + const uint64_t off = static_cast(i) * 8192; + ranges.push_back({off, 4}); + const std::string chunk = data.substr(off, 4); + expected.emplace_back(chunk.begin(), chunk.end()); + } + + std::vector> serial_outs; + { + IoPoolSeamGuard seam(nullptr); // no injected pool + auto status = reader.read_batch(ranges, &serial_outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + std::unique_ptr pool; + auto pool_st = + ThreadPoolBuilder("snii_batch_test").set_min_threads(2).set_max_threads(4).build(&pool); + ASSERT_TRUE(pool_st.ok()) << pool_st.to_string(); + + std::vector> parallel_outs; + { + IoPoolSeamGuard seam(pool.get()); + auto status = reader.read_batch(ranges, ¶llel_outs); + ASSERT_TRUE(status.ok()) << status.to_string(); + } + + EXPECT_EQ(serial_outs, expected); + EXPECT_EQ(parallel_outs, expected); + EXPECT_EQ(serial_outs, parallel_outs); +} + +} // namespace doris::segment_v2::snii_doris diff --git a/be/test/storage/index/snii_frq_direct_encode_test.cpp b/be/test/storage/index/snii_frq_direct_encode_test.cpp new file mode 100644 index 00000000000000..a8dfb02a8b5b85 --- /dev/null +++ b/be/test/storage/index/snii_frq_direct_encode_test.cpp @@ -0,0 +1,411 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/crc32c.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/frq_pod.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii { +namespace { + +ByteSink legacy_raw_dd(const std::vector& docs, uint64_t win_base) { + ByteSink out; + out.put_varint32(static_cast(docs.size())); + std::vector deltas(docs.size()); + uint64_t previous = win_base; + for (size_t i = 0; i < docs.size(); ++i) { + deltas[i] = static_cast(docs[i] - previous); + previous = docs[i]; + } + for (size_t begin = 0; begin < deltas.size(); begin += format::kFrqBaseUnit) { + const size_t count = + std::min(deltas.size() - begin, static_cast(format::kFrqBaseUnit)); + pfor_encode(deltas.data() + begin, count, &out); + } + return out; +} + +ByteSink legacy_raw_freq(const std::vector& freqs) { + ByteSink out; + for (size_t begin = 0; begin < freqs.size(); begin += format::kFrqBaseUnit) { + const size_t count = + std::min(freqs.size() - begin, static_cast(format::kFrqBaseUnit)); + pfor_encode(freqs.data() + begin, count, &out); + } + return out; +} + +std::vector make_docs(size_t count, uint32_t win_base, uint32_t salt) { + std::vector docs; + docs.reserve(count); + uint32_t doc = win_base; + for (size_t i = 0; i < count; ++i) { + doc += 1 + static_cast((i + salt) % 7); + docs.push_back(doc); + } + return docs; +} + +std::vector make_doc_deltas(const std::vector& docs, uint64_t win_base) { + std::vector deltas; + deltas.reserve(docs.size()); + uint64_t previous = win_base; + for (uint32_t doc : docs) { + deltas.push_back(static_cast(static_cast(doc) - previous)); + previous = doc; + } + return deltas; +} + +std::vector make_freqs(size_t count, uint32_t salt) { + std::vector freqs(count); + for (size_t i = 0; i < count; ++i) { + freqs[i] = 1 + static_cast((i * 5 + salt) % 11); + } + return freqs; +} + +void expect_raw_meta(const format::FrqRegionMeta& meta, Slice expected) { + EXPECT_FALSE(meta.zstd); + EXPECT_EQ(meta.uncomp_len, expected.size()); + EXPECT_EQ(meta.disk_len, expected.size()); + EXPECT_EQ(meta.crc, crc32c(expected)); + EXPECT_TRUE(meta.verify_crc); +} + +void expect_meta_equal(const format::FrqRegionMeta& actual, const format::FrqRegionMeta& expected) { + EXPECT_EQ(actual.zstd, expected.zstd); + EXPECT_EQ(actual.uncomp_len, expected.uncomp_len); + EXPECT_EQ(actual.disk_len, expected.disk_len); + EXPECT_EQ(actual.crc, expected.crc); + EXPECT_EQ(actual.verify_crc, expected.verify_crc); +} + +Slice appended_region(const ByteSink& sink, size_t offset, uint64_t length) { + return Slice(sink.buffer().data() + offset, static_cast(length)); +} + +constexpr std::array kBoundaryCounts {0, 1, 255, 256, 257, 512}; + +TEST(SniiFrqDirectEncodeTest, RawDdMatchesReferenceAcrossBoundariesAndConsecutiveAppends) { + for (size_t count : kBoundaryCounts) { + SCOPED_TRACE(count); + constexpr uint32_t kFirstBase = 7; + const std::vector first_docs = make_docs(count, kFirstBase, 1); + const uint32_t second_base = first_docs.empty() ? kFirstBase : first_docs.back(); + const std::vector second_docs = make_docs(count, second_base, 3); + const ByteSink first_reference = legacy_raw_dd(first_docs, kFirstBase); + const ByteSink second_reference = legacy_raw_dd(second_docs, second_base); + + ByteSink actual; + actual.put_fixed32(0x12345678); + const size_t first_offset = actual.size(); + format::FrqRegionMeta first_meta; + ASSERT_TRUE(format::build_dd_region(first_docs, kFirstBase, 0, &actual, &first_meta).ok()); + const size_t second_offset = actual.size(); + format::FrqRegionMeta second_meta; + ASSERT_TRUE( + format::build_dd_region(second_docs, second_base, 0, &actual, &second_meta).ok()); + + ByteSink expected; + expected.put_fixed32(0x12345678); + expected.put_bytes(first_reference.view()); + expected.put_bytes(second_reference.view()); + EXPECT_EQ(actual.buffer(), expected.buffer()); + expect_raw_meta(first_meta, first_reference.view()); + expect_raw_meta(second_meta, second_reference.view()); + + std::vector decoded; + ASSERT_TRUE( + format::decode_dd_region(appended_region(actual, first_offset, first_meta.disk_len), + first_meta, kFirstBase, &decoded) + .ok()); + EXPECT_EQ(decoded, first_docs); + ASSERT_TRUE(format::decode_dd_region( + appended_region(actual, second_offset, second_meta.disk_len), + second_meta, second_base, &decoded) + .ok()); + EXPECT_EQ(decoded, second_docs); + } +} + +TEST(SniiFrqDirectEncodeTest, RawDdDeltasMatchAbsoluteAcrossBoundariesAndWindows) { + const std::vector first_docs {0, 3, 10}; + const std::vector first_deltas = make_doc_deltas(first_docs, /*win_base=*/0); + constexpr uint32_t kSecondBase = 10; + + for (size_t count : kBoundaryCounts) { + SCOPED_TRACE(count); + const std::vector second_docs = make_docs(count, kSecondBase, 3); + const std::vector second_deltas = make_doc_deltas(second_docs, kSecondBase); + + ByteSink absolute; + absolute.put_fixed32(0x12345678); + const size_t first_offset = absolute.size(); + format::FrqRegionMeta absolute_first_meta; + ASSERT_TRUE(format::build_dd_region(first_docs, /*win_base=*/0, /*level=*/0, &absolute, + &absolute_first_meta) + .ok()); + const size_t second_offset = absolute.size(); + format::FrqRegionMeta absolute_second_meta; + ASSERT_TRUE(format::build_dd_region(second_docs, kSecondBase, /*level=*/0, &absolute, + &absolute_second_meta) + .ok()); + + ByteSink direct; + direct.put_fixed32(0x12345678); + format::FrqRegionMeta direct_first_meta; + ASSERT_TRUE(format::build_dd_region_from_deltas(std::span(first_deltas), + /*level=*/0, &direct, &direct_first_meta) + .ok()); + EXPECT_EQ(direct.size(), second_offset); + format::FrqRegionMeta direct_second_meta; + ASSERT_TRUE(format::build_dd_region_from_deltas(std::span(second_deltas), + /*level=*/0, &direct, &direct_second_meta) + .ok()); + + EXPECT_EQ(direct.buffer(), absolute.buffer()); + expect_meta_equal(direct_first_meta, absolute_first_meta); + expect_meta_equal(direct_second_meta, absolute_second_meta); + + std::vector decoded; + ASSERT_TRUE(format::decode_dd_region( + appended_region(direct, first_offset, direct_first_meta.disk_len), + direct_first_meta, /*win_base=*/0, &decoded) + .ok()); + EXPECT_EQ(decoded, first_docs); + ASSERT_TRUE(format::decode_dd_region( + appended_region(direct, second_offset, direct_second_meta.disk_len), + direct_second_meta, kSecondBase, &decoded) + .ok()); + EXPECT_EQ(decoded, second_docs); + } +} + +TEST(SniiFrqDirectEncodeTest, RawFreqMatchesReferenceAcrossBoundariesAndConsecutiveAppends) { + for (size_t count : kBoundaryCounts) { + SCOPED_TRACE(count); + const std::vector first_freqs = make_freqs(count, 2); + const std::vector second_freqs = make_freqs(count, 7); + const ByteSink first_reference = legacy_raw_freq(first_freqs); + const ByteSink second_reference = legacy_raw_freq(second_freqs); + + ByteSink actual; + actual.put_fixed32(0x87654321); + const size_t first_offset = actual.size(); + format::FrqRegionMeta first_meta; + ASSERT_TRUE(format::build_freq_region(first_freqs, 0, &actual, &first_meta).ok()); + const size_t second_offset = actual.size(); + format::FrqRegionMeta second_meta; + ASSERT_TRUE(format::build_freq_region(second_freqs, 0, &actual, &second_meta).ok()); + + ByteSink expected; + expected.put_fixed32(0x87654321); + expected.put_bytes(first_reference.view()); + expected.put_bytes(second_reference.view()); + EXPECT_EQ(actual.buffer(), expected.buffer()); + expect_raw_meta(first_meta, first_reference.view()); + expect_raw_meta(second_meta, second_reference.view()); + + std::vector decoded; + ASSERT_TRUE(format::decode_freq_region( + appended_region(actual, first_offset, first_meta.disk_len), first_meta, + first_freqs.size(), &decoded) + .ok()); + EXPECT_EQ(decoded, first_freqs); + ASSERT_TRUE(format::decode_freq_region( + appended_region(actual, second_offset, second_meta.disk_len), + second_meta, second_freqs.size(), &decoded) + .ok()); + EXPECT_EQ(decoded, second_freqs); + } +} + +TEST(SniiFrqDirectEncodeTest, DescendingDocInLaterRunLeavesOutputsUnchanged) { + std::vector docs(300); + std::iota(docs.begin(), docs.end(), 100); + docs[280] = docs[279] - 1; + ByteSink out; + out.put_fixed32(0x12345678); + const std::vector original_bytes = out.buffer(); + format::FrqRegionMeta meta; + meta.zstd = true; + meta.uncomp_len = 17; + meta.disk_len = 11; + meta.crc = 0x87654321; + meta.verify_crc = false; + const format::FrqRegionMeta original_meta = meta; + + const Status status = format::build_dd_region(docs, 100, 0, &out, &meta); + + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("ascending"), std::string::npos); + EXPECT_EQ(out.buffer(), original_bytes); + EXPECT_EQ(meta.zstd, original_meta.zstd); + EXPECT_EQ(meta.uncomp_len, original_meta.uncomp_len); + EXPECT_EQ(meta.disk_len, original_meta.disk_len); + EXPECT_EQ(meta.crc, original_meta.crc); + EXPECT_EQ(meta.verify_crc, original_meta.verify_crc); +} + +TEST(SniiFrqDirectEncodeTest, WindowedRegionsRoundTripWithAndWithoutFreq) { + for (bool write_freq : {false, true}) { + SCOPED_TRACE(write_freq); + std::vector expected_docs; + std::vector expected_freqs; + writer::TermPostings term; + term.term = "windowed"; + for (uint32_t ordinal = 0; ordinal < format::kSlimDfThreshold; ++ordinal) { + const uint32_t doc = ordinal * 2 + 1; + const uint32_t freq = ordinal % 3 + 1; + expected_docs.push_back(doc); + expected_freqs.push_back(freq); + term.docids.push_back(doc); + term.freqs.push_back(freq); + for (uint32_t position = 0; position < freq; ++position) { + term.positions_flat.push_back(position); + } + } + writer::SniiIndexInput input; + input.index_id = 1; + input.index_suffix = "body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = format::kSlimDfThreshold * 2; + input.write_freq = write_freq; + input.terms.push_back(std::move(term)); + + testing::reset_frq_raw_encode_work(); + snii_test::MemoryFile file; + writer::SniiCompoundWriter compound_writer(&file); + ASSERT_TRUE(compound_writer.add_logical_index(input).ok()); + ASSERT_TRUE(compound_writer.finish().ok()); + + reader::SniiSegmentReader segment_reader; + ASSERT_TRUE(reader::SniiSegmentReader::open(&file, &segment_reader).ok()); + reader::LogicalIndexReader index_reader; + ASSERT_TRUE( + segment_reader.open_index(input.index_id, input.index_suffix, &index_reader).ok()); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + ASSERT_TRUE(index_reader.lookup("windowed", &found, &entry, &frq_base, &prx_base).ok()); + ASSERT_TRUE(found); + ASSERT_EQ(entry.enc, format::DictEntryEnc::kWindowed); + + const uint64_t frq_offset = + index_reader.section_refs().posting_region.offset + frq_base + entry.frq_off_delta; + std::vector frq_bytes; + ASSERT_TRUE(file.read_at(frq_offset, entry.frq_len, &frq_bytes).ok()); + format::FrqPreludeReader prelude; + ASSERT_TRUE( + format::FrqPreludeReader::open( + Slice(frq_bytes.data(), static_cast(entry.prelude_len)), &prelude) + .ok()); + ASSERT_EQ(prelude.n_windows(), 2); + EXPECT_EQ(prelude.has_freq(), write_freq); + EXPECT_TRUE(prelude.has_prx()); + EXPECT_EQ(entry.frq_len, + entry.prelude_len + prelude.dd_block_len() + prelude.freq_block_len()); + + size_t doc_begin = 0; + for (uint32_t window = 0; window < prelude.n_windows(); ++window) { + format::WindowMeta window_meta; + ASSERT_TRUE(prelude.window(window, &window_meta).ok()); + format::FrqRegionMeta dd_meta; + dd_meta.zstd = window_meta.dd_zstd; + dd_meta.uncomp_len = window_meta.dd_uncomp_len; + dd_meta.disk_len = window_meta.dd_disk_len; + dd_meta.crc = window_meta.crc_dd; + const size_t dd_offset = static_cast(entry.prelude_len + window_meta.dd_off); + std::vector decoded_docs; + ASSERT_TRUE( + format::decode_dd_region(Slice(frq_bytes.data() + dd_offset, + static_cast(window_meta.dd_disk_len)), + dd_meta, window_meta.win_base, &decoded_docs) + .ok()); + const std::vector window_docs( + expected_docs.begin() + doc_begin, + expected_docs.begin() + doc_begin + window_meta.doc_count); + EXPECT_EQ(decoded_docs, window_docs); + + if (write_freq) { + format::FrqRegionMeta freq_meta; + freq_meta.zstd = window_meta.freq_zstd; + freq_meta.uncomp_len = window_meta.freq_uncomp_len; + freq_meta.disk_len = window_meta.freq_disk_len; + freq_meta.crc = window_meta.crc_freq; + const size_t freq_offset = static_cast( + entry.prelude_len + prelude.dd_block_len() + window_meta.freq_off); + std::vector decoded_freqs; + ASSERT_TRUE(format::decode_freq_region( + Slice(frq_bytes.data() + freq_offset, + static_cast(window_meta.freq_disk_len)), + freq_meta, window_meta.doc_count, &decoded_freqs) + .ok()); + const std::vector window_freqs( + expected_freqs.begin() + doc_begin, + expected_freqs.begin() + doc_begin + window_meta.doc_count); + EXPECT_EQ(decoded_freqs, window_freqs); + if (window == 1) { + std::vector corrupted( + frq_bytes.begin() + freq_offset, + frq_bytes.begin() + freq_offset + window_meta.freq_disk_len); + corrupted.front() ^= 1; + EXPECT_FALSE(format::decode_freq_region(Slice(corrupted), freq_meta, + window_meta.doc_count, &decoded_freqs) + .ok()); + } + } + if (window == 1) { + std::vector corrupted( + frq_bytes.begin() + dd_offset, + frq_bytes.begin() + dd_offset + window_meta.dd_disk_len); + corrupted.front() ^= 1; + EXPECT_FALSE(format::decode_dd_region(Slice(corrupted), dd_meta, + window_meta.win_base, &decoded_docs) + .ok()); + } + doc_begin += window_meta.doc_count; + } + EXPECT_EQ(doc_begin, expected_docs.size()); + EXPECT_EQ(testing::frq_dd_validation_doc_visits(), 0); + EXPECT_EQ(testing::frq_dd_materialized_values(), 0); + EXPECT_EQ(testing::frq_raw_region_copy_bytes(), 0); + } +} + +} // namespace +} // namespace doris::snii diff --git a/be/test/storage/index/snii_frq_prelude_test.cpp b/be/test/storage/index/snii_frq_prelude_test.cpp new file mode 100644 index 00000000000000..b45802f6646a6a --- /dev/null +++ b/be/test/storage/index/snii_frq_prelude_test.cpp @@ -0,0 +1,520 @@ +// 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. + +// T18 (FP-01..FP-09): FrqPrelude window-row trim. The trimmed on-disk row drops +// the three running-sum offsets (dd_off/freq_off/prx_off) and the two raw-region +// uncomp_len fields; the reader DERIVES them. These tests pin round-trip +// equivalence of the derived WindowMeta, the exact/shrunken byte size, the +// cross-super-block accumulation, the conditional zstd uncomp_len, corruption on +// truncation / sum-overflow, and end-to-end query equivalence with a smaller +// prelude fetch. + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/frq_prelude.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii_query_test_util.h" + +using doris::Status; +using doris::snii::ByteSink; +using doris::snii::Slice; +using doris::snii::varint_len; +using doris::snii::format::build_frq_prelude; +using doris::snii::format::DictEntry; +using doris::snii::format::FrqPreludeColumns; +using doris::snii::format::FrqPreludeReader; +using doris::snii::format::WindowMeta; +using doris::snii::query::term_query; +using doris::snii::reader::LogicalIndexReader; +using doris::snii::reader::SniiSegmentReader; +using doris::snii::snii_test::assert_ok; +using doris::snii::snii_test::build_reader; +using doris::snii::snii_test::MemoryFile; + +namespace { + +// Writer constant (logical_index_writer.cpp kPreludeGroupSize): windows per +// super-block. Used both by the make_windows fixture and to reconstruct the E2E +// term's columns for the size model. +constexpr uint32_t kPreludeGroupSize = 64; + +// --------------------------------------------------------------------------- +// Exact byte-size model of the prelude. row_bytes/prelude_size mirror +// encode_window_row / build_frq_prelude field-for-field so prelude_size(kNew) +// reproduces the real build output exactly (cross-checked by EXPECT_EQ against +// sink.size() in every round-trip test), while prelude_size(kOld) models the +// pre-T18 layout (three extra offset varints + unconditional uncomp_len). +// --------------------------------------------------------------------------- +enum class Layout { kNew, kOld }; + +size_t row_bytes(const WindowMeta& m, uint64_t last_docid_delta, bool has_freq, bool has_prx, + Layout layout) { + const bool old_layout = layout == Layout::kOld; + size_t n = 0; + n += varint_len(last_docid_delta); + n += varint_len(m.doc_count); + n += 1; // win_mode (u8) + if (old_layout) { + n += varint_len(m.dd_off); + } + n += varint_len(m.dd_disk_len); + if (old_layout || m.dd_zstd) { + n += varint_len(m.dd_uncomp_len); + } + n += sizeof(uint32_t); // crc_dd + if (has_freq) { + if (old_layout) { + n += varint_len(m.freq_off); + } + n += varint_len(m.freq_disk_len); + if (old_layout || m.freq_zstd) { + n += varint_len(m.freq_uncomp_len); + } + n += sizeof(uint32_t); // crc_freq + } + if (has_prx) { + if (old_layout) { + n += varint_len(m.prx_off); + } + n += varint_len(m.prx_len); + } + n += varint_len(m.max_freq); + n += 1; // max_norm (u8) + return n; +} + +size_t prelude_size(const FrqPreludeColumns& cols, Layout layout) { + const uint64_t g = cols.group_size; + const size_t n = cols.windows.size(); + const size_t n_super = (g == 0) ? 0 : (n + static_cast(g) - 1) / static_cast(g); + + // Per-super-block window-block byte length + absolute last docid. + std::vector block_len(n_super, 0); + std::vector block_last(n_super, 0); + uint64_t prev_last = 0; + size_t s = 0; + for (size_t start = 0; start < n; start += static_cast(g), ++s) { + const size_t end = std::min(n, start + static_cast(g)); + size_t bl = 0; + for (size_t w = start; w < end; ++w) { + const uint64_t delta = cols.windows[w].last_docid - prev_last; + bl += row_bytes(cols.windows[w], delta, cols.has_freq, cols.has_prx, layout); + prev_last = cols.windows[w].last_docid; + } + block_len[s] = bl; + block_last[s] = prev_last; + } + + // super_block_dir: VInt last_docid_delta + VInt block_off + VInt block_len. + size_t sbdir_len = 0; + uint64_t dir_prev_last = 0; + uint64_t block_off = 0; + for (size_t i = 0; i < n_super; ++i) { + sbdir_len += varint_len(block_last[i] - dir_prev_last); + sbdir_len += varint_len(block_off); + sbdir_len += varint_len(block_len[i]); + dir_prev_last = block_last[i]; + block_off += block_len[i]; + } + + size_t window_region = 0; + for (size_t i = 0; i < n_super; ++i) { + window_region += block_len[i]; + } + + // covered = u8 flags + VInt N + VInt G + VInt n_super + VInt sbdir_len + sbdir. + const size_t covered = 1 + varint_len(n) + varint_len(g) + varint_len(n_super) + + varint_len(sbdir_len) + sbdir_len; + return covered + sizeof(uint32_t) /*crc*/ + window_region; +} + +// Builds n contiguous stride-100 windows with deterministic, all-raw metadata. +// dd/freq/prx offsets are the CONTIGUOUS running prefix sums the reader derives; +// raw uncomp_len == disk_len. Region disk lengths are constant so the NEW row is a +// fixed 16 B (docs+positions) / 10 B (docs-only), while the OLD row's extra offset +// varints grow past one byte for most windows (~24 B / ~13 B). +FrqPreludeColumns make_windows(uint32_t n, uint32_t group_size, bool has_freq, bool has_prx) { + FrqPreludeColumns cols; + cols.has_freq = has_freq; + cols.has_prx = has_prx; + cols.group_size = group_size; + uint64_t dd_run = 0; + uint64_t freq_run = 0; + uint64_t prx_run = 0; + for (uint32_t w = 0; w < n; ++w) { + WindowMeta m; + m.last_docid = (w + 1) * 100 - 1; // window w covers docids [w*100, w*100+99] + m.doc_count = 10; + m.dd_zstd = false; + m.dd_off = dd_run; + m.dd_disk_len = 64; + m.dd_uncomp_len = 64; // raw => == disk_len + m.crc_dd = 0xDD000000U + w; + dd_run += m.dd_disk_len; + if (has_freq) { + m.freq_zstd = false; + m.freq_off = freq_run; + m.freq_disk_len = 48; + m.freq_uncomp_len = 48; + m.crc_freq = 0xEE000000U + w; + freq_run += m.freq_disk_len; + } + if (has_prx) { + m.prx_off = prx_run; + m.prx_len = 96; + prx_run += m.prx_len; + } + m.max_freq = 5; + m.max_norm = 42; + cols.windows.push_back(m); + } + return cols; +} + +// Round-trips columns through build + open, asserting every decoded WindowMeta +// field (incl. the DERIVED dd_off/freq_off/prx_off/uncomp_len) matches the input, +// and that the real build size equals the exact NEW model. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +void expect_round_trip(const FrqPreludeColumns& cols) { + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + ASSERT_EQ(reader.n_windows(), cols.windows.size()); + EXPECT_EQ(reader.has_freq(), cols.has_freq); + EXPECT_EQ(reader.has_prx(), cols.has_prx); + + uint64_t expect_win_base = 0; + uint64_t dd_run = 0; + uint64_t freq_run = 0; + uint64_t prx_run = 0; + for (uint32_t w = 0; w < reader.n_windows(); ++w) { + WindowMeta got; + ASSERT_TRUE(reader.window(w, &got).ok()) << "w=" << w; + const WindowMeta& exp = cols.windows[w]; + EXPECT_EQ(got.last_docid, exp.last_docid) << "w=" << w; + EXPECT_EQ(got.win_base, expect_win_base) << "w=" << w; + EXPECT_EQ(got.doc_count, exp.doc_count) << "w=" << w; + EXPECT_EQ(got.dd_zstd, exp.dd_zstd) << "w=" << w; + EXPECT_EQ(got.dd_off, exp.dd_off) << "w=" << w; // derived == explicit sum + EXPECT_EQ(got.dd_off, dd_run) << "w=" << w; // ...and == our running sum + EXPECT_EQ(got.dd_disk_len, exp.dd_disk_len) << "w=" << w; + EXPECT_EQ(got.dd_uncomp_len, exp.dd_uncomp_len) << "w=" << w; + EXPECT_EQ(got.crc_dd, exp.crc_dd) << "w=" << w; + dd_run += got.dd_disk_len; + if (cols.has_freq) { + EXPECT_EQ(got.freq_zstd, exp.freq_zstd) << "w=" << w; + EXPECT_EQ(got.freq_off, exp.freq_off) << "w=" << w; + EXPECT_EQ(got.freq_off, freq_run) << "w=" << w; + EXPECT_EQ(got.freq_disk_len, exp.freq_disk_len) << "w=" << w; + EXPECT_EQ(got.freq_uncomp_len, exp.freq_uncomp_len) << "w=" << w; + EXPECT_EQ(got.crc_freq, exp.crc_freq) << "w=" << w; + freq_run += got.freq_disk_len; + } + if (cols.has_prx) { + EXPECT_EQ(got.prx_off, exp.prx_off) << "w=" << w; + EXPECT_EQ(got.prx_off, prx_run) << "w=" << w; + EXPECT_EQ(got.prx_len, exp.prx_len) << "w=" << w; + prx_run += got.prx_len; + } + EXPECT_EQ(got.max_freq, exp.max_freq) << "w=" << w; + EXPECT_EQ(got.max_norm, exp.max_norm) << "w=" << w; + expect_win_base = exp.last_docid; + } + EXPECT_EQ(reader.dd_block_len(), dd_run); + if (cols.has_freq) { + EXPECT_EQ(reader.freq_block_len(), freq_run); + } + EXPECT_EQ(sink.size(), prelude_size(cols, Layout::kNew)); // model tracks the real encoder +} + +} // namespace + +// FP-01: N=200 windows (has_freq+has_prx, all raw) round-trip; the derived +// offsets/uncomp_lens equal the explicit running sums the columns carried. +TEST(SniiFrqPreludeTest, FP01DerivedOffsetsMatchExplicit) { + expect_round_trip(make_windows(/*n=*/200, kPreludeGroupSize, /*has_freq=*/true, + /*has_prx=*/true)); +} + +// FP-02: the trimmed prelude hits an EXACT byte size and is strictly smaller than +// the pre-T18 layout, for both docs+positions and docs-only rows. +TEST(SniiFrqPreludeTest, FP02RowTrimShrinksPreludeBytes) { + // docs+positions: drops 5 varints/row (dd_off, dd_uncomp_len, freq_off, + // freq_uncomp_len, prx_off). + const FrqPreludeColumns dp = make_windows(200, kPreludeGroupSize, /*has_freq=*/true, + /*has_prx=*/true); + ByteSink dp_sink; + ASSERT_TRUE(build_frq_prelude(dp, &dp_sink).ok()); + const size_t dp_new = prelude_size(dp, Layout::kNew); + const size_t dp_old = prelude_size(dp, Layout::kOld); + EXPECT_EQ(dp_sink.size(), dp_new); // exact new size + EXPECT_LT(dp_new, dp_old); // strictly < pre-T18 + EXPECT_GE(dp_old - dp_new, 5U * dp.windows.size()); // >= 5 bytes/row removed + EXPECT_GT((dp_old - dp_new) * 100, dp_old * 25); // ~1/3 saved (30-40% band) + + // docs-only: no freq / no positions -> drops only dd_off + dd_uncomp_len. + const FrqPreludeColumns doc = make_windows(200, kPreludeGroupSize, /*has_freq=*/false, + /*has_prx=*/false); + ByteSink doc_sink; + ASSERT_TRUE(build_frq_prelude(doc, &doc_sink).ok()); + const size_t doc_new = prelude_size(doc, Layout::kNew); + const size_t doc_old = prelude_size(doc, Layout::kOld); + EXPECT_EQ(doc_sink.size(), doc_new); + EXPECT_LT(doc_new, doc_old); + EXPECT_GE(doc_old - doc_new, 2U * doc.windows.size()); // >= 2 bytes/row removed +} + +// FP-03: degenerate empty prelude (N=0) builds + opens with zero windows. +TEST(SniiFrqPreludeTest, FP03EmptyWindows) { + FrqPreludeColumns cols; + cols.has_freq = true; + cols.has_prx = true; + cols.group_size = kPreludeGroupSize; + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + EXPECT_EQ(reader.n_windows(), 0U); + EXPECT_EQ(reader.n_super_blocks(), 0U); + EXPECT_EQ(reader.dd_block_len(), 0U); + EXPECT_EQ(reader.freq_block_len(), 0U); + EXPECT_EQ(sink.size(), prelude_size(cols, Layout::kNew)); +} + +// FP-04: single-window round-trip (win_base=0, one super-block). +TEST(SniiFrqPreludeTest, FP04SingleWindow) { + expect_round_trip( + make_windows(/*n=*/1, kPreludeGroupSize, /*has_freq=*/true, /*has_prx=*/true)); +} + +// FP-05: a df~20000 term (adaptive unit=1024 -> ~20 windows) split into multiple +// super-blocks; the derived dd/freq/prx offsets must accumulate ACROSS block +// boundaries (not reset per block). +TEST(SniiFrqPreludeTest, FP05CrossSuperBlockAccumulation) { + constexpr uint32_t kN = 20; // ceil(20000 / kAdaptiveWindowDocs=1024) + constexpr uint32_t kG = 8; // force 3 super-blocks (8, 8, 4) + const FrqPreludeColumns cols = make_windows(kN, kG, /*has_freq=*/true, /*has_prx=*/true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + ASSERT_EQ(reader.n_windows(), kN); + ASSERT_GT(reader.n_super_blocks(), 1U); // genuinely multi-block + + uint64_t dd_run = 0; + uint64_t freq_run = 0; + uint64_t prx_run = 0; + for (uint32_t w = 0; w < kN; ++w) { + WindowMeta got; + ASSERT_TRUE(reader.window(w, &got).ok()) << "w=" << w; + EXPECT_EQ(got.dd_off, dd_run) << "w=" << w; // continuous over all blocks + EXPECT_EQ(got.freq_off, freq_run) << "w=" << w; + EXPECT_EQ(got.prx_off, prx_run) << "w=" << w; + dd_run += got.dd_disk_len; + freq_run += got.freq_disk_len; + prx_run += got.prx_len; + } + EXPECT_EQ(reader.dd_block_len(), dd_run); + EXPECT_EQ(reader.freq_block_len(), freq_run); +} + +// FP-06: a zstd window stores its uncomp_len (!= disk_len) and reads it back, while +// a raw window in the same prelude derives uncomp_len == disk_len (no stored field). +TEST(SniiFrqPreludeTest, FP06ZstdWindowStoresConditionalUncompLen) { + FrqPreludeColumns cols; + cols.has_freq = true; + cols.has_prx = false; + cols.group_size = kPreludeGroupSize; + + WindowMeta z; // window 0: zstd dd + zstd freq, uncomp != disk + z.last_docid = 99; + z.doc_count = 10; + z.dd_zstd = true; + z.dd_off = 0; + z.dd_disk_len = 40; + z.dd_uncomp_len = 137; // meaningful only because dd_zstd + z.crc_dd = 0x0000ABCDU; + z.freq_zstd = true; + z.freq_off = 0; + z.freq_disk_len = 20; + z.freq_uncomp_len = 71; + z.crc_freq = 0x00001234U; + z.max_freq = 3; + z.max_norm = 7; + + WindowMeta r; // window 1: raw dd + raw freq, uncomp == disk + r.last_docid = 199; + r.doc_count = 10; + r.dd_zstd = false; + r.dd_off = z.dd_disk_len; + r.dd_disk_len = 55; + r.dd_uncomp_len = 55; + r.crc_dd = 0x0000BEEFU; + r.freq_zstd = false; + r.freq_off = z.freq_disk_len; + r.freq_disk_len = 25; + r.freq_uncomp_len = 25; + r.crc_freq = 0x00005678U; + r.max_freq = 4; + r.max_norm = 9; + cols.windows = {z, r}; + + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + FrqPreludeReader reader; + ASSERT_TRUE(FrqPreludeReader::open(sink.view(), &reader).ok()); + + WindowMeta g0; + WindowMeta g1; + ASSERT_TRUE(reader.window(0, &g0).ok()); + ASSERT_TRUE(reader.window(1, &g1).ok()); + EXPECT_TRUE(g0.dd_zstd); + EXPECT_EQ(g0.dd_disk_len, 40U); + EXPECT_EQ(g0.dd_uncomp_len, 137U); // stored + read back + EXPECT_NE(g0.dd_uncomp_len, g0.dd_disk_len); + EXPECT_TRUE(g0.freq_zstd); + EXPECT_EQ(g0.freq_uncomp_len, 71U); + EXPECT_NE(g0.freq_uncomp_len, g0.freq_disk_len); + EXPECT_FALSE(g1.dd_zstd); + EXPECT_EQ(g1.dd_uncomp_len, g1.dd_disk_len); // raw: derived, not stored + EXPECT_EQ(g1.freq_uncomp_len, g1.freq_disk_len); + EXPECT_EQ(sink.size(), prelude_size(cols, Layout::kNew)); +} + +// FP-07: a prelude truncated into its window blocks fails to open. +TEST(SniiFrqPreludeTest, FP07TruncatedPreludeIsCorruption) { + const FrqPreludeColumns cols = make_windows(8, 4, /*has_freq=*/true, /*has_prx=*/true); + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); + std::vector bytes = sink.buffer(); + ASSERT_GT(bytes.size(), 6U); + bytes.resize(bytes.size() - 6); // chop into the last window block + FrqPreludeReader reader; + const Status s = FrqPreludeReader::open(Slice(bytes), &reader); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// FP-08: two raw windows whose dd_disk_len values overflow the derived dd-block +// offset accumulator (checked_add_u64) are rejected on open. +TEST(SniiFrqPreludeTest, FP08DdDiskLenSumOverflowIsCorruption) { + FrqPreludeColumns cols; + cols.has_freq = false; + cols.has_prx = false; + cols.group_size = kPreludeGroupSize; + + WindowMeta a; + a.last_docid = 10; + a.doc_count = 1; + a.dd_disk_len = std::numeric_limits::max() / 2 + 100; + a.dd_uncomp_len = a.dd_disk_len; // raw + a.crc_dd = 1; + a.max_freq = 1; + a.max_norm = 0; + WindowMeta b = a; + b.last_docid = 20; // second window's derivation overflows: sum > UINT64_MAX + b.crc_dd = 2; + cols.windows = {a, b}; + + ByteSink sink; + ASSERT_TRUE(build_frq_prelude(cols, &sink).ok()); // builder does not sum-check + FrqPreludeReader reader; + const Status s = FrqPreludeReader::open(sink.view(), &reader); + EXPECT_TRUE(s.is()) << s.to_string(); +} + +// FP-09 (E2E): a high-df windowed term ("failed", 9000 docs) queried through the +// real reader returns the full docid set (new path == old semantics), the docid +// path fetches exactly the (now-smaller) prelude range once, and the on-disk +// prelude_len matches the NEW model and is strictly smaller than the pre-T18 one. +TEST(SniiFrqPreludeTest, FP09WindowedEntryPreludeShrinksAndQueryEquivalent) { + MemoryFile file; + SniiSegmentReader segment_reader; + LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + bool found = false; + DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup("failed", &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found); + ASSERT_GT(entry.prelude_len, 0U) << "expected a windowed entry with a prelude"; + + const uint64_t prelude_abs = + index_reader.section_refs().posting_region.offset + frq_base + entry.frq_off_delta; + + // Reconstruct the term's windows from the (new) prelude and model both layouts: + // the on-disk length must equal the NEW model and be strictly below the OLD one. + std::vector prelude_bytes; + assert_ok(file.read_at(prelude_abs, entry.prelude_len, &prelude_bytes)); + FrqPreludeReader prelude; + assert_ok(FrqPreludeReader::open(Slice(prelude_bytes), &prelude)); + ASSERT_GT(prelude.n_windows(), 1U); // genuinely windowed + + FrqPreludeColumns cols; + cols.has_freq = prelude.has_freq(); + cols.has_prx = prelude.has_prx(); + cols.group_size = kPreludeGroupSize; + for (uint32_t w = 0; w < prelude.n_windows(); ++w) { + WindowMeta m; + assert_ok(prelude.window(w, &m)); + cols.windows.push_back(m); + } + EXPECT_EQ(entry.prelude_len, prelude_size(cols, Layout::kNew)); // on-disk == new model + EXPECT_LT(entry.prelude_len, prelude_size(cols, Layout::kOld)); // shrinks vs pre-T18 + + // Functional equivalence + deterministic IO on the docid path. + file.clear_reads(); + std::vector docids; + assert_ok(term_query(index_reader, "failed", &docids)); + + std::vector expected(9000); + std::iota(expected.begin(), expected.end(), 0U); + EXPECT_EQ(docids, expected); + + // The docid path issues ONE coalesced read starting at the prelude and spanning the + // prelude + dd-block (docid-only: the freq-block is not fetched), i.e. a single + // round-trip whose length lies within [prelude_len, frq_len]. The prelude — now the + // shrunk new-model size (prelude_len == new < old, asserted above) — is the head of + // that read, so the end-to-end docid fetch is strictly smaller than pre-T18. + const auto frq_region_reads = + std::ranges::count_if(file.reads(), [&](const MemoryFile::Read& r) { + return r.offset == prelude_abs && r.len >= entry.prelude_len && + r.len <= entry.frq_len; + }); + EXPECT_EQ(frq_region_reads, 1) + << "docid path must fetch the frq region in a single round-trip at the prelude; " + << "prelude_abs=" << prelude_abs << " frq_len=" << entry.frq_len + << " prelude_len=" << entry.prelude_len; +} diff --git a/be/test/storage/index/snii_query_test.cpp b/be/test/storage/index/snii_query_test.cpp new file mode 100644 index 00000000000000..2df6cd87af3c33 --- /dev/null +++ b/be/test/storage/index/snii_query_test.cpp @@ -0,0 +1,2260 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/encoding/byte_source.h" +#include "storage/index/snii/encoding/pfor.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/prx_pod.h" +#include "storage/index/snii/format/tail_pointer.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/query/internal/regex_prefix.h" +#include "storage/index/snii/query/internal/resolved_phrase_plan.h" +#include "storage/index/snii/query/internal/term_expansion.h" +#include "storage/index/snii/query/phrase_query.h" +#include "storage/index/snii/query/prefix_query.h" +#include "storage/index/snii/query/regexp_query.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/query/wildcard_query.h" +// T24 query op-count seam. Define the gate before the include so QueryTestCounters +// is visible in this TU; it is also auto-enabled library-wide by BE_TEST, so the +// phrase_query.cpp increments and the reads below share the same singleton. +#define SNII_QUERY_TEST_COUNTERS +#include "storage/index/snii/query/internal/query_test_counters.h" +#include "storage/index/snii/reader/dict_block_cache.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "storage/index/snii_query_test_util.h" + +namespace doris::snii::query { +using doris::Status; // RETURN_IF_ERROR expands to bare Status +namespace { + +// Shared reader/writer fixtures live in snii_query_test_util.h so other SNII test +// files can reuse them; pull them into this suite's scope unqualified. +using snii_test::assert_ok; +using snii_test::build_reader; +using snii_test::make_term; +using snii_test::MemoryFile; +using snii_test::PostingDoc; +using snii_test::ScopedEnv; + +template +concept CanExecuteResolvedPhrasePlan = requires(const reader::LogicalIndexReader& idx, Plan&& plan, + std::vector* docids) { + internal::execute_resolved_phrase_plan(idx, std::forward(plan), docids); +}; + +static_assert(CanExecuteResolvedPhrasePlan); +static_assert(!CanExecuteResolvedPhrasePlan); + +class RecordingDocIdSink final : public DocIdSink { +public: + Status append_sorted(std::span docids) override { + out.insert(out.end(), docids.begin(), docids.end()); + return Status::OK(); + } + + Status append_range(uint32_t first, uint64_t last_exclusive) override { + ++range_calls; + for (uint64_t docid = first; docid < last_exclusive; ++docid) { + out.push_back(static_cast(docid)); + } + return Status::OK(); + } + + std::vector out; + size_t range_calls = 0; +}; + +struct PrxRange { + uint64_t offset = 0; + uint64_t len = 0; +}; + +void assert_selective_prx_matches_constant_positions(Slice window, + const std::vector& selected_docs, + uint32_t expected_position) { + std::vector selected_positions; + std::vector selected_offsets; + ByteSource selected_source(window); + assert_ok(format::read_prx_window_csr_selective(&selected_source, selected_docs, + &selected_positions, &selected_offsets)); + + ASSERT_EQ(selected_offsets.size(), selected_docs.size() + 1); + ASSERT_EQ(selected_positions.size(), selected_docs.size()); + for (size_t i = 0; i < selected_docs.size(); ++i) { + EXPECT_EQ(selected_offsets[i], i); + EXPECT_EQ(selected_positions[i], expected_position); + } + EXPECT_EQ(selected_offsets.back(), selected_positions.size()); +} + +// A FileReader decorator that counts read_batch() invocations. Each BatchRangeFetcher +// fetch() barrier issues exactly one read_batch (== one batched/remote serial round), +// so this isolates the number of I/O rounds a query plan emits. Single read_at() calls +// (dict-block / BSBF loads) delegate straight through and are intentionally not counted. +class BatchRoundCountingReader final : public doris::snii::io::FileReader { +public: + explicit BatchRoundCountingReader(doris::snii::io::FileReader* inner) : inner_(inner) {} + + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + return inner_->read_at(offset, len, out); + } + Status read_batch(const std::vector& ranges, + std::vector>* outs) override { + ++batch_rounds_; + return inner_->read_batch(ranges, outs); + } + uint64_t size() const override { return inner_->size(); } + + size_t batch_rounds() const { return batch_rounds_; } + void reset_rounds() { batch_rounds_ = 0; } + +private: + doris::snii::io::FileReader* inner_; + size_t batch_rounds_ = 0; +}; + +// 480 docids with irregular gaps (deterministic LCG): the slim docs region PFOR +// then exceeds the 256B inline threshold so each term becomes a pod_ref, while +// df < 512 keeps it slim (not windowed). This is the layout that forced one PRX +// fetch() per term before T02. +std::vector slim_pod_ref_docids() { + std::vector ids; + ids.reserve(480); + uint32_t cur = 0; + uint32_t state = 0x9e3779b9U; + for (int i = 0; i < 480; ++i) { + ids.push_back(cur); + state = state * 1664525U + 1013904223U; + cur += 1U + (state >> 23) % 250U; // gap in [1, 250] + } + return ids; +} + +// Builds an index with three overlapping slim pod_ref terms ("paa"/"pbb"/"pcc") +// sharing one docid set, each with a single position (5/6/7) so the phrase +// "paa pbb pcc" matches every shared doc. The index is opened through +// `read_through` (e.g. a counting decorator wrapping `file`). `shared_docids` +// returns the docid set, which equals the expected phrase result. +Status build_slim_pod_ref_phrase_reader(MemoryFile* file, doris::snii::io::FileReader* read_through, + reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader, + std::vector* shared_docids) { + const std::vector docids = slim_pod_ref_docids(); + *shared_docids = docids; + auto make_pos_term = [&](std::string term, uint32_t position) { + std::vector docs; + docs.reserve(docids.size()); + for (uint32_t docid : docids) { + docs.push_back({docid, {position}}); + } + return make_term(std::move(term), std::move(docs)); + }; + + writer::SniiIndexInput input; + input.index_id = 11; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = docids.back() + 1; + input.terms = {make_pos_term("paa", 5), make_pos_term("pbb", 6), make_pos_term("pcc", 7)}; + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + + RETURN_IF_ERROR(reader::SniiSegmentReader::open(read_through, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +writer::SniiIndexInput make_many_term_input(uint64_t index_id, std::string suffix, + uint32_t n_terms) { + writer::SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = std::move(suffix); + input.config = format::IndexConfig::kDocsOnly; + input.doc_count = n_terms + 1; + input.target_dict_block_bytes = 128; + input.terms.reserve(n_terms); + for (uint32_t i = 0; i < n_terms; ++i) { + input.terms.push_back( + make_term("term_" + std::to_string(1000000 + i), {{.docid = i, .positions = {0}}})); + } + return input; +} + +format::TailPointer read_tail_pointer(MemoryFile* file) { + std::vector bytes; + assert_ok(file->read_at(file->size() - format::tail_pointer_size(), format::tail_pointer_size(), + &bytes)); + format::TailPointer tp; + assert_ok(format::decode_tail_pointer(Slice(bytes), &tp)); + return tp; +} + +TEST(SniiSegmentReaderTest, OpenReadsOnlyFooterAndDirectory) { + MemoryFile file; + writer::SniiCompoundWriter writer(&file); + assert_ok(writer.add_logical_index(make_many_term_input(7, "Body", 4096))); + assert_ok(writer.finish()); + + const format::TailPointer tp = read_tail_pointer(&file); + file.clear_reads(); + + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + + ASSERT_EQ(file.reads().size(), 2U); + EXPECT_EQ(file.reads()[0].offset, file.size() - format::tail_pointer_size()); + EXPECT_EQ(file.reads()[0].len, format::tail_pointer_size()); + EXPECT_EQ(file.reads()[1].offset, tp.directory_offset); + EXPECT_EQ(file.reads()[1].len, tp.directory_length); + EXPECT_EQ(file.read_bytes(), format::tail_pointer_size() + tp.directory_length); +} + +TEST(SniiSegmentReaderTest, IndexExistsUsesCachedTailDirectory) { + MemoryFile file; + writer::SniiCompoundWriter writer(&file); + writer::SniiIndexInput input = make_many_term_input(7, "Body", 4096); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + + file.clear_reads(); + bool exists = false; + assert_ok(segment_reader.index_exists(input.index_id, input.index_suffix, &exists)); + EXPECT_TRUE(exists); + EXPECT_EQ(file.read_bytes(), 0); + + assert_ok(segment_reader.index_exists(input.index_id + 1, input.index_suffix, &exists)); + EXPECT_FALSE(exists); + EXPECT_EQ(file.read_bytes(), 0); +} + +// F35: memory_usage() feeds the InvertedIndexSearcherCache charge and must now +// account for the resident sampled term index + DICT block directory (previously +// omitted -> under-charge -> over-commit). Exact per-field equality would need +// the reader's private members; the observable public properties asserted here +// are a charge floor (>= sizeof) and monotonic growth with the vocabulary (more +// DICT blocks -> larger sti_/dbd_ heap). The exact heap_bytes() formula the fix +// sums is pinned deterministically by the SampledTermIndex / DictBlockDirectory / +// DictBlock unit tests in snii_writer_test.cpp. +TEST(SniiSegmentReaderTest, MemoryUsageGrowsWithVocabulary) { + auto open_many_term_reader = [](uint32_t n_terms, MemoryFile* file, + reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader) { + writer::SniiCompoundWriter writer(file); + // target_dict_block_bytes == 128 (make_many_term_input) -> many small DICT + // blocks, so a larger vocabulary yields more sampled terms + block refs. + assert_ok(writer.add_logical_index(make_many_term_input(7, "Body", n_terms))); + assert_ok(writer.finish()); + assert_ok(reader::SniiSegmentReader::open(file, segment_reader)); + assert_ok(segment_reader->open_index(7, "Body", index_reader)); + }; + + MemoryFile small_file; + reader::SniiSegmentReader small_segment; + reader::LogicalIndexReader small_reader; + open_many_term_reader(64, &small_file, &small_segment, &small_reader); + + MemoryFile large_file; + reader::SniiSegmentReader large_segment; + reader::LogicalIndexReader large_reader; + open_many_term_reader(4096, &large_file, &large_segment, &large_reader); + + EXPECT_GE(small_reader.memory_usage(), sizeof(reader::LogicalIndexReader)); + EXPECT_GT(large_reader.memory_usage(), small_reader.memory_usage()); +} + +// G13 end-to-end: a many-term segment's adjacent Core/STI/DBD metadata group +// (the first serial fetch of a cold open, dominated by the STI + DBD tables) +// must SHRINK on disk once those sections are zstd-compressed, the open must +// therefore fetch fewer bytes, and every lookup / prefix enumeration must stay +// equal to an uncompressed control built from the identical input. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST(SniiSegmentReaderTest, MetaCompressionShrinksOpenFetchAndKeepsLookupsEqual) { + constexpr uint32_t kTerms = 4096; + auto build = [](MemoryFile* file) { + writer::SniiCompoundWriter writer(file); + assert_ok(writer.add_logical_index(make_many_term_input(7, "Body", kTerms))); + assert_ok(writer.finish()); + }; + + MemoryFile compressed_file; + build(&compressed_file); // default: G13 compression active + MemoryFile control_file; + { + ScopedEnv off("SNII_META_COMPRESS_MIN", "18446744073709551615"); + build(&control_file); // pre-G13 layout from the identical input + } + + reader::SniiSegmentReader compressed_segment; + assert_ok(reader::SniiSegmentReader::open(&compressed_file, &compressed_segment)); + reader::SniiSegmentReader control_segment; + assert_ok(reader::SniiSegmentReader::open(&control_file, &control_segment)); + + // Cold index open fetches the metadata group once, and fewer bytes end to + // end (the DICT blocks / BSBF bytes are identical). + compressed_file.clear_reads(); + reader::LogicalIndexReader compressed_reader; + assert_ok(compressed_segment.open_index(7, "Body", &compressed_reader)); + const size_t compressed_open_bytes = compressed_file.read_bytes(); + ASSERT_FALSE(compressed_file.reads().empty()); + const size_t compressed_metadata_bytes = compressed_file.reads()[0].len; + + control_file.clear_reads(); + reader::LogicalIndexReader control_reader; + assert_ok(control_segment.open_index(7, "Body", &control_reader)); + const size_t control_open_bytes = control_file.read_bytes(); + ASSERT_FALSE(control_file.reads().empty()); + const size_t control_metadata_bytes = control_file.reads()[0].len; + EXPECT_LT(compressed_metadata_bytes, control_metadata_bytes); + EXPECT_LT(compressed_open_bytes, control_open_bytes); + std::cout << "[G13] metadata group: raw=" << control_metadata_bytes + << "B zstd=" << compressed_metadata_bytes + << "B; open fetch: raw=" << control_open_bytes << "B zstd=" << compressed_open_bytes + << "B\n"; + + // Lookups stay equal across the two layouts: present terms resolve to the + // same entry essentials, absent terms miss on both. + for (uint32_t i = 0; i < kTerms; i += 97) { + const std::string term = "term_" + std::to_string(1000000 + i); + bool found_a = false; + bool found_b = false; + format::DictEntry entry_a; + format::DictEntry entry_b; + uint64_t frq_a = 0; + uint64_t prx_a = 0; + uint64_t frq_b = 0; + uint64_t prx_b = 0; + assert_ok(compressed_reader.lookup(term, &found_a, &entry_a, &frq_a, &prx_a)); + assert_ok(control_reader.lookup(term, &found_b, &entry_b, &frq_b, &prx_b)); + ASSERT_TRUE(found_a) << term; + ASSERT_TRUE(found_b) << term; + EXPECT_EQ(entry_a.term, entry_b.term); + EXPECT_EQ(entry_a.kind, entry_b.kind); + EXPECT_EQ(entry_a.df, entry_b.df); + EXPECT_EQ(frq_a, frq_b); + EXPECT_EQ(prx_a, prx_b); + } + bool found_absent = true; + format::DictEntry absent_entry; + uint64_t frq = 0; + uint64_t prx = 0; + assert_ok( + compressed_reader.lookup("zzzz_not_indexed", &found_absent, &absent_entry, &frq, &prx)); + EXPECT_FALSE(found_absent); + + // Ordered prefix enumeration walks sti + dbd + dict blocks on both layouts + // and must produce the identical term sequence. + std::vector hits_a; + std::vector hits_b; + assert_ok(compressed_reader.prefix_terms("term_10001", &hits_a)); + assert_ok(control_reader.prefix_terms("term_10001", &hits_b)); + ASSERT_EQ(hits_a.size(), hits_b.size()); + ASSERT_EQ(hits_a.size(), 100U); // term_1000100 .. term_1000199 + for (size_t i = 0; i < hits_a.size(); ++i) { + EXPECT_EQ(hits_a[i].term, hits_b[i].term); + EXPECT_EQ(hits_a[i].entry.df, hits_b[i].entry.df); + } +} + +// P1 cold-read fix: a NON-resident bloom is skipped entirely. open() must not +// read the 28B header (which the old L1 path cached) and lookup() must not issue a +// 32B body probe; absent / present terms still resolve correctly via sti -> dict. +TEST(SniiSegmentReaderTest, NonResidentBsbfIsSkippedNotProbed) { + ScopedEnv disable_resident_bsbf("SNII_BSBF_RESIDENT_MAX", "0"); + + MemoryFile file; + writer::SniiCompoundWriter writer(&file); + writer::SniiIndexInput input = make_many_term_input(7, "Body", 1024); + assert_ok(writer.add_logical_index(input)); + assert_ok(writer.finish()); + + reader::SniiSegmentReader segment_reader; + assert_ok(reader::SniiSegmentReader::open(&file, &segment_reader)); + format::SectionRefs refs; + assert_ok(segment_reader.section_refs_for_index(input.index_id, input.index_suffix, &refs)); + ASSERT_GT(refs.bsbf.length, format::kBsbfHeaderSize); // a real filter exists on disk + const uint64_t bsbf_end = refs.bsbf.offset + refs.bsbf.length; + auto touches_bsbf = [&](uint64_t offset, size_t len) { + return offset < bsbf_end && refs.bsbf.offset < offset + len; + }; + + // open() must NOT touch the bsbf section at all on the non-resident path. + file.clear_reads(); + reader::LogicalIndexReader index_reader; + assert_ok(segment_reader.open_index(input.index_id, input.index_suffix, &index_reader)); + for (const auto& read : file.reads()) { + EXPECT_FALSE(touches_bsbf(read.offset, read.len)) + << "non-resident open must skip the bsbf header"; + } + + // An absent term still resolves to empty via sti -> dict, and the lookup must + // NOT probe the bsbf section. + file.clear_reads(); + std::vector docids; + assert_ok(term_query(index_reader, "absent_term", &docids)); + EXPECT_TRUE(docids.empty()); + for (const auto& read : file.reads()) { + EXPECT_FALSE(touches_bsbf(read.offset, read.len)) + << "non-resident lookup must not probe the bsbf section"; + } + + // A present term is still found via sti -> dict (correctness with no bloom). + std::vector present; + assert_ok(term_query(index_reader, "term_1000000", &present)); + EXPECT_FALSE(present.empty()); +} + +TEST(SniiSegmentReaderTest, LogicalIndexOpenCachesResidentMetadataAndSmallHeaders) { + ScopedEnv resident_bsbf("SNII_BSBF_RESIDENT_MAX", "1048576"); + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + file.clear_reads(); + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup("failed", &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + EXPECT_EQ(file.read_bytes(), 0); + + std::vector hits; + assert_ok(index_reader.prefix_terms("ord", &hits, 10)); + ASSERT_EQ(hits.size(), 2); + EXPECT_EQ(hits[0].term, "order"); + EXPECT_EQ(hits[1].term, "ordinal"); + EXPECT_EQ(file.read_bytes(), 0); +} + +TEST(SniiPhraseQueryTest, WindowedPhraseQueryKeepsCorrectCandidateOrdinals) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"failed", "order"}, &docids)); + + const std::vector expected {5000, 7000, 8000}; + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseQueryTest, WindowedPhrasePrefixQueryKeepsCorrectCandidateOrdinals) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"failed", "ord"}, &docids, 10)); + + const std::vector expected {5000, 6000, 7000, 8000}; + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseQueryTest, SingleTailPhrasePrefixUsesStreamingPhrasePath) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"failed", "orde"}, &docids, 10)); + + const std::vector expected {5000, 7000, 8000}; + EXPECT_EQ(docids, expected); +} + +// PV1 (deterministic perf): a 3-term slim pod_ref phrase issues exactly one batched +// PRX round. read_batch count = 1 (round1 docs) + 1 (shared PRX fetch) = 2; the +// docid conjunction decodes slim docs from round1 (no extra round) and dict/BSBF +// loads use uncounted read_at. Before T02 this was 4 (round1 + one PRX fetch per +// term). +TEST(SniiPhraseQueryTest, PhraseQueryIssuesSinglePrxBatchRound) { + MemoryFile file; + BatchRoundCountingReader counting(&file); + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + std::vector shared_docids; + assert_ok(build_slim_pod_ref_phrase_reader(&file, &counting, &segment_reader, &index_reader, + &shared_docids)); + + // Guard: every phrase term must be a slim pod_ref (not inline, not windowed), + // otherwise the per-term PRX fetch this test isolates would not occur. + for (std::string_view term : {"paa", "pbb", "pcc"}) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup(term, &found, &entry, &frq_base, &prx_base)); + ASSERT_TRUE(found) << term; + EXPECT_EQ(entry.kind, format::DictEntryKind::kPodRef) << term; + EXPECT_EQ(entry.enc, format::DictEntryEnc::kSlim) << term; + } + + counting.reset_rounds(); + std::vector docids; + assert_ok(phrase_query(index_reader, {"paa", "pbb", "pcc"}, &docids)); + + EXPECT_EQ(docids, shared_docids); + EXPECT_EQ(counting.batch_rounds(), 2U); +} + +// FV3 (functional equivalence): the shared single-batch PRX path returns the result +// dictated by the position layout across three overlapping slim pod_ref terms. A +// backfill mistake (wrong plan/chunk/handle mapping) would misread a term's +// positions and drop docs. +TEST(SniiPhraseQueryTest, ThreeTermPhraseMatchesAcrossSharedPrxFetch) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + std::vector shared_docids; + assert_ok(build_slim_pod_ref_phrase_reader(&file, &file, &segment_reader, &index_reader, + &shared_docids)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"paa", "pbb", "pcc"}, &docids)); + EXPECT_EQ(docids, shared_docids); + + // Same terms, non-consecutive order (positions 7,6,5): the candidate set is + // identical but no doc satisfies the phrase, so the shared PRX fetch must yield + // an empty result. + std::vector reversed; + assert_ok(phrase_query(index_reader, {"pcc", "pbb", "paa"}, &reversed)); + EXPECT_TRUE(reversed.empty()); +} + +TEST(SniiPhraseQueryTest, TwoTermPhraseUsesUnigramPositions) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + const auto original_prx_span = [&](std::string_view term) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + assert_ok(index_reader.lookup(term, &found, &entry, &frq_base, &prx_base)); + EXPECT_TRUE(found); + return PrxRange {.offset = index_reader.section_refs().posting_region.offset + prx_base + + entry.prx_off_delta, + .len = entry.prx_len}; + }; + const std::vector original_prx { + original_prx_span("failed"), + original_prx_span("order"), + }; + + file.clear_reads(); + std::vector docids; + assert_ok(phrase_query(index_reader, {"failed", "order"}, &docids)); + + const std::vector expected {5000, 7000, 8000}; + EXPECT_EQ(docids, expected); + for (const PrxRange& prx : original_prx) { + const bool original_prx_read = std::ranges::any_of(file.reads(), [&](const auto& read) { + return read.offset < prx.offset + prx.len && prx.offset < read.offset + read.len; + }); + EXPECT_TRUE(original_prx_read); + } +} + +TEST(SniiPhraseQueryTest, TwoTermPhraseWithNonIndexableTermFallsBackToPositions) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"trace", "123"}, &docids)); + + const std::vector expected {42}; + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseQueryTest, TwoTermPhrasePrefixUsesUnigramPositions) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"failed", "ord"}, &docids, 10)); + + const std::vector expected {5000, 6000, 7000, 8000}; + EXPECT_EQ(docids, expected); +} + +std::vector all_docids_0_to(uint32_t end_exclusive) { + std::vector docids(end_exclusive); + std::iota(docids.begin(), docids.end(), 0U); + return docids; +} + +// RQ-01: a plain literal pattern is anchored at both ends (RE2::FullMatch), so it +// matches exactly the "order" term and returns its full docid range. +TEST(SniiRegexpQueryTest, MatchesAnchoredLiteralTerm) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, "order", &docids)); + EXPECT_EQ(docids, all_docids_0_to(9000)); +} + +// RQ-03: alternation under a shared prefix expands to "order" and "ordinal"; both +// span the full docid range, so the union must dedup back to [0..8999]. +TEST(SniiRegexpQueryTest, CharClassMatchesMultipleTermsDeduped) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, "ord(er|inal)", &docids)); + EXPECT_EQ(docids, all_docids_0_to(9000)); +} + +// Alternation across non-adjacent terms with disjoint docids yields a sorted, +// deduplicated union. +TEST(SniiRegexpQueryTest, AlternationUnionsDistinctTerms) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, "(123|needle)", &docids)); + const std::vector expected {42, 100, 101, 102, 6000}; + EXPECT_EQ(docids, expected); +} + +// RQ-04: a non-existent prefix enumerates nothing and returns an empty result. +TEST(SniiRegexpQueryTest, NoTermMatchesEmpty) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, "zzz.*", &docids)); + EXPECT_TRUE(docids.empty()); +} + +// RQ-05: a character-class pattern matches only the numeric "123" term. +TEST(SniiRegexpQueryTest, NumericTermMatch) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, "[0-9]+", &docids)); + const std::vector expected {42}; + EXPECT_EQ(docids, expected); +} + +// RQ-06: backreferences are unsupported by RE2 (but accepted by std::regex); +// re.ok() is false so the call returns InvalidArgument without throwing. +TEST(SniiRegexpQueryTest, InvalidPatternReturnsInvalidArgument) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + const Status status = regexp_query(index_reader, "(a)\\1", &docids); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +// RQ-07: a syntactically invalid pattern also returns InvalidArgument. +TEST(SniiRegexpQueryTest, UnbalancedParenReturnsInvalidArgument) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + const Status status = regexp_query(index_reader, "(order", &docids); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()) << status.to_string(); +} + +// RQ-08: null output / null sink return InvalidArgument (no crash, no throw). +TEST(SniiRegexpQueryTest, NullOutputReturnsInvalidArgument) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector* const null_docids = nullptr; + EXPECT_TRUE(regexp_query(index_reader, "order", null_docids) + .is()); + + DocIdSink* const null_sink = nullptr; + EXPECT_TRUE(regexp_query(index_reader, "order", null_sink) + .is()); +} + +// RQ-09: max_expansions caps the number of expanded terms. Terms enumerate in +// sorted order, so the first term "123" is the only expansion. +TEST(SniiRegexpQueryTest, MaxExpansionsCapsExpansion) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(regexp_query(index_reader, ".*", &docids, /*max_expansions=*/1)); + const std::vector expected {42}; + EXPECT_EQ(docids, expected); +} + +// RQ-10: a match-all pattern returns the union of all ordinary-term docids. +TEST(SniiRegexpQueryTest, MatchAllReturnsAllDocids) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector match_all; + assert_ok(regexp_query(index_reader, ".*", &match_all)); + EXPECT_EQ(match_all, all_docids_0_to(9000)); +} + +// RQ-11: the RE2 path reproduces the std::regex_match golden result set exactly. +TEST(SniiRegexpQueryTest, EquivalenceMatchesStdRegexGolden) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + struct GoldenCase { + std::string_view pattern; + std::vector expected; + }; + const std::vector cases { + {.pattern = "order", .expected = all_docids_0_to(9000)}, + {.pattern = "ord(er|inal)", .expected = all_docids_0_to(9000)}, + {.pattern = "zzz.*", .expected = {}}, + {.pattern = "[0-9]+", .expected = {42}}, + {.pattern = "(123|needle)", .expected = {42, 100, 101, 102, 6000}}, + }; + for (const GoldenCase& c : cases) { + std::vector docids; + assert_ok(regexp_query(index_reader, c.pattern, &docids)); + EXPECT_EQ(docids, c.expected) << "pattern=" << c.pattern; + } +} + +// Deterministic perf (golden prefix): RE2::PossibleMatchRange tightens the +// enumeration prefix for left-anchored patterns whose literal scan stops early. +TEST(SniiRegexpQueryTest, AnchoredPrefixIsTightened) { + auto prefix_of = [](std::string_view pattern) -> std::string { + re2::RE2::Options options; + options.set_log_errors(false); + const re2::RE2 re(re2::StringPiece(pattern.data(), pattern.size()), options); + EXPECT_TRUE(re.ok()) << pattern; + return internal::regex_enum_prefix(pattern, re); + }; + + // Tightened beyond the naive literal scan (which would yield ""). + EXPECT_EQ(prefix_of("^(order)"), "order"); + EXPECT_EQ(prefix_of("^(order|ordinal)"), "ord"); + EXPECT_EQ(prefix_of("^ord[ei]"), "ord"); + // Anchored literal already maximal. + EXPECT_EQ(prefix_of("^order"), "order"); + // Non-anchored patterns fall back to the conservative literal scan. + EXPECT_EQ(prefix_of("ord.*"), "ord"); + EXPECT_EQ(prefix_of(".*failed.*order.*"), ""); + EXPECT_EQ(prefix_of("[0-9]+"), ""); +} + +// Deterministic perf (op-count): the tightened "^(order)" prefix reaches a single +// dictionary term, so the matcher is invoked once instead of once per term. +TEST(SniiRegexpQueryTest, AnchoredPrefixEnumeratesSingleTerm) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + constexpr std::string_view kPattern = "^(order)"; + re2::RE2::Options options; + options.set_log_errors(false); + const re2::RE2 re(re2::StringPiece(kPattern.data(), kPattern.size()), options); + ASSERT_TRUE(re.ok()); + + const std::string enum_prefix = internal::regex_enum_prefix(kPattern, re); + EXPECT_EQ(enum_prefix, "order"); + + auto count_matcher = [&](std::string_view prefix) { + int calls = 0; + std::vector docids; + VectorDocIdSink sink(docids); + assert_ok(internal::emit_expanded_docid_union( + index_reader, prefix, + [&](std::string_view term) { + ++calls; + return re2::RE2::FullMatch(re2::StringPiece(term.data(), term.size()), re); + }, + &sink)); + return std::pair> {calls, std::move(docids)}; + }; + + // Tightened prefix enumerates only "order" -> exactly one matcher call. + auto [tight_calls, tight_docids] = count_matcher(enum_prefix); + EXPECT_EQ(tight_calls, 1); + + // The baseline empty prefix (old behavior) scans all 11 dictionary terms. + auto [full_calls, full_docids] = count_matcher(""); + EXPECT_EQ(full_calls, 11); + + // Narrowing is a pure optimization: identical result either way. + EXPECT_EQ(tight_docids, all_docids_0_to(9000)); + EXPECT_EQ(full_docids, all_docids_0_to(9000)); +} + +TEST(SniiPhraseQueryTest, MultiTailPhrasePrefixFiltersTailPrxByExpectedDocs) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector tail_hits; + assert_ok(index_reader.prefix_terms("ord", &tail_hits, 10)); + ASSERT_EQ(tail_hits.size(), 2); + + struct PrxRange { + uint64_t offset = 0; + uint64_t len = 0; + }; + std::vector full_tail_prx_ranges; + for (const auto& hit : tail_hits) { + full_tail_prx_ranges.push_back({index_reader.section_refs().posting_region.offset + + hit.prx_base + hit.entry.prx_off_delta, + hit.entry.prx_len}); + } + + file.clear_reads(); + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"needle", "ord"}, &docids, 10)); + + const std::vector expected {6000}; + EXPECT_EQ(docids, expected); + for (const PrxRange& prx : full_tail_prx_ranges) { + const bool full_tail_prx_read = std::ranges::any_of(file.reads(), [&](const auto& read) { + return read.offset == prx.offset && read.len == prx.len; + }); + EXPECT_FALSE(full_tail_prx_read); + } +} + +// --------------------------------------------------------------------------- +// T24: phrase-prefix micro-opt (sparsest anchor + expected_docids hoist). +// +// build_reader is shared across many SNII suites, so instead of mutating it we +// build small, self-contained kDocsPositions indexes here. Each scenario uses a +// distinct tail prefix so prefix_terms() expansion stays isolated per test, and +// the deterministic op-counters in query_test_counters.h are read directly. +// --------------------------------------------------------------------------- +Status build_positions_reader(MemoryFile* file, reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader, + std::vector terms, uint32_t doc_count) { + writer::SniiIndexInput input; + input.index_id = 21; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = doc_count; + input.terms = std::move(terms); + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + + RETURN_IF_ERROR(reader::SniiSegmentReader::open(file, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +struct OpaqueMatcherPlanCase { + std::string label; + std::vector plain_terms; + std::vector gram_plan; + std::vector common; +}; + +constexpr std::string_view kReservedCommonGramPrefix = + "\x1f" + "DORIS_COMMON_GRAM_V1" + "\x1f"; + +std::string opaque_gram(std::string_view left, std::string_view right) { + return "opaque:CG(" + std::string(left) + "," + std::string(right) + ")"; +} + +std::string reserved_marker_opaque_gram(std::string_view left, std::string_view right) { + return std::string(kReservedCommonGramPrefix) + "opaque:CG(" + std::string(left) + "," + + std::string(right) + ")"; +} + +std::vector opaque_matcher_plan_cases() { + const std::vector nnn = {"nnn:n0", "nnn:n1", "nnn:n2"}; + const std::vector nns = {"nns:n0", "nns:n1", "nns:s"}; + const std::vector nsn = {"nsn:n0", "nsn:s", "nsn:n1"}; + const std::vector nss = {"nss:n", "nss:s", "nss:s"}; + const std::vector snn = {"snn:s", "snn:n0", "snn:n1"}; + const std::vector sns = {"sns:s", "sns:n", "sns:s"}; + const std::vector ssn = {"ssn:s", "ssn:s", "ssn:n"}; + const std::vector sss = {"sss:s", "sss:s", "sss:s"}; + + return {{"NNN", nnn, nnn, {false, false, false}}, + {"NNS", + nns, + {nns[0], reserved_marker_opaque_gram(nns[1], nns[2])}, + {false, false, true}}, + {"NSN", + nsn, + {opaque_gram(nsn[0], nsn[1]), opaque_gram(nsn[1], nsn[2])}, + {false, true, false}}, + {"NSS", + nss, + {opaque_gram(nss[0], nss[1]), opaque_gram(nss[1], nss[2])}, + {false, true, true}}, + {"SNN", snn, {opaque_gram(snn[0], snn[1]), snn[1], snn[2]}, {true, false, false}}, + {"SNS", + sns, + {opaque_gram(sns[0], sns[1]), opaque_gram(sns[1], sns[2])}, + {true, false, true}}, + {"SSN", + ssn, + {opaque_gram(ssn[0], ssn[1]), opaque_gram(ssn[1], ssn[2])}, + {true, true, false}}, + {"SSS", + sss, + {opaque_gram(sss[0], sss[1]), opaque_gram(sss[1], sss[2])}, + {true, true, true}}}; +} + +Status resolve_opaque_matcher_plan(const reader::LogicalIndexReader& idx, + const std::vector& terms, + internal::ResolvedPhrasePlan* plan) { + plan->unique_terms.clear(); + plan->phrase_plan_index.clear(); + plan->position_offsets.resize(terms.size()); + std::iota(plan->position_offsets.begin(), plan->position_offsets.end(), 0U); + + std::vector unique_terms; + for (const std::string& term : terms) { + if (std::ranges::find(unique_terms, term) == unique_terms.end()) { + unique_terms.push_back(term); + } + } + std::ranges::reverse(unique_terms); + + for (const std::string& term : unique_terms) { + internal::ResolvedQueryTerm resolved; + bool found = false; + RETURN_IF_ERROR(internal::resolve_query_term(idx, term, &resolved, &found)); + if (!found) { + return Status::Error( + "opaque matcher test term is absent"); + } + plan->unique_terms.push_back(std::move(resolved)); + } + for (const std::string& term : terms) { + const auto it = std::ranges::find(unique_terms, term); + DCHECK(it != unique_terms.end()); + plan->phrase_plan_index.push_back(static_cast(it - unique_terms.begin())); + } + return Status::OK(); +} + +Status build_opaque_matcher_plan_reader(MemoryFile* file, reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader, + const std::vector& cases) { + writer::SpimiTermBuffer buffer(/*has_positions=*/true); + uint32_t next_docid = 0; + for (const OpaqueMatcherPlanCase& test_case : cases) { + const auto is_common = [&](std::string_view term) { + for (size_t i = 0; i < test_case.plain_terms.size(); ++i) { + if (test_case.common[i] && term == test_case.plain_terms[i]) { + return true; + } + } + return false; + }; + const auto add_doc = [&](const std::vector& tokens, uint32_t docid) { + for (uint32_t pos = 0; pos < tokens.size(); ++pos) { + buffer.add_token(tokens[pos], docid, pos); + if (pos + 1 < tokens.size() && + (is_common(tokens[pos]) || is_common(tokens[pos + 1]))) { + const std::string gram = + test_case.label == "NNS" + ? reserved_marker_opaque_gram(tokens[pos], tokens[pos + 1]) + : opaque_gram(tokens[pos], tokens[pos + 1]); + buffer.add_token(gram, docid, pos); + } + } + }; + + add_doc(test_case.plain_terms, next_docid++); + add_doc({test_case.label + ":prefix", test_case.plain_terms[0], test_case.plain_terms[1], + test_case.plain_terms[2]}, + next_docid++); + add_doc({test_case.plain_terms[0], test_case.label + ":gap", test_case.plain_terms[1], + test_case.plain_terms[2]}, + next_docid++); + add_doc({test_case.plain_terms[2], test_case.plain_terms[1], test_case.plain_terms[0]}, + next_docid++); + } + + writer::SniiIndexInput input; + input.index_id = 31; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = next_docid; + input.terms = buffer.finalize_sorted(); + + writer::SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + + RETURN_IF_ERROR(reader::SniiSegmentReader::open(file, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +std::vector cursor_lifetime_terms() { + std::vector remaining_positions; + remaining_positions.reserve(8192); + remaining_positions.push_back(12); + for (uint32_t i = 0; i < 8191; ++i) { + remaining_positions.push_back(1000 + i); + } + std::vector repeated_remaining_positions = remaining_positions; + return {make_term("cursor_left", {{100, {0, 10}}}), make_term("cursor_right", {{100, {1, 11}}}), + make_term("cursor_tail", {{100, std::move(remaining_positions)}, {101, {0}}}), + make_term("repeat_cursor", {{300, {0, 1, 10, 11}}}), + make_term("repeat_tail", {{300, std::move(repeated_remaining_positions)}, {301, {0}}})}; +} + +// "lead mid axt*": the leading exact term is high-frequency (5 positions/doc) while +// the middle exact term is a single position/doc -> "mid" is the sparsest exact +// term and becomes the anchor. doc400 is a valid "lead mid" candidate (expected +// tail@8) with NO tail term, so it must be filtered out of the final result. +// NOLINTBEGIN(modernize-use-designated-initializers): positional aggregate init of test posting data +std::vector anchor_scenario_terms() { + return {make_term("lead", {{100, {0, 3, 6, 9, 12}}, + {200, {0, 3, 6, 9, 12}}, + {300, {0, 3, 6, 9, 12}}, + {400, {0, 3, 6, 9, 12}}}), + make_term("mid", {{100, {1}}, {200, {7}}, {300, {7}}, {400, {7}}}), + make_term("axta", {{100, {2}}}), // doc100: lead@0, mid@1 -> tail@2 + make_term("axtb", {{200, {8}}}), // doc200: lead@6, mid@7 -> tail@8 + make_term("axtc", {{300, {8}}})}; // doc300: lead@6, mid@7 -> tail@8 +} + +// "dlead dmid dxt*": dlead is a single position/doc (sparsest), dmid is +// high-frequency. The anchor therefore stays at phrase position 0 (dlead), which +// is exactly the old span[0] behavior -- the degenerate no-change case. +std::vector leading_sparse_scenario_terms() { + return {make_term("dlead", {{500, {3}}, {600, {3}}}), + make_term("dmid", {{500, {0, 2, 4, 6, 8}}, {600, {0, 2, 4, 6, 8}}}), + make_term("dxta", {{500, {5}}}), // dlead@3, dmid@4 -> tail@5 + make_term("dxtb", {{600, {5}}})}; +} + +// "ulead umid uxt*": umid (single position) is the anchor. In doc800 umid sits at +// position 0, which is < its phrase offset (1), so a general anchor would underflow +// `start`; the underflow guard skips it and doc800 is correctly excluded. +std::vector anchor_underflow_scenario_terms() { + return {make_term("ulead", {{800, {5, 9}}, {810, {5, 9}}, {820, {5, 9}}}), + make_term("umid", {{800, {0}}, {810, {6}}, {820, {6}}}), + make_term("uxta", {{810, {7}}}), // ulead@5, umid@6 -> tail@7 + make_term("uxtb", {{820, {7}}})}; +} + +// Two exact positions per doc make both the anchor and first checked span too +// small to amortize selecting a forward scan over the existing binary search. +std::vector small_span_scenario_terms() { + return {make_term("slead", {{1000, {0, 3}}, {1010, {0, 3}}}), + make_term("smid", {{1000, {1, 4}}, {1010, {1, 4}}}), + make_term("sxta", {{1000, {2, 5}}}), make_term("sxtb", {{1010, {2, 5}}})}; +} + +// The scan cost would otherwise be favorable, but bmid@0 cannot produce a +// phrase start at offset 1. The O(1) endpoint gate must keep this boundary shape +// on the binary-search path without probing for a viable subrange. +std::vector invalid_anchor_boundary_terms() { + std::vector leading_positions(64); + std::iota(leading_positions.begin(), leading_positions.end(), 0); + std::vector anchor_positions(48); + std::iota(anchor_positions.begin(), anchor_positions.end(), 0); + return {make_term("blead", {{1020, std::move(leading_positions)}}), + make_term("bmid", {{1020, std::move(anchor_positions)}}), + make_term("bxta", {{1020, {2}}}), make_term("bxtb", {{1020, {3}}})}; +} +// NOLINTEND(modernize-use-designated-initializers) + +// FUNC-1: the new sparsest-anchor path produces the correct result set. The anchor +// (mid) is not phrase-position 0, exercising the general anchor formula + underflow +// guard; doc400 (candidate, no tail term) is filtered out. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixAnchorsOnSparsestTerm) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, anchor_scenario_terms(), + /*doc_count=*/500)); + + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"lead", "mid", "axt"}, &docids, 10)); + + const std::vector expected {100, 200, 300}; + EXPECT_EQ(docids, expected); +} + +// Perf (deterministic): the outer anchor enumeration size == docs x min_span. With +// mid as the anchor (1 position/doc) over 4 candidate docs the count is 4, strictly +// below the old span[0] baseline of Sum|lead| == 4 x 5 == 20. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixAnchorIterationsMinimal) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, anchor_scenario_terms(), + /*doc_count=*/500)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"lead", "mid", "axt"}, &docids, 10)); + + constexpr uint64_t kCandidateDocs = 4; // {100,200,300,400} + constexpr uint64_t kMinSpan = 1; // mid has one position per doc + constexpr uint64_t kOldLeadSpanTotal = kCandidateDocs * 5; // Sum|span[0]| (lead@5pos) + EXPECT_EQ(internal::query_test_counters().anchor_iterations, kCandidateDocs * kMinSpan); + EXPECT_LT(internal::query_test_counters().anchor_iterations, kOldLeadSpanTotal); + EXPECT_EQ(internal::query_test_counters().monotonic_position_scans, 0U); +} + +// FUNC-3: when the leading exact term is already the sparsest, the anchor stays at +// phrase-position 0, so anchor_iterations == Sum|span[0]| exactly (no regression / +// no change vs. the old hardcoded anchor). Result is still correct. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixLeadingTermSparsestIsDegenerate) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, + leading_sparse_scenario_terms(), /*doc_count=*/700)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"dlead", "dmid", "dxt"}, &docids, 10)); + + const std::vector expected {500, 600}; + EXPECT_EQ(docids, expected); + // dlead has one position/doc over 2 docs -> Sum|span[0]| == 2, unchanged. + EXPECT_EQ(internal::query_test_counters().anchor_iterations, 2U); + EXPECT_EQ(internal::query_test_counters().monotonic_position_scans, 0U); +} + +// FUNC-4: a doc whose anchor position is smaller than the anchor's phrase offset +// (umid@0, offset 1) is skipped by the underflow guard and never false-matches; the +// two well-formed docs still match via distinct tails. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixSkipsAnchorUnderflowDoc) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, + anchor_underflow_scenario_terms(), /*doc_count=*/900)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"ulead", "umid", "uxt"}, &docids, 10)); + + const std::vector expected {810, 820}; // doc800 excluded (underflow) + EXPECT_EQ(docids, expected); + // 3 candidates {800,810,820}, umid is the single-position anchor -> 3 x 1. + EXPECT_EQ(internal::query_test_counters().anchor_iterations, 3U); + EXPECT_EQ(internal::query_test_counters().monotonic_position_scans, 0U); +} + +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixSmallSpansKeepBinarySearch) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, + small_span_scenario_terms(), /*doc_count=*/1100)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"slead", "smid", "sxt"}, &docids, 10)); + + const std::vector expected {1000, 1010}; + EXPECT_EQ(docids, expected); + EXPECT_EQ(internal::query_test_counters().anchor_iterations, 4U); + EXPECT_EQ(internal::query_test_counters().monotonic_position_scans, 0U); +} + +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixInvalidAnchorBoundaryKeepsBinarySearch) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, + invalid_anchor_boundary_terms(), /*doc_count=*/1100)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"blead", "bmid", "bxt"}, &docids, 10)); + + const std::vector expected {1020}; + EXPECT_EQ(docids, expected); + EXPECT_EQ(internal::query_test_counters().anchor_iterations, 48U); + EXPECT_EQ(internal::query_test_counters().monotonic_position_scans, 0U); +} + +// Perf (deterministic): the multi-tail branch materializes expected_docids exactly +// once per query (hoisted out of the per-tail loop). Here prefix "axt" expands to 3 +// tails; the old per-tail rebuild would have counted 3. +TEST(SniiPhraseQueryTest, MultiTailPhrasePrefixBuildsExpectedDocidsOnce) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, anchor_scenario_terms(), + /*doc_count=*/500)); + + std::vector tail_hits; + assert_ok(index_reader.prefix_terms("axt", &tail_hits, 10)); + ASSERT_EQ(tail_hits.size(), 3); // axta, axtb, axtc -> multi-tail branch + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"lead", "mid", "axt"}, &docids, 10)); + + const std::vector expected {100, 200, 300}; + EXPECT_EQ(docids, expected); + EXPECT_EQ(internal::query_test_counters().expected_docids_build, 1U); +} + +// FUNC-6: a single tail expansion takes the streaming ExecuteResolvedPhraseTerms +// path, never the multi-tail branch, so expected_docids_build stays 0. Result is +// unchanged from SingleTailPhrasePrefixUsesStreamingPhrasePath. +TEST(SniiPhraseQueryTest, SingleTailPhrasePrefixDoesNotBuildExpectedDocids) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"failed", "orde"}, &docids, 10)); + + const std::vector expected {5000, 7000, 8000}; + EXPECT_EQ(docids, expected); + EXPECT_EQ(internal::query_test_counters().expected_docids_build, 0U); +} + +// FUNC-5: an empty tail expansion returns OK with an empty result before the +// multi-tail branch, so no expected_docids vector is built. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixEmptyTailExpansionReturnsEmpty) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, anchor_scenario_terms(), + /*doc_count=*/500)); + + internal::query_test_counters() = internal::QueryTestCounters {}; + std::vector docids; + assert_ok(phrase_prefix_query(index_reader, {"lead", "mid", "zzz"}, &docids, 10)); + + EXPECT_TRUE(docids.empty()); + EXPECT_EQ(internal::query_test_counters().expected_docids_build, 0U); +} + +// FUNC-7: a null output pointer returns InvalidArgument (no crash, no throw). +TEST(SniiPhraseQueryTest, PhrasePrefixQueryNullOutReturnsInvalidArgument) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, anchor_scenario_terms(), + /*doc_count=*/500)); + + std::vector* const null_docids = nullptr; + EXPECT_TRUE(phrase_prefix_query(index_reader, {"lead", "mid", "axt"}, null_docids, 10) + .is()); +} + +// FUNC-2: byte-for-byte result equivalence on the existing fixture across the three +// canonical phrase-prefix shapes (multi-tail single-exact, single-tail, multi-tail +// single-doc). Locks the sparsest-anchor + hoist changes as pure optimizations. +TEST(SniiPhraseQueryTest, MultiTailPhrasePrefixEquivalenceRegression) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector failed_ord; + assert_ok(phrase_prefix_query(index_reader, {"failed", "ord"}, &failed_ord, 10)); + EXPECT_EQ(failed_ord, (std::vector {5000, 6000, 7000, 8000})); + + std::vector failed_orde; + assert_ok(phrase_prefix_query(index_reader, {"failed", "orde"}, &failed_orde, 10)); + EXPECT_EQ(failed_orde, (std::vector {5000, 7000, 8000})); + + std::vector needle_ord; + assert_ok(phrase_prefix_query(index_reader, {"needle", "ord"}, &needle_ord, 10)); + EXPECT_EQ(needle_ord, (std::vector {6000})); +} + +// Perf (deterministic, corroborating): the CPU-only anchor reorder + expected_docids +// hoist do not change which bytes are fetched (spans/candidates are materialized +// before anchor selection). Two independently-built identical readers therefore +// issue an identical, non-zero number of physical reads for the same query. +TEST(SniiPhraseQueryTest, MultiTermPhrasePrefixDoesNotRegressReadCount) { + MemoryFile file_a; + reader::SniiSegmentReader segment_a; + reader::LogicalIndexReader index_a; + assert_ok(build_positions_reader(&file_a, &segment_a, &index_a, anchor_scenario_terms(), + /*doc_count=*/500)); + file_a.clear_reads(); + std::vector docids_a; + assert_ok(phrase_prefix_query(index_a, {"lead", "mid", "axt"}, &docids_a, 10)); + const size_t reads_a = file_a.reads().size(); + + MemoryFile file_b; + reader::SniiSegmentReader segment_b; + reader::LogicalIndexReader index_b; + assert_ok(build_positions_reader(&file_b, &segment_b, &index_b, anchor_scenario_terms(), + /*doc_count=*/500)); + file_b.clear_reads(); + std::vector docids_b; + assert_ok(phrase_prefix_query(index_b, {"lead", "mid", "axt"}, &docids_b, 10)); + const size_t reads_b = file_b.reads().size(); + + EXPECT_EQ(docids_a, docids_b); + // T24 is a CPU-only change (sparsest-term anchor + expected_docids hoist), so the + // query's physical IO must be unchanged. The small anchor-scenario index is fully + // resident, so the read count is a deterministic constant across identical readers. + EXPECT_EQ(reads_a, reads_b); +} + +TEST(SniiPhraseQueryTest, MultiTermPhraseUsesPairPrefilter) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"failed", "order", "ordinal"}, &docids)); + + const std::vector expected {5000, 7000}; + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseQueryTest, RepeatedTermPhraseUsesCachedPostingSpan) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"repeat", "repeat", "repeat"}, &docids)); + + std::vector expected(9000); + std::iota(expected.begin(), expected.end(), 0); + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseMatcherInvariantTest, OpaqueCommonGramPlansMatchPlainPositionOracle) { + const std::vector cases = opaque_matcher_plan_cases(); + ASSERT_TRUE(cases[1].gram_plan[1].starts_with(kReservedCommonGramPrefix)); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + for (const OpaqueMatcherPlanCase& test_case : cases) { + std::vector plain; + assert_ok(phrase_query(index_reader, test_case.plain_terms, &plain)); + ASSERT_FALSE(plain.empty()) << test_case.label; + + internal::ResolvedPhrasePlan resolved_plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, test_case.gram_plan, &resolved_plan)); + std::vector gram; + assert_ok(internal::execute_resolved_phrase_plan(index_reader, std::move(resolved_plan), + &gram)); + EXPECT_EQ(gram, plain) << test_case.label; + } +} + +TEST(SniiPhraseMatcherInvariantTest, RepeatedOpaqueGramUsesOnePlanAtTwoPositions) { + const std::vector cases = opaque_matcher_plan_cases(); + const OpaqueMatcherPlanCase& sss = cases.back(); + ASSERT_EQ(sss.label, "SSS"); + ASSERT_EQ(sss.gram_plan.size(), 2U); + ASSERT_EQ(sss.gram_plan[0], sss.gram_plan[1]); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + std::vector plain; + assert_ok(phrase_query(index_reader, sss.plain_terms, &plain)); + internal::ResolvedPhrasePlan resolved_plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, sss.gram_plan, &resolved_plan)); + format::PrxDecodeStats stats; + format::PrxDecodeContext decode_context {.stats = &stats}; + std::vector gram; + assert_ok(internal::execute_resolved_phrase_plan(index_reader, std::move(resolved_plan), &gram, + &decode_context)); + std::vector gram_candidates; + assert_ok(term_query(index_reader, sss.gram_plan[0], &gram_candidates)); + + EXPECT_EQ(gram, plain); + EXPECT_EQ(stats.selected_docs, gram_candidates.size()) + << "the repeated clause must decode its shared plan once"; +} + +TEST(SniiPhraseMatcherInvariantTest, SingleOpaqueGramMatchesTwoTermPlainPhrase) { + const std::vector cases = opaque_matcher_plan_cases(); + const OpaqueMatcherPlanCase& sss = cases.back(); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + std::vector plain; + assert_ok(phrase_query(index_reader, {sss.plain_terms[0], sss.plain_terms[1]}, &plain)); + internal::ResolvedPhrasePlan resolved_plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, {sss.gram_plan[0]}, &resolved_plan)); + format::PrxDecodeStats stats; + format::PrxDecodeContext decode_context {.stats = &stats}; + std::vector gram; + assert_ok(internal::execute_resolved_phrase_plan(index_reader, std::move(resolved_plan), &gram, + &decode_context)); + EXPECT_EQ(gram, plain); + EXPECT_EQ(stats.frame_count(), 0U) << "a one-clause plan must use the doc-posting term route"; +} + +TEST(SniiPhraseMatcherInvariantTest, ResolvedPlanMovesInlinePayloadIntoTermPlans) { + const std::vector cases = opaque_matcher_plan_cases(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + internal::ResolvedPhrasePlan plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, cases[1].gram_plan, &plan)); + const size_t unique_term_count = plan.unique_terms.size(); + size_t inline_payload_count = 0; + for (const internal::ResolvedQueryTerm& term : plan.unique_terms) { + inline_payload_count += !term.entry.frq_bytes.empty(); + inline_payload_count += !term.entry.prx_bytes.empty(); + } + ASSERT_GT(inline_payload_count, 0U); + + internal::query_test_counters() = {}; + std::vector docids; + assert_ok(internal::execute_resolved_phrase_plan(index_reader, std::move(plan), &docids)); + EXPECT_EQ(internal::query_test_counters().resolved_term_entry_copies, 0U); + EXPECT_EQ(internal::query_test_counters().resolved_term_entry_moves, unique_term_count); + EXPECT_EQ(internal::query_test_counters().resolved_term_payload_pointer_reuses, + inline_payload_count); +} + +TEST(SniiPhraseMatcherInvariantTest, ResolvedPlanRejectsStructuralMismatches) { + const std::vector cases = opaque_matcher_plan_cases(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + internal::ResolvedPhrasePlan plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, cases[1].gram_plan, &plan)); + ASSERT_TRUE(plan.is_valid()); + + internal::ResolvedPhrasePlan mismatched_offsets = plan; + mismatched_offsets.position_offsets.pop_back(); + EXPECT_FALSE(mismatched_offsets.is_valid()); + + internal::ResolvedPhrasePlan invalid_mapping = plan; + invalid_mapping.phrase_plan_index[0] = invalid_mapping.unique_terms.size(); + EXPECT_FALSE(invalid_mapping.is_valid()); + + internal::ResolvedPhrasePlan empty_mismatch = plan; + empty_mismatch.phrase_plan_index.clear(); + empty_mismatch.position_offsets.clear(); + EXPECT_FALSE(empty_mismatch.is_valid()); +} + +TEST(SniiPhraseMatcherInvariantTest, ResolvedPlanRejectsInvalidCoverageAndOffsets) { + const std::vector cases = opaque_matcher_plan_cases(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + internal::ResolvedPhrasePlan plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, cases[1].gram_plan, &plan)); + ASSERT_TRUE(plan.is_valid()); + + internal::ResolvedPhrasePlan unused_unique = plan; + unused_unique.unique_terms.push_back(unused_unique.unique_terms.front()); + EXPECT_FALSE(unused_unique.is_valid()); + + // Since the selective CommonGrams postings change, position offsets are only + // required to start at 0 and ascend STRICTLY -- gaps are legal (a gram + // covers two positions, so the next clause's offset jumps past it). The + // fatal invariants are a nonzero first offset and a non-ascending step. + internal::ResolvedPhrasePlan nonzero_first = plan; + ++nonzero_first.position_offsets.front(); + EXPECT_FALSE(nonzero_first.is_valid()); + + internal::ResolvedPhrasePlan non_ascending = plan; + ASSERT_GE(non_ascending.position_offsets.size(), 2U); + non_ascending.position_offsets.back() = + non_ascending.position_offsets[non_ascending.position_offsets.size() - 2]; + EXPECT_FALSE(non_ascending.is_valid()); +} + +TEST(SniiPhraseMatcherInvariantTest, ResolvedPlanNullOutputIsInvalidArgument) { + const std::vector cases = opaque_matcher_plan_cases(); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_opaque_matcher_plan_reader(&file, &segment_reader, &index_reader, cases)); + + internal::ResolvedPhrasePlan plan; + assert_ok(resolve_opaque_matcher_plan(index_reader, cases[1].gram_plan, &plan)); + EXPECT_TRUE(internal::execute_resolved_phrase_plan(index_reader, std::move(plan), nullptr) + .is()); +} + +TEST(SniiPhraseMatcherInvariantTest, ThreeClauseCursorSurvivesRemainingPlanDecode) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, cursor_lifetime_terms(), + /*doc_count=*/400)); + + internal::query_test_counters() = {}; + std::vector docids; + assert_ok(phrase_query(index_reader, {"cursor_left", "cursor_right", "cursor_tail"}, &docids)); + EXPECT_EQ(docids, (std::vector {100})); + EXPECT_EQ(internal::query_test_counters().phrase_position_epoch_cache_hits, 0U); + EXPECT_EQ(internal::query_test_counters().phrase_position_epoch_cache_misses, 3U); +} + +TEST(SniiPhraseMatcherInvariantTest, ThreeClauseRepeatedPlanKeepsPairCursorAlive) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_positions_reader(&file, &segment_reader, &index_reader, cursor_lifetime_terms(), + /*doc_count=*/400)); + + internal::query_test_counters() = {}; + std::vector docids; + assert_ok( + phrase_query(index_reader, {"repeat_cursor", "repeat_cursor", "repeat_tail"}, &docids)); + EXPECT_EQ(docids, (std::vector {300})); + EXPECT_EQ(internal::query_test_counters().phrase_position_epoch_cache_hits, 1U); + EXPECT_EQ(internal::query_test_counters().phrase_position_epoch_cache_misses, 2U); +} + +TEST(SniiPhraseQueryTest, DenseTermWithMissingDocKeepsCandidateOrdinals) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector driver_docids; + assert_ok(term_query(index_reader, "driver", &driver_docids)); + EXPECT_EQ(driver_docids.size(), 8000); + + std::vector almost_docids; + assert_ok(term_query(index_reader, "almost", &almost_docids)); + EXPECT_EQ(almost_docids.size(), 8999); + ASSERT_GT(almost_docids.size(), 6144); + EXPECT_EQ(almost_docids[3999], 3999); + EXPECT_EQ(almost_docids[4000], 4001); + EXPECT_EQ(almost_docids[6143], 6144); + EXPECT_EQ(almost_docids[6144], 6145); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"driver", "almost"}, &docids)); + + std::vector expected; + expected.reserve(7999); + for (uint32_t docid = 0; docid < 8000; ++docid) { + if (docid != 4000) { + expected.push_back(docid); + } + } + EXPECT_EQ(docids, expected); +} + +TEST(SniiPhraseQueryTest, SparseWindowBitsetKeepsCandidateOrdinals) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + std::vector docids; + assert_ok(phrase_query(index_reader, {"sparse_left", "sparse_right"}, &docids)); + + std::vector expected; + for (uint32_t docid = 0; docid < 9000; ++docid) { + if (docid % 3 == 0 && docid % 4 != 1) { + expected.push_back(docid); + } + } + EXPECT_EQ(docids, expected); +} + +TEST(SniiTermQueryTest, WindowedDenseTermEmitsRangesToSink) { + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + RecordingDocIdSink sink; + assert_ok(term_query(index_reader, "failed", &sink)); + + std::vector expected(9000); + std::iota(expected.begin(), expected.end(), 0); + EXPECT_EQ(sink.out, expected); + EXPECT_GT(sink.range_calls, 0); +} + +TEST(SniiPrxPodTest, SelectivePforCsrMatchesFullCsrAcrossRuns) { + std::vector freqs; + std::vector positions; + freqs.reserve(320); + for (uint32_t doc = 0; doc < 320; ++doc) { + const uint32_t freq = (doc % 5 == 0) ? 2 : 1; + freqs.push_back(freq); + positions.push_back(doc * 3); + if (freq == 2) { + positions.push_back(doc * 3 + 2); + } + } + + ByteSink sink; + assert_ok(format::build_prx_window_flat(positions, freqs, -1, &sink)); + + std::vector full_positions; + std::vector full_offsets; + ByteSource full_source(sink.view()); + assert_ok(format::read_prx_window_csr(&full_source, &full_positions, &full_offsets)); + + auto assert_selected_matches_full = [&](const std::vector& selected_docs) { + std::vector selected_positions; + std::vector selected_offsets; + ByteSource selected_source(sink.view()); + assert_ok(format::read_prx_window_csr_selective(&selected_source, selected_docs, + &selected_positions, &selected_offsets)); + + ASSERT_EQ(selected_offsets.size(), selected_docs.size() + 1); + for (size_t i = 0; i < selected_docs.size(); ++i) { + const uint32_t doc = selected_docs[i]; + const std::vector expected(full_positions.begin() + full_offsets[doc], + full_positions.begin() + full_offsets[doc + 1]); + const std::vector actual( + selected_positions.begin() + selected_offsets[i], + selected_positions.begin() + selected_offsets[i + 1]); + EXPECT_EQ(actual, expected); + } + }; + + assert_selected_matches_full({0, 1, 2}); + assert_selected_matches_full({0, 1, 127, 128, 129, 255, 256, 319}); +} + +TEST(SniiPrxPodTest, SelectivePforCsrHandlesDocsSpanningPforRuns) { + const std::vector freqs {300, 1, 260, 2, 1}; + std::vector positions; + positions.reserve(564); + auto append_doc = [&](uint32_t count, uint32_t base, uint32_t seed) { + uint32_t pos = base; + uint32_t state = seed; + for (uint32_t i = 0; i < count; ++i) { + state = state * 1664525 + 1013904223; + pos += 1 + (state & 0xFFFF); + positions.push_back(pos); + } + }; + append_doc(freqs[0], 3, 11); + append_doc(freqs[1], 11, 13); + append_doc(freqs[2], 19, 17); + append_doc(freqs[3], 29, 19); + append_doc(freqs[4], 37, 23); + + ByteSink sink; + assert_ok(format::build_prx_window_flat(positions, freqs, -1, &sink)); + ASSERT_FALSE(sink.buffer().empty()); + ASSERT_EQ(sink.buffer().front(), static_cast(format::PrxCodec::kPfor)); + + std::vector full_positions; + std::vector full_offsets; + ByteSource full_source(sink.view()); + assert_ok(format::read_prx_window_csr(&full_source, &full_positions, &full_offsets)); + + const std::vector selected_docs {0, 2, 3}; + std::vector selected_positions; + std::vector selected_offsets; + ByteSource selected_source(sink.view()); + assert_ok(format::read_prx_window_csr_selective(&selected_source, selected_docs, + &selected_positions, &selected_offsets)); + + ASSERT_EQ(selected_offsets.size(), selected_docs.size() + 1); + for (size_t i = 0; i < selected_docs.size(); ++i) { + const uint32_t doc = selected_docs[i]; + const std::vector expected(full_positions.begin() + full_offsets[doc], + full_positions.begin() + full_offsets[doc + 1]); + const std::vector actual(selected_positions.begin() + selected_offsets[i], + selected_positions.begin() + selected_offsets[i + 1]); + EXPECT_EQ(actual, expected); + } +} + +TEST(SniiPrxPodTest, AutoCodecUsesZstdWhenItIsSmaller) { + std::vector freqs(1024, 1); + std::vector positions(1024, 7); + + ByteSink sink; + assert_ok(format::build_prx_window_flat(positions, freqs, -1, &sink)); + + ASSERT_FALSE(sink.buffer().empty()); + EXPECT_EQ(sink.buffer().front(), static_cast(format::PrxCodec::kZstd)); + + std::vector decoded_positions; + std::vector decoded_offsets; + ByteSource source(sink.view()); + assert_ok(format::read_prx_window_csr(&source, &decoded_positions, &decoded_offsets)); + + ASSERT_EQ(decoded_offsets.size(), freqs.size() + 1); + ASSERT_EQ(decoded_positions.size(), positions.size()); + for (size_t i = 0; i < freqs.size(); ++i) { + EXPECT_EQ(decoded_offsets[i], i); + EXPECT_EQ(decoded_positions[i], 7); + } + EXPECT_EQ(decoded_offsets.back(), positions.size()); + + const std::vector selected_docs {0, 17, 511, 1023}; + assert_selective_prx_matches_constant_positions(sink.view(), selected_docs, 7); +} + +TEST(SniiPrxPodTest, AutoCodecKeepsPforForTinyWindows) { + const std::vector freqs {1, 2}; + const std::vector positions {3, 5, 8}; + + ByteSink sink; + assert_ok(format::build_prx_window_flat(positions, freqs, -1, &sink)); + + ASSERT_FALSE(sink.buffer().empty()); + EXPECT_EQ(sink.buffer().front(), static_cast(format::PrxCodec::kPfor)); +} + +TEST(SniiPforTest, LowBitWidthFastPathsRoundTrip) { + auto assert_round_trip = [](const std::vector& values, uint8_t expected_width) { + ByteSink sink; + doris::snii::pfor_encode(values.data(), values.size(), &sink); + ASSERT_FALSE(sink.buffer().empty()); + EXPECT_EQ(sink.buffer().front(), expected_width); + + std::vector decoded(values.size(), 0xFFFFFFFF); + ByteSource source(sink.view()); + assert_ok(doris::snii::pfor_decode(&source, values.size(), decoded.data())); + EXPECT_TRUE(source.eof()); + EXPECT_EQ(decoded, values); + }; + + std::vector one_bit(128); + for (size_t i = 0; i < one_bit.size(); ++i) { + one_bit[i] = static_cast(i & 1); + } + assert_round_trip(one_bit, 1); + + one_bit[17] = 1000; + assert_round_trip(one_bit, 1); + + std::vector two_bit(128); + for (size_t i = 0; i < two_bit.size(); ++i) { + two_bit[i] = static_cast(i & 3); + } + assert_round_trip(two_bit, 2); + + std::vector three_bit(131); + for (size_t i = 0; i < three_bit.size(); ++i) { + three_bit[i] = static_cast(i & 7); + } + assert_round_trip(three_bit, 3); + + std::vector four_bit(128); + for (size_t i = 0; i < four_bit.size(); ++i) { + four_bit[i] = static_cast(i & 15); + } + assert_round_trip(four_bit, 4); + + std::vector five_bit(129); + for (size_t i = 0; i < five_bit.size(); ++i) { + five_bit[i] = static_cast(i & 31); + } + assert_round_trip(five_bit, 5); + + std::vector six_bit(130); + for (size_t i = 0; i < six_bit.size(); ++i) { + six_bit[i] = static_cast(i & 63); + } + assert_round_trip(six_bit, 6); + + std::vector seven_bit(131); + for (size_t i = 0; i < seven_bit.size(); ++i) { + seven_bit[i] = static_cast(i & 127); + } + assert_round_trip(seven_bit, 7); + + std::vector eight_bit(256); + for (size_t i = 0; i < eight_bit.size(); ++i) { + eight_bit[i] = static_cast(i); + } + assert_round_trip(eight_bit, 8); +} + +// =========================================================================== +// T04 -- DICT block request-scoped cache (MRU) + resident single-range read. +// =========================================================================== + +namespace t04 { + +std::string many_term_key(uint32_t i) { + return "term_" + std::to_string(1000000 + i); +} + +// Builds an index of `n_terms` tiny docs-only terms with a small target block +// size so the dictionary spans several DICT blocks, opened through `read_through`. +Status build_multi_block_reader(MemoryFile* file, doris::snii::io::FileReader* read_through, + reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader, uint32_t n_terms) { + writer::SniiIndexInput input = make_many_term_input(21, "Body", n_terms); + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + writer::SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + RETURN_IF_ERROR(reader::SniiSegmentReader::open(read_through, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +// Serializes read_at calls so the recording MemoryFile can be shared by the +// concurrency test without racing on its own bookkeeping. Test infra only -- the +// production FileReader (Doris IO / S3) is itself concurrent-read safe; this lock +// is NOT part of the reader under test and never wraps a decode. +class LockedFileReader final : public doris::snii::io::FileReader { +public: + explicit LockedFileReader(doris::snii::io::FileReader* inner) : inner_(inner) {} + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + std::lock_guard guard(mu_); + return inner_->read_at(offset, len, out); + } + uint64_t size() const override { return inner_->size(); } + +private: + doris::snii::io::FileReader* inner_; + std::mutex mu_; +}; + +// Captures everything lookup() returns so resident / on-demand / cached paths can +// be asserted byte-identical. +struct LookupResult { + bool found = false; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + std::string term; + uint64_t frq_off_delta = 0; + uint64_t frq_len = 0; + uint64_t df = 0; + bool operator==(const LookupResult&) const = default; +}; + +LookupResult do_lookup(const reader::LogicalIndexReader& idx, std::string_view term, + reader::DictBlockCache* cache) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + const Status st = idx.lookup(term, &found, &entry, &frq_base, &prx_base, cache); + EXPECT_TRUE(st.ok()) << st.to_string(); + LookupResult r; + r.found = found; + if (found) { + r.frq_base = frq_base; + r.prx_base = prx_base; + r.term = entry.term; + r.frq_off_delta = entry.frq_off_delta; + r.frq_len = entry.frq_len; + r.df = static_cast(entry.df); + } + return r; +} + +size_t count_reads_in_region(const MemoryFile& file, const format::RegionRef& region, + bool* one_covers_region) { + size_t count = 0; + *one_covers_region = false; + for (const MemoryFile::Read& r : file.reads()) { + if (r.offset >= region.offset && r.offset < region.offset + region.length) { + ++count; + if (r.offset == region.offset && r.len == region.length) { + *one_covers_region = true; + } + } + } + return count; +} + +} // namespace t04 + +// F10: a resident dictionary that spans several blocks is loaded with ONE range +// read over dict_region (was one read_at per block -> up to ~4 serial S3 rounds). +TEST(SniiLogicalReaderTest, ResidentDictLoadIssuesSingleRangeRead) { + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + ScopedEnv resident_bsbf("SNII_BSBF_RESIDENT_MAX", "1048576"); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(t04::build_multi_block_reader(&file, &file, &segment_reader, &index_reader, 64)); + + const format::RegionRef& dict_region = index_reader.section_refs().dict_region; + ASSERT_GT(dict_region.length, 0U); + bool one_covers_region = false; + const size_t region_reads = t04::count_reads_in_region(file, dict_region, &one_covers_region); + EXPECT_EQ(region_reads, 1U); + EXPECT_TRUE(one_covers_region); +} + +// F08/F20 + perf gate: forced on-demand, a block hit K times by repeated lookups +// AND a block shared by several distinct terms of one query each decode ONCE when +// a single request-scoped cache is threaded through -- dict_decode_counter() == +// unique_blocks (1 here; build_reader's small dictionary is a single block). +TEST(SniiLogicalReaderTest, OnDemandLookupDecompressesBlockOncePerUniqueBlock) { + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + doris::snii::testing::reset_dict_decode_counter(); + reader::DictBlockCache cache; + for (int i = 0; i < 5; ++i) { + const t04::LookupResult r = t04::do_lookup(index_reader, "failed", &cache); + EXPECT_TRUE(r.found); + } + EXPECT_EQ(doris::snii::testing::dict_decode_counter(), 1U); +} + +// The headline gate: a multi-term query whose terms fall in the same DICT block +// decodes it ONCE with a shared cache, and once-per-term without one. +TEST(SniiLogicalReaderTest, SharedDictBlockDecodesOncePerQuery) { + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + const std::vector terms = {"failed", "order", "driver"}; + + // With one request-scoped cache: the shared block decodes once. + doris::snii::testing::reset_dict_decode_counter(); + reader::DictBlockCache cache; + for (std::string_view t : terms) { + EXPECT_TRUE(t04::do_lookup(index_reader, t, &cache).found); + } + EXPECT_EQ(doris::snii::testing::dict_decode_counter(), 1U); // == unique_blocks + + // Baseline (no cache): each term re-decodes the same block. + doris::snii::testing::reset_dict_decode_counter(); + for (std::string_view t : terms) { + EXPECT_TRUE(t04::do_lookup(index_reader, t, nullptr).found); + } + EXPECT_EQ(doris::snii::testing::dict_decode_counter(), terms.size()); +} + +// New/old equivalence: the on-demand + cache path returns exactly what the +// resident path does -- lookups and full term_query docid sets are identical. +TEST(SniiLogicalReaderTest, OnDemandResultsMatchResidentBaseline) { + const std::vector present = {"failed", "order", "ordinal", + "driver", "needle", "almost"}; + const std::string_view absent = "definitely_absent_term"; + + std::vector resident_lookups; + std::vector> resident_docids; + { + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + for (std::string_view t : present) { + resident_lookups.push_back(t04::do_lookup(index_reader, t, nullptr)); + std::vector docids; + assert_ok(term_query(index_reader, t, &docids)); + resident_docids.push_back(std::move(docids)); + } + EXPECT_FALSE(t04::do_lookup(index_reader, absent, nullptr).found); + } + + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + MemoryFile file; + reader::SniiSegmentReader segment_reader; + reader::LogicalIndexReader index_reader; + assert_ok(build_reader(&file, &segment_reader, &index_reader)); + + reader::DictBlockCache cache; + for (size_t i = 0; i < present.size(); ++i) { + // cached vs uncached on-demand both equal the resident baseline. + EXPECT_EQ(t04::do_lookup(index_reader, present[i], &cache), resident_lookups[i]); + EXPECT_EQ(t04::do_lookup(index_reader, present[i], nullptr), resident_lookups[i]); + std::vector docids; + assert_ok(term_query(index_reader, present[i], &docids)); + EXPECT_EQ(docids, resident_docids[i]); + } + EXPECT_FALSE(t04::do_lookup(index_reader, absent, &cache).found); +} + +// F-05: prefix enumeration across many on-demand blocks is correct and a second +// pass over the same cache adds no decodes (cross-call request-scoped reuse). +TEST(SniiLogicalReaderTest, PrefixEnumerationReusesCachedBlocks) { + constexpr uint32_t kTerms = 64; + + // Resident baseline ordering. + std::vector resident_terms; + { + ScopedEnv resident_dict("SNII_DICT_RESIDENT_MAX", "1048576"); + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(t04::build_multi_block_reader(&file, &file, &seg, &idx, kTerms)); + std::vector hits; + assert_ok(idx.prefix_terms("term_", &hits, 0)); + for (auto& h : hits) { + resident_terms.push_back(h.term); + } + } + ASSERT_EQ(resident_terms.size(), kTerms); + + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(t04::build_multi_block_reader(&file, &file, &seg, &idx, kTerms)); + + reader::DictBlockCache cache(/*max_entries=*/128); + doris::snii::testing::reset_dict_decode_counter(); + + std::vector hits1; + assert_ok(idx.prefix_terms("term_", &hits1, 0, &cache)); + const uint64_t blocks = doris::snii::testing::dict_decode_counter(); + EXPECT_GE(blocks, 2U); // genuinely multi-block (else the reuse gate is vacuous) + + std::vector hits2; + assert_ok(idx.prefix_terms("term_", &hits2, 0, &cache)); + EXPECT_EQ(doris::snii::testing::dict_decode_counter(), blocks); // second pass: no re-decode + + std::vector got1; + std::vector got2; + for (auto& h : hits1) { + got1.push_back(h.term); + } + for (auto& h : hits2) { + got2.push_back(h.term); + } + EXPECT_EQ(got1, resident_terms); + EXPECT_EQ(got2, resident_terms); +} + +// Request-scoped cache unit: MRU promotion, LRU eviction, a hard size bound, and +// pins that keep an evicted block alive -- no file IO involved. +TEST(SniiLogicalReaderTest, RequestCacheEvictsLruStaysBoundedAndPinsSurvive) { + reader::DictBlockCache cache(/*max_entries=*/2); + int loads = 0; + auto loader_for = [&](uint8_t tag) { + return [&loads, tag](std::shared_ptr* out) -> Status { + ++loads; + auto block = std::make_shared(); + block->bytes.assign(4, tag); + *out = block; + return Status::OK(); + }; + }; + + std::shared_ptr pin0; + assert_ok(cache.get_or_load(0, loader_for(0), &pin0)); + EXPECT_EQ(loads, 1); + std::shared_ptr tmp; + assert_ok(cache.get_or_load(1, loader_for(1), &tmp)); // {1,0} + EXPECT_EQ(loads, 2); + assert_ok(cache.get_or_load(0, loader_for(0), &tmp)); // hit -> {0,1} + EXPECT_EQ(loads, 2); + assert_ok(cache.get_or_load(2, loader_for(2), &tmp)); // miss -> evict LRU(1) -> {2,0} + EXPECT_EQ(loads, 3); + EXPECT_EQ(cache.size(), 2U); + + assert_ok(cache.get_or_load(0, loader_for(0), &tmp)); // 0 still resident -> hit + EXPECT_EQ(loads, 3); + assert_ok(cache.get_or_load(1, loader_for(1), &tmp)); // 1 was evicted -> reload + EXPECT_EQ(loads, 4); + EXPECT_LE(cache.size(), cache.max_entries()); + + // pin0 was taken before 0 ever cycled; even across evictions its bytes stay live. + ASSERT_TRUE(pin0); + ASSERT_EQ(pin0->bytes.size(), 4U); + EXPECT_EQ(pin0->bytes[0], 0); +} + +// F-06: with a 1-entry cache, alternating two different blocks forces reloads; +// a reloaded block still decodes correctly (Slice not corrupted by eviction). +TEST(SniiLogicalReaderTest, SmallCacheReloadsEvictedBlockCorrectly) { + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + constexpr uint32_t kTerms = 64; + + MemoryFile file; + reader::SniiSegmentReader seg; + reader::LogicalIndexReader idx; + assert_ok(t04::build_multi_block_reader(&file, &file, &seg, &idx, kTerms)); + + const std::string first = t04::many_term_key(0); // smallest -> block 0 + const std::string last = t04::many_term_key(kTerms - 1); // largest -> last block + + reader::DictBlockCache cache(/*max_entries=*/1); + doris::snii::testing::reset_dict_decode_counter(); + + const t04::LookupResult first_a = t04::do_lookup(idx, first, &cache); + const uint64_t after_first = doris::snii::testing::dict_decode_counter(); + EXPECT_TRUE(t04::do_lookup(idx, last, &cache).found); // evicts block 0 + const t04::LookupResult first_b = t04::do_lookup(idx, first, &cache); // reload block 0 + const uint64_t after_reload = doris::snii::testing::dict_decode_counter(); + + EXPECT_TRUE(first_a.found); + EXPECT_EQ(first_a, first_b); // reload produced the identical entry + EXPECT_GT(after_reload, after_first); // a reload actually happened (evicted) + EXPECT_LE(cache.size(), cache.max_entries()); +} + +// Concurrency: N queries share the const reader, each with its OWN request-scoped +// cache. No shared mutable state is added to the reader, so every thread decodes +// the on-demand block once -> dict_decode_counter() == thread count, results are +// correct, and the run is TSAN-clean (BUILD_TYPE_UT=TSAN). +TEST(SniiLogicalReaderConcurrencyTest, ConcurrentQueriesUseIndependentRequestCaches) { + ScopedEnv on_demand("SNII_DICT_RESIDENT_MAX", "0"); + + MemoryFile file; + reader::SniiSegmentReader seg0; + reader::LogicalIndexReader idx0; + assert_ok(build_reader(&file, &seg0, &idx0)); // populate `file` + + t04::LockedFileReader locked(&file); + reader::SniiSegmentReader seg; + assert_ok(reader::SniiSegmentReader::open(&locked, &seg)); + reader::LogicalIndexReader idx; + assert_ok(seg.open_index(7, "Body", &idx)); + + doris::snii::testing::reset_dict_decode_counter(); + constexpr int kThreads = 8; + constexpr int kIters = 16; + std::atomic failures {0}; + auto worker = [&]() { + reader::DictBlockCache cache; // request-scoped: one per query/thread + for (int i = 0; i < kIters; ++i) { + bool found = false; + format::DictEntry entry; + uint64_t frq_base = 0; + uint64_t prx_base = 0; + const Status st = idx.lookup("failed", &found, &entry, &frq_base, &prx_base, &cache); + if (!st.ok() || !found) { + failures.fetch_add(1, std::memory_order_relaxed); + } + } + }; + + std::vector threads; + threads.reserve(kThreads); + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back(worker); + } + for (std::thread& t : threads) { + t.join(); + } + + EXPECT_EQ(failures.load(), 0); + // Each thread's cache decodes the shared block exactly once. + EXPECT_EQ(doris::snii::testing::dict_decode_counter(), static_cast(kThreads)); +} + +} // namespace +} // namespace doris::snii::query diff --git a/be/test/storage/index/snii_query_test_util.h b/be/test/storage/index/snii_query_test_util.h new file mode 100644 index 00000000000000..e6af5217417760 --- /dev/null +++ b/be/test/storage/index/snii_query_test_util.h @@ -0,0 +1,258 @@ +// 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. + +#pragma once + +// Shared gtest fixtures for the SNII reader/writer test suites. Other test files +// reuse these by including this header and pulling the symbols in with +// `using namespace doris::snii::snii_test;` (see snii_query_test.cpp). + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/io/file_reader.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/reader/logical_index_reader.h" +#include "storage/index/snii/reader/snii_segment_reader.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" + +namespace doris::snii::snii_test { +using doris::Status; // RETURN_IF_ERROR / Status::OK() expand to a bare Status below. + +// An in-memory FileReader+FileWriter that records every read() (offset/len) so +// snii-layer round-trips and exact-range assertions can be made against it. +class MemoryFile final : public doris::snii::io::FileReader, public doris::snii::io::FileWriter { +public: + struct Read { + uint64_t offset = 0; + size_t len = 0; + }; + + Status append(Slice data) override { + data_.insert(data_.end(), data.data(), data.data() + data.size()); + return Status::OK(); + } + + Status finalize() override { + finalized_ = true; + return Status::OK(); + } + + uint64_t bytes_written() const override { return data_.size(); } + + // NOLINTBEGIN(readability-non-const-parameter): FileReader interface writes into out. + Status read_at(uint64_t offset, size_t len, std::vector* out) override { + if (offset > data_.size() || len > data_.size() - offset) { + return Status::Corruption("memory file read past eof"); + } + reads_.push_back({offset, len}); + read_bytes_ += len; + out->resize(len); + if (len != 0) { + std::memcpy(out->data(), data_.data() + offset, len); + } + return Status::OK(); + } + // NOLINTEND(readability-non-const-parameter) + + uint64_t size() const override { return data_.size(); } + const std::vector& data() const { return data_; } + bool finalized() const { return finalized_; } + const std::vector& reads() const { return reads_; } + size_t read_bytes() const { return read_bytes_; } + void clear_reads() { + reads_.clear(); + read_bytes_ = 0; + } + +private: + std::vector data_; + std::vector reads_; + size_t read_bytes_ = 0; + bool finalized_ = false; +}; + +// RAII helper that sets an environment variable for the duration of a test and +// restores (or clears) the previous value on destruction. +class ScopedEnv { +public: + ScopedEnv(const char* key, const char* value) : key_(key) { + if (const char* old = std::getenv(key); old != nullptr) { + old_value_ = old; + } + setenv(key, value, 1); + } + + ~ScopedEnv() { + if (old_value_.has_value()) { + setenv(key_, old_value_->c_str(), 1); + } else { + unsetenv(key_); + } + } + +private: + const char* key_; + std::optional old_value_; +}; + +// A single document's postings (docid + its in-doc positions) used to drive the +// writer fixtures below. +struct PostingDoc { + uint32_t docid = 0; + std::vector positions; +}; + +inline writer::TermPostings make_term(std::string term, std::vector docs) { + std::ranges::sort(docs, [](const PostingDoc& lhs, const PostingDoc& rhs) { + return lhs.docid < rhs.docid; + }); + + writer::TermPostings posting; + posting.term = std::move(term); + posting.docids.reserve(docs.size()); + posting.freqs.reserve(docs.size()); + for (const PostingDoc& doc : docs) { + posting.docids.push_back(doc.docid); + posting.freqs.push_back(static_cast(doc.positions.size())); + posting.positions_flat.insert(posting.positions_flat.end(), doc.positions.begin(), + doc.positions.end()); + } + return posting; +} + +inline std::vector docs_with_one_position(uint32_t begin, uint32_t end, + uint32_t position) { + std::vector docs; + docs.reserve(end - begin); + for (uint32_t docid = begin; docid < end; ++docid) { + docs.push_back({docid, {position}}); + } + return docs; +} + +inline void assert_ok(const Status& status) { + ASSERT_TRUE(status.ok()) << status.to_string(); +} + +inline segment_v2::inverted_index::CommonGramsSegmentMetadata make_plain_scoring_metadata( + uint64_t doc_count, uint64_t token_count) { + using namespace segment_v2::inverted_index; + CommonGramsSegmentMetadata metadata; + metadata.plain_term_key_version = PlainTermKeyVersion::kRawNoInternal; + metadata.common_grams_coverage = CommonGramsCoverage::kNone; + metadata.base_analyzer_fingerprint = "snii-test-plain-v1"; + metadata.scoring_coverage = ScoringCoverage::kComplete; + metadata.scoring_stats_version = COMMON_GRAMS_SCORING_STATS_VERSION_V1; + metadata.norm_semantics_version = COMMON_GRAMS_NORM_SEMANTICS_VERSION_V1; + metadata.scoring_doc_count = doc_count; + metadata.scoring_token_count = token_count; + return metadata; +} + +// The standard reader-side fixture: a 9000-doc kDocsPositions index whose +// ordinary terms exercise dense/sparse/windowed/tail postings. Opens +// `segment_reader`/`index_reader` over `file`. +inline Status build_reader(MemoryFile* file, reader::SniiSegmentReader* segment_reader, + reader::LogicalIndexReader* index_reader) { + constexpr uint32_t kDocCount = 9000; + auto failed_docs = docs_with_one_position(0, kDocCount, 0); + auto order_docs = docs_with_one_position(0, kDocCount, 2); + auto ordinal_docs = docs_with_one_position(0, kDocCount, 2); + auto driver_docs = docs_with_one_position(0, 8000, 0); + auto almost_docs = docs_with_one_position(0, kDocCount, 1); + std::vector needle_docs {{.docid = 100, .positions = {0}}, + {.docid = 101, .positions = {0}}, + {.docid = 102, .positions = {0}}, + {.docid = 6000, .positions = {0}}}; + std::vector numeric_tail_docs {{.docid = 42, .positions = {1}}}; + std::vector sparse_left_docs; + std::vector sparse_right_docs; + std::vector repeat_docs; + std::vector trace_docs {{.docid = 42, .positions = {0}}}; + sparse_left_docs.reserve(kDocCount / 3 + 1); + sparse_right_docs.reserve(kDocCount); + repeat_docs.reserve(kDocCount); + for (uint32_t docid = 0; docid < kDocCount; ++docid) { + if (docid % 3 == 0) { + sparse_left_docs.push_back({docid, {0}}); + } + if (docid % 4 != 1) { + sparse_right_docs.push_back({docid, {1}}); + } + repeat_docs.push_back({docid, {0, 1, 2}}); + } + almost_docs.erase(almost_docs.begin() + 4000); + failed_docs[8000].positions = {0, 4}; + for (PostingDoc& doc : order_docs) { + if (doc.docid == 5000 || doc.docid == 7000) { + doc.positions = {1}; + } else if (doc.docid == 8000) { + doc.positions = {5}; + } + } + for (PostingDoc& doc : ordinal_docs) { + if (doc.docid == 6000) { + doc.positions = {1}; + } + } + + writer::SniiIndexInput input; + input.index_id = 7; + input.index_suffix = "Body"; + input.config = format::IndexConfig::kDocsPositions; + input.doc_count = kDocCount; + input.terms = {make_term("almost", std::move(almost_docs)), + make_term("123", std::move(numeric_tail_docs)), + make_term("driver", std::move(driver_docs)), + make_term("failed", std::move(failed_docs)), + make_term("needle", std::move(needle_docs)), + make_term("order", std::move(order_docs)), + make_term("ordinal", std::move(ordinal_docs)), + make_term("repeat", std::move(repeat_docs)), + make_term("sparse_left", std::move(sparse_left_docs)), + make_term("sparse_right", std::move(sparse_right_docs)), + make_term("trace", std::move(trace_docs))}; + std::ranges::sort(input.terms, + [](const writer::TermPostings& lhs, const writer::TermPostings& rhs) { + return lhs.term < rhs.term; + }); + + writer::SniiCompoundWriter writer(file); + RETURN_IF_ERROR(writer.add_logical_index(input)); + RETURN_IF_ERROR(writer.finish()); + EXPECT_TRUE(file->finalized()); + + RETURN_IF_ERROR(reader::SniiSegmentReader::open(file, segment_reader)); + return segment_reader->open_index(input.index_id, input.index_suffix, index_reader); +} + +} // namespace doris::snii::snii_test diff --git a/be/test/storage/index/snii_spill_io_test.cpp b/be/test/storage/index/snii_spill_io_test.cpp new file mode 100644 index 00000000000000..e8ec2239a3abe3 --- /dev/null +++ b/be/test/storage/index/snii_spill_io_test.cpp @@ -0,0 +1,226 @@ +// 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. + +// R10-io-local-s3 spill backend tests. +// +// SpillableByteBuffer now spills to disk via Doris IO (io::global_local_filesystem() +// create_file / appendv / close + read_at on read-back) instead of the standalone +// doris::snii::io::Local{File}Writer/Reader. These tests pin the externally observable contract: +// whatever lands on the spill scratch file, streamed back in append order, is byte-for-byte +// identical to the concatenation of every append -- across the 0-byte branch, the +// cap-crossing spill trigger, post-spill appends, and a >256 KiB chunk that also spans the +// stream_into() copy window. The scratch file is a process-private intermediate (NOT the +// published on-disk format v2), so this is an equivalence/round-trip guarantee, not a golden +// byte test. Everything is deterministic: fixed cap_bytes and a fixed byte pattern. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/status.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/io/file_writer.h" +#include "storage/index/snii/writer/spillable_byte_buffer.h" + +namespace doris::snii::writer { +using doris::Status; +namespace { + +// Deterministic, position- and seed-dependent byte pattern so a reorder, truncation, or +// in-place corruption of any chunk is caught by a full-buffer comparison. +std::vector make_bytes(uint32_t seed, size_t n) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = static_cast((seed * 1103515245U + static_cast(i) * 12345U + 7U) & + 0xFFU); + } + return v; +} + +// In-memory doris::snii::io::FileWriter that records the exact bytes (and order) streamed into it. +class CapturingSniiFileWriter final : public doris::snii::io::FileWriter { +public: + doris::Status append(doris::snii::Slice data) override { + bytes_.insert(bytes_.end(), data.data(), data.data() + data.size()); + return doris::Status::OK(); + } + doris::Status finalize() override { return doris::Status::OK(); } + uint64_t bytes_written() const override { return bytes_.size(); } + + const std::vector& bytes() const { return bytes_; } + +private: + std::vector bytes_; +}; + +void expect_bytes_eq(const std::vector& actual, const std::vector& expected) { + // Compare sizes first, then memcmp, so a mismatch on a multi-MiB buffer does not dump the + // whole vector into the test log. + ASSERT_EQ(actual.size(), expected.size()); + if (!expected.empty()) { + EXPECT_EQ(std::memcmp(actual.data(), expected.data(), expected.size()), 0); + } +} + +// Points SNII_TEMP_DIR (resolve_temp_dir()'s first choice) at a throwaway directory so the +// spill scratch files are hermetic to the test and cleaned up afterwards. +class SniiSpillIoTest : public ::testing::Test { +protected: + void SetUp() override { + std::error_code ec; + std::filesystem::path base = std::filesystem::temp_directory_path(ec); + if (ec) { + base = "/tmp"; + } + dir_ = base / ("snii_spill_io_test_" + std::to_string(::getpid())); + std::filesystem::create_directories(dir_, ec); + ::setenv("SNII_TEMP_DIR", dir_.c_str(), 1); + } + + void TearDown() override { + ::unsetenv("SNII_TEMP_DIR"); + std::error_code ec; + std::filesystem::remove_all(dir_, ec); + } + + std::filesystem::path dir_; +}; + +// Append a mix of pre-spill, cap-crossing, and post-spill chunks (including a 0-byte append on +// each side of the spill), then assert the streamed-back bytes equal the concatenation in order. +TEST_F(SniiSpillIoTest, SpillRoundTripPreservesBytesAndOrder) { + SpillableByteBuffer buf(/*cap_bytes=*/16, "roundtrip"); + + const std::vector> chunks = { + make_bytes(1, 5), // pre-spill (ram=5) + make_bytes(2, 0), // 0-byte append, pre-spill (must be a no-op, not a chunk) + make_bytes(3, 9), // pre-spill (ram=14, still under cap) + make_bytes(4, 7), // crosses cap (ram=21 >= 16) -> spill flushes [c0,c2,c3] + make_bytes(5, 0), // 0-byte append, post-spill (len==0 appendv no-op) + make_bytes(6, 4), // post-spill, routed straight to the scratch file + make_bytes(7, 33) // post-spill + }; + + std::vector expected; + for (const auto& c : chunks) { + ASSERT_TRUE(buf.append(doris::snii::Slice(c)).ok()); + expected.insert(expected.end(), c.begin(), c.end()); + } + + EXPECT_TRUE(buf.spilled()); + EXPECT_EQ(buf.size(), static_cast(expected.size())); + + ASSERT_TRUE(buf.seal().ok()); + + CapturingSniiFileWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + expect_bytes_eq(out.bytes(), expected); +} + +// Exercises the 0-byte branch and a >256 KiB chunk appended after the spill (the old standalone +// writer's "direct write" path); the total also exceeds stream_into()'s 1 MiB copy window so the +// read-back loop runs more than once. +TEST_F(SniiSpillIoTest, EmptyAndLargeChunk) { + SpillableByteBuffer buf(/*cap_bytes=*/16, "largechunk"); + + constexpr size_t kLarge = 1300U * 1024U; // > 256 KiB and > the 1 MiB stream window + + std::vector> chunks; + chunks.push_back(make_bytes(10, 20)); // 20 >= 16 -> spill immediately, flushes [c0] + chunks.push_back(make_bytes(11, 0)); // post-spill 0-byte append (len==0 branch) + chunks.push_back(make_bytes(12, kLarge)); // post-spill large append straight to disk + chunks.push_back(make_bytes(13, 5)); // post-spill tail + + std::vector expected; + for (const auto& c : chunks) { + ASSERT_TRUE(buf.append(doris::snii::Slice(c)).ok()); + expected.insert(expected.end(), c.begin(), c.end()); + } + + EXPECT_TRUE(buf.spilled()); + EXPECT_EQ(buf.size(), static_cast(expected.size())); + + ASSERT_TRUE(buf.seal().ok()); + + CapturingSniiFileWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + expect_bytes_eq(out.bytes(), expected); +} + +// A buffer that never crosses the cap stays RAM-resident; stream_into() must still reproduce the +// exact bytes from the in-memory chunk chain (no scratch file is ever created). +TEST_F(SniiSpillIoTest, RamOnlyRoundTripNoSpill) { + SpillableByteBuffer buf(/*cap_bytes=*/UINT64_MAX, "ramonly"); // spilling disabled + + const std::vector> chunks = {make_bytes(30, 7), make_bytes(31, 0), + make_bytes(32, 13), make_bytes(33, 50)}; + + std::vector expected; + for (const auto& c : chunks) { + ASSERT_TRUE(buf.append(doris::snii::Slice(c)).ok()); + expected.insert(expected.end(), c.begin(), c.end()); + } + + EXPECT_FALSE(buf.spilled()); + EXPECT_EQ(buf.size(), static_cast(expected.size())); + + ASSERT_TRUE(buf.seal().ok()); // no-op for a RAM-resident buffer + + CapturingSniiFileWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + expect_bytes_eq(out.bytes(), expected); +} + +// The move-append path must spill and round-trip identically to the copy-append path. +TEST_F(SniiSpillIoTest, AppendMoveSpillRoundTrip) { + SpillableByteBuffer buf(/*cap_bytes=*/16, "movespill"); + + const std::vector> spec = { + {40, 10}, // pre-spill (ram=10) + {41, 0}, // empty, pre-spill + {42, 12}, // crosses cap (ram=22) -> spill flushes [c0,c2] + {43, 0}, // empty, post-spill + {44, 7} // post-spill + }; + + std::vector expected; + for (const auto& [seed, n] : spec) { + std::vector c = make_bytes(seed, n); + expected.insert(expected.end(), c.begin(), c.end()); + ASSERT_TRUE(buf.append_move(std::move(c)).ok()); + } + + EXPECT_TRUE(buf.spilled()); + EXPECT_EQ(buf.size(), static_cast(expected.size())); + + ASSERT_TRUE(buf.seal().ok()); + + CapturingSniiFileWriter out; + ASSERT_TRUE(buf.stream_into(&out).ok()); + expect_bytes_eq(out.bytes(), expected); +} + +} // namespace +} // namespace doris::snii::writer diff --git a/be/test/storage/index/snii_spimi_intern_test.cpp b/be/test/storage/index/snii_spimi_intern_test.cpp new file mode 100644 index 00000000000000..603b0ed606741f --- /dev/null +++ b/be/test/storage/index/snii_spimi_intern_test.cpp @@ -0,0 +1,1364 @@ +// 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. + +// T05: SPIMI vocab transparent-hash interning + single-string storage. +// +// These tests pin the writer-side SpimiTermBuffer owned-mode interning to its new +// shape: each distinct vocab string is materialized into owned_vocab_ EXACTLY ONCE +// (no double-store, no per-token temporary probe std::string), and the term-id +// assignment / finalize output is byte-identical to the prior behavior. Writer-only, +// no reader fixture (build_reader) needed -- the buffer is driven directly via +// add_token(string_view) and drained via finalize_sorted(). + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/exception.h" +#include "common/status.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" +#include "util/md5.h" + +using doris::snii::writer::ClassifiedPlainTerm; +using doris::snii::writer::PlainTermId; +using doris::snii::writer::SpimiTermBuffer; +using doris::snii::writer::StreamedTermPostings; +using doris::snii::writer::TermPostingBuffer; +using doris::snii::writer::TermPostings; +using doris::Status; +// Alias the writer's TEST-ONLY counter namespace so it never collides with gtest's +// own ::testing namespace. +namespace stb_testing = doris::snii::writer::testing; + +namespace { + +// An ordinary >SSO ASCII term with a prefix shared by several cases below. +std::string MixedLongAscii() { + return "ordinary-prefix-sharing-quick-brown-term"; +} + +// A >SSO multibyte (CJK) token (a direct UTF-8 string literal): 18 bytes (6 chars x +// 3 bytes), well past libstdc++'s 15-byte SSO, exercising a long non-ASCII vocab key. +std::string MixedCjk() { + return "中文长词条目"; +} + +Status MaterializeInIrregularWindows(StreamedTermPostings&& source, TermPostings* output) { + constexpr std::array kChunkSizes = {1, 3, 17, 64, 127, 251, 509}; + if (source.source == nullptr || output == nullptr) { + return Status::Error( + "test posting materializer: invalid arguments"); + } + *output = TermPostings(); + output->term = std::move(source.term); + output->retain_positions = source.retain_positions; + TermPostingBuffer buffer(nullptr); + size_t chunk_index = 0; + bool exhausted = false; + while (!exhausted) { + buffer.clear_reuse(); + const uint32_t target_docs = + static_cast(kChunkSizes[chunk_index % kChunkSizes.size()]); + RETURN_IF_ERROR(source.source->fill(target_docs, &buffer, &exhausted)); + if (!exhausted && buffer.document_count() != target_docs) { + return Status::Error( + "posting source returned a short non-terminal fill"); + } + output->docids.insert(output->docids.end(), buffer.docids().begin(), buffer.docids().end()); + output->freqs.insert(output->freqs.end(), buffer.freqs().begin(), buffer.freqs().end()); + output->positions_flat.insert(output->positions_flat.end(), buffer.positions_flat().begin(), + buffer.positions_flat().end()); + ++chunk_index; + } + return Status::OK(); +} + +std::vector AddSortedStatlessCommonGram(SpimiTermBuffer* common_grams, + uint32_t document_count) { + common_grams->enable_common_gram_pair_keys(); + const PlainTermId left = common_grams->intern_plain_term("of"); + const PlainTermId right = common_grams->intern_plain_term("world"); + std::vector expected; + expected.reserve(document_count); + for (uint32_t ordinal = 0; ordinal < document_count; ++ordinal) { + const uint32_t docid = 3 * ordinal + 1; + expected.push_back(docid); + common_grams->add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + return expected; +} + +// Feeds a fixed, implementation-independent token script mixing a long ASCII +// term, a short ASCII term, and a >SSO CJK token across several docids (some +// re-touched). Two buffers fed this script must finalize identically. +void FeedMixedScript(SpimiTermBuffer& b) { + const std::string long_ascii = MixedLongAscii(); + const std::string cjk = MixedCjk(); + b.add_token(std::string_view("alpha"), 0, 0); + b.add_token(std::string_view(long_ascii), 0, 1); + b.add_token(std::string_view(cjk), 0, 2); + b.add_token(std::string_view("alpha"), 1, 0); + b.add_token(std::string_view(long_ascii), 1, 1); + b.add_token(std::string_view("alpha"), 5, 3); + b.add_token(std::string_view(cjk), 5, 4); + b.add_token(std::string_view("alpha"), 9, 0); +} + +void ExpectPostingsEqual(const std::vector& a, const std::vector& b) { + ASSERT_EQ(a.size(), b.size()); + for (size_t i = 0; i < a.size(); ++i) { + EXPECT_EQ(a[i].term, b[i].term); + EXPECT_EQ(a[i].docids, b[i].docids); + EXPECT_EQ(a[i].freqs, b[i].freqs); + EXPECT_EQ(a[i].positions_flat, b[i].positions_flat); + EXPECT_EQ(a[i].retain_positions, b[i].retain_positions); + } +} + +namespace inverted_index = doris::segment_v2::inverted_index; + +enum class CommonGramsPostingEventKind { kPlain, kGram }; + +struct CommonGramsPostingEvent { + CommonGramsPostingEventKind kind; + std::string left; + std::string right; + uint32_t docid; + uint32_t position; + bool retain_positions; +}; + +std::string PhysicalPlainTerm(std::string_view logical_term) { + if (logical_term.front() != inverted_index::PLAIN_ESCAPE_PREFIX && + logical_term.front() != '\x1f') { + return std::string(logical_term); + } + std::string physical_term; + EXPECT_TRUE(inverted_index::try_encode_escaped_plain_term_prevalidated(logical_term, + physical_term)); + return physical_term; +} + +size_t PlainHotCacheSet(std::string_view term) { + constexpr size_t kPlainHotCacheSetCount = 1024; + const size_t raw_hash = SpimiTermBuffer::hash_term_bytes_for_test(term); + const size_t mixed_hash = phmap::phmap_mix()(raw_hash); + return mixed_hash & (kPlainHotCacheSetCount - 1); +} + +std::array FindPlainHotCacheColliders(std::string_view target) { + std::array result; + size_t found = 0; + for (uint64_t candidate = 0; found < result.size(); ++candidate) { + std::string term = "plain-cache-collider-" + std::to_string(candidate); + if (PlainHotCacheSet(term) == PlainHotCacheSet(target)) { + result[found++] = std::move(term); + } + } + return result; +} + +size_t CommonGramPairL0Index(uint32_t left, uint32_t right) { + constexpr uint64_t kHashMultiplier = 11400714819323198485ULL; + const uint64_t pair = (static_cast(left) << 32) | right; + return static_cast((pair * kHashMultiplier) >> 54); +} + +std::optional FindCommonGramPairL0RightCollider(uint32_t left, uint32_t target_right) { + const size_t target_index = CommonGramPairL0Index(left, target_right); + for (uint32_t candidate = target_right + 1; candidate < 4096; ++candidate) { + if (CommonGramPairL0Index(left, candidate) == target_index) { + return candidate; + } + } + return std::nullopt; +} + +std::string NativePairPlainTerm(uint32_t id) { + return "native-pair-plain-" + std::to_string(id); +} + +const std::vector>& CommonGramsLogicalDocuments() { + static const std::string internal = std::string(1, '\x1f') + "literal"; + static const std::string escaped = + std::string(1, inverted_index::PLAIN_ESCAPE_PREFIX) + "literal"; + static const std::vector> documents = { + {internal, "of", "the", escaped, internal, "of"}, + {internal, "of", "the"}, + }; + return documents; +} + +void FeedSeparateCommonGramAndPlainAdds(SpimiTermBuffer* buffer) { + buffer->enable_common_gram_pair_keys(); + const auto& common_words = inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + const auto& documents = CommonGramsLogicalDocuments(); + for (uint32_t docid = 0; docid < documents.size(); ++docid) { + std::optional previous; + bool previous_is_common = false; + for (uint32_t position = 0; position < documents[docid].size(); ++position) { + const std::string& logical = documents[docid][position]; + const PlainTermId current = + buffer->intern_plain_term(PhysicalPlainTerm(logical), logical); + const bool current_is_common = common_words.contains(logical); + if (previous.has_value() && (previous_is_common || current_is_common)) { + buffer->add_common_gram(*previous, current, docid, position, + previous_is_common && current_is_common); + } + buffer->add_plain_token(current, docid, position + 1); + previous = current; + previous_is_common = current_is_common; + } + } +} + +void FeedFusedCommonGramAndPlainAdds(SpimiTermBuffer* buffer) { + buffer->enable_common_gram_pair_keys(); + const auto& common_words = inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + const auto& documents = CommonGramsLogicalDocuments(); + for (uint32_t docid = 0; docid < documents.size(); ++docid) { + std::optional previous; + for (uint32_t position = 0; position < documents[docid].size(); ++position) { + const std::string& logical = documents[docid][position]; + const ClassifiedPlainTerm current = buffer->intern_classified_plain_term( + PhysicalPlainTerm(logical), logical, common_words); + if (previous.has_value() && (previous->is_common || current.is_common)) { + buffer->add_common_gram_and_plain(previous->id, current.id, docid, position, + position + 1, + previous->is_common && current.is_common); + } else { + buffer->add_plain_token(current.id, docid, position + 1); + } + previous = current; + } + } +} + +std::vector CommonGramsPostingScript() { + const std::string internal_plain = std::string(1, '\x1f') + "literal"; + const std::string escaped_plain = + std::string(1, inverted_index::PLAIN_ESCAPE_PREFIX) + "literal"; + const std::string utf8 = "中文长词条目"; + return { + {CommonGramsPostingEventKind::kPlain, internal_plain, {}, 0, 0, true}, + {CommonGramsPostingEventKind::kGram, internal_plain, "of", 0, 0, false}, + {CommonGramsPostingEventKind::kPlain, "of", {}, 0, 1, true}, + {CommonGramsPostingEventKind::kGram, "of", "the", 0, 1, true}, + {CommonGramsPostingEventKind::kPlain, "the", {}, 0, 2, true}, + {CommonGramsPostingEventKind::kGram, "the", escaped_plain, 0, 2, false}, + {CommonGramsPostingEventKind::kPlain, escaped_plain, {}, 0, 3, true}, + {CommonGramsPostingEventKind::kGram, internal_plain, "of", 0, 8, false}, + {CommonGramsPostingEventKind::kGram, "of", "the", 0, 9, true}, + {CommonGramsPostingEventKind::kPlain, internal_plain, {}, 1, 10, true}, + {CommonGramsPostingEventKind::kGram, internal_plain, "of", 1, 10, false}, + {CommonGramsPostingEventKind::kPlain, "of", {}, 1, 11, true}, + {CommonGramsPostingEventKind::kGram, "of", "the", 1, 11, true}, + {CommonGramsPostingEventKind::kPlain, "the", {}, 1, 12, true}, + {CommonGramsPostingEventKind::kGram, "the", utf8, 2, 0, false}, + {CommonGramsPostingEventKind::kPlain, utf8, {}, 2, 1, true}, + }; +} + +std::vector AdversarialPairSortScript() { + const std::string control_before_escape = std::string(1, '\x1d') + "control"; + const std::string escape_prefix = + std::string(1, inverted_index::PLAIN_ESCAPE_PREFIX) + "escaped"; + const std::string internal_prefix = std::string(1, '\x1f') + "internal"; + constexpr size_t kCommonGramLengthAndSeparatorBytes = 9; + const std::string max_left(64, 'm'); + const size_t max_right_size = inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES - + inverted_index::CG_V1_MARKER.size() - + kCommonGramLengthAndSeparatorBytes - max_left.size(); + const std::string max_right(max_right_size, 'r'); + + return { + // A control-byte component pins physical/logical rank equivalence + // immediately below the EscapedV1 namespace. + {CommonGramsPostingEventKind::kGram, control_before_escape, "tail", 0, 1, true}, + // The physical common-gram order starts with encoded left length, not + // left lexical order: one-byte "z" must precede two-byte "aa". + {CommonGramsPostingEventKind::kGram, "aa", "right", 0, 2, true}, + {CommonGramsPostingEventKind::kGram, "z", "right", 0, 3, true}, + // Equal left components fall through to logical right-byte order. + {CommonGramsPostingEventKind::kGram, "same", "z", 0, 4, true}, + {CommonGramsPostingEventKind::kGram, "same", "aa", 0, 5, true}, + // EscapedV1 physical bytes must be decoded before comparing logical + // components; 0x1e and 0x1f are both reserved physical prefixes. + {CommonGramsPostingEventKind::kGram, escape_prefix, "tail", 0, 6, true}, + {CommonGramsPostingEventKind::kGram, internal_prefix, "tail", 0, 7, true}, + {CommonGramsPostingEventKind::kGram, "same", escape_prefix, 0, 8, true}, + {CommonGramsPostingEventKind::kGram, "same", internal_prefix, 0, 9, true}, + // Exact maximum encodable physical key size exercises the trusted + // materializer without weakening the writer's analyzer-side limit. + {CommonGramsPostingEventKind::kGram, max_left, max_right, 0, 10, true}, + }; +} + +void MaybeRequestForcedSpill(SpimiTermBuffer* buffer, bool force_spill, size_t event_index) { + if (force_spill && (event_index == 3 || event_index == 7 || event_index == 12)) { + buffer->request_global_spill_for_test(); + } +} + +void FeedPhysicalCommonGrams(SpimiTermBuffer* buffer, + const std::vector& events, bool force_spill) { + for (size_t i = 0; i < events.size(); ++i) { + const auto& event = events[i]; + if (event.kind == CommonGramsPostingEventKind::kPlain) { + buffer->add_token(PhysicalPlainTerm(event.left), event.docid, event.position, + event.retain_positions); + } else { + std::string physical_gram; + EXPECT_TRUE(inverted_index::try_encode_common_gram_prevalidated(event.left, event.right, + physical_gram)); + buffer->add_token(physical_gram, event.docid, event.position, event.retain_positions); + } + MaybeRequestForcedSpill(buffer, force_spill, i); + } +} + +void FeedPairKeyCommonGrams(SpimiTermBuffer* buffer, + const std::vector& events, bool force_spill) { + buffer->enable_common_gram_pair_keys(); + for (size_t i = 0; i < events.size(); ++i) { + const auto& event = events[i]; + const PlainTermId left = buffer->intern_plain_term(PhysicalPlainTerm(event.left)); + if (event.kind == CommonGramsPostingEventKind::kPlain) { + buffer->add_plain_token(left, event.docid, event.position); + } else { + const PlainTermId right = buffer->intern_plain_term(PhysicalPlainTerm(event.right)); + buffer->add_common_gram(left, right, event.docid, event.position, + event.retain_positions); + } + MaybeRequestForcedSpill(buffer, force_spill, i); + } +} + +void ExpectPairKeysMatchPhysicalTerms(const std::vector& events, + bool force_spill) { + SpimiTermBuffer physical(/*has_positions=*/true); + SpimiTermBuffer pair_keys(/*has_positions=*/true); + if (force_spill) { + for (SpimiTermBuffer* buffer : {&physical, &pair_keys}) { + buffer->set_forced_spill_min_arena_bytes(0); + buffer->set_max_run_files(1); + } + } + + FeedPhysicalCommonGrams(&physical, events, force_spill); + FeedPairKeyCommonGrams(&pair_keys, events, force_spill); + ASSERT_TRUE(physical.status().ok()) << physical.status(); + ASSERT_TRUE(pair_keys.status().ok()) << pair_keys.status(); + if (force_spill) { + EXPECT_GE(physical.run_count_for_test(), 1U); + EXPECT_GE(pair_keys.run_count_for_test(), 1U); + } + + const std::vector expected = physical.finalize_sorted(); + const std::vector actual = pair_keys.finalize_sorted(); + ASSERT_TRUE(physical.status().ok()) << physical.status(); + ASSERT_TRUE(pair_keys.status().ok()) << pair_keys.status(); + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(actual[i].term, expected[i].term); + EXPECT_EQ(actual[i].docids, expected[i].docids); + EXPECT_EQ(actual[i].retain_positions, expected[i].retain_positions); + EXPECT_EQ(actual[i].positions_flat, expected[i].positions_flat); + if (expected[i].retain_positions) { + EXPECT_EQ(actual[i].freqs, expected[i].freqs); + } else { + EXPECT_TRUE(actual[i].freqs.empty()); + } + } +} + +} // namespace + +TEST(SniiSpimiTermBufferTest, InternHashUsesFastStringViewHash) { + const std::vector terms = { + "", + "short", + MixedLongAscii(), + MixedCjk(), + std::string("\x1F\x01\0\x03gram-key", 12), + std::string("embedded\0nul", 12), + }; + + for (const std::string& term : terms) { + const std::string_view view(term); + EXPECT_EQ(SpimiTermBuffer::hash_term_bytes_for_test(view), + std::hash {}(view)); + } +} + +TEST(SniiSpimiTermBufferTest, PhysicalCommonGramsRemainPhysicalOnFinalDrain) { + namespace inverted_index = doris::segment_v2::inverted_index; + const std::string first = inverted_index::encode_common_gram("of", "the").value(); + const std::string second = inverted_index::encode_common_gram("the", "world").value(); + + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token(second, 0, 1); + buf.add_token(std::string_view("plain"), 0, 0); + buf.add_token(first, 0, 0); + buf.add_token(first, 1, 2); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 3U); + EXPECT_EQ(terms[0].term, first); + EXPECT_EQ(terms[1].term, second); + EXPECT_EQ(terms[2].term, "plain"); + EXPECT_EQ(terms[0].docids, (std::vector {0U, 1U})); + EXPECT_EQ(terms[0].positions_flat, (std::vector {0U, 2U})); +} + +TEST(SniiSpimiTermBufferTest, PairKeyModeRejectsGenericStringTokenEntryPoint) { + EXPECT_DEATH( + { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + buffer.add_token(std::string_view("plain"), /*docid=*/0, /*pos=*/0); + }, + ".*"); +} + +TEST(SniiSpimiTermBufferTest, PairKeysMatchPhysicalTermsForMixedRepeatedEscapedAndUtf8Grams) { + ExpectPairKeysMatchPhysicalTerms(CommonGramsPostingScript(), /*force_spill=*/false); +} + +TEST(SniiSpimiTermBufferTest, FusedCommonGramAndPlainAddsMatchSeparatePostingAdds) { + SpimiTermBuffer separate(/*has_positions=*/true); + SpimiTermBuffer fused(/*has_positions=*/true); + FeedSeparateCommonGramAndPlainAdds(&separate); + FeedFusedCommonGramAndPlainAdds(&fused); + + const std::vector expected = separate.finalize_sorted(); + const std::vector actual = fused.finalize_sorted(); + ASSERT_TRUE(separate.status().ok()) << separate.status(); + ASSERT_TRUE(fused.status().ok()) << fused.status(); + ExpectPostingsEqual(actual, expected); + + const std::string internal = std::string(1, '\x1f') + "literal"; + const std::string escaped = std::string(1, inverted_index::PLAIN_ESCAPE_PREFIX) + "literal"; + const std::string repeated_mixed_gram = + inverted_index::encode_common_gram(internal, "of").value(); + const std::string positioned_gram = inverted_index::encode_common_gram("of", "the").value(); + const auto mixed = std::ranges::find(actual, repeated_mixed_gram, &TermPostings::term); + const auto positioned = std::ranges::find(actual, positioned_gram, &TermPostings::term); + const auto escaped_plain = + std::ranges::find(actual, PhysicalPlainTerm(escaped), &TermPostings::term); + ASSERT_NE(mixed, actual.end()); + ASSERT_NE(positioned, actual.end()); + ASSERT_NE(escaped_plain, actual.end()); + EXPECT_FALSE(mixed->retain_positions); + EXPECT_EQ(mixed->docids, (std::vector {0U, 1U})); + EXPECT_TRUE(mixed->freqs.empty()); + EXPECT_TRUE(positioned->retain_positions); + EXPECT_EQ(positioned->docids, (std::vector {0U, 1U})); +} + +TEST(SniiSpimiTermBufferTest, FusedAllCommonTokensCheckSpillGateOncePerInputToken) { + const auto& common_words = inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + std::vector terms; + for (std::string_view logical : {"of", "the", "and", "to"}) { + terms.push_back(buffer.intern_classified_plain_term(logical, logical, common_words)); + ASSERT_TRUE(terms.back().is_common); + } + + stb_testing::reset_spill_gate_check_count(); + buffer.add_plain_token(terms.front().id, /*docid=*/0, /*pos=*/1); + for (uint32_t i = 1; i < terms.size(); ++i) { + buffer.add_common_gram_and_plain(terms[i - 1].id, terms[i].id, /*docid=*/0, + /*gram_pos=*/i, /*plain_pos=*/i + 1, + /*retain_positions=*/true); + } + + EXPECT_EQ(buffer.total_tokens(), 7U); + EXPECT_EQ(stb_testing::spill_gate_check_count(), 4U); +} + +TEST(SniiSpimiTermBufferTest, PairKeysMatchPhysicalTermsAcrossForcedSpills) { + ExpectPairKeysMatchPhysicalTerms(CommonGramsPostingScript(), /*force_spill=*/true); +} + +TEST(SniiSpimiTermBufferTest, PairRankSortPreservesOrderWithLinearPairDecodes) { + stb_testing::reset_common_gram_pair_fast_path_counts(); + + SpimiTermBuffer physical(/*has_positions=*/true); + SpimiTermBuffer pair_keys(/*has_positions=*/true); + const auto events = AdversarialPairSortScript(); + FeedPhysicalCommonGrams(&physical, events, /*force_spill=*/false); + FeedPairKeyCommonGrams(&pair_keys, events, /*force_spill=*/false); + + const std::vector expected = physical.finalize_sorted(); + const std::vector actual = pair_keys.finalize_sorted(); + ASSERT_TRUE(physical.status().ok()) << physical.status(); + ASSERT_TRUE(pair_keys.status().ok()) << pair_keys.status(); + ExpectPostingsEqual(actual, expected); + ASSERT_EQ(actual.size(), events.size()); + EXPECT_EQ(actual.back().term.size(), inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES); + // Each pair is decoded once while building integer component ranks and once + // while materializing the final physical term. Sorting must not decode pairs. + EXPECT_EQ(stb_testing::common_gram_pair_unchecked_decode_count(), 2U * events.size()); + EXPECT_EQ(stb_testing::common_gram_trusted_plain_decode_count(), 2U * events.size()); +} + +TEST(SniiSpimiTermBufferTest, PairRankAdversarialOrderSurvivesForcedSpills) { + stb_testing::reset_run_compactions(); + ExpectPairKeysMatchPhysicalTerms(AdversarialPairSortScript(), /*force_spill=*/true); + EXPECT_GT(stb_testing::run_compactions(), 0U); +} + +TEST(SniiSpimiTermBufferTest, StatlessCommonGramOmitsFrequenciesAcrossSpillAndCompaction) { + stb_testing::reset_run_compactions(); + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + common_grams.set_forced_spill_min_arena_bytes(0); + common_grams.set_max_run_files(1); + const PlainTermId left = common_grams.intern_plain_term("of"); + const PlainTermId right = common_grams.intern_plain_term("world"); + + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/0, + /*retain_positions=*/false); + common_grams.request_global_spill_for_test(); + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/1, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/2, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/9, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/9, /*pos=*/1, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/13, /*pos=*/0, + /*retain_positions=*/false); + common_grams.request_global_spill_for_test(); + common_grams.add_common_gram(left, right, /*docid=*/13, /*pos=*/1, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/17, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/21, /*pos=*/0, + /*retain_positions=*/false); + common_grams.request_global_spill_for_test(); + common_grams.add_common_gram(left, right, /*docid=*/25, /*pos=*/0, + /*retain_positions=*/false); + + EXPECT_EQ(common_grams.total_tokens(), 10U); + EXPECT_GE(common_grams.run_count_for_test(), 1U); + EXPECT_GT(stb_testing::run_compactions(), 0U); + const std::vector gram_terms = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(gram_terms.size(), 1U); + EXPECT_FALSE(gram_terms[0].retain_positions); + EXPECT_EQ(gram_terms[0].docids, (std::vector {5U, 9U, 13U, 17U, 21U, 25U})); + EXPECT_TRUE(gram_terms[0].freqs.empty()); + EXPECT_TRUE(gram_terms[0].positions_flat.empty()); + + SpimiTermBuffer ordinary_docs(/*has_positions=*/false); + ordinary_docs.add_token("ordinary", /*docid=*/5, /*pos=*/0); + ordinary_docs.add_token("ordinary", /*docid=*/5, /*pos=*/0); + ordinary_docs.add_token("ordinary", /*docid=*/9, /*pos=*/0); + const std::vector ordinary_terms = ordinary_docs.finalize_sorted(); + ASSERT_EQ(ordinary_terms.size(), 1U); + EXPECT_EQ(ordinary_terms[0].docids, (std::vector {5U, 9U})); + EXPECT_EQ(ordinary_terms[0].freqs, (std::vector {1U, 1U})); +} + +TEST(SniiSpimiTermBufferTest, PlainTermHotCacheSurvivesForcedSpill) { + stb_testing::reset_vocab_string_materialization_count(); + stb_testing::reset_owned_term_full_byte_comparison_count(); + stb_testing::reset_common_gram_plain_cache_counts(); + + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + common_grams.set_forced_spill_min_arena_bytes(0); + common_grams.set_max_run_files(1); + + constexpr std::array terms = {"the", "database", "the", "database", + "the", "database", "the", "database"}; + for (size_t i = 0; i < terms.size(); ++i) { + if (i == 2) { + common_grams.request_global_spill_for_test(); + stb_testing::reset_owned_term_full_byte_comparison_count(); + } + const PlainTermId term = common_grams.intern_plain_term(terms[i], terms[i]); + common_grams.add_plain_token(term, static_cast(i / 2), + static_cast(i % 2)); + } + + EXPECT_GE(common_grams.run_count_for_test(), 1U); + const std::vector postings = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(postings.size(), 2U); + EXPECT_EQ(postings[0].term, "database"); + EXPECT_EQ(postings[0].docids, (std::vector {0U, 1U, 2U, 3U})); + EXPECT_EQ(postings[0].freqs, (std::vector {1U, 1U, 1U, 1U})); + EXPECT_EQ(postings[0].positions_flat, (std::vector {1U, 1U, 1U, 1U})); + EXPECT_EQ(postings[1].term, "the"); + EXPECT_EQ(postings[1].docids, (std::vector {0U, 1U, 2U, 3U})); + EXPECT_EQ(postings[1].freqs, (std::vector {1U, 1U, 1U, 1U})); + EXPECT_EQ(postings[1].positions_flat, (std::vector {0U, 0U, 0U, 0U})); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 2U); + EXPECT_EQ(stb_testing::common_gram_plain_cache_probes(), terms.size()); + EXPECT_EQ(stb_testing::common_gram_plain_cache_hits(), terms.size() - 2); + EXPECT_EQ(stb_testing::common_gram_plain_intern_table_probes(), 2U); + EXPECT_EQ(stb_testing::owned_term_full_byte_comparison_count(), 0U); +} + +TEST(SniiSpimiTermBufferTest, PlainTermHotCacheUsesFullBytesForHashCollisions) { + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + buffer.set_owned_term_hash_mask_for_test(0); + + const PlainTermId alpha = buffer.intern_plain_term("alpha", "alpha"); + const PlainTermId beta = buffer.intern_plain_term("beta", "beta"); + const PlainTermId gamma = buffer.intern_plain_term("gamma", "gamma"); + EXPECT_NE(alpha.value, beta.value); + EXPECT_NE(alpha.value, gamma.value); + EXPECT_NE(beta.value, gamma.value); + + EXPECT_EQ(buffer.intern_plain_term("beta", "beta").value, beta.value); + EXPECT_EQ(buffer.intern_plain_term("alpha", "alpha").value, alpha.value); + EXPECT_EQ(buffer.intern_plain_term("gamma", "gamma").value, gamma.value); +} + +TEST(SniiSpimiTermBufferTest, PlainTermHotCacheDoesNotPublishFailedInsertion) { + stb_testing::reset_vocab_string_materialization_count(); + stb_testing::reset_common_gram_plain_cache_counts(); + + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + stb_testing::fail_next_owned_term_reserve(); + EXPECT_THROW(buffer.intern_plain_term("reserve-recoverable", "reserve-recoverable"), + std::bad_alloc); + stb_testing::fail_next_owned_term_emplace(); + EXPECT_THROW(buffer.intern_plain_term("emplace-recoverable", "emplace-recoverable"), + std::bad_alloc); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 0U); + EXPECT_EQ(stb_testing::common_gram_plain_cache_hits(), 0U); + + stb_testing::reset_common_gram_plain_cache_counts(); + const PlainTermId reserve_retry = + buffer.intern_plain_term("reserve-recoverable", "reserve-recoverable"); + const PlainTermId emplace_retry = + buffer.intern_plain_term("emplace-recoverable", "emplace-recoverable"); + EXPECT_EQ(reserve_retry.value, 0U); + EXPECT_EQ(emplace_retry.value, 1U); + EXPECT_EQ(stb_testing::common_gram_plain_cache_probes(), 2U); + EXPECT_EQ(stb_testing::common_gram_plain_cache_hits(), 0U); + EXPECT_EQ(stb_testing::common_gram_plain_intern_table_probes(), 2U); + + EXPECT_EQ(buffer.intern_plain_term("reserve-recoverable", "reserve-recoverable").value, + reserve_retry.value); + EXPECT_EQ(buffer.intern_plain_term("emplace-recoverable", "emplace-recoverable").value, + emplace_retry.value); + EXPECT_EQ(stb_testing::common_gram_plain_cache_probes(), 4U); + EXPECT_EQ(stb_testing::common_gram_plain_cache_hits(), 2U); + EXPECT_EQ(stb_testing::common_gram_plain_intern_table_probes(), 2U); +} + +TEST(SniiSpimiTermBufferTest, ClassifiedPlainTermMembershipSurvivesHotCacheEviction) { + const auto& common_words = inverted_index::CommonWordSet::builtin_english_stop_words_v1(); + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + + inverted_index::common_grams_testing::reset_common_word_membership_lookup_count(); + const ClassifiedPlainTerm first = + buffer.intern_classified_plain_term("the", "the", common_words); + ASSERT_TRUE(first.is_common); + for (const std::string& term : FindPlainHotCacheColliders("the")) { + const ClassifiedPlainTerm classified = + buffer.intern_classified_plain_term(term, term, common_words); + EXPECT_FALSE(classified.is_common); + } + EXPECT_EQ(inverted_index::common_grams_testing::common_word_membership_lookup_count(), 3U); + + stb_testing::reset_common_gram_plain_cache_counts(); + const ClassifiedPlainTerm after_eviction = + buffer.intern_classified_plain_term("the", "the", common_words); + + EXPECT_EQ(after_eviction.id.value, first.id.value); + EXPECT_TRUE(after_eviction.is_common); + EXPECT_EQ(stb_testing::common_gram_plain_cache_hits(), 0U); + EXPECT_EQ(stb_testing::common_gram_plain_intern_table_probes(), 1U); + EXPECT_EQ(inverted_index::common_grams_testing::common_word_membership_lookup_count(), 3U); +} + +TEST(SniiSpimiTermBufferTest, ClassifiedPlainTermUsesLogicalBytesForMembership) { + const std::string logical = std::string(1, '\x1f') + "common"; + const std::string physical = PhysicalPlainTerm(logical); + ASSERT_NE(physical, logical); + const std::string wordset_content = logical + "\n"; + auto parsed = inverted_index::CommonWordSet::parse_words(wordset_content); + ASSERT_TRUE(parsed.has_value()) << parsed.error(); + + SpimiTermBuffer buffer(/*has_positions=*/true); + buffer.enable_common_gram_pair_keys(); + inverted_index::common_grams_testing::reset_common_word_membership_lookup_count(); + const ClassifiedPlainTerm classified = + buffer.intern_classified_plain_term(physical, logical, *parsed); + + EXPECT_TRUE(classified.is_common); + EXPECT_EQ(inverted_index::common_grams_testing::common_word_membership_lookup_count(), 1U); +} + +TEST(SniiSpimiTermBufferTest, DocsOnlyPairCacheSkipsHashAndSurvivesSpill) { + stb_testing::reset_common_gram_pair_cache_counts(); + + SpimiTermBuffer common_grams(/*has_positions=*/true); + const uint64_t bytes_before_cache = common_grams.resident_bytes_for_test(); + common_grams.enable_common_gram_pair_keys(); + EXPECT_EQ(common_grams.resident_bytes_for_test() - bytes_before_cache, 32U * 1024U); + common_grams.set_forced_spill_min_arena_bytes(0); + common_grams.set_max_run_files(1); + const PlainTermId left = common_grams.intern_plain_term("of"); + const PlainTermId right = common_grams.intern_plain_term("world"); + + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/0, + /*retain_positions=*/false); + common_grams.request_global_spill_for_test(); + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/1, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/5, /*pos=*/2, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/9, /*pos=*/0, + /*retain_positions=*/false); + + EXPECT_EQ(common_grams.total_tokens(), 4U); + EXPECT_GE(common_grams.run_count_for_test(), 1U); + const std::vector gram_terms = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(gram_terms.size(), 1U); + EXPECT_EQ(gram_terms[0].docids, (std::vector {5U, 9U})); + EXPECT_TRUE(gram_terms[0].freqs.empty()); + EXPECT_EQ(stb_testing::common_gram_pair_cache_probes(), 4U); + EXPECT_EQ(stb_testing::common_gram_pair_cache_pair_hits(), 3U); + EXPECT_EQ(stb_testing::common_gram_pair_cache_same_doc_hits(), 2U); +} + +TEST(SniiSpimiTermBufferTest, NativePairInternerSurvivesL0CollisionAndForcedSpill) { + constexpr uint32_t kLeft = 0; + constexpr uint32_t kTargetRight = 1; + const std::optional collider_right = + FindCommonGramPairL0RightCollider(kLeft, kTargetRight); + ASSERT_TRUE(collider_right.has_value()); + std::optional second_collider_right; + for (uint32_t candidate = *collider_right + 1; candidate < 4096; ++candidate) { + if (CommonGramPairL0Index(kLeft, candidate) == CommonGramPairL0Index(kLeft, kTargetRight)) { + second_collider_right = candidate; + break; + } + } + ASSERT_TRUE(second_collider_right.has_value()); + ASSERT_NE(*collider_right, kTargetRight); + ASSERT_EQ(CommonGramPairL0Index(kLeft, *collider_right), + CommonGramPairL0Index(kLeft, kTargetRight)); + ASSERT_EQ(CommonGramPairL0Index(kLeft, *second_collider_right), + CommonGramPairL0Index(kLeft, kTargetRight)); + + stb_testing::reset_common_gram_pair_cache_counts(); + stb_testing::reset_common_gram_native_pair_intern_counts(); + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + common_grams.set_forced_spill_min_arena_bytes(0); + common_grams.set_max_run_files(1); + + std::vector ids; + ids.reserve(static_cast(*second_collider_right) + 1); + for (uint32_t id = 0; id <= *second_collider_right; ++id) { + const PlainTermId interned = common_grams.intern_plain_term(NativePairPlainTerm(id)); + ASSERT_EQ(interned.value, id); + ids.push_back(interned); + } + + common_grams.add_common_gram(ids[kLeft], ids[kTargetRight], /*docid=*/1, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(ids[kLeft], ids[kTargetRight], /*docid=*/2, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(ids[kLeft], ids[*collider_right], /*docid=*/3, /*pos=*/0, + /*retain_positions=*/false); + common_grams.request_global_spill_for_test(); + common_grams.add_common_gram(ids[kLeft], ids[*second_collider_right], /*docid=*/4, /*pos=*/0, + /*retain_positions=*/false); + ASSERT_GE(common_grams.run_count_for_test(), 1U); + common_grams.add_common_gram(ids[kLeft], ids[kTargetRight], /*docid=*/5, /*pos=*/0, + /*retain_positions=*/false); + + EXPECT_EQ(stb_testing::common_gram_pair_cache_probes(), 5U); + EXPECT_EQ(stb_testing::common_gram_pair_cache_pair_hits(), 1U); + EXPECT_EQ(stb_testing::common_gram_native_pair_probes(), 4U); + EXPECT_EQ(stb_testing::common_gram_native_pair_hits(), 1U); + EXPECT_EQ(stb_testing::common_gram_native_pair_inserts(), 3U); + const std::vector terms = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(terms.size(), 3U); + + auto target_key = inverted_index::encode_common_gram(NativePairPlainTerm(kLeft), + NativePairPlainTerm(kTargetRight)); + ASSERT_TRUE(target_key.has_value()); + const auto target = std::ranges::find(terms, target_key.value(), &TermPostings::term); + ASSERT_NE(target, terms.end()); + EXPECT_EQ(target->docids, (std::vector {1U, 2U, 5U})); + EXPECT_EQ(common_grams.resident_bytes_for_test(), 0U); +} + +TEST(SniiSpimiTermBufferTest, PositionedPairCacheSurvivesForcedSpill) { + stb_testing::reset_common_gram_pair_cache_counts(); + SpimiTermBuffer positioned(/*has_positions=*/true); + positioned.enable_common_gram_pair_keys(); + positioned.set_forced_spill_min_arena_bytes(0); + positioned.set_max_run_files(1); + const PlainTermId positioned_left = positioned.intern_plain_term("of"); + const PlainTermId positioned_right = positioned.intern_plain_term("the"); + positioned.request_global_spill_for_test(); + positioned.add_common_gram(positioned_left, positioned_right, /*docid=*/7, /*pos=*/3, + /*retain_positions=*/true); + EXPECT_GE(positioned.run_count_for_test(), 1U); + positioned.add_common_gram(positioned_left, positioned_right, /*docid=*/7, /*pos=*/4, + /*retain_positions=*/true); + const std::vector positioned_terms = positioned.finalize_sorted(); + ASSERT_TRUE(positioned.status().ok()) << positioned.status(); + ASSERT_EQ(positioned_terms.size(), 1U); + EXPECT_EQ(positioned_terms[0].freqs, (std::vector {2U})); + EXPECT_EQ(positioned_terms[0].positions_flat, (std::vector {3U, 4U})); + EXPECT_EQ(stb_testing::common_gram_pair_cache_probes(), 2U); + EXPECT_EQ(stb_testing::common_gram_pair_cache_pair_hits(), 1U); + EXPECT_EQ(stb_testing::common_gram_pair_cache_same_doc_hits(), 0U); +} + +TEST(SniiSpimiTermBufferTest, CommonGramSpillGateObservesResidentGrowthImmediately) { + doris::snii::writer::MemoryReporter reporter( + nullptr, /*cap_bytes=*/1, + doris::snii::writer::MemoryReporter::CapPolicy::kSpillThreshold); + SpimiTermBuffer common_grams(/*has_positions=*/true, /*spill_threshold_bytes=*/0, &reporter); + common_grams.enable_common_gram_pair_keys(); + const PlainTermId term = common_grams.intern_plain_term("common"); + + common_grams.add_plain_token(term, /*docid=*/0, /*pos=*/0); + EXPECT_EQ(common_grams.run_count_for_test(), 1U); + const std::vector terms = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].freqs, (std::vector {1U})); +} + +TEST(SniiSpimiTermBufferTest, StatlessCommonGramStoresOnlyDocumentDeltas) { + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + const PlainTermId left = common_grams.intern_plain_term("of"); + const PlainTermId right = common_grams.intern_plain_term("world"); + const uint64_t resident_before_postings = common_grams.resident_bytes_for_test(); + + constexpr uint32_t kDocumentCount = 20000; + for (uint32_t docid = 1; docid <= kDocumentCount; ++docid) { + common_grams.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + + // One one-byte zigzag delta per document fits in a single 32 KiB arena block. + // The old tagged representation wrote a second, constant byte and needed two. + EXPECT_LT(common_grams.resident_bytes_for_test() - resident_before_postings, 60U * 1024U); + const std::vector terms = common_grams.finalize_sorted(); + ASSERT_TRUE(common_grams.status().ok()) << common_grams.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_FALSE(terms[0].retain_positions); + ASSERT_EQ(terms[0].docids.size(), kDocumentCount); + EXPECT_EQ(terms[0].docids.front(), 1U); + EXPECT_EQ(terms[0].docids.back(), kDocumentCount); + EXPECT_TRUE(terms[0].freqs.empty()); +} + +TEST(SniiSpimiTermBufferTest, SortedWideStatlessCommonGramStreamsSourceWindows) { + for (const uint32_t document_count : {doris::snii::format::kSlimDfThreshold, uint32_t {777}}) { + SCOPED_TRACE(document_count); + SpimiTermBuffer common_grams(/*has_positions=*/true); + const std::vector expected = + AddSortedStatlessCommonGram(&common_grams, document_count); + + size_t callback_count = 0; + const doris::Status status = common_grams.for_each_term_sorted( + [&](StreamedTermPostings&& source) -> doris::Status { + ++callback_count; + TermPostings postings; + RETURN_IF_ERROR(MaterializeInIrregularWindows(std::move(source), &postings)); + EXPECT_FALSE(postings.retain_positions); + EXPECT_TRUE(postings.freqs.empty()); + EXPECT_TRUE(postings.positions_flat.empty()); + EXPECT_EQ(postings.document_count(), document_count); + EXPECT_EQ(postings.docids, expected); + return doris::Status::OK(); + }); + + ASSERT_TRUE(status.ok()) << status; + EXPECT_EQ(callback_count, 1U); + } +} + +TEST(SniiSpimiTermBufferTest, StreamedStatlessCommonGramRequiresSynchronousConsumption) { + SpimiTermBuffer common_grams(/*has_positions=*/true); + constexpr uint32_t kDocumentCount = 777; + AddSortedStatlessCommonGram(&common_grams, kDocumentCount); + + const doris::Status status = common_grams.for_each_term_sorted( + [&](StreamedTermPostings&&) { return doris::Status::OK(); }); + + EXPECT_TRUE(status.is()) << status; +} + +TEST(SniiSpimiTermBufferTest, ForcedSpillStatlessCommonGramSourceCoalescesSeamDuplicates) { + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + common_grams.set_forced_spill_min_arena_bytes(0); + const PlainTermId left = common_grams.intern_plain_term("of"); + const PlainTermId right = common_grams.intern_plain_term("world"); + + std::vector expected; + expected.reserve(768); + for (uint32_t docid = 0; docid < 256; ++docid) { + expected.push_back(docid); + common_grams.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + common_grams.request_global_spill_for_test(); + expected.push_back(256); + common_grams.add_common_gram(left, right, /*docid=*/256, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/256, /*pos=*/0, + /*retain_positions=*/false); + for (uint32_t docid = 257; docid < 512; ++docid) { + expected.push_back(docid); + common_grams.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + common_grams.request_global_spill_for_test(); + expected.push_back(512); + common_grams.add_common_gram(left, right, /*docid=*/512, /*pos=*/0, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/512, /*pos=*/0, + /*retain_positions=*/false); + for (uint32_t docid = 513; docid < 768; ++docid) { + expected.push_back(docid); + common_grams.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + + ASSERT_EQ(common_grams.run_count_for_test(), 2U); + size_t callback_count = 0; + const doris::Status status = + common_grams.for_each_term_sorted([&](StreamedTermPostings&& source) { + ++callback_count; + TermPostings postings; + RETURN_IF_ERROR(MaterializeInIrregularWindows(std::move(source), &postings)); + EXPECT_EQ(postings.document_count(), expected.size()); + EXPECT_EQ(postings.docids, expected); + EXPECT_TRUE(std::ranges::is_sorted(postings.docids)); + EXPECT_EQ(std::ranges::adjacent_find(postings.docids), postings.docids.end()); + return doris::Status::OK(); + }); + + ASSERT_TRUE(status.ok()) << status; + EXPECT_EQ(callback_count, 1U); +} + +TEST(SniiSpimiTermBufferTest, OutOfOrderWideStatlessCommonGramStaysMaterialized) { + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + const PlainTermId left = common_grams.intern_plain_term("of"); + const PlainTermId right = common_grams.intern_plain_term("world"); + constexpr uint32_t kDocumentCount = 600; + for (uint32_t docid = 0; docid < kDocumentCount; ++docid) { + common_grams.add_common_gram(left, right, docid, /*pos=*/0, + /*retain_positions=*/false); + } + common_grams.add_common_gram(left, right, /*docid=*/100, /*pos=*/1, + /*retain_positions=*/false); + common_grams.add_common_gram(left, right, /*docid=*/50, /*pos=*/2, + /*retain_positions=*/false); + + size_t callback_count = 0; + const doris::Status status = + common_grams.for_each_term_sorted([&](StreamedTermPostings&& source) { + ++callback_count; + TermPostings postings; + RETURN_IF_ERROR(MaterializeInIrregularWindows(std::move(source), &postings)); + EXPECT_EQ(postings.document_count(), kDocumentCount); + if (postings.docids.size() != kDocumentCount) { + return doris::Status::Error( + "unexpected materialized document count"); + } + for (uint32_t docid = 0; docid < kDocumentCount; ++docid) { + EXPECT_EQ(postings.docids[docid], docid); + } + EXPECT_TRUE(postings.freqs.empty()); + EXPECT_TRUE(postings.positions_flat.empty()); + return doris::Status::OK(); + }); + + ASSERT_TRUE(status.ok()) << status; + EXPECT_EQ(callback_count, 1U); +} + +TEST(SniiSpimiTermBufferTest, OrdinaryDocsOnlyMarkerTermRetainsFrequency) { + const std::string literal_marker_term = std::string(inverted_index::CG_V1_MARKER) + "literal"; + SpimiTermBuffer ordinary(/*has_positions=*/false); + ordinary.add_token(literal_marker_term, /*docid=*/3, /*pos=*/0, + /*retain_positions=*/false); + + const std::vector terms = ordinary.finalize_sorted(); + ASSERT_TRUE(ordinary.status().ok()) << ordinary.status(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, literal_marker_term); + EXPECT_EQ(terms[0].docids, (std::vector {3U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U})); +} + +TEST(SniiSpimiTermBufferTest, CommonGramLogicalValidationRunsOncePerDistinctPlainTerm) { + stb_testing::reset_common_gram_logical_validation_count(); + stb_testing::reset_vocab_string_materialization_count(); + + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + const std::string internal = std::string(1, '\x1f') + "literal"; + const std::string physical_internal = PhysicalPlainTerm(internal); + + for (int repeat = 0; repeat < 8; ++repeat) { + common_grams.intern_plain_term("ordinary", "ordinary"); + common_grams.intern_plain_term("中文长词条目", "中文长词条目"); + common_grams.intern_plain_term(physical_internal, internal); + } + + EXPECT_EQ(stb_testing::common_gram_logical_validation_count(), 3U); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 3U); +} + +TEST(SniiSpimiTermBufferTest, CommonGramLogicalValidationFailureDoesNotMutateVocabulary) { + stb_testing::reset_common_gram_logical_validation_count(); + stb_testing::reset_vocab_string_materialization_count(); + + SpimiTermBuffer common_grams(/*has_positions=*/true); + common_grams.enable_common_gram_pair_keys(); + for (const std::string& invalid : {std::string("bad\0term", 8), std::string("\xc3", 1)}) { + try { + common_grams.intern_plain_term(invalid, invalid); + FAIL() << "expected analyzer error"; + } catch (const doris::Exception& error) { + EXPECT_EQ(error.code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + } + } + + EXPECT_EQ(stb_testing::common_gram_logical_validation_count(), 2U); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 0U); + EXPECT_EQ(common_grams.intern_plain_term("valid", "valid").value, 0U); +} + +TEST(SniiSpimiTermBufferTest, CompactPostingPoolVarintFastPathCrossesSliceBoundary) { + doris::snii::writer::CompactPostingPool pool; + doris::snii::writer::CompactPostingPool::SliceWriter writer; + uint8_t level = 0; + const uint32_t head = pool.start_chain(&writer, &level); + for (uint32_t i = 1; i < doris::snii::writer::CompactPostingPool::kSliceSizes_level0(); ++i) { + pool.append_byte(&writer, &level, 0); + } + + const std::array values = { + 300, 0, 127, 128, static_cast(std::numeric_limits::max())}; + for (const uint64_t value : values) { + pool.append_varint(&writer, &level, value); + } + + auto cursor = pool.cursor(head, writer.cur); + for (uint32_t i = 1; i < doris::snii::writer::CompactPostingPool::kSliceSizes_level0(); ++i) { + EXPECT_EQ(cursor.next(), 0U); + } + for (const uint64_t expected : values) { + EXPECT_EQ(cursor.read_varint(), expected); + } +} + +TEST(SniiSpimiTermBufferTest, PerTermDocsOnlyDropsSameDocOccurrencesBeforeDecode) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token("docs", /*docid=*/5, /*pos=*/10, /*retain_positions=*/false); + buf.add_token("docs", /*docid=*/5, /*pos=*/11, /*retain_positions=*/false); + buf.add_token("docs", /*docid=*/1, /*pos=*/3, /*retain_positions=*/false); + buf.add_token("docs", /*docid=*/5, /*pos=*/12, /*retain_positions=*/false); + buf.add_token("positioned", /*docid=*/5, /*pos=*/20, /*retain_positions=*/true); + buf.add_token("positioned", /*docid=*/5, /*pos=*/21, /*retain_positions=*/true); + + EXPECT_EQ(buf.total_tokens(), 6U); + const std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].term, "docs"); + EXPECT_FALSE(terms[0].retain_positions); + EXPECT_EQ(terms[0].docids, (std::vector {1U, 5U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U, 1U})); + EXPECT_TRUE(terms[0].positions_flat.empty()); + EXPECT_EQ(terms[1].term, "positioned"); + EXPECT_TRUE(terms[1].retain_positions); + EXPECT_EQ(terms[1].docids, (std::vector {5U})); + EXPECT_EQ(terms[1].freqs, (std::vector {2U})); + EXPECT_EQ(terms[1].positions_flat, (std::vector {20U, 21U})); +} + +// --------------------------------------------------------------------------------- +// Functional verification (FV1-FV9) +// --------------------------------------------------------------------------------- + +// FV1: ids are assigned in first-seen order (b=0,a=1,c=2) but the emitted order is +// lexicographic (a,b,c); docids/freqs are recovered correctly. +TEST(SniiSpimiTermBufferTest, VocabAssignsIdsInFirstSeenOrder) { + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token(std::string_view("b"), 0, 0); + buf.add_token(std::string_view("a"), 1, 0); + buf.add_token(std::string_view("b"), 2, 0); + buf.add_token(std::string_view("c"), 3, 0); + buf.add_token(std::string_view("a"), 4, 0); + + EXPECT_EQ(buf.unique_terms(), 3U); + EXPECT_EQ(buf.total_tokens(), 5U); + EXPECT_TRUE(buf.status().ok()); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 3U); + EXPECT_EQ(terms[0].term, "a"); + EXPECT_EQ(terms[1].term, "b"); + EXPECT_EQ(terms[2].term, "c"); + EXPECT_EQ(terms[0].docids, (std::vector {1U, 4U})); + EXPECT_EQ(terms[1].docids, (std::vector {0U, 2U})); + EXPECT_EQ(terms[2].docids, (std::vector {3U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U, 1U})); + EXPECT_EQ(terms[1].freqs, (std::vector {1U, 1U})); + EXPECT_EQ(terms[2].freqs, (std::vector {1U})); +} + +// FV2: the same >SSO ordinary term fed 1000 times reuses ONE id (heterogeneous hit +// path), yielding a single term with 1000 ascending docids and freq 1 each. +TEST(SniiSpimiTermBufferTest, RepeatedTermReusesSingleId) { + SpimiTermBuffer buf(/*has_positions=*/false); + const std::string term = MixedLongAscii(); + ASSERT_GT(term.size(), 15U); // exceeds libstdc++ SSO: the OLD probe heap-allocated + + constexpr uint32_t kRepeats = 1000; + for (uint32_t d = 0; d < kRepeats; ++d) { + buf.add_token(std::string_view(term), d, 0); + } + EXPECT_EQ(buf.unique_terms(), 1U); + EXPECT_EQ(buf.total_tokens(), kRepeats); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, term); + ASSERT_EQ(terms[0].docids.size(), kRepeats); + for (uint32_t d = 0; d < kRepeats; ++d) { + EXPECT_EQ(terms[0].docids[d], d); + EXPECT_EQ(terms[0].freqs[d], 1U); + } +} + +// FV3 (also the byte-identity perf gate): two independent buffers fed the same mixed +// script (long ASCII + plain + CJK token) finalize to element-identical postings. +TEST(SniiSpimiTermBufferTest, FinalizeIsByteIdenticalAcrossRuns) { + SpimiTermBuffer a(/*has_positions=*/true); + SpimiTermBuffer b(/*has_positions=*/true); + FeedMixedScript(a); + FeedMixedScript(b); + + std::vector ra = a.finalize_sorted(); + std::vector rb = b.finalize_sorted(); + ExpectPostingsEqual(ra, rb); + EXPECT_TRUE(a.status().ok()); + EXPECT_TRUE(b.status().ok()); + + bool saw_long_ascii = false; + bool saw_cjk = false; + for (const auto& tp : ra) { + if (tp.term == MixedLongAscii()) { + saw_long_ascii = true; + } + if (tp.term == MixedCjk()) { + saw_cjk = true; + } + } + EXPECT_TRUE(saw_long_ascii); + EXPECT_TRUE(saw_cjk); +} + +// FV4: no tokens -> empty result, status stays OK. +TEST(SniiSpimiTermBufferTest, EmptyVocabProducesNoTerms) { + SpimiTermBuffer buf(/*has_positions=*/false); + EXPECT_EQ(buf.unique_terms(), 0U); + std::vector terms = buf.finalize_sorted(); + EXPECT_TRUE(terms.empty()); + EXPECT_TRUE(buf.status().ok()); +} + +// FV5: a single token yields a single term, single docid, freq 1. +TEST(SniiSpimiTermBufferTest, SingleTokenProducesSingleTerm) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token(std::string_view("solo"), 7, 3); + EXPECT_EQ(buf.unique_terms(), 1U); + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].term, "solo"); + EXPECT_EQ(terms[0].docids, (std::vector {7U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U})); + EXPECT_EQ(terms[0].positions_flat, (std::vector {3U})); +} + +// FV6: the empty string is a valid distinct term; the heterogeneous equality functor +// matches "" against a stored "" so a repeat empty token reuses the same id. +TEST(SniiSpimiTermBufferTest, EmptyStringIsAValidDistinctTerm) { + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token(std::string_view(""), 0, 0); // empty term, first occurrence + buf.add_token(std::string_view("x"), 1, 0); + buf.add_token(std::string_view(""), 2, 0); // empty term reused via transparent eq + + EXPECT_EQ(buf.unique_terms(), 2U); + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 2U); + EXPECT_EQ(terms[0].term, ""); // "" sorts before "x" + EXPECT_EQ(terms[0].docids, (std::vector {0U, 2U})); + EXPECT_EQ(terms[1].term, "x"); + EXPECT_EQ(terms[1].docids, (std::vector {1U})); + EXPECT_TRUE(buf.status().ok()); +} + +// FV7: add_token(string_view) on a BORROWED-vocab buffer is rejected (latches +// InvalidArgument, token ignored). The interning functors hold &owned_vocab_ but are +// never dereferenced on this path (reject happens before the find), so empty +// owned_vocab_ is never indexed out of bounds. +TEST(SniiSpimiTermBufferTest, BorrowedModeRejectsStringView) { + const std::vector vocab = {"a", "b"}; + SpimiTermBuffer buf(&vocab, /*has_positions=*/false); + buf.add_token(0U, 0, 0); // valid id-path token + buf.add_token(std::string_view("a"), 1, 0); // illegal on a borrowed-vocab buffer + + EXPECT_FALSE(buf.status().ok()); + EXPECT_EQ(buf.total_tokens(), 1U); // the string-view token was ignored + EXPECT_EQ(buf.unique_terms(), 1U); +} + +// FV8: long ordinary terms sharing a prefix remain distinct and round-trip with +// exact bytes. +TEST(SniiSpimiTermBufferTest, PrefixSharingOrdinaryTermsRoundTrip) { + SpimiTermBuffer buf(/*has_positions=*/true); + const std::string first = "ordinary-prefix-sharing-alpha-suffix"; + const std::string second = "ordinary-prefix-sharing-beta-suffix"; + const std::string third = "ordinary-prefix-sharing-gamma-suffix"; + buf.add_token(std::string_view(first), 0, 0); + buf.add_token(std::string_view(second), 0, 1); + buf.add_token(std::string_view(third), 0, 2); + + EXPECT_EQ(buf.unique_terms(), 3U); + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 3U); + + bool saw_first = false; + bool saw_second = false; + bool saw_third = false; + for (const auto& tp : terms) { + if (tp.term == first) { + saw_first = true; + } else if (tp.term == second) { + saw_second = true; + } else if (tp.term == third) { + saw_third = true; + } + } + EXPECT_TRUE(saw_first); + EXPECT_TRUE(saw_second); + EXPECT_TRUE(saw_third); +} + +// FV9: out-of-order / revisited docids for one term coalesce into one strictly +// ascending entry per docid (orthogonal to the intern change; guards no regression). +TEST(SniiSpimiTermBufferTest, OutOfOrderDocidCoalesces) { + SpimiTermBuffer buf(/*has_positions=*/true); + buf.add_token(std::string_view("t"), 5, 50); + buf.add_token(std::string_view("t"), 1, 10); + buf.add_token(std::string_view("t"), 5, 52); // revisit doc 5 + + std::vector terms = buf.finalize_sorted(); + ASSERT_EQ(terms.size(), 1U); + EXPECT_EQ(terms[0].docids, (std::vector {1U, 5U})); + EXPECT_EQ(terms[0].freqs, (std::vector {1U, 2U})); + EXPECT_EQ(terms[0].positions_flat, (std::vector {10U, 50U, 52U})); + EXPECT_TRUE(buf.status().ok()); +} + +// --------------------------------------------------------------------------------- +// Deterministic performance verification (allocation seam counts) +// --------------------------------------------------------------------------------- + +// The same >SSO term fed M times materializes its string EXACTLY ONCE: no per-token +// temporary probe std::string (F21) and no second owned-string map key (F03). The +// OLD instrumented baseline would have been M temporaries + 1 emplace = M+1. +TEST(SniiSpimiTermBufferTest, VocabInterningMaterializesEachStringOnce) { + stb_testing::reset_vocab_string_materialization_count(); + SpimiTermBuffer buf(/*has_positions=*/false); + const std::string term = MixedLongAscii(); + ASSERT_GT(term.size(), 15U); + + constexpr uint32_t kRepeats = 1000; + for (uint32_t d = 0; d < kRepeats; ++d) { + buf.add_token(std::string_view(term), d, 0); + } + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 1U); + EXPECT_EQ(buf.unique_terms(), 1U); +} + +// N distinct >SSO terms, each fed twice, materialize exactly N strings: the count +// tracks DISTINCT terms (one owned_vocab_.emplace_back each), not total tokens -- the +// repeat of an already-seen term allocates nothing on the heterogeneous hit path. +TEST(SniiSpimiTermBufferTest, VocabMaterializesOncePerDistinctTerm) { + stb_testing::reset_vocab_string_materialization_count(); + SpimiTermBuffer buf(/*has_positions=*/false); + + constexpr uint32_t kDistinct = 500; + for (uint32_t i = 0; i < kDistinct; ++i) { + const std::string term = + "ordinary-prefix-sharing-distinct-term-number-" + std::to_string(i); + buf.add_token(std::string_view(term), i, 0); + buf.add_token(std::string_view(term), i + kDistinct, 0); // repeat: zero materialization + } + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), static_cast(kDistinct)); + EXPECT_EQ(buf.unique_terms(), static_cast(kDistinct)); +} + +// The seam resets cleanly between measurements (guards the reset_/count_ contract the +// two perf tests above rely on for determinism in a shared process). +TEST(SniiSpimiTermBufferTest, MaterializationCounterResets) { + stb_testing::reset_vocab_string_materialization_count(); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 0U); + SpimiTermBuffer buf(/*has_positions=*/false); + buf.add_token(std::string_view("one"), 0, 0); + buf.add_token(std::string_view("two"), 0, 0); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 2U); + stb_testing::reset_vocab_string_materialization_count(); + EXPECT_EQ(stb_testing::vocab_string_materialization_count(), 0U); +} diff --git a/be/test/storage/index/snii_writer_test.cpp b/be/test/storage/index/snii_writer_test.cpp new file mode 100644 index 00000000000000..c282a3f102e32c --- /dev/null +++ b/be/test/storage/index/snii_writer_test.cpp @@ -0,0 +1,526 @@ +// 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. + +// SNII writer and reader regression tests: +// - heap_bytes() accessors on the resident format readers +// (SampledTermIndexReader / DictBlockDirectoryReader / DictBlockReader) that +// LogicalIndexReader::memory_usage() sums so the searcher-cache charge stops +// under-counting. Exact hand-computed equality for SSO terms; the string-heap +// accumulation is exercised with an over-15-byte term. +// - geometric null-docid accumulation growth. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/status.h" +#include "gen_cpp/AgentService_types.h" +#include "runtime/exec_env.h" +#include "runtime/index_policy/index_policy_mgr.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/analyzer/custom_analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_grams_segment_metadata.h" +#include "storage/index/inverted/common_grams/common_word_set.h" +#include "storage/index/snii/common/slice.h" +#include "storage/index/snii/encoding/byte_sink.h" +#include "storage/index/snii/format/dict_block.h" +#include "storage/index/snii/format/dict_block_directory.h" +#include "storage/index/snii/format/dict_entry.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/format/phrase_bigram.h" +#include "storage/index/snii/format/sampled_term_index.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/docid_sink.h" +#include "storage/index/snii/snii_index_writer.h" +#include "storage/tablet/tablet_schema.h" +#include "util/slice.h" + +namespace doris::segment_v2 { +bool snii_effective_write_freq(snii::format::IndexConfig index_config); +} + +namespace { + +using doris::snii::ByteSink; +using doris::snii::Slice; +using namespace doris::snii::format; // NOLINT(google-build-using-namespace) + +class ScopedCommonGramsPolicies { +public: + ScopedCommonGramsPolicies() { + auto* exec_env = doris::ExecEnv::GetInstance(); + previous_ = exec_env->index_policy_mgr(); + exec_env->_index_policy_mgr = &manager_; + + doris::TIndexPolicy tokenizer; + tokenizer.id = 920010; + tokenizer.name = "snii_writer_cg_tokenizer"; + tokenizer.type = doris::TIndexPolicyType::TOKENIZER; + tokenizer.properties["type"] = "char_group"; + tokenizer.properties["tokenize_on_chars"] = "[\\u0020]"; + tokenizer.properties["max_token_length"] = "16383"; + + doris::TIndexPolicy common_grams; + common_grams.id = 920011; + common_grams.name = "snii_writer_cg_filter"; + common_grams.type = doris::TIndexPolicyType::TOKEN_FILTER; + common_grams.properties["type"] = "common_grams"; + + doris::TIndexPolicy analyzer; + analyzer.id = 920012; + analyzer.name = analyzer_name(); + analyzer.type = doris::TIndexPolicyType::ANALYZER; + analyzer.properties["tokenizer"] = tokenizer.name; + analyzer.properties["token_filter"] = "lowercase," + common_grams.name; + manager_.apply_policy_changes({tokenizer, common_grams, analyzer}, {}); + } + + ~ScopedCommonGramsPolicies() { doris::ExecEnv::GetInstance()->_index_policy_mgr = previous_; } + + static std::string analyzer_name() { return "snii_writer_cg_analyzer"; } + + // Derive a metadata seed from the analyzer's real identity, the same + // way production derives it for a fresh segment; hand-rolled fingerprints can + // never match the provider identity and fail init's consistency check. + doris::segment_v2::inverted_index::CommonGramsSegmentMetadata metadata_seed() { + auto provider = manager_.get_analyzer_provider_by_name(analyzer_name()); + DORIS_CHECK(provider != nullptr); + const auto* identity = provider->common_grams_identity(); + DORIS_CHECK(identity != nullptr); + return doris::segment_v2::inverted_index::make_common_grams_segment_metadata(*identity); + } + +private: + doris::IndexPolicyMgr manager_; + doris::IndexPolicyMgr* previous_ = nullptr; +}; + +// A fatal assertion inside a helper FUNCTION only aborts the helper; the calling +// test keeps running and may dereference state that failed to initialize (this +// bit us as a null-analyzer SEGV). A macro expands in the test body, so the +// fatal assertion aborts the test itself. +bool status_is_ok(const doris::Status& status) { + return status.ok(); +} + +#define ASSERT_OK(status) ASSERT_PRED1(status_is_ok, status) + +int assert_ok_evaluation_count = 0; + +doris::Status counted_failure_status() { + ++assert_ok_evaluation_count; + return doris::Status::InternalError("counted failure"); +} + +TEST(AssertOkMacroTest, FailingExpressionIsEvaluatedOnce) { + assert_ok_evaluation_count = 0; + EXPECT_FATAL_FAILURE(ASSERT_OK(counted_failure_status()), "counted failure"); + EXPECT_EQ(assert_ok_evaluation_count, 1); +} + +void init_failure_index_meta(doris::TabletIndex* index_meta, int64_t index_id) { + doris::TabletIndexPB index_pb; + index_pb.set_index_type(doris::IndexType::INVERTED); + index_pb.set_index_id(index_id); + index_pb.set_index_name("common_grams_failure_latch"); + index_pb.add_col_unique_id(0); + index_pb.mutable_properties()->insert({"parser", "english"}); + index_pb.mutable_properties()->insert({"support_phrase", "true"}); + index_meta->init_from_pb(index_pb); +} + +void init_common_grams_index_meta(doris::TabletIndex* index_meta, int64_t index_id) { + doris::TabletIndexPB index_pb; + index_pb.set_index_type(doris::IndexType::INVERTED); + index_pb.set_index_id(index_id); + index_pb.set_index_name("common_grams_typed_writer"); + index_pb.add_col_unique_id(0); + index_pb.mutable_properties()->insert({"analyzer", ScopedCommonGramsPolicies::analyzer_name()}); + index_pb.mutable_properties()->insert({"support_phrase", "true"}); + index_meta->init_from_pb(index_pb); +} + +std::string encoded_gram(std::string_view left, std::string_view right) { + auto encoded = doris::segment_v2::inverted_index::encode_common_gram(left, right); + EXPECT_TRUE(encoded.has_value()) << encoded.error(); + return encoded.value(); +} + +std::shared_ptr create_failure_analyzer() { + using namespace doris::segment_v2::inverted_index; // NOLINT + Settings tokenizer_settings; + tokenizer_settings.set("tokenize_on_chars", "[whitespace]"); + CustomAnalyzerConfig::Builder builder; + builder.with_tokenizer_config("char_group", tokenizer_settings); + builder.add_token_filter_config("lowercase", {}); + builder.add_token_filter_config("common_grams", {}); + return CustomAnalyzer::build_custom_analyzer(builder.build(), AnalysisPurpose::kIndex); +} + +doris::Slice malformed_value_after_valid_token() { + static const std::string input = + std::string("VALID B") + static_cast(0xFF) + std::string("AD"); + return doris::Slice(input); +} + +TEST(SniiWriterTest, EffectiveWriteFreqResolver) { + const bool saved = doris::config::snii_positions_index_write_freq; + doris::config::snii_positions_index_write_freq = false; + EXPECT_FALSE(doris::segment_v2::snii_effective_write_freq(IndexConfig::kDocsPositions)); + EXPECT_TRUE(doris::segment_v2::snii_effective_write_freq(IndexConfig::kDocsPositionsScoring)); + doris::config::snii_positions_index_write_freq = true; + EXPECT_TRUE(doris::segment_v2::snii_effective_write_freq(IndexConfig::kDocsPositions)); + doris::config::snii_positions_index_write_freq = saved; +} + +// SampledTermIndexReader::heap_bytes(): all-SSO sample terms have no per-string +// heap, so the charge is exactly n_blocks * sizeof(std::string) (reserve-exact +// backing buffer). +TEST(SniiSegmentReaderTest, SampledTermIndexHeapBytesMatchesFormula) { + const std::vector terms = {"s000", "s001", "s002", "s003", "s004", "s005"}; + SampledTermIndexBuilder builder; + for (const auto& term : terms) { + builder.add_block_first_term(term); // strictly ascending, SSO + } + ByteSink sink; + builder.finish(&sink); + + SampledTermIndexReader reader; + ASSERT_OK(SampledTermIndexReader::open(sink.view(), &reader)); + ASSERT_EQ(reader.n_blocks(), terms.size()); + EXPECT_EQ(reader.heap_bytes(), terms.size() * sizeof(std::string)); +} + +// The std_string_heap_bytes accumulation: an over-15-byte sample term adds its +// heap buffer on top of the vector buffer. +TEST(SniiSegmentReaderTest, SampledTermIndexHeapBytesCountsLongTerms) { + const std::string long_term = "b_this_is_a_long_sample_term_well_over_15_bytes"; + ASSERT_GT(long_term.size(), 15U); + const std::vector terms = {"a_short", long_term, "c_short"}; + SampledTermIndexBuilder builder; + for (const auto& term : terms) { + builder.add_block_first_term(term); + } + ByteSink sink; + builder.finish(&sink); + + SampledTermIndexReader reader; + ASSERT_OK(SampledTermIndexReader::open(sink.view(), &reader)); + const size_t vector_only = terms.size() * sizeof(std::string); + EXPECT_GT(reader.heap_bytes(), vector_only); + // capacity() >= size(), so the long term contributes >= size()+1 heap bytes. + EXPECT_GE(reader.heap_bytes(), vector_only + long_term.size() + 1); + // Cross-check the shared helper on an SSO vs non-SSO string. + EXPECT_EQ(std_string_heap_bytes(std::string("short")), 0U); + EXPECT_GT(std_string_heap_bytes(long_term), 0U); +} + +// DictBlockDirectoryReader::heap_bytes(): BlockRef is trivially copyable, so the +// charge is exactly n_blocks * sizeof(BlockRef). +TEST(SniiSegmentReaderTest, DictBlockDirectoryHeapBytesMatchesFormula) { + DictBlockDirectoryBuilder builder; + constexpr uint32_t kBlocks = 5; + for (uint32_t i = 0; i < kBlocks; ++i) { + BlockRef ref; + ref.offset = 100000ULL * (i + 1); // multi-byte varints -> each ref > 8 bytes + ref.length = 640; + ref.n_entries = 3; + ref.flags = 0; + ref.checksum = 0xDEAD0000U + i; + builder.add(ref); + } + ByteSink sink; + builder.finish(&sink); + + DictBlockDirectoryReader reader; + ASSERT_OK(DictBlockDirectoryReader::open(sink.view(), &reader)); + ASSERT_EQ(reader.n_blocks(), kBlocks); + EXPECT_EQ(reader.heap_bytes(), static_cast(kBlocks) * sizeof(BlockRef)); +} + +// A minimal slim pod_ref entry that round-trips at tier T1 (extra tier>=T2 fields +// are ignored on encode). Terms are supplied by the caller in ascending order. +DictEntry make_pod_ref(std::string term) { + DictEntry e; + e.term = std::move(term); + e.kind = DictEntryKind::kPodRef; + e.enc = DictEntryEnc::kSlim; + e.df = 3; + e.ttf_delta = 6; + e.max_freq = 9; + e.frq_off_delta = 0; + e.frq_len = 128; + e.frq_docs_len = 64; // dd region on-disk length (<= frq_len) + e.dd_meta.uncomp_len = 70; + e.dd_meta.crc = 0xABCD1234U; + e.freq_meta.uncomp_len = 40; + e.freq_meta.crc = 0x55AA00FFU; + e.prx_off_delta = 0; + e.prx_len = 64; + return e; +} + +std::vector build_dict_block(const std::vector& terms, + uint32_t anchor_interval) { + DictBlockBuilder builder(IndexTier::kT1, /*has_positions=*/false, /*frq_base=*/0, + /*prx_base=*/0, anchor_interval); + for (const auto& term : terms) { + builder.add_entry(make_pod_ref(term)); + } + ByteSink sink; + builder.finish(&sink); + return sink.buffer(); +} + +// DictBlockReader::heap_bytes(): with anchor_interval 16 and 20 SSO entries there +// are two anchors (indices 0 and 16), each with an SSO anchor term, so the charge +// is exactly n_anchors * (sizeof(uint32_t) + sizeof(std::string)). +TEST(SniiSegmentReaderTest, DictBlockAnchorHeapBytesMatchesFormula) { + constexpr uint32_t kEntries = 20; + constexpr uint32_t kAnchorInterval = 16; + std::vector terms; + terms.reserve(kEntries); + for (uint32_t i = 0; i < kEntries; ++i) { + // "dt_00".."dt_19": strictly ascending, 5 bytes (SSO). + terms.push_back("dt_" + std::string(1, static_cast('0' + i / 10)) + + std::string(1, static_cast('0' + i % 10))); + } + const std::vector bytes = build_dict_block(terms, kAnchorInterval); + + DictBlockReader reader; + ASSERT_OK( + DictBlockReader::open(Slice(bytes), IndexTier::kT1, /*has_positions=*/false, &reader)); + ASSERT_EQ(reader.n_entries(), kEntries); + + const size_t n_anchors = (kEntries + kAnchorInterval - 1) / kAnchorInterval; // == 2 + EXPECT_EQ(reader.heap_bytes(), n_anchors * (sizeof(uint32_t) + sizeof(std::string))); +} + +// A long (> 15 byte) anchor term adds its heap buffer beyond the anchor vectors. +TEST(SniiSegmentReaderTest, DictBlockAnchorHeapBytesCountsLongAnchor) { + // One entry -> one anchor (entry 0), whose term is > 15 bytes. + const std::string long_term = "a_dict_anchor_term_well_over_15_bytes"; + ASSERT_GT(long_term.size(), 15U); + const std::vector bytes = build_dict_block({long_term}, /*anchor_interval=*/16); + + DictBlockReader reader; + ASSERT_OK( + DictBlockReader::open(Slice(bytes), IndexTier::kT1, /*has_positions=*/false, &reader)); + ASSERT_EQ(reader.n_entries(), 1U); + const size_t vector_only = sizeof(uint32_t) + sizeof(std::string); // one anchor + EXPECT_GT(reader.heap_bytes(), vector_only); + EXPECT_GE(reader.heap_bytes(), vector_only + long_term.size() + 1); +} + +// ==================== null-docids growth-policy regression pins ==================== +// +// append_nullable feeds add_nulls once per NULL RUN -- millions of calls on a +// large interleaved-null compaction segment. An exact reserve(size()+count) +// inside add_nulls capped capacity at "just enough", so EVERY subsequent call +// reallocated + memcpy'd the whole array: O(runs x N) total memcpy (the +// agentlogs full-compaction pathology: ~TBs of memcpy per tablet, 8x+ slower +// than V3). These pins count capacity changes across many small appends: with +// geometric growth that is O(log n); with the exact-reserve bug it was one per +// call. add_nulls touches only _null_docids/_rid, so a scaffold-free writer +// (null collaborators, no init()) exercises the real production code path. + +TEST(SniiWriterNullDocids, AddNullsGrowsGeometricallyNotQuadratically) { + doris::segment_v2::SniiIndexColumnWriter writer(nullptr, nullptr, /*single_field=*/true); + constexpr uint32_t kRuns = 4096; + size_t capacity_changes = 0; + size_t last_cap = writer.null_docids_for_test().capacity(); + for (uint32_t i = 0; i < kRuns; ++i) { + ASSERT_OK(writer.add_nulls(1)); + const size_t cap = writer.null_docids_for_test().capacity(); + if (cap != last_cap) { + ++capacity_changes; + last_cap = cap; + } + } + // Geometric growth reallocates O(log n) times (libstdc++ doubling: ~13 for + // 4096); the exact-reserve bug reallocated on every call (4096). The bound + // leaves generous headroom for any sane growth policy while still failing + // a per-call realloc by two orders of magnitude. + EXPECT_LE(capacity_changes, 64U) << "add_nulls reallocates per call again"; + // Content unchanged by the policy fix: docids 0..kRuns-1 in order. + const auto& nulls = writer.null_docids_for_test(); + ASSERT_EQ(nulls.size(), kRuns); + EXPECT_EQ(nulls.front(), 0U); + EXPECT_EQ(nulls.back(), kRuns - 1); + EXPECT_TRUE(std::ranges::is_sorted(nulls)); +} + +TEST(SniiWriterFailureLatch, AnalyzerFailureDiscardsStateAndBlocksFinish) { + doris::TabletIndex index_meta; + init_failure_index_meta(&index_meta, 91); + + doris::segment_v2::SniiIndexColumnWriter writer(nullptr, &index_meta, + /*single_field=*/true); + ASSERT_OK(writer.init()); + + writer.set_analysis_for_test( + doris::segment_v2::inverted_index::InvertedIndexAnalyzer::create_reader({}), + create_failure_analyzer()); + + const doris::Slice value = malformed_value_after_valid_token(); + auto add_status = writer.add_values("", &value, 1); + ASSERT_EQ(add_status.code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_EQ(writer.term_buffer_for_test(), nullptr); + EXPECT_EQ(writer.memory_reporter_for_test(), nullptr); + + EXPECT_EQ(writer.add_nulls(1).code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_EQ(writer.finish().code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + EXPECT_EQ(writer.term_buffer_for_test(), nullptr); + EXPECT_EQ(writer.memory_reporter_for_test(), nullptr); +} + +TEST(SniiCommonGramsWriter, EncodesTypedTermsAndBuildsSemanticScoringInput) { + using doris::segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX; + ScopedCommonGramsPolicies policies; + doris::TabletIndex index_meta; + init_common_grams_index_meta(&index_meta, 93); + + doris::segment_v2::SniiIndexColumnWriter writer(nullptr, &index_meta, /*single_field=*/true, + policies.metadata_seed()); + ASSERT_OK(writer.init()); + + const std::string marker_plain = std::string("\x1f") + "raw"; + const std::string first_text = marker_plain + " of the"; + const std::string second_text = "plain terms"; + const std::vector values = {doris::Slice(first_text), doris::Slice(second_text)}; + ASSERT_OK(writer.add_values("", values.data(), values.size())); + ASSERT_OK(writer.add_nulls(1)); + + auto postings = writer.term_buffer_for_test()->finalize_sorted(); + std::map by_term; + for (auto& posting : postings) { + EXPECT_FALSE(posting.term.starts_with(doris::snii::format::kPhraseBigramTermMarker)); + const std::string term = posting.term; + by_term.emplace(term, std::move(posting)); + } + + const std::string escaped_plain = std::string(1, PLAIN_ESCAPE_PREFIX) + "Graw"; + ASSERT_TRUE(by_term.contains(escaped_plain)); + EXPECT_EQ(by_term.at(escaped_plain).docids, (std::vector {0})); + EXPECT_EQ(by_term.at(escaped_plain).positions_flat, (std::vector {1})); + ASSERT_TRUE(by_term.contains(encoded_gram(marker_plain, "of"))); + EXPECT_EQ(by_term.at(encoded_gram(marker_plain, "of")).docids, (std::vector {0})); + // A gram keeps positions only when BOTH sides are common words + // (add_common_gram_and_plain's retain flag); a mixed pair such as + // escaped-raw + common keeps co-occurrence docids but never participates + // in phrase-position semantics. + EXPECT_TRUE(by_term.at(encoded_gram(marker_plain, "of")).positions_flat.empty()); + ASSERT_TRUE(by_term.contains(encoded_gram("of", "the"))); + EXPECT_EQ(by_term.at(encoded_gram("of", "the")).positions_flat, (std::vector {2})); + EXPECT_EQ(by_term.at("of").positions_flat, (std::vector {2})); + EXPECT_EQ(by_term.at("the").positions_flat, (std::vector {3})); + EXPECT_EQ(by_term.at("plain").positions_flat, (std::vector {1})); + EXPECT_EQ(by_term.at("terms").positions_flat, (std::vector {2})); + + EXPECT_EQ(writer.encoded_norms_for_test(), + (std::vector {doris::snii::query::encode_norm(3), + doris::snii::query::encode_norm(2), + doris::snii::query::encode_norm(0)})); + EXPECT_EQ(writer.scoring_token_count_for_test(), 5U); + const auto metadata = writer.common_grams_metadata_for_test(); + EXPECT_EQ(metadata.scoring_doc_count, 3U); + EXPECT_EQ(metadata.scoring_token_count, 5U); +} + +TEST(SniiCommonGramsWriter, EscapedPlainOverflowPoisonsAllAccumulatedState) { + ScopedCommonGramsPolicies policies; + doris::TabletIndex index_meta; + init_common_grams_index_meta(&index_meta, 94); + + doris::segment_v2::SniiIndexColumnWriter writer(nullptr, &index_meta, /*single_field=*/true, + policies.metadata_seed()); + ASSERT_OK(writer.init()); + + const std::string valid_text = "man of the year"; + const doris::Slice valid_value(valid_text); + ASSERT_OK(writer.add_values("", &valid_value, 1)); + ASSERT_OK(writer.add_nulls(1)); + ASSERT_NE(writer.term_buffer_for_test(), nullptr); + EXPECT_FALSE(writer.encoded_norms_for_test().empty()); + EXPECT_FALSE(writer.null_docids_for_test().empty()); + EXPECT_GT(writer.scoring_token_count_for_test(), 0U); + + std::string oversized_after_escape( + doris::segment_v2::inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + oversized_after_escape.front() = '\x1f'; + const doris::Slice value(oversized_after_escape); + auto status = writer.add_values("", &value, 1); + ASSERT_EQ(status.code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); + ASSERT_EQ(writer.term_buffer_for_test(), nullptr); + ASSERT_EQ(writer.memory_reporter_for_test(), nullptr); + EXPECT_TRUE(writer.encoded_norms_for_test().empty()); + EXPECT_TRUE(writer.null_docids_for_test().empty()); + EXPECT_EQ(writer.scoring_token_count_for_test(), 0U); + EXPECT_EQ(writer.finish().code(), doris::ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); +} + +TEST(SniiCommonGramsWriter, PlainControlKeepsRawTermsAndLegacyConfig) { + doris::TabletIndex index_meta; + init_failure_index_meta(&index_meta, 95); + doris::segment_v2::SniiIndexColumnWriter writer(nullptr, &index_meta, + /*single_field=*/true); + ASSERT_OK(writer.init()); + + const std::string raw = std::string("\x1f") + "literal"; + const doris::Slice value(raw); + ASSERT_OK(writer.add_values("", &value, 1)); + auto postings = writer.term_buffer_for_test()->finalize_sorted(); + ASSERT_FALSE(postings.empty()); + EXPECT_TRUE(std::ranges::none_of(postings, [](const auto& posting) { + return posting.term.starts_with(doris::segment_v2::inverted_index::PLAIN_ESCAPE_PREFIX); + })); + EXPECT_EQ(writer.config_for_test(), IndexConfig::kDocsPositions); + EXPECT_TRUE(writer.encoded_norms_for_test().empty()); + EXPECT_FALSE(writer.has_common_grams_metadata_seed_for_test()); +} + +TEST(SniiDocIdSinkGrowth, AppendRangeGrowsGeometrically) { + std::vector docids; + doris::snii::query::VectorDocIdSink sink(docids); + constexpr uint32_t kRuns = 4096; + size_t capacity_changes = 0; + size_t last_cap = docids.capacity(); + for (uint32_t i = 0; i < kRuns; ++i) { + ASSERT_OK(sink.append_range(i, static_cast(i) + 1)); // one docid per run + const size_t cap = docids.capacity(); + if (cap != last_cap) { + ++capacity_changes; + last_cap = cap; + } + } + EXPECT_LE(capacity_changes, 64U) << "append_range reallocates per call again"; + ASSERT_EQ(docids.size(), kRuns); + EXPECT_EQ(docids.front(), 0U); + EXPECT_EQ(docids.back(), kRuns - 1); + EXPECT_TRUE(std::ranges::is_sorted(docids)); +} + +} // namespace diff --git a/be/test/storage/segment/count_on_index_fastpath_test.cpp b/be/test/storage/segment/count_on_index_fastpath_test.cpp new file mode 100644 index 00000000000000..e529aa7c5a6bdd --- /dev/null +++ b/be/test/storage/segment/count_on_index_fastpath_test.cpp @@ -0,0 +1,313 @@ +// 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 "storage/segment/count_on_index_fastpath.h" + +#include + +// Truth table for the G02 count-only fast-path caller guard. SegmentIterator +// fills CountOnIndexFastpathFacts from its state right before the index apply; +// this test pins that the ONLY admitted configuration is: COUNT_ON_INDEX agg + +// exactly one pushed-down MATCH expr + zero other filters + zero deletes + +// full row bitmap + zero row-id consumers + the no-read-data contract active. +// Each single deviation must veto the fast path (fall through to the +// row-accurate bitmap), because the fabricated [0, df) bitmap is only +// count-equivalent, never row-equivalent. +namespace doris::segment_v2 { + +namespace { + +// The one configuration that admits the fast path. +CountOnIndexFastpathFacts safe_facts() { + CountOnIndexFastpathFacts f; + f.is_count_on_index_agg = true; + f.has_column_predicates = false; + f.common_expr_count = 1; + f.single_expr_is_match_pred = true; + f.has_virtual_column_exprs = false; + f.has_delete_predicates = false; + f.segment_delete_bitmap_empty = true; + f.has_col_id_predicates = false; + f.has_topn_filters = false; + f.has_external_row_ranges = false; + f.row_bitmap_is_full = true; + f.record_rowids = false; + f.has_ann_topn = false; + f.has_score_runtime = false; + f.no_need_read_data_opt_enabled = true; + f.keys_type_supported = true; + return f; +} + +} // namespace + +TEST(CountOnIndexFastpath, AllGuardsSatisfiedAdmits) { + EXPECT_TRUE(count_on_index_fastpath_safe(safe_facts())); +} + +// Non-count context: no COUNT_ON_INDEX pushdown (the "context flag absent" +// case) -> the reader must take the normal bitmap path. +TEST(CountOnIndexFastpath, NotCountOnIndexVetoes) { + auto f = safe_facts(); + f.is_count_on_index_agg = false; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); +} + +// Extra predicates beside the single MATCH veto: the count of one predicate +// is not the count of a conjunction. +TEST(CountOnIndexFastpath, ExtraPredicatesVeto) { + { + auto f = safe_facts(); + f.has_column_predicates = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.common_expr_count = 2; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.common_expr_count = 0; + f.single_expr_is_match_pred = false; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + // One expr, but a compound / non-MATCH root: leaf-level fabricated + // bitmaps would be combined with sibling bitmaps. + auto f = safe_facts(); + f.single_expr_is_match_pred = false; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_col_id_predicates = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_topn_filters = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_virtual_column_exprs = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } +} + +// Deletes veto: COUNT_ON_INDEX counts matching rows MINUS deleted rows, and a +// fabricated id range cannot participate in the delete-bitmap subtraction or +// in delete-predicate filtering (mirror of the V3 _lazy_init handling). +TEST(CountOnIndexFastpath, DeletesVeto) { + { + auto f = safe_facts(); + f.segment_delete_bitmap_empty = false; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_delete_predicates = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } +} + +// A pre-pruned or later-restricted row space vetoes: [0, df) only counts +// correctly against the full [0, num_rows) bitmap. +TEST(CountOnIndexFastpath, PrunedRowSpaceVetoes) { + { + auto f = safe_facts(); + f.row_bitmap_is_full = false; // condition cache / key ranges pruned + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_external_row_ranges = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } +} + +// Consumers of REAL row ids veto. +TEST(CountOnIndexFastpath, RowIdConsumersVeto) { + { + auto f = safe_facts(); + f.record_rowids = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_ann_topn = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.has_score_runtime = true; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } +} + +// The COUNT_ON_INDEX no-read-data contract must be active, otherwise column +// data would be materialized at fabricated row ids. +TEST(CountOnIndexFastpath, DataReadContractVetoes) { + { + auto f = safe_facts(); + f.no_need_read_data_opt_enabled = false; + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } + { + auto f = safe_facts(); + f.keys_type_supported = false; // e.g. AGG keys / MOR without MOW + EXPECT_FALSE(count_on_index_fastpath_safe(f)); + } +} + +// --- G03 count-emission shortcut guard truth table -------------------------- +// The shortcut replaces the per-rowid batch loop with a defaults-fill +// countdown; the ONLY admitted configuration is: reader answered from df + +// zero surviving evaluation stages + zero row-id/value consumers + pure +// countdown batch accounting + a schema-shaped block whose every column is +// defaults-fillable. Each single deviation must refuse (keep today's batch +// loop, which is always count-exact). + +namespace { + +CountEmitShortcutFacts safe_emit_facts() { + CountEmitShortcutFacts f; + f.count_fastpath_hit = true; + f.needs_vec_eval = false; + f.needs_short_eval = false; + f.needs_expr_eval = false; + f.has_remaining_col_predicates = false; + f.has_remaining_common_exprs = false; + f.has_delete_predicates = false; + f.lazy_materialization_read = false; + f.has_virtual_columns = false; + f.record_rowids = false; + f.has_read_limit = false; + f.read_orderby_key_reverse = false; + f.has_condition_cache_digest = false; + f.block_shape_matches_schema = true; + f.all_columns_emit_defaults = true; + return f; +} + +} // namespace + +TEST(CountEmitShortcut, AllGuardsSatisfiedAdmits) { + EXPECT_TRUE(count_emit_shortcut_safe(safe_emit_facts())); +} + +// The reader must have ANSWERED via the fabricated count bitmap. A mere G02 +// guard pass whose reader fell through to a row-accurate decode (multi-term, +// pruned bigram, CLucene index, query-cache hit) keeps today's emission. +TEST(CountEmitShortcut, NoReaderHitRefuses) { + auto f = safe_emit_facts(); + f.count_fastpath_hit = false; + EXPECT_FALSE(count_emit_shortcut_safe(f)); +} + +// Any surviving evaluation stage refuses: the shortcut emits unconditionally +// and cannot re-apply filters (e.g. an index-eval downgrade kept the expr). +TEST(CountEmitShortcut, SurvivingEvaluationRefuses) { + { + auto f = safe_emit_facts(); + f.needs_vec_eval = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.needs_short_eval = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.needs_expr_eval = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.has_remaining_col_predicates = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.has_remaining_common_exprs = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.has_delete_predicates = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.lazy_materialization_read = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } +} + +// Consumers of real row ids or per-row values refuse. +TEST(CountEmitShortcut, RowIdOrValueConsumersRefuse) { + { + auto f = safe_emit_facts(); + f.has_virtual_columns = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.record_rowids = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } +} + +// Batch accounting must be a pure countdown: limits, reverse key-ordered +// reads, and condition-cache writes all keep today's loop. +TEST(CountEmitShortcut, NonCountdownBatchAccountingRefuses) { + { + auto f = safe_emit_facts(); + f.has_read_limit = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.read_orderby_key_reverse = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.has_condition_cache_digest = true; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } +} + +// The emitted block must be exactly the read schema and every column must be +// one today's path fills with defaults (no real read, no storage->schema +// cast, no version/lsn/tso rewrite). +TEST(CountEmitShortcut, BlockShapeOrColumnFillMismatchRefuses) { + { + auto f = safe_emit_facts(); + f.block_shape_matches_schema = false; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } + { + auto f = safe_emit_facts(); + f.all_columns_emit_defaults = false; + EXPECT_FALSE(count_emit_shortcut_safe(f)); + } +} + +} // namespace doris::segment_v2 diff --git a/be/test/storage/segment/inverted_index_file_reader_test.cpp b/be/test/storage/segment/inverted_index_file_reader_test.cpp index ba14fbb8359622..4aecc7848a22a2 100644 --- a/be/test/storage/segment/inverted_index_file_reader_test.cpp +++ b/be/test/storage/segment/inverted_index_file_reader_test.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include "io/fs/local_file_system.h" #include "runtime/exec_env.h" @@ -31,6 +32,9 @@ #include "storage/index/inverted/inverted_index_cache.h" #include "storage/index/inverted/inverted_index_desc.h" #include "storage/index/inverted/inverted_index_fs_directory.h" +#include "storage/index/snii/snii_doris_adapter.h" +#include "storage/index/snii/writer/snii_compound_writer.h" +#include "storage/index/snii/writer/spimi_term_buffer.h" #include "storage/options.h" #include "storage/storage_engine.h" #include "storage/tablet/tablet_schema.h" @@ -214,6 +218,44 @@ class InvertedIndexFileReaderTest : public testing::Test { create_v2_file_with_null_bitmap(file_path, index_id, index_suffix, false); } + void create_snii_file_with_null_bitmap(const std::string& file_path, uint64_t index_id, + const std::string& index_suffix, bool has_null_bitmap) { + std::filesystem::path parent_path = std::filesystem::path(file_path).parent_path(); + if (!std::filesystem::exists(parent_path)) { + std::filesystem::create_directories(parent_path); + } + + io::FileWriterPtr file_writer; + io::FileWriterOptions opts; + Status st = io::global_local_filesystem()->create_file(file_path, &file_writer, &opts); + ASSERT_TRUE(st.ok()) << st; + + snii_doris::DorisSniiFileWriter snii_file_writer(file_writer.get()); + doris::snii::writer::SniiCompoundWriter writer(&snii_file_writer); + doris::snii::writer::TermPostings term; + term.term = "apple"; + term.docids = {0}; + term.freqs = {1}; + term.positions_flat = {0}; + + doris::snii::writer::SniiIndexInput input; + input.index_id = index_id; + input.index_suffix = index_suffix; + input.config = doris::snii::format::IndexConfig::kDocsPositions; + input.doc_count = 2; + input.terms = {std::move(term)}; + if (has_null_bitmap) { + input.null_docids = {1}; + } + + st = writer.add_logical_index(input); + ASSERT_TRUE(st.ok()) << st; + st = writer.finish(); + ASSERT_TRUE(st.ok()) << st; + st = file_writer->close(false); + ASSERT_TRUE(st.ok()) << st; + } + private: StorageEngine* _engine_ref = nullptr; std::unique_ptr _data_dir = nullptr; @@ -376,6 +418,42 @@ TEST_F(InvertedIndexFileReaderTest, TestHasNullV2WithSmallNullBitmap) { EXPECT_FALSE(res); // Small bitmap should return false } +TEST_F(InvertedIndexFileReaderTest, TestHasNullSniiWithNullBitmap) { + std::string index_path = kTestDir + "/has_null_snii"; + create_snii_file_with_null_bitmap(index_path + ".idx", 1, "test", true); + + InvertedIndexFileInfo file_info; + IndexFileReader reader(io::global_local_filesystem(), index_path, + InvertedIndexStorageFormatPB::SNII, file_info); + + Status init_status = reader.init(4096); + EXPECT_TRUE(init_status.ok()); + + MockTabletIndex tablet_index(1, "test"); + bool res = false; + Status status = reader.has_null(&tablet_index, &res); + EXPECT_TRUE(status.ok()); + EXPECT_TRUE(res); +} + +TEST_F(InvertedIndexFileReaderTest, TestHasNullSniiWithoutNullBitmap) { + std::string index_path = kTestDir + "/has_null_snii_without_nulls"; + create_snii_file_with_null_bitmap(index_path + ".idx", 1, "test", false); + + InvertedIndexFileInfo file_info; + IndexFileReader reader(io::global_local_filesystem(), index_path, + InvertedIndexStorageFormatPB::SNII, file_info); + + Status init_status = reader.init(4096); + EXPECT_TRUE(init_status.ok()); + + MockTabletIndex tablet_index(1, "test"); + bool res = true; + Status status = reader.has_null(&tablet_index, &res); + EXPECT_TRUE(status.ok()); + EXPECT_FALSE(res); +} + // Test case for has_null method with V2 format stream nullptr TEST_F(InvertedIndexFileReaderTest, TestHasNullV2StreamNullptr) { std::string index_path = kTestDir + "/has_null_stream_null"; diff --git a/be/test/storage/segment/inverted_index_file_writer_test.cpp b/be/test/storage/segment/inverted_index_file_writer_test.cpp index f0a7680d25152b..49f40ea85584b5 100644 --- a/be/test/storage/segment/inverted_index_file_writer_test.cpp +++ b/be/test/storage/segment/inverted_index_file_writer_test.cpp @@ -32,6 +32,33 @@ using namespace doris; class IndexFileWriterTest : public ::testing::Test { protected: + class TrackingFileWriter final : public io::FileWriter { + public: + Status close(bool non_block = false) override { + close_calls.push_back(non_block); + _state = non_block ? State::ASYNC_CLOSING : State::CLOSED; + return Status::OK(); + } + + Status appendv(const Slice* data, size_t data_cnt) override { + for (size_t i = 0; i < data_cnt; ++i) { + _bytes_appended += data[i].size; + } + return Status::OK(); + } + + const io::Path& path() const override { return _path; } + size_t bytes_appended() const override { return _bytes_appended; } + State state() const override { return _state; } + + std::vector close_calls; + + private: + io::Path _path {"tracking.idx"}; + size_t _bytes_appended = 0; + State _state = State::OPENED; + }; + class MockDorisFSDirectoryFileLength : public DorisFSDirectory { public: //MOCK_METHOD(lucene::store::IndexOutput*, createOutput, (const char* name), (override)); @@ -2181,6 +2208,21 @@ TEST_F(IndexFileWriterTest, EmptyIndexV2Test) { ASSERT_TRUE(close_status.ok()); } +TEST_F(IndexFileWriterTest, SniiBeginCloseStartsUnderlyingFileClose) { + auto file_writer = std::make_unique(); + auto* tracking_writer = file_writer.get(); + IndexFileWriter writer(_fs, _index_path_prefix, _rowset_id, _seg_id, + InvertedIndexStorageFormatPB::SNII, std::move(file_writer)); + + ASSERT_TRUE(writer.begin_close().ok()); + EXPECT_EQ(tracking_writer->state(), io::FileWriter::State::ASYNC_CLOSING); + EXPECT_EQ(tracking_writer->close_calls, std::vector({true})); + + ASSERT_TRUE(writer.finish_close().ok()); + EXPECT_EQ(tracking_writer->state(), io::FileWriter::State::CLOSED); + EXPECT_EQ(tracking_writer->close_calls, std::vector({true, false})); +} + // Test for StreamSinkFileWriter path in close() TEST_F(IndexFileWriterTest, StreamSinkFileWriterCloseTest) { // Create a writer without providing a file_writer to trigger the StreamSinkFileWriter path diff --git a/be/test/storage/segment/inverted_index_fs_directory_test.cpp b/be/test/storage/segment/inverted_index_fs_directory_test.cpp index 4109dba52beaff..ab9ad886530640 100644 --- a/be/test/storage/segment/inverted_index_fs_directory_test.cpp +++ b/be/test/storage/segment/inverted_index_fs_directory_test.cpp @@ -288,6 +288,58 @@ TEST_F(DorisFSDirectoryTest, FSIndexInputReadInternalWithBytesReadError) { _CLDELETE(input); } +TEST_F(DorisFSDirectoryTest, FSIndexInputReadInternalRecordsIndexIOStatsAndContext) { + std::filesystem::path test_file = _tmp_dir / "test_file_with_stats"; + std::ofstream ofs(test_file); + ofs << "test content for stats"; + ofs.close(); + + lucene::store::IndexInput* input = nullptr; + CLuceneError error; + + bool result = + DorisFSDirectory::FSIndexInput::open(_fs, test_file.string().c_str(), input, error); + EXPECT_TRUE(result); + + io::FileCacheStatistics stats; + io::IOContext io_ctx; + io_ctx.is_disposable = true; + io_ctx.is_index_data = false; + io_ctx.read_file_cache = false; + io_ctx.file_cache_stats = &stats; + + input->setIoContext(&io_ctx); + input->setIndexFile(true); + + uint8_t buffer[6]; + input->readBytes(buffer, 6, false); + EXPECT_EQ(std::string(reinterpret_cast(buffer), 6), "test c"); + + const auto* captured = static_cast(input->getIoContext()); + EXPECT_TRUE(captured->is_inverted_index); + EXPECT_TRUE(captured->is_index_data); + EXPECT_FALSE(captured->read_file_cache); + EXPECT_TRUE(captured->is_disposable); + EXPECT_EQ(captured->file_cache_stats, &stats); + + EXPECT_EQ(stats.inverted_index_request_bytes, 6); + EXPECT_EQ(stats.inverted_index_read_bytes, 6); + EXPECT_EQ(stats.inverted_index_range_read_count, 1); + EXPECT_EQ(stats.inverted_index_serial_read_rounds, 1); + + input->setIoContext(nullptr); + captured = static_cast(input->getIoContext()); + EXPECT_TRUE(captured->is_inverted_index); + EXPECT_TRUE(captured->is_index_data); + EXPECT_EQ(captured->file_cache_stats, nullptr); + + input->setIndexFile(false); + captured = static_cast(input->getIoContext()); + EXPECT_FALSE(captured->is_index_data); + + _CLDELETE(input); +} + // Test 19: FSIndexOutput init error TEST_F(DorisFSDirectoryTest, FSIndexOutputInitError) { DebugPoints::instance()->add( diff --git a/be/test/storage/segment/inverted_index_iterator_test.cpp b/be/test/storage/segment/inverted_index_iterator_test.cpp index cbbb910f65c549..4313c1035627ad 100644 --- a/be/test/storage/segment/inverted_index_iterator_test.cpp +++ b/be/test/storage/segment/inverted_index_iterator_test.cpp @@ -23,6 +23,9 @@ #include #include +#include "common/exception.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_string.h" #include "storage/index/inverted/inverted_index_parser.h" #include "storage/index/inverted/inverted_index_reader.h" @@ -140,9 +143,10 @@ TEST_F(InvertedIndexIteratorTest, AddReader_SingleReader) { TEST_F(InvertedIndexIteratorTest, AddReader_MultipleReadersWithDifferentKeys) { InvertedIndexIterator iterator; - auto reader1 = create_mock_reader("chinese"); - auto reader2 = create_mock_reader("english"); - auto reader3 = create_mock_reader(""); // empty key stays empty (no properties set) + auto reader1 = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 1); + auto reader2 = create_mock_reader("english", InvertedIndexReaderType::FULLTEXT, 2); + auto reader3 = create_mock_reader("", InvertedIndexReaderType::FULLTEXT, + 3); // empty key stays empty (no properties set) iterator.add_reader(InvertedIndexReaderType::FULLTEXT, reader1); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, reader2); @@ -165,11 +169,39 @@ TEST_F(InvertedIndexIteratorTest, AddReader_MultipleReadersWithDifferentKeys) { // Don't assert specific reader - fallback mode returns first available } +TEST_F(InvertedIndexIteratorTest, AddReader_DuplicateIndexIdFails) { + auto first_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 7); + auto duplicate_reader = create_mock_reader("english", InvertedIndexReaderType::FULLTEXT, 7); + +#ifndef NDEBUG + EXPECT_DEATH( + { + InvertedIndexIterator iterator; + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, first_reader); + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, duplicate_reader); + }, + "Duplicate inverted index id 7 in one field"); +#else + InvertedIndexIterator iterator; + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, first_reader); + try { + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, duplicate_reader); + FAIL() << "Expected doris::Exception"; + } catch (const Exception& e) { + const auto message = e.to_string(); + EXPECT_NE(message.find("Duplicate inverted index id 7 in one field"), std::string::npos) + << message; + } +#endif +} + // Test that "none" is treated as a distinct analyzer key (not empty string) TEST_F(InvertedIndexIteratorTest, AddReader_NoneAnalyzerIsDistinct) { InvertedIndexIterator iterator; - auto empty_reader = create_mock_reader(""); // empty key (no properties) - auto none_reader = create_mock_reader("none"); // explicit "none" key + auto empty_reader = create_mock_reader("", InvertedIndexReaderType::FULLTEXT, + 1); // empty key (no properties) + auto none_reader = create_mock_reader("none", InvertedIndexReaderType::FULLTEXT, + 2); // explicit "none" key iterator.add_reader(InvertedIndexReaderType::FULLTEXT, empty_reader); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, none_reader); @@ -251,8 +283,8 @@ TEST_F(InvertedIndexIteratorTest, EmptyString_NoSpecifiedAnalyzer_ReturnsAny) { // select_best_reader with column_type tests TEST_F(InvertedIndexIteratorTest, SelectBestReader_MatchQuerySelectsFulltext) { InvertedIndexIterator iterator; - auto fulltext_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT); - auto string_reader = create_mock_reader("chinese", InvertedIndexReaderType::STRING_TYPE); + auto fulltext_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 1); + auto string_reader = create_mock_reader("chinese", InvertedIndexReaderType::STRING_TYPE, 2); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, fulltext_reader); iterator.add_reader(InvertedIndexReaderType::STRING_TYPE, string_reader); @@ -266,8 +298,8 @@ TEST_F(InvertedIndexIteratorTest, SelectBestReader_MatchQuerySelectsFulltext) { TEST_F(InvertedIndexIteratorTest, SelectBestReader_EqualQuerySelectsStringType) { InvertedIndexIterator iterator; - auto fulltext_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT); - auto string_reader = create_mock_reader("chinese", InvertedIndexReaderType::STRING_TYPE); + auto fulltext_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 1); + auto string_reader = create_mock_reader("chinese", InvertedIndexReaderType::STRING_TYPE, 2); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, fulltext_reader); iterator.add_reader(InvertedIndexReaderType::STRING_TYPE, string_reader); @@ -279,13 +311,43 @@ TEST_F(InvertedIndexIteratorTest, SelectBestReader_EqualQuerySelectsStringType) EXPECT_EQ(result.value(), string_reader); } +TEST_F(InvertedIndexIteratorTest, SelectBestReader_ArrayUsesLeafStringType) { + InvertedIndexIterator iterator; + auto string_reader = create_mock_reader("chinese", InvertedIndexReaderType::STRING_TYPE, 10); + auto fulltext_reader = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 20); + + iterator.add_reader(InvertedIndexReaderType::STRING_TYPE, string_reader); + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, fulltext_reader); + + DataTypePtr column_type = std::make_shared(make_nullable( + std::make_shared(make_nullable(std::make_shared())))); + auto result = iterator.select_best_reader(column_type, InvertedIndexQueryType::MATCH_ANY_QUERY, + "chinese"); + ASSERT_TRUE(result.has_value()) << result.error(); + EXPECT_EQ(result.value(), fulltext_reader); +} + +TEST_F(InvertedIndexIteratorTest, SelectAnyReaderIsDeterministicByIndexId) { + InvertedIndexIterator iterator; + auto reader_id_100 = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 100); + auto reader_id_50 = create_mock_reader("english", InvertedIndexReaderType::FULLTEXT, 50); + + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, reader_id_100); + iterator.add_reader(InvertedIndexReaderType::FULLTEXT, reader_id_50); + + auto result = iterator.select_any_reader(); + ASSERT_TRUE(result.has_value()) << result.error(); + EXPECT_EQ(result.value(), reader_id_50); +} + // Index lookup performance test TEST_F(InvertedIndexIteratorTest, IndexLookup_ManyReadersStillFast) { InvertedIndexIterator iterator; std::vector> readers; for (int i = 0; i < 100; i++) { - auto reader = create_mock_reader("analyzer_" + std::to_string(i)); + auto reader = create_mock_reader("analyzer_" + std::to_string(i), + InvertedIndexReaderType::FULLTEXT, i + 1); readers.push_back(reader); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, reader); } @@ -326,8 +388,8 @@ TEST_F(InvertedIndexIteratorTest, EdgeCase_CaseInsensitiveQuery) { TEST_F(InvertedIndexIteratorTest, EdgeCase_GetReaderByType) { InvertedIndexIterator iterator; - auto fulltext = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT); - auto string_type = create_mock_reader("english", InvertedIndexReaderType::STRING_TYPE); + auto fulltext = create_mock_reader("chinese", InvertedIndexReaderType::FULLTEXT, 1); + auto string_type = create_mock_reader("english", InvertedIndexReaderType::STRING_TYPE, 2); iterator.add_reader(InvertedIndexReaderType::FULLTEXT, fulltext); iterator.add_reader(InvertedIndexReaderType::STRING_TYPE, string_type); diff --git a/be/test/storage/segment/inverted_index_reader_test.cpp b/be/test/storage/segment/inverted_index_reader_test.cpp index 9004e7f1476984..ad20a3d438b17b 100644 --- a/be/test/storage/segment/inverted_index_reader_test.cpp +++ b/be/test/storage/segment/inverted_index_reader_test.cpp @@ -34,6 +34,7 @@ #include "runtime/runtime_state.h" #include "storage/index/index_file_reader.h" #include "storage/index/index_file_writer.h" +#include "storage/index/inverted/analyzer/analyzer.h" #include "storage/index/inverted/inverted_index_desc.h" #include "storage/index/inverted/inverted_index_iterator.h" #include "storage/index/inverted/inverted_index_writer.h" @@ -70,12 +71,19 @@ class InvertedIndexReaderTest : public testing::Test { _inverted_index_query_cache = std::unique_ptr( InvertedIndexQueryCache::create_global_cache(inverted_index_cache_limit, 1)); + // Both caches are owned by this fixture, so the previous globals must come back in + // TearDown -- otherwise ExecEnv keeps pointing at them after the fixture is destroyed and + // the next test that reaches InvertedIndexQueryCache::instance() reads freed memory. + _previous_searcher_cache = ExecEnv::GetInstance()->get_inverted_index_searcher_cache(); + _previous_query_cache = ExecEnv::GetInstance()->get_inverted_index_query_cache(); ExecEnv::GetInstance()->set_inverted_index_searcher_cache( _inverted_index_searcher_cache.get()); - ExecEnv::GetInstance()->_inverted_index_query_cache = _inverted_index_query_cache.get(); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_inverted_index_query_cache.get()); } void TearDown() override { + ExecEnv::GetInstance()->set_inverted_index_searcher_cache(_previous_searcher_cache); + ExecEnv::GetInstance()->set_inverted_index_query_cache(_previous_query_cache); ASSERT_TRUE(io::global_local_filesystem()->delete_directory(kTestDir).ok()); } @@ -703,11 +711,14 @@ class InvertedIndexReaderTest : public testing::Test { EXPECT_EQ(0, stats.inverted_index_searcher_cache_miss); } - // Query cache disabled (miss) while searcher cache should still hit. + // Query cache disabled while searcher cache should still hit. { OlapReaderStatistics stats; run_match(false, true, kImages, &stats); - EXPECT_EQ(1, stats.inverted_index_query_cache_miss); + EXPECT_EQ(0, stats.inverted_index_query_cache_lookup); + EXPECT_EQ(0, stats.inverted_index_query_cache_hit); + EXPECT_EQ(0, stats.inverted_index_query_cache_miss); + EXPECT_EQ(0, stats.inverted_index_query_cache_insert); EXPECT_EQ(1, stats.inverted_index_searcher_cache_hit); EXPECT_EQ(0, stats.inverted_index_searcher_cache_miss); } @@ -722,12 +733,14 @@ class InvertedIndexReaderTest : public testing::Test { EXPECT_EQ(1, stats.inverted_index_searcher_cache_miss); } - // Both caches disabled should report misses. + // Both caches disabled should not touch the query cache. { OlapReaderStatistics stats; run_match(false, false, kUnique, &stats); - EXPECT_EQ(1, stats.inverted_index_query_cache_miss); + EXPECT_EQ(0, stats.inverted_index_query_cache_lookup); EXPECT_EQ(0, stats.inverted_index_query_cache_hit); + EXPECT_EQ(0, stats.inverted_index_query_cache_miss); + EXPECT_EQ(0, stats.inverted_index_query_cache_insert); EXPECT_EQ(0, stats.inverted_index_searcher_cache_hit); EXPECT_EQ(1, stats.inverted_index_searcher_cache_miss); } @@ -874,6 +887,7 @@ class InvertedIndexReaderTest : public testing::Test { OlapReaderStatistics stats; RuntimeState runtime_state; TQueryOptions query_options; + query_options.enable_inverted_index_query_cache = false; query_options.enable_inverted_index_searcher_cache = false; runtime_state.set_query_options(query_options); @@ -909,6 +923,10 @@ class InvertedIndexReaderTest : public testing::Test { EXPECT_TRUE(query_status.ok()) << query_status; EXPECT_EQ(bitmap->cardinality(), 600) << "V3: Should find 600 documents matching 'common_term'"; + EXPECT_EQ(0, stats.inverted_index_query_cache_lookup); + EXPECT_EQ(0, stats.inverted_index_query_cache_hit); + EXPECT_EQ(0, stats.inverted_index_query_cache_miss); + EXPECT_EQ(0, stats.inverted_index_query_cache_insert); // Verify first and last document IDs EXPECT_TRUE(bitmap->contains(0)) << "V3: First document should match 'common_term'"; @@ -943,6 +961,51 @@ class InvertedIndexReaderTest : public testing::Test { InvertedIndexQueryType::MATCH_ANY_QUERY, bitmap); EXPECT_TRUE(query_status.ok()) << query_status; EXPECT_EQ(bitmap->cardinality(), 0) << "V3: Should find 0 documents matching 'noexist'"; + + InvertedIndexAnalyzerConfig analyzer_config; + analyzer_config.parser_type = InvertedIndexParserType::PARSER_ENGLISH; + analyzer_config.lower_case = INVERTED_INDEX_PARSER_TRUE; + auto analyzer_provider = + inverted_index::InvertedIndexAnalyzer::create_analyzer_provider( + &analyzer_config); + InvertedIndexAnalyzerCtx analyzer_ctx; + analyzer_ctx.parser_type = analyzer_config.parser_type; + analyzer_ctx.analyzer_provider = std::move(analyzer_provider); + + RuntimeState cached_runtime_state; + TQueryOptions cached_query_options; + cached_query_options.enable_inverted_index_query_cache = true; + cached_query_options.enable_inverted_index_searcher_cache = true; + cached_runtime_state.set_query_options(cached_query_options); + OlapReaderStatistics cached_stats; + auto cached_context = std::make_shared(); + cached_context->io_ctx = &io_ctx; + cached_context->stats = &cached_stats; + cached_context->runtime_state = &cached_runtime_state; + const Field cached_query = Field::create_field("common_term"); + + std::shared_ptr first_bitmap; + query_status = str_reader->query(cached_context, field_name, cached_query, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + first_bitmap, &analyzer_ctx); + EXPECT_TRUE(query_status.ok()) << query_status; + ASSERT_NE(first_bitmap, nullptr); + EXPECT_EQ(first_bitmap->cardinality(), 600); + EXPECT_EQ(cached_stats.inverted_index_query_cache_lookup, 1); + EXPECT_EQ(cached_stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(cached_stats.inverted_index_query_cache_insert, 1); + + std::shared_ptr second_bitmap; + query_status = str_reader->query(cached_context, field_name, cached_query, + InvertedIndexQueryType::MATCH_PHRASE_QUERY, + second_bitmap, &analyzer_ctx); + EXPECT_TRUE(query_status.ok()) << query_status; + ASSERT_NE(second_bitmap, nullptr); + EXPECT_EQ(*second_bitmap, *first_bitmap); + EXPECT_EQ(cached_stats.inverted_index_query_cache_lookup, 2); + EXPECT_EQ(cached_stats.inverted_index_query_cache_hit, 1); + EXPECT_EQ(cached_stats.inverted_index_query_cache_miss, 1); + EXPECT_EQ(cached_stats.inverted_index_query_cache_insert, 1); } { TabletIndex idx_meta; @@ -4048,6 +4111,8 @@ class InvertedIndexReaderTest : public testing::Test { } private: + InvertedIndexSearcherCache* _previous_searcher_cache = nullptr; + InvertedIndexQueryCache* _previous_query_cache = nullptr; std::unique_ptr _inverted_index_searcher_cache; std::unique_ptr _inverted_index_query_cache; }; diff --git a/be/test/storage/segment/inverted_index_writer_test.cpp b/be/test/storage/segment/inverted_index_writer_test.cpp index 8b7f56221cf8cf..e4e5d5f3aadfc8 100644 --- a/be/test/storage/segment/inverted_index_writer_test.cpp +++ b/be/test/storage/segment/inverted_index_writer_test.cpp @@ -28,23 +28,35 @@ #include #include #include +#include #include #include +#include "common/config.h" #include "core/block/block.h" #include "core/data_type/data_type_factory.hpp" #include "core/data_type/data_type_number.h" #include "core/field.h" #include "io/fs/local_file_system.h" +#include "runtime/index_policy/index_policy_mgr.h" #include "runtime/runtime_state.h" #include "storage/index/index_file_reader.h" #include "storage/index/index_file_writer.h" +#include "storage/index/inverted/analyzer/analyzer.h" +#include "storage/index/inverted/analyzer/custom_analyzer.h" +#include "storage/index/inverted/common_grams/common_grams_key_codec.h" +#include "storage/index/inverted/common_grams/common_word_set.h" #include "storage/index/inverted/inverted_index_desc.h" #include "storage/index/inverted/inverted_index_fs_directory.h" #include "storage/index/inverted/inverted_index_reader.h" +#include "storage/index/snii/query/bm25_scorer.h" +#include "storage/index/snii/query/term_query.h" +#include "storage/index/snii/snii_index_writer.h" +#include "storage/index/snii/stats/snii_stats_provider.h" #include "storage/iterator/olap_data_convertor.h" #include "storage/tablet/tablet_schema.h" #include "storage/types.h" +#include "util/defer_op.h" #include "util/faststring.h" #include "util/slice.h" @@ -57,6 +69,101 @@ namespace doris::segment_v2 { using InvertedIndexDirectoryMap = std::map, std::shared_ptr>; +class CommonGramsBuildFlagRestorer { +public: + CommonGramsBuildFlagRestorer() : _original(config::enable_common_grams_index_build) {} + ~CommonGramsBuildFlagRestorer() { config::enable_common_grams_index_build = _original; } + +private: + const bool _original; +}; + +class GappedTokenStream final : public lucene::analysis::TokenStream { +public: + lucene::analysis::Token* next(lucene::analysis::Token* token) override { + if (_emitted) { + return nullptr; + } + _emitted = true; + token->clear(); + token->setTextNoCopy(_term.data(), static_cast(_term.size())); + token->setPositionIncrement(2); + return token; + } + + void close() override {} + void reset() override { _emitted = false; } + +private: + std::string _term = "gapped"; + bool _emitted = false; +}; + +class GappedTokenAnalyzer final : public lucene::analysis::Analyzer { +public: + bool isSDocOpt() override { return true; } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, lucene::util::Reader*) override { + return new GappedTokenStream(); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + lucene::util::Reader*) override { + _reusable = std::make_unique(); + return _reusable.get(); + } + + lucene::analysis::TokenStream* tokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + return new GappedTokenStream(); + } + + lucene::analysis::TokenStream* reusableTokenStream(const TCHAR*, + const inverted_index::ReaderPtr&) override { + _reusable = std::make_unique(); + return _reusable.get(); + } + +private: + std::unique_ptr _reusable; +}; + +void install_common_grams_policy(IndexPolicyMgr* policy_mgr, int64_t policy_id, + std::string_view name, int64_t index_id, TabletIndex* index_meta) { + TIndexPolicy tokenizer; + tokenizer.id = policy_id; + tokenizer.name = fmt::format("{}_tokenizer", name); + tokenizer.type = TIndexPolicyType::TOKENIZER; + tokenizer.properties["type"] = "char_group"; + tokenizer.properties["tokenize_on_chars"] = "[\\u0020]"; + tokenizer.properties["max_token_length"] = + std::to_string(inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES); + + TIndexPolicy common_grams; + common_grams.id = policy_id + 1; + common_grams.name = fmt::format("{}_filter", name); + common_grams.type = TIndexPolicyType::TOKEN_FILTER; + common_grams.properties["type"] = "common_grams"; + // Preparation authority: a CommonGrams policy is only prepared when FE + + TIndexPolicy analyzer; + analyzer.id = policy_id + 2; + analyzer.name = fmt::format("{}_analyzer", name); + analyzer.type = TIndexPolicyType::ANALYZER; + analyzer.properties["tokenizer"] = tokenizer.name; + analyzer.properties["token_filter"] = "lowercase," + common_grams.name; + policy_mgr->apply_policy_changes({tokenizer, common_grams, analyzer}, {}); + + TabletIndexPB index_pb; + index_pb.set_index_type(IndexType::INVERTED); + index_pb.set_index_id(index_id); + index_pb.set_index_name(fmt::format("{}_index", name)); + index_pb.add_col_unique_id(1); + index_pb.mutable_properties()->insert({"analyzer", analyzer.name}); + index_pb.mutable_properties()->insert({"support_phrase", "true"}); + index_meta->init_from_pb(index_pb); +} + class InvertedIndexWriterTest : public testing::Test { using ExpectedDocMap = std::map>; @@ -1478,6 +1585,269 @@ TEST_F(InvertedIndexWriterTest, FileCreationAndOutputErrorHandling) { // but it should not crash } +TEST_F(InvertedIndexWriterTest, CommonGramsBuildSwitchAppliesOnlyToSnii) { + CommonGramsBuildFlagRestorer flag_restorer; + auto* exec_env = ExecEnv::GetInstance(); + auto* previous_policy_mgr = exec_env->index_policy_mgr(); + IndexPolicyMgr policy_mgr; + exec_env->_index_policy_mgr = &policy_mgr; + DEFER(exec_env->_index_policy_mgr = previous_policy_mgr); + + TabletIndex index_meta; + install_common_grams_policy(&policy_mgr, 920040, "common_grams_enabled_snapshot", 98, + &index_meta); + + const std::string v3_rowset_id = "v3_common_grams_enabled_snapshot"; + const std::string v3_prefix {InvertedIndexDescriptor::get_index_file_path_prefix( + local_segment_path(kTestDir, v3_rowset_id, 0))}; + const std::string snii_rowset_id = "snii_common_grams_enabled_snapshot"; + const std::string snii_prefix {InvertedIndexDescriptor::get_index_file_path_prefix( + local_segment_path(kTestDir, snii_rowset_id, 0))}; + io::FileWriterOptions opts; + auto fs = io::global_local_filesystem(); + io::FileWriterPtr v3_compound_file; + ASSERT_TRUE(fs->create_file(InvertedIndexDescriptor::get_index_file_path_v2(v3_prefix), + &v3_compound_file, &opts) + .ok()); + io::FileWriterPtr snii_compound_file; + ASSERT_TRUE(fs->create_file(InvertedIndexDescriptor::get_index_file_path_v2(snii_prefix), + &snii_compound_file, &opts) + .ok()); + IndexFileWriter v3_file_writer(fs, v3_prefix, v3_rowset_id, 0, InvertedIndexStorageFormatPB::V3, + std::move(v3_compound_file)); + IndexFileWriter snii_file_writer(fs, snii_prefix, snii_rowset_id, 0, + InvertedIndexStorageFormatPB::SNII, + std::move(snii_compound_file)); + + config::enable_common_grams_index_build = true; + InvertedIndexColumnWriter v3_writer( + "1", &v3_file_writer, &index_meta, /*single_field=*/true); + SniiIndexColumnWriter snii_writer(&snii_file_writer, &index_meta, /*single_field=*/true); + config::enable_common_grams_index_build = false; + + ASSERT_TRUE(v3_writer.init().ok()); + ASSERT_TRUE(snii_writer.init().ok()); + std::string marker_leading_value(inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + marker_leading_value.front() = '\x1f'; + const Slice marker_value(marker_leading_value); + const Slice phrase_value("man of the year"); + ASSERT_TRUE(v3_writer.add_values("", &marker_value, 1).ok()); + ASSERT_TRUE(v3_writer.add_values("", &phrase_value, 1).ok()); + ASSERT_TRUE(v3_writer.finish().ok()); + ASSERT_TRUE(v3_file_writer.begin_close().ok()); + ASSERT_TRUE(v3_file_writer.finish_close().ok()); + + IndexFileReader v3_file_reader(fs, v3_prefix, InvertedIndexStorageFormatPB::V3); + ASSERT_TRUE(v3_file_reader.init().ok()); + auto directory_result = v3_file_reader.open(&index_meta); + ASSERT_TRUE(directory_result.has_value()) << directory_result.error(); + auto directory = std::move(directory_result.value()); + EXPECT_FALSE(directory->fileExists("doris_common_grams.meta")); + + const std::string encoded_gram = inverted_index::encode_common_gram("man", "of").value(); + std::set indexed_terms; + std::unique_ptr reader( + lucene::index::IndexReader::open(directory.get()), [](lucene::index::IndexReader* ptr) { + ptr->close(); + _CLLDELETE(ptr); + }); + std::unique_ptr terms(reader->terms(), [](TermEnum* ptr) { + ptr->close(); + _CLLDELETE(ptr); + }); + while (terms->next()) { + indexed_terms.emplace(lucene_wcstoutf8string(terms->term(false)->text(), + terms->term(false)->textLength())); + } + EXPECT_TRUE(indexed_terms.contains(marker_leading_value)); + EXPECT_TRUE(indexed_terms.contains("man")); + EXPECT_TRUE(indexed_terms.contains("of")); + EXPECT_TRUE(indexed_terms.contains("the")); + EXPECT_TRUE(indexed_terms.contains("year")); + EXPECT_FALSE(indexed_terms.contains(encoded_gram)); + + const Status snii_add_status = snii_writer.add_values("", &marker_value, 1); + EXPECT_EQ(snii_add_status.code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR) << snii_add_status; + EXPECT_NE(snii_add_status.to_string().find("enable_common_grams_index_build=false"), + std::string::npos); + EXPECT_EQ(snii_writer.finish().code(), ErrorCode::INVERTED_INDEX_ANALYZER_ERROR); +} + +TEST_F(InvertedIndexWriterTest, CommonGramsDisabledBuildSwitchSnapshotWritesPlainSnii) { + CommonGramsBuildFlagRestorer flag_restorer; + auto* exec_env = ExecEnv::GetInstance(); + auto* previous_policy_mgr = exec_env->index_policy_mgr(); + IndexPolicyMgr policy_mgr; + exec_env->_index_policy_mgr = &policy_mgr; + DEFER(exec_env->_index_policy_mgr = previous_policy_mgr); + + TabletIndex index_meta; + install_common_grams_policy(&policy_mgr, 920050, "common_grams_disabled_snapshot", 99, + &index_meta); + + const std::string rowset_id = "snii_common_grams_disabled_snapshot"; + const std::string index_path_prefix {InvertedIndexDescriptor::get_index_file_path_prefix( + local_segment_path(kTestDir, rowset_id, 0))}; + io::FileWriterOptions opts; + auto fs = io::global_local_filesystem(); + io::FileWriterPtr compound_file; + ASSERT_TRUE(fs->create_file(InvertedIndexDescriptor::get_index_file_path_v2(index_path_prefix), + &compound_file, &opts) + .ok()); + IndexFileWriter file_writer(fs, index_path_prefix, rowset_id, 0, + InvertedIndexStorageFormatPB::SNII, std::move(compound_file)); + + const auto stale_seed = inverted_index::make_common_grams_segment_metadata( + {.common_grams_dictionary_identity = "stale-dictionary", + .base_analyzer_fingerprint = "stale-base-analyzer", + .common_grams_fingerprint = "stale-common-grams"}); + + config::enable_common_grams_index_build = false; + SniiIndexColumnWriter writer(&file_writer, &index_meta, /*single_field=*/true, stale_seed); + config::enable_common_grams_index_build = true; + + ASSERT_TRUE(writer.init().ok()); + std::string marker_leading_value(inverted_index::COMMON_GRAM_MAX_ENCODED_BYTES, 'x'); + marker_leading_value.front() = '\x1f'; + const Slice marker_value(marker_leading_value); + const Slice phrase_value("man of"); + ASSERT_TRUE(writer.add_values("", &marker_value, 1).ok()); + ASSERT_TRUE(writer.add_values("", &phrase_value, 1).ok()); + writer.set_analysis_for_test(inverted_index::InvertedIndexAnalyzer::create_reader({}), + std::make_shared()); + const Slice gapped_value("ignored"); + ASSERT_TRUE(writer.add_values("", &gapped_value, 1).ok()); + ASSERT_TRUE(writer.finish().ok()); + ASSERT_TRUE(file_writer.begin_close().ok()); + ASSERT_TRUE(file_writer.finish_close().ok()); + + const std::string encoded_gram = inverted_index::encode_common_gram("man", "of").value(); + IndexFileReader file_reader(fs, index_path_prefix, InvertedIndexStorageFormatPB::SNII); + ASSERT_TRUE(file_reader.init().ok()); + auto logical_result = file_reader.open_snii_index(&index_meta); + ASSERT_TRUE(logical_result.has_value()) << logical_result.error(); + auto logical = std::move(logical_result.value()); + EXPECT_EQ(logical->common_grams_metadata(), nullptr); + EXPECT_EQ(logical->tier(), snii::format::tier_of(snii::format::IndexConfig::kDocsPositions)); + std::vector docids; + ASSERT_TRUE(snii::query::term_query(*logical, marker_leading_value, &docids).ok()); + EXPECT_EQ(docids, (std::vector {0})); + ASSERT_TRUE(snii::query::term_query(*logical, encoded_gram, &docids).ok()); + EXPECT_TRUE(docids.empty()); + ASSERT_TRUE(snii::query::term_query(*logical, "gapped", &docids).ok()); + EXPECT_EQ(docids, (std::vector {2})); +} + +TEST_F(InvertedIndexWriterTest, CommonGramsSniiPersistsTypedTermsAndSemanticScoringInputs) { + auto* exec_env = ExecEnv::GetInstance(); + auto* previous_policy_mgr = exec_env->index_policy_mgr(); + IndexPolicyMgr policy_mgr; + exec_env->_index_policy_mgr = &policy_mgr; + DEFER(exec_env->_index_policy_mgr = previous_policy_mgr); + + TIndexPolicy tokenizer; + tokenizer.id = 920030; + tokenizer.name = "snii_cg_writer_tokenizer"; + tokenizer.type = TIndexPolicyType::TOKENIZER; + tokenizer.properties["type"] = "char_group"; + tokenizer.properties["tokenize_on_chars"] = "[\\u0020]"; + TIndexPolicy common_grams; + common_grams.id = 920031; + common_grams.name = "snii_cg_writer_filter"; + common_grams.type = TIndexPolicyType::TOKEN_FILTER; + common_grams.properties["type"] = "common_grams"; + TIndexPolicy analyzer; + analyzer.id = 920032; + analyzer.name = "snii_cg_writer_analyzer"; + analyzer.type = TIndexPolicyType::ANALYZER; + analyzer.properties["tokenizer"] = tokenizer.name; + analyzer.properties["token_filter"] = "lowercase," + common_grams.name; + policy_mgr.apply_policy_changes({tokenizer, common_grams, analyzer}, {}); + + TabletIndexPB index_pb; + index_pb.set_index_type(IndexType::INVERTED); + index_pb.set_index_id(95); + index_pb.set_index_name("common_grams_snii_typed_writer"); + index_pb.add_col_unique_id(1); + index_pb.mutable_properties()->insert({"analyzer", analyzer.name}); + index_pb.mutable_properties()->insert({"support_phrase", "true"}); + TabletIndex index_meta; + index_meta.init_from_pb(index_pb); + + const std::string rowset_id = "snii_common_grams_typed_writer"; + const std::string index_path_prefix {InvertedIndexDescriptor::get_index_file_path_prefix( + local_segment_path(kTestDir, rowset_id, 0))}; + const std::string index_path = + InvertedIndexDescriptor::get_index_file_path_v2(index_path_prefix); + io::FileWriterPtr compound_file; + io::FileWriterOptions opts; + auto fs = io::global_local_filesystem(); + ASSERT_TRUE(fs->create_file(index_path, &compound_file, &opts).ok()); + IndexFileWriter file_writer(fs, index_path_prefix, rowset_id, 0, + InvertedIndexStorageFormatPB::SNII, std::move(compound_file)); + + SniiIndexColumnWriter writer(&file_writer, &index_meta, /*single_field=*/true); + ASSERT_TRUE(writer.init().ok()); + const std::string marker_plain = std::string(1, '\x1f') + "literal"; + const std::string marker_phrase = marker_plain + " of"; + const std::vector values {Slice("man of the year"), Slice(marker_phrase)}; + ASSERT_TRUE(writer.add_values("", values.data(), values.size()).ok()); + ASSERT_TRUE(writer.add_nulls(1).ok()); + ASSERT_TRUE(writer.finish().ok()); + ASSERT_TRUE(file_writer.begin_close().ok()); + ASSERT_TRUE(file_writer.finish_close().ok()); + + IndexFileReader file_reader(fs, index_path_prefix, InvertedIndexStorageFormatPB::SNII); + ASSERT_TRUE(file_reader.init().ok()); + auto logical_result = file_reader.open_snii_index(&index_meta); + ASSERT_TRUE(logical_result.has_value()) << logical_result.error(); + auto logical = std::move(logical_result.value()); + ASSERT_NE(logical->common_grams_metadata(), nullptr); + EXPECT_EQ(logical->common_grams_metadata()->common_grams_dictionary_identity, + inverted_index::CommonWordSet::default_word_set()->identity()); + EXPECT_EQ(logical->common_grams_metadata()->base_analyzer_fingerprint.size(), 64U); + EXPECT_EQ(logical->common_grams_metadata()->common_grams_fingerprint.size(), 64U); + EXPECT_EQ(logical->common_grams_metadata()->scoring_doc_count, 3); + EXPECT_EQ(logical->common_grams_metadata()->scoring_token_count, 6); + + snii::stats::SniiStatsProvider stats; + ASSERT_TRUE(snii::stats::SniiStatsProvider::open(logical.get(), &stats).ok()); + EXPECT_EQ(stats.doc_count(), 3); + EXPECT_EQ(stats.sum_total_term_freq(), 6); + uint8_t encoded_norm = 0; + ASSERT_TRUE(stats.encoded_norm(0, &encoded_norm).ok()); + EXPECT_EQ(encoded_norm, snii::query::encode_norm(4)); + ASSERT_TRUE(stats.encoded_norm(1, &encoded_norm).ok()); + EXPECT_EQ(encoded_norm, snii::query::encode_norm(2)); + ASSERT_TRUE(stats.encoded_norm(2, &encoded_norm).ok()); + EXPECT_EQ(encoded_norm, snii::query::encode_norm(0)); + + const std::string expected_gram = inverted_index::encode_common_gram("man", "of").value(); + const std::string expected_plain = + inverted_index::encode_plain_term(marker_plain, + inverted_index::PlainTermKeyVersion::kEscapedV1) + .value(); + const std::string expected_marker_gram = + inverted_index::encode_common_gram(marker_plain, "of").value(); + const std::string unexpected_double_plain = + inverted_index::encode_plain_term(expected_plain, + inverted_index::PlainTermKeyVersion::kEscapedV1) + .value(); + const std::string unexpected_physical_gram = + inverted_index::encode_common_gram(expected_plain, "of").value(); + std::vector docids; + ASSERT_TRUE(snii::query::term_query(*logical, expected_gram, &docids).ok()); + EXPECT_EQ(docids, (std::vector {0})); + ASSERT_TRUE(snii::query::term_query(*logical, expected_plain, &docids).ok()); + EXPECT_EQ(docids, (std::vector {1})); + ASSERT_TRUE(snii::query::term_query(*logical, expected_marker_gram, &docids).ok()); + EXPECT_EQ(docids, (std::vector {1})); + ASSERT_TRUE(snii::query::term_query(*logical, unexpected_double_plain, &docids).ok()); + EXPECT_TRUE(docids.empty()); + ASSERT_TRUE(snii::query::term_query(*logical, unexpected_physical_gram, &docids).ok()); + EXPECT_TRUE(docids.empty()); +} + // Test case to verify .nrm file creation behavior with different tokenization settings // This test verifies the change in inverted_index_writer.cpp lines 165-171 // where .nrm file is only created when field requires tokenization (_should_analyzer == true) diff --git a/be/test/storage/segment/segment_iterator_count_emit_shortcut_test.cpp b/be/test/storage/segment/segment_iterator_count_emit_shortcut_test.cpp new file mode 100644 index 00000000000000..d141637c48b6d0 --- /dev/null +++ b/be/test/storage/segment/segment_iterator_count_emit_shortcut_test.cpp @@ -0,0 +1,419 @@ +// 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. + +// G03 count-emission shortcut white-box tests. Given a post-apply _row_bitmap +// of cardinality N (the exact count G02 fabricated), the shortcut must emit +// EXACTLY N rows shaped as NOT-NULL defaults across VStatisticsIterator-sized +// batches, then EOF -- byte-for-byte what today's per-rowid batch loop emits +// for a count-fastpath scan, minus the per-rowid work. The engage decision +// must admit only the provably emission-only configuration and refuse on any +// deviation (falling through to today's loop). Uses the established +// `#define private public` convention of segment_iterator_limit_opt_test.cpp +// over a bare SegmentIterator (no real segment needed: the shortcut never +// touches segment data). +#include + +#include +#include +#include +#include + +#include "common/status.h" +#include "core/assert_cast.h" +#include "core/block/block.h" +#include "core/block/column_with_type_and_name.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "storage/index/index_query_context.h" +#include "storage/olap_common.h" +#include "storage/segment/count_on_index_fastpath.h" +#include "storage/tablet/tablet_schema.h" + +#if defined(__clang__) +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wkeyword-macro" +#endif +#define private public +#define protected public +#include "storage/segment/segment_iterator.h" +#undef private +#undef protected +#if defined(__clang__) +#pragma clang diagnostic pop +#endif + +namespace doris::segment_v2 { + +namespace { + +constexpr uint64_t kBatchRows = SegmentIterator::kCountEmitBatchRows; // 65535 + +void reset_emit_counters() { + internal::count_emit_test_counters() = internal::CountEmitTestCounters {}; +} + +uint64_t emit_hits() { + return internal::count_emit_test_counters().count_emit_shortcut_hits; +} + +uint64_t emit_batches() { + return internal::count_emit_test_counters().count_emit_shortcut_batches; +} + +// DUP-key schema shaped like a COUNT_ON_INDEX scan: an INT key plus the +// (nullable STRING) match column. Both columns end up defaults-filled by +// today's path once the index consumed the MATCH predicate. +TabletSchemaSPtr make_tablet_schema() { + TabletSchemaPB schema_pb; + schema_pb.set_keys_type(KeysType::DUP_KEYS); + auto* key_col = schema_pb.add_column(); + key_col->set_unique_id(0); + key_col->set_name("k1"); + key_col->set_type("INT"); + key_col->set_is_key(true); + key_col->set_is_nullable(true); + auto* match_col = schema_pb.add_column(); + match_col->set_unique_id(1); + match_col->set_name("content"); + match_col->set_type("STRING"); + match_col->set_is_key(false); + match_col->set_is_nullable(true); + auto tablet_schema = std::make_shared(); + tablet_schema->init_from_pb(schema_pb); + return tablet_schema; +} + +SchemaSPtr make_read_schema(const TabletSchemaSPtr& tablet_schema) { + std::vector read_column_ids(tablet_schema->num_columns()); + for (uint32_t cid = 0; cid < read_column_ids.size(); ++cid) { + read_column_ids[cid] = cid; + } + return std::make_shared(tablet_schema->columns(), read_column_ids); +} + +Block make_block_for(const Schema& schema) { + Block block; + for (size_t i = 0; i < schema.num_column_ids(); ++i) { + const auto* col_desc = schema.column(schema.column_id(i)); + auto data_type = Schema::get_data_type_ptr(*col_desc); + block.insert( + ColumnWithTypeAndName(data_type->create_column(), data_type, col_desc->name())); + } + return block; +} + +struct Fixture { + TabletSchemaSPtr tablet_schema; + SchemaSPtr read_schema; + std::unique_ptr iter; + OlapReaderStatistics stats; + + Fixture() { + tablet_schema = make_tablet_schema(); + read_schema = make_read_schema(tablet_schema); + iter = std::make_unique(nullptr, read_schema); + iter->_opts.tablet_schema = tablet_schema; + iter->_opts.push_down_agg_type_opt = TPushAggOp::COUNT_ON_INDEX; + iter->_opts.stats = &stats; + // State _lazy_init/_vec_init_lazy_materialization would have produced + // for a count-fastpath scan: no predicate columns, no lazy + // materialization, index fully answered every column's conditions. + iter->_is_pred_column.resize(read_schema->columns().size(), false); + iter->_lazy_materialization_read = false; + iter->_storage_name_and_type.resize(read_schema->columns().size()); + for (size_t i = 0; i < read_schema->num_column_ids(); ++i) { + ColumnId cid = read_schema->column_id(i); + iter->_storage_name_and_type[cid] = + std::make_pair(read_schema->column(cid)->name(), + Schema::get_data_type_ptr(*read_schema->column(cid))); + iter->_need_read_data_indices[cid] = false; + } + } + + // Simulates the reader having fabricated a count bitmap (G02 hit). + void set_hit() { iter->_count_fastpath_hit = true; } + + // Simulates the _lazy_init engage step for a fabricated bitmap of + // cardinality `count`. + void engage_with_count(uint64_t count) { + iter->_count_emit_shortcut = true; + iter->_count_emit_rows_remaining = count; + } + + // Drives _emit_count_shortcut_batch until EOF; returns per-batch rows. + std::vector drain(Block* block) { + std::vector batch_rows; + while (true) { + Status st = iter->_emit_count_shortcut_batch(block); + if (st.is()) { + EXPECT_EQ(block->rows(), 0U) << "EOF must deliver an empty block"; + break; + } + EXPECT_TRUE(st.ok()) << st.to_string(); + batch_rows.push_back(block->rows()); + } + return batch_rows; + } +}; + +uint64_t total_rows(const std::vector& batches) { + uint64_t total = 0; + for (size_t rows : batches) { + total += rows; + } + return total; +} + +} // namespace + +// The emitted total must EQUAL the fabricated count for counts spanning +// several batches plus a tail: this is the count-equality contract with +// today's per-rowid loop (which emits exactly |_row_bitmap| rows). +TEST(CountEmitShortcut, EmitsExactCountAcrossBatchesThenEof) { + reset_emit_counters(); + Fixture fx; + const uint64_t count = 3 * kBatchRows + 3395; // 200000 + fx.engage_with_count(count); + + Block block = make_block_for(*fx.read_schema); + const auto batches = fx.drain(&block); + + ASSERT_EQ(batches.size(), 4U); + EXPECT_EQ(batches[0], kBatchRows); + EXPECT_EQ(batches[1], kBatchRows); + EXPECT_EQ(batches[2], kBatchRows); + EXPECT_EQ(batches[3], 3395U); + EXPECT_EQ(total_rows(batches), count); + EXPECT_EQ(fx.stats.raw_rows_read, static_cast(count)); + EXPECT_EQ(emit_batches(), 4U); + + // A second call after EOF stays EOF (the scanner may re-poll). + Status st = fx.iter->_emit_count_shortcut_batch(&block); + EXPECT_TRUE(st.is()); + EXPECT_EQ(block.rows(), 0U); +} + +// Batch-boundary counts: exact batch multiples and off-by-one neighbours all +// sum to the requested count with ceil(count / batch) batches. +TEST(CountEmitShortcut, BatchBoundaryCountsAreExact) { + for (const uint64_t count : {uint64_t(1), kBatchRows - 1, kBatchRows, kBatchRows + 1}) { + reset_emit_counters(); + Fixture fx; + fx.engage_with_count(count); + Block block = make_block_for(*fx.read_schema); + const auto batches = fx.drain(&block); + EXPECT_EQ(total_rows(batches), count) << "count: " << count; + const uint64_t expected_batches = (count + kBatchRows - 1) / kBatchRows; + EXPECT_EQ(batches.size(), expected_batches) << "count: " << count; + EXPECT_EQ(emit_batches(), expected_batches) << "count: " << count; + for (size_t rows : batches) { + EXPECT_LE(rows, kBatchRows); + } + } +} + +// Empty result (df == 0, e.g. the term missed the dict): immediate EOF with +// zero emitted rows -- identical to today's empty-bitmap first batch. +TEST(CountEmitShortcut, EmptyResultIsImmediateEof) { + reset_emit_counters(); + Fixture fx; + fx.engage_with_count(0); + Block block = make_block_for(*fx.read_schema); + Status st = fx.iter->_emit_count_shortcut_batch(&block); + EXPECT_TRUE(st.is()); + EXPECT_EQ(block.rows(), 0U); + EXPECT_EQ(fx.stats.raw_rows_read, 0); + EXPECT_EQ(emit_batches(), 0U); +} + +// Nullable columns must be filled with NOT-NULL defaults exactly like the +// _no_need_read_key_data/_prune_column split fill. A raw +// ColumnNullable::insert_many_defaults would insert NULLs, and count(col) +// (planned under COUNT_ON_INDEX for count(match_col)) counts non-null values: +// NULL-filled rows would collapse the count to zero. +TEST(CountEmitShortcut, NullableColumnsEmitNotNullDefaults) { + Fixture fx; + fx.engage_with_count(1000); + Block block = make_block_for(*fx.read_schema); + ASSERT_TRUE(fx.iter->_emit_count_shortcut_batch(&block).ok()); + ASSERT_EQ(block.rows(), 1000U); + + // INT key column: not-null zeros. + const auto* key_col = assert_cast(block.get_by_position(0).column.get()); + EXPECT_FALSE(key_col->has_null()); + const auto& key_data = + assert_cast&>(key_col->get_nested_column()).get_data(); + for (size_t j = 0; j < key_data.size(); ++j) { + ASSERT_EQ(key_data[j], 0) << "row " << j; + } + + // STRING match column: not-null empty strings with valid offsets. + const auto* match_col = + assert_cast(block.get_by_position(1).column.get()); + EXPECT_FALSE(match_col->has_null()); + const auto& match_data = assert_cast(match_col->get_nested_column()); + ASSERT_EQ(match_data.size(), 1000U); + for (size_t j = 0; j < match_data.size(); ++j) { + ASSERT_EQ(match_data.get_data_at(j).size, 0U) << "row " << j; + } +} + +// The shortcut consumes only the CARDINALITY of the post-apply bitmap; the id +// positions are irrelevant. Pins the null-bearing G02 shape (null-disjoint +// fabrication produces a SHIFTED range like [100, 150), not [0, 50)). +TEST(CountEmitShortcut, NullDisjointFabricatedShapeCountsByCardinality) { + Fixture fx; + fx.iter->_row_bitmap = roaring::Roaring {}; + fx.iter->_row_bitmap.addRange(100, 150); + // The exact engage math from _lazy_init. + fx.engage_with_count(fx.iter->_row_bitmap.cardinality()); + Block block = make_block_for(*fx.read_schema); + const auto batches = fx.drain(&block); + EXPECT_EQ(total_rows(batches), 50U); +} + +// Segment iterators keep independent countdowns: a multi-segment scan is the +// sum of per-segment counts, each emitted by its own iterator. +TEST(CountEmitShortcut, MultipleIteratorsEmitIndependently) { + reset_emit_counters(); + Fixture fx1; + Fixture fx2; + fx1.engage_with_count(kBatchRows + 5000); // 2 batches + fx2.engage_with_count(3); // 1 batch + Block block1 = make_block_for(*fx1.read_schema); + Block block2 = make_block_for(*fx2.read_schema); + + const auto batches1 = fx1.drain(&block1); + const auto batches2 = fx2.drain(&block2); + EXPECT_EQ(total_rows(batches1), kBatchRows + 5000); + EXPECT_EQ(total_rows(batches2), 3U); + EXPECT_EQ(batches1.size(), 2U); + EXPECT_EQ(batches2.size(), 1U); + EXPECT_EQ(emit_batches(), 3U); +} + +// G02->G03 handshake teardown: the reply flag is captured into the iterator +// and BOTH context flags are cleared so no later read_from_index call can +// observe or forge them. +TEST(CountEmitShortcut, CaptureHitRecordsReplyAndResetsHandshake) { + Fixture fx; + fx.iter->_index_query_context = std::make_shared(); + fx.iter->_index_query_context->count_on_index_fastpath = true; + fx.iter->_index_query_context->count_on_index_fastpath_hit = true; + + fx.iter->_capture_count_fastpath_hit(); + EXPECT_TRUE(fx.iter->_count_fastpath_hit); + EXPECT_FALSE(fx.iter->_index_query_context->count_on_index_fastpath); + EXPECT_FALSE(fx.iter->_index_query_context->count_on_index_fastpath_hit); + + // No reply (row-accurate decode / cache hit): captured false. + fx.iter->_index_query_context->count_on_index_fastpath = true; + fx.iter->_capture_count_fastpath_hit(); + EXPECT_FALSE(fx.iter->_count_fastpath_hit); + EXPECT_FALSE(fx.iter->_index_query_context->count_on_index_fastpath); + + // Null context (index machinery never initialized): no crash, no hit. + fx.iter->_index_query_context = nullptr; + fx.iter->_capture_count_fastpath_hit(); + EXPECT_FALSE(fx.iter->_count_fastpath_hit); +} + +// Engage admits the one emission-only configuration and counts a seam hit. +TEST(CountEmitShortcut, EngageAdmitsCleanCountFastpathState) { + reset_emit_counters(); + Fixture fx; + fx.set_hit(); + Block block = make_block_for(*fx.read_schema); + EXPECT_TRUE(fx.iter->_should_engage_count_emit_shortcut(&block)); + EXPECT_EQ(emit_hits(), 1U); +} + +// Engage refusals: every deviation keeps today's per-rowid loop. The seam +// stays untouched on refusals. +TEST(CountEmitShortcut, EngageRefusalTruthTable) { + reset_emit_counters(); + // Without the reader hit (row-accurate bitmap): refuse. + { + Fixture fx; + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // A surviving evaluation stage (e.g. index-eval downgrade kept the expr). + { + Fixture fx; + fx.set_hit(); + fx.iter->_is_need_expr_eval = true; + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // A pushed-down read limit must keep the limit-aware loop. + { + Fixture fx; + fx.set_hit(); + fx.iter->_opts.read_limit = 10; + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // Condition-cache writes are only produced by the real batch loop. + { + Fixture fx; + fx.set_hit(); + fx.iter->_opts.condition_cache_digest = 7; + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // Rowid consumers need real row ids. + { + Fixture fx; + fx.set_hit(); + fx.iter->_opts.record_rowids = true; + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // Block shape != read schema (e.g. delete-condition column beyond block). + { + Fixture fx; + fx.set_hit(); + Block block = make_block_for(*fx.read_schema); + block.erase(1); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // A column today's path would REALLY read (index did not answer its + // conditions): values would come from column data, not defaults. + { + Fixture fx; + fx.set_hit(); + fx.iter->_need_read_data_indices.erase(1); + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + // A storage->schema cast column: today's path emits CAST(file-type + // default), which the shortcut does not reproduce. Simulated by giving + // the INT key column the STRING column's storage type. + { + Fixture fx; + fx.set_hit(); + fx.iter->_storage_name_and_type[0].second = + Schema::get_data_type_ptr(*fx.read_schema->column(1)); + Block block = make_block_for(*fx.read_schema); + EXPECT_FALSE(fx.iter->_should_engage_count_emit_shortcut(&block)); + } + EXPECT_EQ(emit_hits(), 0U); +} + +} // namespace doris::segment_v2 diff --git a/build.sh b/build.sh index 8df39074a42867..6f0056814505ad 100755 --- a/build.sh +++ b/build.sh @@ -338,6 +338,7 @@ else ;; --index-tool) BUILD_INDEX_TOOL='ON' + BUILD_BE=1 shift ;; --benchmark) diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/InvertedIndexUtil.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/InvertedIndexUtil.java index 449e0333066f14..378e5643182091 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/analysis/InvertedIndexUtil.java +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/InvertedIndexUtil.java @@ -25,6 +25,7 @@ import org.apache.doris.catalog.info.IndexType; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.DdlException; +import org.apache.doris.indexpolicy.IndexPolicyMgr; import org.apache.doris.nereids.trees.plans.commands.info.IndexDefinition; import org.apache.doris.nereids.types.DataType; import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; @@ -105,6 +106,12 @@ public static void checkInvertedIndexParser(String indexColName, PrimitiveType c checkInvertedIndexProperties(properties, colType, invertedIndexFileStorageFormat); } + if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII + && !colType.isStringType() && !colType.isArrayType()) { + throw new AnalysisException("SNII inverted index storage format only supports string columns, column: " + + indexColName + " type: " + colType); + } + // default is "none" if not set if (parser == null) { parser = INVERTED_INDEX_PARSER_NONE; @@ -229,7 +236,7 @@ private static void checkInvertedIndexProperties(Map properties, + "or 'normalizer' for text normalization without tokenization."); } - checkAnalyzerName(analyzerName, colType); + checkAnalyzerName(analyzerName, colType, invertedIndexFileStorageFormat, supportPhrase); checkNormalizerName(normalizerName, colType); if (parser != null && !parser.matches("none|english|unicode|chinese|standard|icu|basic|ik")) { @@ -318,16 +325,36 @@ private static void normalizeInvertedIndexProperties(Map propert INVERTED_INDEX_PARSER_KEY_ALIAS); } - private static void checkAnalyzerName(String analyzerName, PrimitiveType colType) throws AnalysisException { + private static void checkAnalyzerName(String analyzerName, PrimitiveType colType, + TInvertedIndexFileStorageFormat storageFormat, String supportPhrase) throws AnalysisException { if (analyzerName == null || analyzerName.isEmpty()) { return; } - if (!colType.isStringType() && !colType.isVariantType()) { - throw new AnalysisException("INVERTED index with analyzer: " + analyzerName - + " is not supported for column of type " + colType); - } try { - Env.getCurrentEnv().getIndexPolicyMgr().validateAnalyzerExists(analyzerName); + IndexPolicyMgr indexPolicyMgr = Env.getCurrentEnv().getIndexPolicyMgr(); + if (indexPolicyMgr.validateAnalyzerUsesCommonGrams(analyzerName)) { + if (colType.isArrayType()) { + throw new AnalysisException("CommonGrams analyzer '" + analyzerName + + "' does not support ARRAY columns"); + } + if (!colType.isCharFamily()) { + throw new AnalysisException("CommonGrams analyzer '" + analyzerName + + "' is supported only on scalar CHAR, VARCHAR, or STRING columns"); + } + if (storageFormat != TInvertedIndexFileStorageFormat.SNII) { + throw new AnalysisException("CommonGrams analyzer '" + analyzerName + + "' is supported only by SNII inverted indexes"); + } + if (!"true".equals(supportPhrase)) { + throw new AnalysisException("CommonGrams analyzer '" + analyzerName + + "' requires support_phrase=true"); + } + return; + } + if (!colType.isStringType() && !colType.isVariantType()) { + throw new AnalysisException("INVERTED index with analyzer: " + analyzerName + + " is not supported for column of type " + colType); + } } catch (DdlException e) { throw new AnalysisException("Invalid custom analyzer: " + e.getMessage()); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/cloud/datasource/CloudInternalCatalog.java b/fe/fe-core/src/main/java/org/apache/doris/cloud/datasource/CloudInternalCatalog.java index ee3091ec8e1751..51049bd28c8a94 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/cloud/datasource/CloudInternalCatalog.java +++ b/fe/fe-core/src/main/java/org/apache/doris/cloud/datasource/CloudInternalCatalog.java @@ -384,11 +384,15 @@ public OlapFile.TabletMetaCloudPB.Builder createTabletMetaBuilder(long tableId, schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.V2); } else if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V3) { schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.V3); + } else if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII) { + schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.SNII); } else if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.DEFAULT) { if (Config.inverted_index_storage_format.equalsIgnoreCase("V1")) { schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.V1); } else if (Config.inverted_index_storage_format.equalsIgnoreCase("V2")) { schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.V2); + } else if (Config.inverted_index_storage_format.equalsIgnoreCase("SNII")) { + schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.SNII); } else { schemaBuilder.setInvertedIndexStorageFormat(OlapFile.InvertedIndexStorageFormatPB.V3); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/util/PropertyAnalyzer.java b/fe/fe-core/src/main/java/org/apache/doris/common/util/PropertyAnalyzer.java index 305fcf19064603..424d25d98de6d8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/util/PropertyAnalyzer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/util/PropertyAnalyzer.java @@ -1231,6 +1231,8 @@ public static TInvertedIndexFileStorageFormat analyzeInvertedIndexFileStorageFor return TInvertedIndexFileStorageFormat.V1; } else if (Config.inverted_index_storage_format.equalsIgnoreCase("V2")) { return TInvertedIndexFileStorageFormat.V2; + } else if (Config.inverted_index_storage_format.equalsIgnoreCase("SNII")) { + return TInvertedIndexFileStorageFormat.SNII; } else { return TInvertedIndexFileStorageFormat.V3; } @@ -1242,11 +1244,15 @@ public static TInvertedIndexFileStorageFormat analyzeInvertedIndexFileStorageFor return TInvertedIndexFileStorageFormat.V2; } else if (invertedIndexFileStorageFormat.equalsIgnoreCase("v3")) { return TInvertedIndexFileStorageFormat.V3; + } else if (invertedIndexFileStorageFormat.equalsIgnoreCase("snii")) { + return TInvertedIndexFileStorageFormat.SNII; } else if (invertedIndexFileStorageFormat.equalsIgnoreCase("default")) { if (Config.inverted_index_storage_format.equalsIgnoreCase("V1")) { return TInvertedIndexFileStorageFormat.V1; } else if (Config.inverted_index_storage_format.equalsIgnoreCase("V2")) { return TInvertedIndexFileStorageFormat.V2; + } else if (Config.inverted_index_storage_format.equalsIgnoreCase("SNII")) { + return TInvertedIndexFileStorageFormat.SNII; } else { return TInvertedIndexFileStorageFormat.V3; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/AsciiFoldingTokenFilterValidator.java b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/AsciiFoldingTokenFilterValidator.java index 2c0837dd5cfde0..ba83edd53f2457 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/AsciiFoldingTokenFilterValidator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/AsciiFoldingTokenFilterValidator.java @@ -25,7 +25,8 @@ import java.util.Set; public class AsciiFoldingTokenFilterValidator extends BasePolicyValidator { - private static final Set ALLOWED_PROPS = ImmutableSet.of("type"); + private static final Set ALLOWED_PROPS = + ImmutableSet.of("type", "preserve_original"); public AsciiFoldingTokenFilterValidator() { super(ALLOWED_PROPS); diff --git a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicy.java b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicy.java index fb0338c8ccf460..024ea7734b9515 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicy.java +++ b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicy.java @@ -58,6 +58,9 @@ public class IndexPolicy implements Writable, GsonPostProcessable { public static final String PROP_TOKENIZER = "tokenizer"; public static final String PROP_TOKEN_FILTER = "token_filter"; public static final String PROP_CHAR_FILTER = "char_filter"; + // Marks the token filter whose grams only SNII can read. The word list itself is a BE-local + // config, so no policy property names it -- this constant is only used to recognise the type. + public static final String COMMON_GRAMS_TYPE = "common_grams"; public static final Set BUILTIN_TOKENIZERS = ImmutableSet.of( "empty", "ngram", "edge_ngram", "keyword", "standard", "char_group", "basic", "icu", "pinyin"); @@ -127,4 +130,10 @@ public List getShowInfo() { public boolean isInvalid() { return false; } + + public boolean isCommonGramsPolicy() { + return type == IndexPolicyTypeEnum.TOKEN_FILTER + && properties != null + && COMMON_GRAMS_TYPE.equals(properties.get(PROP_TYPE)); + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicyMgr.java b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicyMgr.java index 1a3e2a433d1297..c30f0c18881dd3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicyMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/indexpolicy/IndexPolicyMgr.java @@ -31,6 +31,7 @@ import org.apache.doris.persist.gson.GsonUtils; import org.apache.doris.qe.ShowResultSet; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.gson.annotations.SerializedName; @@ -40,6 +41,7 @@ import java.io.DataInput; import java.io.DataOutput; import java.io.IOException; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Set; @@ -115,6 +117,26 @@ public void validateAnalyzerExists(String analyzerName) throws DdlException { } } + /** + * Validate that {@code analyzerName} resolves to a usable analyzer graph and report whether + * that graph ends in a common_grams token filter. Callers need the flag because CommonGrams + * indexes store gram terms, which only SNII can read back. + */ + public boolean validateAnalyzerUsesCommonGrams(String analyzerName) throws DdlException { + String normalizedName = normalizeKey(analyzerName); + if (IndexPolicy.BUILTIN_ANALYZERS.contains(normalizedName)) { + return false; + } + + readLock(); + try { + IndexPolicy analyzer = requireAnalyzerLocked(analyzerName); + return validateAnalyzerGraphLocked(analyzer.getName(), analyzer.getProperties()); + } finally { + readUnlock(); + } + } + public void validateNormalizerExists(String normalizerName) throws DdlException { String normalizedName = normalizeKey(normalizerName); // Built-in normalizers are stored in lowercase, so use normalized name for comparison @@ -159,11 +181,13 @@ public void createIndexPolicy(boolean ifNotExists, String policyName, throw new DdlException("Policy name '" + policyName + "' conflicts with built-in analyzer name"); } - IndexPolicy indexPolicy = IndexPolicy.create(policyName, type, properties); - writeLock(); try { validatePolicyProperties(type, properties); + if (type == IndexPolicyTypeEnum.ANALYZER) { + validateAnalyzerGraphLocked(policyName, properties); + } + IndexPolicy indexPolicy = IndexPolicy.create(policyName, type, properties); if (nameToIndexPolicy.containsKey(normalizedName)) { if (ifNotExists) { @@ -183,7 +207,7 @@ public void createIndexPolicy(boolean ifNotExists, String policyName, } finally { writeUnlock(); } - LOG.info("Created index policy successfully: {}", indexPolicy); + LOG.info("Created index policy successfully: {}", policyName); } public IndexPolicy getPolicyByName(String name) { @@ -254,6 +278,119 @@ private void validateAnalyzerProperties(Map properties) throws D } } + // CommonGrams pairs each token with its neighbour, so it is only correct when every token + // upstream of it advances the position by exactly one. These are the tokenizers and token + // filters that guarantee that; anything else (a synonym expander, a word splitter, an + // asciifolding that preserves the original) can stack tokens at one position and would make + // the resulting grams span the wrong terms. + private static final Set UNIT_POSITION_TOKENIZERS = + ImmutableSet.of("empty", "char_group", "ngram", "edge_ngram"); + private static final Set UNIT_POSITION_TOKEN_FILTERS = + ImmutableSet.of("empty", "lowercase", "icu_normalizer", "asciifolding"); + + private boolean validateAnalyzerGraphLocked(String analyzerName, + Map properties) throws DdlException { + ResolvedAnalyzerComponent tokenizer = resolveAnalyzerComponentLocked( + properties.get(IndexPolicy.PROP_TOKENIZER), IndexPolicyTypeEnum.TOKENIZER); + List tokenFilters = new ArrayList<>(); + String tokenFilterNames = properties.get(IndexPolicy.PROP_TOKEN_FILTER); + if (tokenFilterNames != null && !tokenFilterNames.isEmpty()) { + for (String tokenFilterName : tokenFilterNames.split(",\\s*")) { + tokenFilters.add(resolveAnalyzerComponentLocked( + tokenFilterName, IndexPolicyTypeEnum.TOKEN_FILTER)); + } + } + + int commonGramsIndex = -1; + int commonGramsCount = 0; + for (int i = 0; i < tokenFilters.size(); i++) { + if (IndexPolicy.COMMON_GRAMS_TYPE.equals(tokenFilters.get(i).componentType)) { + commonGramsIndex = i; + commonGramsCount++; + } + } + if (commonGramsCount == 0) { + return false; + } + if (commonGramsCount != 1 || commonGramsIndex != tokenFilters.size() - 1) { + throw new DdlException("CommonGrams analyzer '" + analyzerName + + "' requires common_grams exactly once as the terminal token filter"); + } + + if (!UNIT_POSITION_TOKENIZERS.contains(tokenizer.componentType)) { + throw new DdlException("CommonGrams analyzer '" + analyzerName + + "' tokenizer '" + tokenizer.referenceName + + "' does not guarantee unit position increments"); + } + for (int i = 0; i < commonGramsIndex; i++) { + ResolvedAnalyzerComponent tokenFilter = tokenFilters.get(i); + boolean unitPositionFilter = + UNIT_POSITION_TOKEN_FILTERS.contains(tokenFilter.componentType); + if (unitPositionFilter && "asciifolding".equals(tokenFilter.componentType)) { + unitPositionFilter = "false".equalsIgnoreCase( + tokenFilter.properties.getOrDefault("preserve_original", "false")); + } + if (!unitPositionFilter) { + throw new DdlException("CommonGrams analyzer '" + analyzerName + + "' token filter '" + tokenFilter.referenceName + + "' does not guarantee unit position increments"); + } + } + return true; + } + + private IndexPolicy requireAnalyzerLocked(String analyzerName) throws DdlException { + IndexPolicy analyzer = nameToIndexPolicy.get(normalizeKey(analyzerName)); + if (analyzer == null) { + throw new DdlException("Analyzer '" + analyzerName + "' does not exist"); + } + if (analyzer.getType() != IndexPolicyTypeEnum.ANALYZER) { + throw new DdlException("Policy '" + analyzerName + "' is not an analyzer"); + } + if (analyzer.isInvalid()) { + throw new DdlException("Analyzer '" + analyzerName + "' is invalid"); + } + return analyzer; + } + + private ResolvedAnalyzerComponent resolveAnalyzerComponentLocked(String referenceName, + IndexPolicyTypeEnum expectedType) throws DdlException { + String normalizedName = normalizeKey(referenceName); + boolean builtin = expectedType == IndexPolicyTypeEnum.TOKENIZER + ? IndexPolicy.BUILTIN_TOKENIZERS.contains(normalizedName) + : IndexPolicy.BUILTIN_TOKEN_FILTERS.contains(normalizedName); + if (builtin) { + return new ResolvedAnalyzerComponent(referenceName, normalizedName, + Map.of(IndexPolicy.PROP_TYPE, normalizedName)); + } + + IndexPolicy policy = nameToIndexPolicy.get(normalizedName); + if (policy == null) { + throw new DdlException("Referenced " + expectedType + " policy '" + + referenceName + "' does not exist"); + } + if (policy.getType() != expectedType) { + throw new DdlException("Referenced policy '" + referenceName + "' is of type " + + policy.getType() + " but expected " + expectedType); + } + return new ResolvedAnalyzerComponent(referenceName, + policy.getProperties().get(IndexPolicy.PROP_TYPE), + policy.getProperties()); + } + + private static final class ResolvedAnalyzerComponent { + private final String referenceName; + private final String componentType; + private final Map properties; + + private ResolvedAnalyzerComponent(String referenceName, String componentType, + Map properties) { + this.referenceName = referenceName; + this.componentType = componentType; + this.properties = properties; + } + } + private void validateNormalizerProperties(Map properties) throws DdlException { if (properties.containsKey(IndexPolicy.PROP_TOKENIZER)) { throw new DdlException("Normalizer cannot contain 'tokenizer' field"); @@ -291,16 +428,17 @@ private void validateNormalizerProperties(Map properties) throws private void validatePolicyReference(String name, IndexPolicyTypeEnum expectedType) throws DdlException { + String normalizedName = normalizeKey(name); if (expectedType == IndexPolicyTypeEnum.TOKENIZER - && IndexPolicy.BUILTIN_TOKENIZERS.contains(name)) { + && IndexPolicy.BUILTIN_TOKENIZERS.contains(normalizedName)) { return; } if (expectedType == IndexPolicyTypeEnum.TOKEN_FILTER - && IndexPolicy.BUILTIN_TOKEN_FILTERS.contains(name)) { + && IndexPolicy.BUILTIN_TOKEN_FILTERS.contains(normalizedName)) { return; } if (expectedType == IndexPolicyTypeEnum.CHAR_FILTER - && IndexPolicy.BUILTIN_CHAR_FILTERS.contains(name)) { + && IndexPolicy.BUILTIN_CHAR_FILTERS.contains(normalizedName)) { return; } @@ -383,6 +521,12 @@ private void validateTokenFilterProperties(Map properties) throw case "icu_normalizer": validator = new ICUNormalizerTokenFilterValidator(); break; + case "common_grams": + // No properties beyond the type: the word list it grams against is a BE-local file + // named by be.conf's common_grams_wordset_path, deliberately not selectable per + // policy, since every replica of a tablet must gram the same terms. + validator = new NoOperationValidator("common_grams token filter"); + break; default: Set userFacingTypes = IndexPolicy.BUILTIN_TOKEN_FILTERS.stream() .filter(t -> !t.equals("empty")) @@ -539,7 +683,7 @@ private void checkPolicyNotReferenced(IndexPolicy policy) throws DdlException { if (policyType == IndexPolicyTypeEnum.TOKENIZER && otherType == IndexPolicyTypeEnum.ANALYZER) { String tokenizer = properties.get(IndexPolicy.PROP_TOKENIZER); - if (policyName.equals(tokenizer)) { + if (normalizeKey(policyName).equals(normalizeKey(tokenizer))) { throw new DdlException("Cannot drop " + policyType + " policy '" + policyName + "' as it is referenced by " + otherType + " policy '" + otherPolicy.getName() + "'"); @@ -559,7 +703,7 @@ private void checkFilterReference(String policyName, IndexPolicyTypeEnum policyT String filterList) throws DdlException { if (filterList != null && !filterList.isEmpty()) { for (String filter : filterList.split(",\\s*")) { - if (policyName.equals(filter)) { + if (normalizeKey(policyName).equals(normalizeKey(filter))) { throw new DdlException("Cannot drop " + policyType + " policy '" + policyName + "' as it is referenced by " + referencingType + " policy '" + referencingPolicy.getName() + "'"); @@ -631,7 +775,10 @@ public static IndexPolicyMgr read(DataInput in) throws IOException { @Override public void gsonPostProcess() throws IOException { // Store with normalized key for case-insensitive lookup + nameToIndexPolicy.clear(); idToIndexPolicy.forEach( - (id, indexPolicy) -> nameToIndexPolicy.put(normalizeKey(indexPolicy.getName()), indexPolicy)); + (id, indexPolicy) -> + nameToIndexPolicy.put(normalizeKey(indexPolicy.getName()), indexPolicy)); } + } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsage.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsage.java index e70a39176697ef..34867f69ce22cb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsage.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsage.java @@ -17,22 +17,35 @@ package org.apache.doris.nereids.rules.rewrite; +import org.apache.doris.analysis.InvertedIndexProperties; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Index; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.common.DdlException; +import org.apache.doris.indexpolicy.IndexPolicyMgr; import org.apache.doris.nereids.exceptions.AnalysisException; 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.Expression; +import org.apache.doris.nereids.trees.expressions.Match; +import org.apache.doris.nereids.trees.expressions.SearchExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; import com.google.common.collect.ImmutableList; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import java.util.List; +import java.util.Map; +import java.util.Optional; /** * Check score function usage in project and aggregate without proper optimization context. @@ -70,6 +83,81 @@ public List buildRules() { ); } + static void checkScoringPolicyAdmission(LogicalFilter filter, + LogicalOlapScan scan, IndexPolicyMgr indexPolicyMgr) { + for (Expression conjunct : filter.getConjuncts()) { + for (Expression expression : conjunct.collect(e -> e instanceof Match)) { + Match match = (Match) expression; + if (!(match.left() instanceof SlotReference)) { + // Wrapped MATCH inputs are resolved later; do not infer a selected index from nested slots. + continue; + } + SlotReference slot = (SlotReference) match.left(); + Index selectedIndex = resolveSelectedInvertedIndex( + slot, match.getAnalyzer().orElse(null)); + checkSelectedIndexPolicyAdmission(selectedIndex, scan, indexPolicyMgr); + } + + for (Expression expression : conjunct.collect(e -> e instanceof SearchExpression)) { + SearchExpression search = (SearchExpression) expression; + for (Expression slotExpression : search.getSlotChildren()) { + if (!(slotExpression instanceof SlotReference)) { + // Variant ElementAt bindings are resolved after this rule or per segment in the BE. + continue; + } + SlotReference slot = (SlotReference) slotExpression; + // SEARCH has no explicit analyzer option. Mirror the implicit index selection in + // ExpressionTranslator.visitSearchExpression so admission checks the same index. + Index selectedIndex = resolveSelectedInvertedIndex(slot, null); + checkSelectedIndexPolicyAdmission(selectedIndex, scan, indexPolicyMgr); + } + } + } + } + + private static Index resolveSelectedInvertedIndex(SlotReference slot, String analyzer) { + Optional originalTable = slot.getOriginalTable() + .filter(OlapTable.class::isInstance) + .map(OlapTable.class::cast); + Optional originalColumn = slot.getOriginalColumn(); + if (!originalTable.isPresent() || !originalColumn.isPresent()) { + // The BE remains the authoritative admission gate when rewrite metadata is lost. + return null; + } + return originalTable.get().getInvertedIndex( + originalColumn.get(), slot.getSubPath(), analyzer); + } + + private static void checkSelectedIndexPolicyAdmission(Index selectedIndex, + LogicalOlapScan scan, IndexPolicyMgr indexPolicyMgr) { + if (selectedIndex == null) { + return; + } + Map properties = selectedIndex.getProperties(); + if (properties == null + || !properties.containsKey( + InvertedIndexProperties.INVERTED_INDEX_ANALYZER_NAME_KEY)) { + return; + } + + String analyzerName = properties.get( + InvertedIndexProperties.INVERTED_INDEX_ANALYZER_NAME_KEY); + boolean usesCommonGrams; + try { + usesCommonGrams = indexPolicyMgr.validateAnalyzerUsesCommonGrams(analyzerName); + } catch (DdlException e) { + throw new AnalysisException("score() cannot use inverted index '" + + selectedIndex.getIndexName() + "': " + e.getMessage(), e); + } + if (usesCommonGrams + && scan.getTable().getInvertedIndexFileStorageFormat() + != TInvertedIndexFileStorageFormat.SNII) { + throw new AnalysisException("score() cannot use CommonGrams analyzer '" + + analyzerName + "' on inverted index '" + selectedIndex.getIndexName() + + "': CommonGrams scoring is supported only by SNII"); + } + } + private boolean hasScoreFunction(LogicalProject project) { return project.getProjects().stream() .anyMatch(projection -> { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownScoreTopNIntoOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownScoreTopNIntoOlapScan.java index 1395c2febb09bd..68d206a45018de 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownScoreTopNIntoOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownScoreTopNIntoOlapScan.java @@ -17,6 +17,7 @@ package org.apache.doris.nereids.rules.rewrite; +import org.apache.doris.catalog.Env; import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.rules.Rule; import org.apache.doris.nereids.rules.RuleType; @@ -128,6 +129,9 @@ private Plan pushDown( + " for score() push down optimization"); } + CheckScoreUsage.checkScoringPolicyAdmission( + filter, scan, Env.getCurrentEnv().getIndexPolicyMgr()); + // 3. Check for score() predicates in WHERE clause and extract score range info List scorePredicates = filter.getConjuncts().stream() .filter(conjunct -> !conjunct.collect(e -> e instanceof Score).isEmpty()) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java index 494e756538b112..bf5aac95225629 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java @@ -31,6 +31,7 @@ import org.apache.doris.common.Config; import org.apache.doris.common.UserException; import org.apache.doris.qe.ConnectContext; +import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; import com.google.common.collect.Maps; import org.apache.commons.lang3.StringUtils; @@ -134,6 +135,10 @@ public void validate(ConnectContext ctx) throws UserException { } IndexType indexType = existedIdx.getIndexType(); + OlapTable olapTable = (OlapTable) table; + if (olapTable.getInvertedIndexFileStorageFormat() == TInvertedIndexFileStorageFormat.SNII) { + throw new AnalysisException("BUILD INDEX is not supported for SNII inverted index storage format yet"); + } if ((Config.isNotCloudMode() && indexType == IndexType.NGRAM_BF) || indexType == IndexType.BLOOMFILTER || (Config.isCloudMode() diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateTableInfo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateTableInfo.java index fd30dacc9d1d8e..a0a2718a232c3e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateTableInfo.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateTableInfo.java @@ -857,8 +857,10 @@ public void validate(ConnectContext ctx) { } if (indexDef.getIndexType() == IndexType.ANN) { if (invertedIndexFileStorageFormat != null - && invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V1) { - throw new AnalysisException("ANN index is not supported in index format V1"); + && (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V1 + || invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII)) { + throw new AnalysisException("ANN index is not supported in index format " + + invertedIndexFileStorageFormat); } } for (String indexColName : indexDef.getColumnNames()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java index 8630d80b7dc0ab..36f256994a7116 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java @@ -164,6 +164,11 @@ public void checkColumn(ColumnDefinition column, KeysType keysType, "ANN index can only be used in DUP_KEYS table or UNIQUE_KEYS table with" + " merge-on-write enabled"); } + if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V1 + || invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII) { + throw new AnalysisException("ANN index is not supported in index format " + + invertedIndexFileStorageFormat); + } return; } @@ -177,6 +182,17 @@ public void checkColumn(ColumnDefinition column, KeysType keysType, throw new AnalysisException(colType + " is not supported in " + indexType.toString() + " index. " + "invalid index: " + name); } + if (indexType == IndexType.INVERTED + && invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII) { + boolean isStringIndex = colType.isStringLikeType() + || (colType.isArrayType() + && ((ArrayType) colType).getItemType().isStringLikeType()); + if (!isStringIndex) { + throw new AnalysisException( + "SNII inverted index storage format does not support BKD index on column: " + + indexColName); + } + } // In inverted index format v1, each subcolumn of a variant has its own index file, leading to high IOPS. // when the subcolumn type changes, it may result in missing files, causing link file failure. @@ -264,8 +280,10 @@ public void checkColumn(Column column, KeysType keysType, boolean enableUniqueKe "ANN index can only be used in DUP_KEYS table or UNIQUE_KEYS table with" + " merge-on-write enabled"); } - if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V1) { - throw new AnalysisException("ANN index is not supported in index format V1"); + if (invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.V1 + || invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII) { + throw new AnalysisException("ANN index is not supported in index format " + + invertedIndexFileStorageFormat); } return; } @@ -280,9 +298,16 @@ public void checkColumn(Column column, KeysType keysType, boolean enableUniqueKe throw new AnalysisException(colType + " is not supported in " + indexType.toString() + " index. " + "invalid index: " + name); } - - if (indexType == IndexType.ANN && !colType.isArrayType()) { - throw new AnalysisException("ANN index column must be array type"); + if (indexType == IndexType.INVERTED + && invertedIndexFileStorageFormat == TInvertedIndexFileStorageFormat.SNII) { + boolean isStringIndex = colType.isStringType() + || (colType.isArrayType() + && ((org.apache.doris.catalog.ArrayType) columnType).getItemType().isStringType()); + if (!isStringIndex) { + throw new AnalysisException( + "SNII inverted index storage format does not support BKD index on column: " + + indexColName); + } } // In inverted index format v1, each subcolumn of a variant has its own index file, leading to high IOPS. diff --git a/fe/fe-core/src/test/java/org/apache/doris/alter/IndexChangeJobTest.java b/fe/fe-core/src/test/java/org/apache/doris/alter/IndexChangeJobTest.java index fa6260d19f7a8d..8a836b6b5d6f2c 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/alter/IndexChangeJobTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/alter/IndexChangeJobTest.java @@ -46,6 +46,7 @@ import org.apache.doris.qe.ConnectContext; import org.apache.doris.task.AgentTask; import org.apache.doris.task.AgentTaskQueue; +import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; import org.apache.doris.thrift.TStatusCode; import org.apache.doris.thrift.TTaskType; import org.apache.doris.transaction.FakeTransactionIDGenerator; @@ -195,6 +196,47 @@ public void testBuildIndexIndexChange() throws UserException { Assert.assertEquals(OlapTableState.NORMAL, olapTable.getState()); } + @Test + public void testBuildIndexRejectedForSniiStorageFormat() throws UserException { + if (fakeEnv != null) { + fakeEnv.close(); + } + fakeEnv = new FakeEnv(); + if (fakeEditLog != null) { + fakeEditLog.close(); + } + fakeEditLog = new FakeEditLog(); + FakeEnv.setEnv(masterEnv); + SchemaChangeHandler schemaChangeHandler = Env.getCurrentEnv().getSchemaChangeHandler(); + ArrayList alterOps = new ArrayList<>(); + Database db = masterEnv.getInternalCatalog().getDbOrDdlException(CatalogTestUtil.testDbId1); + OlapTable olapTable = (OlapTable) db.getTableOrDdlException(CatalogTestUtil.testTableId1); + String indexName = "index1"; + TableNameInfo tableNameInfo = new TableNameInfo(masterEnv.getInternalCatalog().getName(), db.getName(), + olapTable.getName()); + IndexDefinition indexDefinition = new IndexDefinition(indexName, false, + Lists.newArrayList(olapTable.getBaseSchema().get(1).getName()), + "INVERTED", + Maps.newHashMap(), "balabala"); + CreateIndexOp createIndexClause = new CreateIndexOp(tableNameInfo, indexDefinition, false); + ConnectContext connectContext = new ConnectContext(); + createIndexClause.validate(connectContext); + alterOps.add(createIndexClause); + schemaChangeHandler.process(alterOps, db, olapTable); + TInvertedIndexFileStorageFormat originalFormat = olapTable.getInvertedIndexFileStorageFormat(); + try { + olapTable.setInvertedIndexFileStorageFormat(TInvertedIndexFileStorageFormat.SNII); + BuildIndexOp buildIndexClause = new BuildIndexOp(tableNameInfo, indexName, null, false); + buildIndexClause.validate(connectContext); + Assert.fail("BUILD INDEX should be rejected for SNII inverted index storage format."); + } catch (AnalysisException e) { + Assert.assertTrue(e.getMessage().contains( + "BUILD INDEX is not supported for SNII inverted index storage format yet")); + } finally { + olapTable.setInvertedIndexFileStorageFormat(originalFormat); + } + } + @Test public void testDropIndexIndexChange() throws UserException { if (fakeEnv != null) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/analysis/InvertedIndexPropertiesTest.java b/fe/fe-core/src/test/java/org/apache/doris/analysis/InvertedIndexPropertiesTest.java index b2e2e01e278b68..c6075fe9593293 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/analysis/InvertedIndexPropertiesTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/analysis/InvertedIndexPropertiesTest.java @@ -17,12 +17,16 @@ package org.apache.doris.analysis; +import org.apache.doris.catalog.Env; import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.common.AnalysisException; +import org.apache.doris.indexpolicy.IndexPolicyMgr; import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import java.util.HashMap; import java.util.Map; @@ -343,6 +347,52 @@ public void testCheckInvertedIndexParserAllowsDotCharFilterPattern() { PrimitiveType.VARCHAR, props, TInvertedIndexFileStorageFormat.V2)); } + @Test + public void testCommonGramsAnalyzerAcceptsOnlyPhraseEnabledScalarSnii() throws Exception { + IndexPolicyMgr manager = commonGramsManager(); + + withIndexPolicyManager(manager, () -> Assertions.assertDoesNotThrow( + () -> InvertedIndexUtil.checkInvertedIndexParser("c", PrimitiveType.VARCHAR, + new HashMap<>(Map.of("analyzer", "domain_analyzer", + "support_phrase", "true")), + TInvertedIndexFileStorageFormat.SNII))); + } + + @Test + public void testCommonGramsAnalyzerRejectsV3ArrayVariantAndMissingPhrase() throws Exception { + IndexPolicyMgr manager = commonGramsManager(); + + withIndexPolicyManager(manager, () -> { + assertCommonGramsIndexError(manager, PrimitiveType.VARCHAR, + Map.of("analyzer", "domain_analyzer", "support_phrase", "true"), + TInvertedIndexFileStorageFormat.V3, + "supported only by SNII inverted indexes"); + assertCommonGramsIndexError(manager, PrimitiveType.ARRAY, + Map.of("analyzer", "domain_analyzer", "support_phrase", "true"), + TInvertedIndexFileStorageFormat.SNII, + "does not support ARRAY columns"); + assertCommonGramsIndexError(manager, PrimitiveType.VARIANT, + Map.of("analyzer", "domain_analyzer", "support_phrase", "true"), + TInvertedIndexFileStorageFormat.SNII, + "supported only on scalar CHAR, VARCHAR, or STRING columns"); + assertCommonGramsIndexError(manager, PrimitiveType.VARCHAR, + Map.of("analyzer", "domain_analyzer"), + TInvertedIndexFileStorageFormat.SNII, + "requires support_phrase=true"); + }); + } + + @Test + public void testPlainCustomAnalyzerBehaviorRemainsUnchanged() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("plain_analyzer")).thenReturn(false); + + withIndexPolicyManager(manager, () -> Assertions.assertDoesNotThrow( + () -> InvertedIndexUtil.checkInvertedIndexParser("c", PrimitiveType.VARIANT, + new HashMap<>(Map.of("analyzer", "plain_analyzer")), + TInvertedIndexFileStorageFormat.V3))); + } + // --- buildAnalyzerSqlFragment (migrated from InvertedIndexSqlGeneratorTest) --- @Test @@ -368,4 +418,29 @@ public void testBuildAnalyzerSqlFragmentQuoted() { InvertedIndexProperties.buildAnalyzerSqlFragment("O'Reilly")); } + private static IndexPolicyMgr commonGramsManager() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("domain_analyzer")).thenReturn(true); + return manager; + } + + private static void assertCommonGramsIndexError(IndexPolicyMgr manager, + PrimitiveType columnType, Map properties, + TInvertedIndexFileStorageFormat storageFormat, String expectedMessage) { + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> InvertedIndexUtil.checkInvertedIndexParser( + "c", columnType, new HashMap<>(properties), storageFormat)); + Assertions.assertTrue(exception.getMessage().contains(expectedMessage), + exception.getMessage()); + } + + private static void withIndexPolicyManager(IndexPolicyMgr manager, Runnable action) { + Env env = Mockito.mock(Env.class); + Mockito.when(env.getIndexPolicyMgr()).thenReturn(manager); + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + action.run(); + } + } + } diff --git a/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/CommonGramsDdlValidationTest.java b/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/CommonGramsDdlValidationTest.java new file mode 100644 index 00000000000000..a3455e465bed67 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/CommonGramsDdlValidationTest.java @@ -0,0 +1,154 @@ +// 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.indexpolicy; + +import org.apache.doris.common.DdlException; +import org.apache.doris.common.UserException; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +public class CommonGramsDdlValidationTest { + + @Test + public void testAcceptsTerminalCommonGramsGraph() throws Exception { + IndexPolicyMgr manager = managerWithCommonGrams(); + manager.replayCreateIndexPolicy(policy(3, "lower", IndexPolicyTypeEnum.TOKEN_FILTER, + Map.of("type", "lowercase"))); + manager.replayCreateIndexPolicy(policy(4, "domain_analyzer", IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", "token_filter", "lower,domain_grams"))); + + Assertions.assertTrue(manager.validateAnalyzerUsesCommonGrams("domain_analyzer")); + } + + @Test + public void testRejectsDuplicateAndNonTerminalCommonGrams() { + IndexPolicyMgr duplicate = managerWithCommonGrams(); + duplicate.replayCreateIndexPolicy(policy(3, "dup", IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", + "token_filter", "domain_grams,domain_grams"))); + assertAnalyzerError(duplicate, "dup", "exactly once as the terminal token filter"); + + IndexPolicyMgr nonTerminal = managerWithCommonGrams(); + nonTerminal.replayCreateIndexPolicy(policy(3, "lower", IndexPolicyTypeEnum.TOKEN_FILTER, + Map.of("type", "lowercase"))); + nonTerminal.replayCreateIndexPolicy(policy(4, "non_terminal", IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", + "token_filter", "domain_grams,lower"))); + assertAnalyzerError(nonTerminal, "non_terminal", + "exactly once as the terminal token filter"); + } + + @Test + public void testRejectsUnsafePositionFactories() { + IndexPolicyMgr unsafeTokenizer = managerWithCommonGrams(); + unsafeTokenizer.replayCreateIndexPolicy(policy(3, "unsafe_tokenizer", + IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "standard", "token_filter", "domain_grams"))); + assertAnalyzerError(unsafeTokenizer, "unsafe_tokenizer", + "tokenizer 'standard' does not guarantee unit position increments"); + + IndexPolicyMgr unsafeFilter = managerWithCommonGrams(); + unsafeFilter.replayCreateIndexPolicy(policy(3, "word_parts", + IndexPolicyTypeEnum.TOKEN_FILTER, Map.of("type", "word_delimiter"))); + unsafeFilter.replayCreateIndexPolicy(policy(4, "unsafe_filter", + IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", + "token_filter", "word_parts,domain_grams"))); + assertAnalyzerError(unsafeFilter, "unsafe_filter", + "token filter 'word_parts' does not guarantee unit position increments"); + + IndexPolicyMgr stackedAscii = managerWithCommonGrams(); + stackedAscii.replayCreateIndexPolicy(policy(3, "folded", + IndexPolicyTypeEnum.TOKEN_FILTER, + Map.of("type", "asciifolding", "preserve_original", "true"))); + stackedAscii.replayCreateIndexPolicy(policy(4, "stacked_ascii", + IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", "token_filter", "folded,domain_grams"))); + assertAnalyzerError(stackedAscii, "stacked_ascii", + "token filter 'folded' does not guarantee unit position increments"); + + IndexPolicyMgr malformedAscii = managerWithCommonGrams(); + malformedAscii.replayCreateIndexPolicy(policy(3, "malformed_folded", + IndexPolicyTypeEnum.TOKEN_FILTER, + Map.of("type", "asciifolding", "preserve_original", "garbage"))); + malformedAscii.replayCreateIndexPolicy(policy(4, "malformed_ascii", + IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "char_group", + "token_filter", "malformed_folded,domain_grams"))); + assertAnalyzerError(malformedAscii, "malformed_ascii", + "token filter 'malformed_folded' does not guarantee unit position increments"); + } + + @Test + public void testPlainAnalyzerBehaviorRemainsUnchanged() throws Exception { + IndexPolicyMgr manager = new IndexPolicyMgr(); + manager.replayCreateIndexPolicy(policy(1, "plain_analyzer", IndexPolicyTypeEnum.ANALYZER, + Map.of("tokenizer", "standard", "token_filter", "lowercase"))); + + Assertions.assertFalse(manager.validateAnalyzerUsesCommonGrams("plain_analyzer")); + } + + @Test + public void testRejectsCommonGramsPropertiesBeyondType() { + IndexPolicyMgr manager = new IndexPolicyMgr(); + + // The word list common_grams matches against is a BE-local file named by be.conf's + // common_grams_wordset_path, so no policy property can name one. An unknown property has + // to fail the DDL rather than be silently ignored, or the user would believe a word list + // they supplied was in effect. + for (Map properties : List.of( + Map.of("type", "common_grams", "words", "FILE:db/index_common_words/domain.txt"), + Map.of("type", "common_grams", "format", "wordset"), + Map.of("type", "common_grams", "ignore_case", "false"))) { + UserException exception = Assertions.assertThrows(UserException.class, + () -> manager.createIndexPolicy(false, "domain_grams", + IndexPolicyTypeEnum.TOKEN_FILTER, properties)); + Assertions.assertTrue( + exception.getMessage().contains("common_grams token filter does not support"), + exception.getMessage()); + } + } + + private static IndexPolicyMgr managerWithCommonGrams() { + IndexPolicyMgr manager = new IndexPolicyMgr(); + manager.replayCreateIndexPolicy(commonGramsPolicy()); + return manager; + } + + private static IndexPolicy commonGramsPolicy() { + return policy(2, "domain_grams", IndexPolicyTypeEnum.TOKEN_FILTER, + Map.of("type", "common_grams")); + } + + private static IndexPolicy policy(long id, String name, IndexPolicyTypeEnum type, + Map properties) { + return new IndexPolicy(id, name, type, properties); + } + + private static void assertAnalyzerError(IndexPolicyMgr manager, String analyzer, + String expectedMessage) { + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> manager.validateAnalyzerUsesCommonGrams(analyzer)); + Assertions.assertTrue(exception.getMessage().contains(expectedMessage), + exception.getMessage()); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/PolicyValidatorTests.java b/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/PolicyValidatorTests.java index 60ac8a68ff6658..4418d6270a4952 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/PolicyValidatorTests.java +++ b/fe/fe-core/src/test/java/org/apache/doris/indexpolicy/PolicyValidatorTests.java @@ -24,6 +24,10 @@ // import org.junit.jupiter.params.ParameterizedTest; // import org.junit.jupiter.params.provider.ValueSource; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; import java.util.HashMap; import java.util.Map; @@ -49,6 +53,19 @@ public void testAsciiFoldingValidator_InvalidProperty() { Assertions.assertTrue(exception.getMessage().contains("does not support parameter")); } + @Test + public void testAsciiFoldingValidatorAcceptsPreserveOriginal() throws Exception { + AsciiFoldingTokenFilterValidator validator = new AsciiFoldingTokenFilterValidator(); + validator.validate(Map.of("type", "asciifolding", "preserve_original", "true")); + validator.validate(Map.of("type", "asciifolding", "preserve_original", "false")); + } + + private static IndexPolicy roundTrip(IndexPolicy policy) throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + policy.write(new DataOutputStream(bytes)); + return IndexPolicy.read(new DataInputStream(new ByteArrayInputStream(bytes.toByteArray()))); + } + // @ParameterizedTest // @ValueSource(strings = {"yes", "no", "1", "0"}) // public void testAsciiFoldingValidator_InvalidBooleanValue(String value) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsageTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsageTest.java new file mode 100644 index 00000000000000..40fb2c2db23b2c --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CheckScoreUsageTest.java @@ -0,0 +1,471 @@ +// 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.SearchDslParser; +import org.apache.doris.catalog.AggregateType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.Index; +import org.apache.doris.catalog.KeysType; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.PartitionInfo; +import org.apache.doris.catalog.TableIndexes; +import org.apache.doris.catalog.Type; +import org.apache.doris.catalog.info.IndexType; +import org.apache.doris.common.DdlException; +import org.apache.doris.indexpolicy.IndexPolicyMgr; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.properties.OrderKey; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.MatchPhrase; +import org.apache.doris.nereids.trees.expressions.SearchExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.types.StringType; +import org.apache.doris.thrift.TInvertedIndexFileStorageFormat; +import org.apache.doris.thrift.TStorageType; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public class CheckScoreUsageTest { + private static final AtomicLong NEXT_INDEX_ID = new AtomicLong(1); + + @Test + public void testAcceptsCompatibleCommonGramsPolicyForSelectedSniiIndex() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("domain_analyzer")) + .thenReturn(true); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "domain_analyzer")), + "body", null); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("domain_analyzer"); + } + + @Test + public void testScorePushDownRuleInvokesPolicyAdmission() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "missing_analyzer")), + "body", null); + LogicalTopN>> topN = scoreTopN(filter); + Rule rule = new PushDownScoreTopNIntoOlapScan().buildRules().get(0); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getIndexPolicyMgr()).thenReturn(manager); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> rule.transform(topN, Mockito.mock(CascadesContext.class))); + Assertions.assertTrue(exception.getMessage().contains( + "Analyzer 'missing_analyzer' does not exist"), exception.getMessage()); + } + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + } + + @Test + public void testSearchScorePushDownRuleInvokesPolicyAdmission() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalFilter filter = searchScoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + builtInIndex("idx_body", "body"), + index("idx_other", "other", "missing_analyzer")), + "body", "other"); + LogicalTopN>> topN = scoreTopN(filter); + Rule rule = new PushDownScoreTopNIntoOlapScan().buildRules().get(0); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getIndexPolicyMgr()).thenReturn(manager); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> rule.transform(topN, Mockito.mock(CascadesContext.class))); + Assertions.assertTrue(exception.getMessage().contains( + "Analyzer 'missing_analyzer' does not exist"), exception.getMessage()); + } + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + } + + @Test + public void testSearchRejectsCommonGramsPolicyOnSecondFieldForV3() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("domain_analyzer")) + .thenReturn(true); + LogicalFilter filter = searchScoreFilter( + table(TInvertedIndexFileStorageFormat.V3, + builtInIndex("idx_body", "body"), + index("idx_other", "other", "domain_analyzer")), + "body", "other"); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Assertions.assertTrue(exception.getMessage().contains("supported only by SNII"), + exception.getMessage()); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("domain_analyzer"); + } + + @Test + public void testRejectsCommonGramsPolicyForSelectedV3Index() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("domain_analyzer")) + .thenReturn(true); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.V3, + index("idx_body", "body", "domain_analyzer")), + "body", null); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Assertions.assertTrue(exception.getMessage().contains("supported only by SNII"), + exception.getMessage()); + } + + @Test + public void testRejectsMissingSelectedAnalyzerPolicy() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "missing_analyzer")), + "body", null); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Assertions.assertTrue(exception.getMessage().contains( + "Analyzer 'missing_analyzer' does not exist"), exception.getMessage()); + } + + @Test + public void testRejectsIncompatibleSelectedCommonGramsPolicy() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException( + "CommonGrams token-filter 'domain_grams' current state is PREPARING")) + .when(manager).validateAnalyzerUsesCommonGrams("domain_analyzer"); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "domain_analyzer")), + "body", null); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Assertions.assertTrue(exception.getMessage().contains("current state is PREPARING"), + exception.getMessage()); + } + + @Test + public void testIgnoresUnselectedIndexWithMissingPolicy() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("plain_analyzer")) + .thenReturn(false); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.V3, + index("idx_body", "body", "plain_analyzer"), + index("idx_other", "other", "missing_analyzer")), + "body", null); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("plain_analyzer"); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams("missing_analyzer"); + } + + @Test + public void testExplicitAnalyzerSelectsTargetIndexOnSameColumn() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("target_analyzer")) + .thenReturn(true); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body_first", "body", "first_analyzer"), + index("idx_body_target", "body", "target_analyzer")), + "body", "target_analyzer"); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("target_analyzer"); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams("first_analyzer"); + } + + @Test + public void testAnalyzerlessMatchSelectsFirstAnalyzedSiblingIndex() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("first_analyzer")) + .thenReturn(false); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body_first", "body", "first_analyzer"), + index("idx_body_second", "body", "second_analyzer")), + "body", null); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("first_analyzer"); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams("second_analyzer"); + } + + @Test + public void testSearchSelectsFirstAnalyzedSiblingIndex() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.when(manager.validateAnalyzerUsesCommonGrams("first_analyzer")) + .thenReturn(false); + LogicalFilter filter = searchScoreFilter( + table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body_first", "body", "first_analyzer"), + index("idx_body_second", "body", "second_analyzer")), + "body"); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("first_analyzer"); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams("second_analyzer"); + } + + @Test + public void testSelectedIndexWithoutExplicitAnalyzerSkipsPolicyAdmission() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.V3, + builtInIndex("idx_body", "body")), + "body", null); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams(Mockito.anyString()); + } + + @Test + public void testNoSelectedIndexFallsBackWithoutPolicyAdmission() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + LogicalFilter filter = scoreFilter( + table(TInvertedIndexFileStorageFormat.SNII), "body", null); + + Assertions.assertDoesNotThrow(() -> CheckScoreUsage.checkScoringPolicyAdmission( + filter, filter.child(), manager)); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams(Mockito.anyString()); + } + + @Test + public void testMatchWithoutOriginalColumnFallsBackToBeAtRuleLevel() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + LogicalOlapScan scan = newScan(table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "missing_analyzer"))); + SlotReference detachedSlot = new SlotReference("body", StringType.INSTANCE); + LogicalFilter filter = scoreFilter(scan, detachedSlot, "missing_analyzer"); + + assertScorePushDownSucceeds(filter, manager); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams(Mockito.anyString()); + } + + @Test + public void testMatchWithNonSlotLeftFallsBackToBeAtRuleLevel() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalOlapScan scan = newScan(table(TInvertedIndexFileStorageFormat.SNII, + index("idx_body", "body", "missing_analyzer"))); + Alias wrappedSlot = new Alias(findSlot(scan, "body"), "wrapped_body"); + LogicalFilter filter = scoreFilter(scan, wrappedSlot, null); + + assertScorePushDownSucceeds(filter, manager); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams(Mockito.anyString()); + } + + @Test + public void testSkippedMatchDoesNotSuppressAdmissionForLaterMatch() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalOlapScan scan = newScan(table(TInvertedIndexFileStorageFormat.SNII, + index("idx_other", "other", "missing_analyzer"))); + Alias wrappedSlot = new Alias(findSlot(scan, "body"), "wrapped_body"); + MatchPhrase skippedMatch = new MatchPhrase( + wrappedSlot, new StringLiteral("alpha beta"), null); + MatchPhrase admittedMatch = new MatchPhrase( + findSlot(scan, "other"), new StringLiteral("alpha beta"), null); + LogicalFilter filter = new LogicalFilter<>( + ImmutableSet.of(skippedMatch, admittedMatch), scan); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission(filter, scan, manager)); + Assertions.assertTrue(exception.getMessage().contains("missing_analyzer"), + exception.getMessage()); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + } + + @Test + public void testSearchWithNonSlotBindingFallsBackToBeAtRuleLevel() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + LogicalOlapScan scan = newScan(table(TInvertedIndexFileStorageFormat.SNII, + index("idx_variant", "variant_body", "missing_analyzer"))); + ElementAt variantBinding = new ElementAt( + findSlot(scan, "variant_body"), new StringLiteral("path")); + LogicalFilter filter = searchScoreFilter( + scan, "variant_body.path:alpha", ImmutableList.of(variantBinding)); + + assertScorePushDownSucceeds(filter, manager); + Mockito.verify(manager, Mockito.never()).validateAnalyzerUsesCommonGrams(Mockito.anyString()); + } + + @Test + public void testSkippedSearchBindingDoesNotSuppressAdmissionForLaterField() throws Exception { + IndexPolicyMgr manager = Mockito.mock(IndexPolicyMgr.class); + Mockito.doThrow(new DdlException("Analyzer 'missing_analyzer' does not exist")) + .when(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + LogicalOlapScan scan = newScan(table(TInvertedIndexFileStorageFormat.SNII, + index("idx_other", "other", "missing_analyzer"))); + ElementAt variantBinding = new ElementAt( + findSlot(scan, "variant_body"), new StringLiteral("path")); + LogicalFilter filter = searchScoreFilter( + scan, "variant_body.path:alpha AND other:alpha", + ImmutableList.of(variantBinding, findSlot(scan, "other"))); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> CheckScoreUsage.checkScoringPolicyAdmission(filter, scan, manager)); + Assertions.assertTrue(exception.getMessage().contains("missing_analyzer"), + exception.getMessage()); + Mockito.verify(manager).validateAnalyzerUsesCommonGrams("missing_analyzer"); + } + + private static LogicalFilter scoreFilter( + OlapTable table, String columnName, String analyzer) { + LogicalOlapScan scan = newScan(table); + return scoreFilter(scan, findSlot(scan, columnName), analyzer); + } + + private static LogicalFilter scoreFilter( + LogicalOlapScan scan, Expression left, String analyzer) { + MatchPhrase match = new MatchPhrase(left, new StringLiteral("alpha beta"), analyzer); + return new LogicalFilter<>(ImmutableSet.of(match), scan); + } + + private static LogicalFilter searchScoreFilter( + OlapTable table, String... columnNames) { + LogicalOlapScan scan = newScan(table); + List slotChildren = new ArrayList<>(); + List clauses = new ArrayList<>(); + for (String columnName : columnNames) { + slotChildren.add(findSlot(scan, columnName)); + clauses.add(columnName + ":alpha"); + } + return searchScoreFilter(scan, String.join(" AND ", clauses), slotChildren); + } + + private static LogicalFilter searchScoreFilter( + LogicalOlapScan scan, String dsl, List slotChildren) { + SearchExpression search = new SearchExpression( + dsl, SearchDslParser.parseDsl(dsl, null), slotChildren); + return new LogicalFilter<>(ImmutableSet.of(search), scan); + } + + private static LogicalOlapScan newScan(OlapTable table) { + return new LogicalOlapScan( + StatementScopeIdGenerator.newRelationId(), table, ImmutableList.of("db")); + } + + private static SlotReference findSlot(LogicalOlapScan scan, String columnName) { + return (SlotReference) scan.getOutput().stream() + .filter(output -> output.getName().equals(columnName)) + .findFirst() + .orElseThrow(); + } + + private static LogicalTopN>> scoreTopN( + LogicalFilter filter) { + Alias scoreAlias = new Alias(new Score(), "score"); + LogicalProject> project = new LogicalProject<>( + ImmutableList.of(scoreAlias), filter); + return new LogicalTopN<>( + ImmutableList.of(new OrderKey(scoreAlias.toSlot(), false, false)), + 10, 0, project); + } + + private static void assertScorePushDownSucceeds( + LogicalFilter filter, IndexPolicyMgr manager) { + LogicalTopN>> topN = scoreTopN(filter); + Rule rule = new PushDownScoreTopNIntoOlapScan().buildRules().get(0); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getIndexPolicyMgr()).thenReturn(manager); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Assertions.assertFalse( + rule.transform(topN, Mockito.mock(CascadesContext.class)).isEmpty()); + } + } + + private static OlapTable table(TInvertedIndexFileStorageFormat storageFormat, + Index... indexes) { + List columns = ImmutableList.of( + new Column("id", Type.INT, true, AggregateType.NONE, "0", ""), + new Column("body", Type.STRING, false, AggregateType.NONE, "", ""), + new Column("other", Type.STRING, false, AggregateType.NONE, "", ""), + new Column("variant_body", Type.VARIANT, false, AggregateType.NONE, "", "")); + OlapTable table = new OlapTable(10, "score_table", false, columns, + KeysType.DUP_KEYS, new PartitionInfo(), null, + new TableIndexes(ImmutableList.copyOf(indexes))); + table.setIndexMeta(-1, "score_table", table.getFullSchema(), + 0, 0, (short) 0, TStorageType.COLUMN, KeysType.DUP_KEYS); + table.setInvertedIndexFileStorageFormat(storageFormat); + return table; + } + + private static Index index(String name, String column, String analyzer) { + return new Index(NEXT_INDEX_ID.getAndIncrement(), name, ImmutableList.of(column), IndexType.INVERTED, + Map.of("analyzer", analyzer, "support_phrase", "true"), ""); + } + + private static Index builtInIndex(String name, String column) { + return new Index(NEXT_INDEX_ID.getAndIncrement(), name, ImmutableList.of(column), IndexType.INVERTED, + Map.of("parser", "standard", "support_phrase", "true"), ""); + } + +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IndexDefinitionTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IndexDefinitionTest.java index 7b41ddc95cf840..060e687b495242 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IndexDefinitionTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IndexDefinitionTest.java @@ -18,7 +18,9 @@ package org.apache.doris.nereids.trees.plans.commands; import org.apache.doris.catalog.AggregateType; +import org.apache.doris.catalog.Column; import org.apache.doris.catalog.KeysType; +import org.apache.doris.catalog.Type; import org.apache.doris.catalog.info.IndexType; import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.trees.plans.commands.info.ColumnDefinition; @@ -57,6 +59,68 @@ void testVariantIndexFormatV1() throws AnalysisException { } } + @Test + void testSniiInvertedIndexColumnTypes() throws AnalysisException { + IndexDefinition def = new IndexDefinition("snii_index", false, Lists.newArrayList("col1"), + "INVERTED", null, "comment"); + + def.checkColumn(new ColumnDefinition("col1", StringType.INSTANCE, false, AggregateType.NONE, true, + null, "comment"), KeysType.DUP_KEYS, false, TInvertedIndexFileStorageFormat.SNII); + def.checkColumn(new ColumnDefinition("col1", ArrayType.of(StringType.INSTANCE), false, + AggregateType.NONE, true, null, "comment"), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII); + + AnalysisException intException = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new ColumnDefinition("col1", IntegerType.INSTANCE, false, AggregateType.NONE, + true, null, "comment"), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(intException.getMessage().contains("does not support BKD index")); + + AnalysisException arrayIntException = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new ColumnDefinition("col1", ArrayType.of(IntegerType.INSTANCE), false, + AggregateType.NONE, true, null, "comment"), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(arrayIntException.getMessage().contains("does not support BKD index")); + } + + @Test + void testSniiInvertedIndexCatalogColumnTypes() throws AnalysisException { + IndexDefinition def = new IndexDefinition("snii_index", false, Lists.newArrayList("col1"), + "INVERTED", null, "comment"); + + def.checkColumn(new Column("col1", Type.STRING, true), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII); + def.checkColumn(new Column("col1", org.apache.doris.catalog.ArrayType.create(Type.STRING), true), + KeysType.DUP_KEYS, false, TInvertedIndexFileStorageFormat.SNII); + + AnalysisException intException = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new Column("col1", Type.INT, true), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(intException.getMessage().contains("does not support BKD index")); + + AnalysisException arrayIntException = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new Column("col1", org.apache.doris.catalog.ArrayType.create(Type.INT), true), + KeysType.DUP_KEYS, false, TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(arrayIntException.getMessage().contains("does not support BKD index")); + } + + @Test + void testSniiRejectsAnnIndex() { + IndexDefinition def = new IndexDefinition("ann_index", false, Lists.newArrayList("col1"), + "ANN", null, "comment"); + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new ColumnDefinition("col1", ArrayType.of(FloatType.INSTANCE), false, + AggregateType.NONE, false, null, "comment"), KeysType.DUP_KEYS, false, + TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(exception.getMessage().contains("ANN index is not supported in index format SNII")); + + AnalysisException catalogException = Assertions.assertThrows(AnalysisException.class, () -> + def.checkColumn(new Column("col1", org.apache.doris.catalog.ArrayType.create(Type.FLOAT), false), + KeysType.DUP_KEYS, false, TInvertedIndexFileStorageFormat.SNII)); + Assertions.assertTrue(catalogException.getMessage().contains( + "ANN index is not supported in index format SNII")); + } + void testArrayTypeSupport() throws AnalysisException { IndexDefinition def = new IndexDefinition("array_index", false, Lists.newArrayList("col1"), "INVERTED", null, "array test"); diff --git a/gensrc/proto/olap_file.proto b/gensrc/proto/olap_file.proto index 7f1c62a446100b..c3172969dfad09 100644 --- a/gensrc/proto/olap_file.proto +++ b/gensrc/proto/olap_file.proto @@ -446,6 +446,7 @@ enum InvertedIndexStorageFormatPB { V1 = 0; V2 = 1; V3 = 2; + SNII = 3; } // Tablet-level storage format. Values match TStorageFormat (Thrift) integer values so diff --git a/gensrc/proto/snii.proto b/gensrc/proto/snii.proto new file mode 100644 index 00000000000000..cd605a5aed7692 --- /dev/null +++ b/gensrc/proto/snii.proto @@ -0,0 +1,82 @@ +// 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. + +syntax = "proto2"; + +package doris.snii; + +message SniiMetadataDirectoryPB { + repeated uint32 required_features = 1 [packed = true]; + repeated SniiLogicalIndexMetadataPB indexes = 2; +} + +message SniiLogicalIndexMetadataPB { + optional uint64 index_id = 1; + optional bytes index_suffix = 2; + optional SniiBlobRefPB core_metadata = 3; + optional SniiBlobRefPB sampled_term_index = 4; + optional SniiBlobRefPB dict_block_directory = 5; +} + +message SniiBlobRefPB { + optional uint64 offset = 1; + optional uint64 length = 2; +} + +message SniiCoreMetadataPB { + optional uint32 index_config = 1; + optional SniiStatsPB stats = 2; + optional SniiSectionRefsPB section_refs = 3; + optional SniiCommonGramsMetadataPB common_grams = 4; + optional uint32 common_grams_posting_policy = 5; +} + +message SniiStatsPB { + optional uint64 doc_count = 1; + optional uint64 indexed_doc_count = 2; + optional uint64 term_count = 3; + optional uint64 sum_total_term_freq = 4; + optional uint64 null_count = 5; +} + +message SniiRegionRefPB { + optional uint64 offset = 1; + optional uint64 length = 2; +} + +message SniiSectionRefsPB { + optional SniiRegionRefPB dict_region = 1; + optional SniiRegionRefPB posting_region = 2; + optional SniiRegionRefPB norms = 3; + optional SniiRegionRefPB null_bitmap = 4; + optional SniiRegionRefPB bsbf = 5; +} + +message SniiCommonGramsMetadataPB { + optional uint32 plain_term_key_version = 1; + optional uint32 common_grams_coverage = 2; + optional uint32 common_grams_semantics_version = 3; + optional uint32 common_grams_key_version = 4; + optional bytes common_grams_dictionary_identity = 5; + optional bytes base_analyzer_fingerprint = 6; + optional bytes common_grams_fingerprint = 7; + optional uint32 scoring_coverage = 8; + optional uint32 scoring_stats_version = 9; + optional uint32 norm_semantics_version = 10; + optional uint64 scoring_doc_count = 11; + optional uint64 scoring_token_count = 12; +} diff --git a/gensrc/thrift/Types.thrift b/gensrc/thrift/Types.thrift index c6b9c705307380..d088a936b9e05f 100644 --- a/gensrc/thrift/Types.thrift +++ b/gensrc/thrift/Types.thrift @@ -130,7 +130,8 @@ enum TInvertedIndexFileStorageFormat { DEFAULT = 0, // Default format, unspecified storage method. V1 = 1, // Index per idx: Each index is stored separately based on its identifier. V2 = 2, // Segment id per idx: Indexes are organized based on segment identifiers, grouping indexes by their associated segment. - V3 = 3 // Position and dictionary compression + V3 = 3, // Position and dictionary compression + SNII = 4 // SNII native inverted index storage format } struct TScalarType { diff --git a/regression-test/data/inverted_index_p0/storage_format/common_grams_docs.csv b/regression-test/data/inverted_index_p0/storage_format/common_grams_docs.csv new file mode 100644 index 00000000000000..6ca68c00de5706 --- /dev/null +++ b/regression-test/data/inverted_index_p0/storage_format/common_grams_docs.csv @@ -0,0 +1,20 @@ +1|alpha +2|alpha beta +3|alpha beta gamma +4|alpha the beta of gamma and delta in epsilon to +5|alpha the beta of gamma and delta in epsilon to zeta +6|foo bar baz +7|foo bar the +8|foo the bar +9|foo the of +10|the foo bar +11|the foo of +12|the of foo +13|the of and +14|the the the +15|the world +16|the worker +17|foo theater +18|foo of the world +19|the bar bazaar +20|alpha beta gamma delta epsilon zeta eta theta iota kappa diff --git a/regression-test/data/inverted_index_p0/storage_format/test_common_grams_snii.out b/regression-test/data/inverted_index_p0/storage_format/test_common_grams_snii.out new file mode 100644 index 00000000000000..8fab3e9454bc61 --- /dev/null +++ b/regression-test/data/inverted_index_p0/storage_format/test_common_grams_snii.out @@ -0,0 +1,668 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !cg_any_one -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_any_many -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_any_three -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_any_common_many -- +1 +10 +11 +12 +13 +14 +15 +16 +18 +19 +2 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +3 +30 +31 +4 +5 +7 +8 +9 + +-- !cg_all_one -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_all_many -- +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_all_three -- +20 +21 +29 +3 +31 +4 +5 + +-- !cg_all_common_many -- +31 +4 +5 + +-- !cg_all_common_non_adjacent -- +13 +26 +31 +4 +5 + +-- !cg_exact_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_exact_2 -- +2 +20 +22 +29 +3 + +-- !cg_exact_2_common -- +12 +13 +27 +9 + +-- !cg_exact_3 -- +20 +29 +3 + +-- !cg_exact_6 -- +31 +4 +5 + +-- !cg_exact_10 -- +4 +5 + +-- !cg_shape_nnn -- +6 + +-- !cg_shape_nns -- +7 + +-- !cg_shape_nsn -- +8 + +-- !cg_shape_nss -- +9 + +-- !cg_shape_snn -- +10 + +-- !cg_shape_sns -- +11 + +-- !cg_shape_ssn -- +12 + +-- !cg_shape_sss -- +13 + +-- !cg_repeated_common -- +14 + +-- !cg_prefix_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_prefix_2 -- +2 +20 +22 +29 +3 + +-- !cg_prefix_3 -- +20 +29 +3 + +-- !cg_prefix_6 -- +31 +4 +5 + +-- !cg_prefix_10 -- +4 +5 + +-- !cg_prefix_the_wo -- +15 +16 +18 +30 + +-- !cg_prefix_foo_the -- +17 +24 +8 +9 + +-- !cg_prefix_foo_of_th -- +18 +28 + +-- !cg_prefix_the_bar_ba -- +19 + +-- !v3_sloppy_plain_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !snii_plain_oracle_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !v3_sloppy_plain_2 -- +15 +18 +30 + +-- !snii_plain_oracle_2 -- +15 +18 +30 + +-- !v3_sloppy_plain_3 -- +30 + +-- !snii_plain_oracle_3 -- +30 + +-- !v3_sloppy_plain_6 -- +31 +4 +5 + +-- !snii_plain_oracle_6 -- +31 +4 +5 + +-- !v3_sloppy_plain_10 -- +4 +5 + +-- !snii_plain_oracle_10 -- +4 +5 + +-- !cg_regexp_namespace -- +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +20 +24 +25 +26 +27 +29 +30 +31 +4 +5 +7 +8 +9 + +-- !cg_literal_marker_namespace -- + +-- !cg_regexp_leading_namespace -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_search_leading_wildcard -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !v3_search_leading_wildcard -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !cg_search_internal_gram_namespace -- + +-- !v3_search_internal_gram_namespace -- + +-- !cg_search_literal_marker_namespace -- + +-- !cg_search_field_exists -- +1 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +2 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +3 +30 +31 +4 +5 +6 +7 +8 +9 + +-- !cg_search_field_exists_oracle -- +1 +10 +11 +12 +13 +14 +15 +16 +17 +18 +19 +2 +20 +21 +22 +23 +24 +25 +26 +27 +28 +29 +3 +30 +31 +4 +5 +6 +7 +8 +9 + +-- !explicit_plain_index_routing -- +2 +20 +22 +29 +3 + +-- !safety_off_plain -- +13 + +-- !no_index_exact -- +20 +29 +3 + +-- !plain_exact_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !plain_exact_2 -- +2 +20 +22 +29 +3 + +-- !plain_exact_2_common -- +12 +13 +27 +9 + +-- !plain_exact_3 -- +20 +29 +3 + +-- !plain_exact_6 -- +31 +4 +5 + +-- !plain_exact_10 -- +4 +5 + +-- !plain_shape_nnn -- +6 + +-- !plain_shape_nns -- +7 + +-- !plain_shape_nsn -- +8 + +-- !plain_shape_nss -- +9 + +-- !plain_shape_snn -- +10 + +-- !plain_shape_sns -- +11 + +-- !plain_shape_ssn -- +12 + +-- !plain_shape_sss -- +13 + +-- !plain_repeated_common -- +14 + +-- !plain_prefix_1 -- +1 +2 +20 +21 +22 +29 +3 +31 +4 +5 + +-- !plain_prefix_2 -- +2 +20 +22 +29 +3 + +-- !plain_prefix_3 -- +20 +29 +3 + +-- !plain_prefix_6 -- +31 +4 +5 + +-- !plain_prefix_10 -- +4 +5 + +-- !plain_prefix_the_wo -- +15 +16 +18 +30 + +-- !plain_prefix_foo_the -- +17 +24 +8 +9 + +-- !plain_prefix_foo_of_th -- +18 +28 + +-- !plain_prefix_the_bar_ba -- +19 + +-- !cg_cache_reuse_across_plan_toggle -- +13 + +-- !cg_authoritative_empty -- + +-- !cg_authoritative_empty_prefix -- + +-- !plain_authoritative_empty -- + +-- !plain_authoritative_empty_prefix -- + +-- !score_snii_exact_limit_10 -- +2 6.162817 +3 4.848449 +22 4.848449 +20 1.944889 +29 1.944889 + +-- !score_v3_exact_limit_10 -- +2 6.162818 +3 4.848449 +22 4.848449 +20 1.944889 +29 1.944889 + +-- !score_snii_exact_limit_100 -- +2 6.162817 +3 4.848449 +22 4.848449 +20 1.944889 +29 1.944889 + +-- !score_v3_exact_limit_100 -- +2 6.162818 +3 4.848449 +22 4.848449 +20 1.944889 +29 1.944889 + +-- !score_snii_prefix_limit_10 -- +2 3.042638 +3 2.393722 +22 2.393722 +20 0.960209 +29 0.960209 + +-- !score_snii_prefix_limit_100 -- +2 3.042638 +3 2.393722 +22 2.393722 +20 0.960209 +29 0.960209 + +-- !v3_prefix_limit_10 -- +2 +20 +22 +29 +3 + +-- !v3_prefix_limit_100 -- +2 +20 +22 +29 +3 + +-- !score_v3_missing_control -- +1 + +-- !missing_plain_fallback -- +1 + +-- !missing_prefix_plain_fallback -- +1 + +-- !score_v3_mixed_control -- +1 +2 + +-- !mixed_plain_fallback -- +1 +2 + +-- !mixed_prefix_plain_fallback -- +1 +2 + +-- !mixed_recovered_phrase -- +1 +2 + +-- !mixed_recovered_prefix -- +1 +2 + +-- !mixed_recovered_score -- +1 0.364643 +2 0.364643 + +-- !after_full_compaction_phrase -- +31 +4 +5 + +-- !after_full_compaction_score -- +2 6.162817 +3 4.848449 +22 4.848449 +20 1.944889 +29 1.944889 + diff --git a/regression-test/data/inverted_index_p0/storage_format/test_storage_format_snii.out b/regression-test/data/inverted_index_p0/storage_format/test_storage_format_snii.out new file mode 100644 index 00000000000000..33e05cf4214d2f --- /dev/null +++ b/regression-test/data/inverted_index_p0/storage_format/test_storage_format_snii.out @@ -0,0 +1,16 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !match_any -- +1 +2 + +-- !match_all -- +1 + +-- !match_phrase -- +5 + +-- !null_bitmap -- +4 + +-- !array_contains -- +1 diff --git a/regression-test/suites/inverted_index_p0/storage_format/test_common_grams_snii.groovy b/regression-test/suites/inverted_index_p0/storage_format/test_common_grams_snii.groovy new file mode 100644 index 00000000000000..51939d6bdc0c33 --- /dev/null +++ b/regression-test/suites/inverted_index_p0/storage_format/test_common_grams_snii.groovy @@ -0,0 +1,1043 @@ +// 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. + +import java.util.regex.Pattern + +import org.apache.doris.regression.action.ProfileAction + +suite("test_common_grams_snii", "p0,nonConcurrent") { + def backendIdToIp = [:] + def backendIdToHttpPort = [:] + getBackendIpHttpPort(backendIdToIp, backendIdToHttpPort) + + def readBeConfig = { backendId, String key -> + def (code, out, err) = show_be_config( + backendIdToIp.get(backendId), backendIdToHttpPort.get(backendId)) + assertEquals(0, code, err) + def row = parseJson(out.trim()).find { it[0] == key } + assertNotNull(row, "BE config ${key} is missing on backend ${backendId}") + return row[2].toString() + } + + def originalBeConfigs = [:] + ["enable_common_grams_query_plan", "enable_common_grams_index_build", + "common_grams_plan_cost_ratio_percent"].each { key -> + backendIdToIp.keySet().each { backendId -> + originalBeConfigs.computeIfAbsent(backendId) { [:] }[key] = + readBeConfig(backendId, key) + } + } + + def setBeConfig = { String key, String value -> + backendIdToIp.keySet().each { backendId -> + def (code, out, err) = update_be_config( + backendIdToIp.get(backendId), backendIdToHttpPort.get(backendId), key, value) + assertEquals(0, code, "update ${key}=${value} failed on backend ${backendId}: ${out} ${err}") + assertEquals(value, readBeConfig(backendId, key)) + } + } + + def createAnalyzer = { String name, String tokenizer, String tokenFilters -> + sql """ + CREATE INVERTED INDEX ANALYZER ${name} + PROPERTIES ( + "tokenizer" = "${tokenizer}", + "token_filter" = "${tokenFilters}" + ) + """ + } + + def waitAnalyzerInstalled = { String name -> + def deadline = System.currentTimeMillis() + 180_000 + Exception lastNotFound = null + while (System.currentTimeMillis() < deadline) { + try { + sql """SELECT TOKENIZE('probe', '"analyzer"="${name}"')""" + return + } catch (Exception e) { + if (!e.message.contains("Policy not found")) { + throw e + } + lastNotFound = e + sleep(1000) + } + } + throw new IllegalStateException("analyzer ${name} was not installed on BE", lastNotFound) + } + + def assertProfileCounterPositive = { String label, String query, String counter -> + def profileId = "test_common_grams_${label}_${System.nanoTime()}" + String profileString + sql "SET enable_profile = true" + try { + sql """ + /* ${profileId} */ + ${query} + """ + profileString = new ProfileAction(context).getProfileBySql( + profileId, [counter]) + def matcher = Pattern.compile(Pattern.quote(counter) + ":\\s*(\\d+)") + .matcher(profileString) + assertTrue(matcher.find(), "${counter} is missing from profile") + assertTrue(Long.parseLong(matcher.group(1)) > 0, "${counter} must be positive") + } finally { + sql "SET enable_profile = false" + } + return profileString + } + + try { + sql "DROP TABLE IF EXISTS test_common_grams_snii" + sql "DROP TABLE IF EXISTS test_common_grams_v3_plain" + sql "DROP TABLE IF EXISTS test_common_grams_mixed_v3_plain" + sql "DROP TABLE IF EXISTS test_common_grams_missing" + sql "DROP TABLE IF EXISTS test_common_grams_mixed" + sql "DROP TABLE IF EXISTS test_common_grams_no_index" + try_sql "DROP INVERTED INDEX ANALYZER IF EXISTS cg_default_analyzer" + try_sql "DROP INVERTED INDEX ANALYZER IF EXISTS cg_plain_analyzer" + try_sql "DROP INVERTED INDEX TOKEN_FILTER IF EXISTS cg_default_grams" + try_sql "DROP INVERTED INDEX TOKENIZER IF EXISTS cg_char_group_tokenizer" + + setBeConfig("enable_common_grams_query_plan", "true") + setBeConfig("enable_common_grams_index_build", "true") + // 100 is the validator maximum: grams win whenever not estimated more + // expensive than plain. The filler rows below make the probed stopword + // postings genuinely dominate so the gate stays decisively open. + setBeConfig("common_grams_plan_cost_ratio_percent", "100") + + sql """ + CREATE INVERTED INDEX TOKENIZER cg_char_group_tokenizer + PROPERTIES ( + "type" = "char_group", + "tokenize_on_chars" = "[whitespace], [punctuation]" + ) + """ + sql """ + CREATE INVERTED INDEX ANALYZER cg_plain_analyzer + PROPERTIES ( + "tokenizer" = "cg_char_group_tokenizer", + "token_filter" = "lowercase" + ) + """ + sql """ + CREATE INVERTED INDEX TOKEN_FILTER cg_default_grams + PROPERTIES ("type" = "common_grams") + """ + // The word list lives on the BE (/common_grams/default_words.txt) + // and is deliberately not selectable per policy, so naming one has to fail the DDL rather + // than be silently ignored. + test { + sql """ + CREATE INVERTED INDEX TOKEN_FILTER cg_words_rejected + PROPERTIES ( + "type" = "common_grams", + "words" = "FILE:db/cg_words/common_grams_words.txt" + ) + """ + exception "does not support parameter 'words'" + } + + createAnalyzer("cg_default_analyzer", "cg_char_group_tokenizer", + "lowercase,cg_default_grams") + + test { + sql """ + CREATE INVERTED INDEX ANALYZER cg_unsafe_tokenizer + PROPERTIES ( + "tokenizer" = "standard", + "token_filter" = "cg_default_grams" + ) + """ + exception "does not guarantee unit position increments" + } + test { + sql """ + CREATE INVERTED INDEX ANALYZER cg_non_terminal + PROPERTIES ( + "tokenizer" = "cg_char_group_tokenizer", + "token_filter" = "cg_default_grams,lowercase" + ) + """ + exception "exactly once as the terminal token filter" + } + + waitAnalyzerInstalled("cg_plain_analyzer") + waitAnalyzerInstalled("cg_default_analyzer") + + sql """ + CREATE TABLE test_common_grams_snii ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer", + "support_phrase" = "true" + ), + INDEX idx_body_plain (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_plain_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "SNII" + ) + """ + + test { + sql """ + CREATE TABLE test_common_grams_v3_rejected ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "inverted_index_storage_format" = "V3" + ) + """ + exception "supported only by SNII inverted indexes" + } + test { + sql """ + CREATE TABLE test_common_grams_array_rejected ( + id INT NOT NULL, + body ARRAY NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "inverted_index_storage_format" = "SNII" + ) + """ + exception "does not support ARRAY columns" + } + test { + sql """ + CREATE TABLE test_common_grams_phrase_rejected ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "inverted_index_storage_format" = "SNII" + ) + """ + exception "requires support_phrase=true" + } + + sql """ + CREATE TABLE test_common_grams_v3_plain ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_plain_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "V3" + ) + """ + sql """ + CREATE TABLE test_common_grams_no_index ( + id INT NOT NULL, + body STRING NULL + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + sql """ + CREATE TABLE test_common_grams_mixed_v3_plain ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_plain_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "V3" + ) + """ + + streamLoad { + table "test_common_grams_snii" + set "column_separator", "|" + set "columns", "id,body" + file "common_grams_docs.csv" + time 30_000 + check { result, exception, startTime, endTime -> + if (exception != null) { + throw exception + } + def json = parseJson(result) + assertEquals("success", json.Status.toLowerCase()) + assertEquals(20, json.NumberLoadedRows) + } + } + // Cost-gate ballast: the plain-plan estimate is driven by the candidate + // df of the probed terms (rarest term included), while the gram side is + // driven by the pair-term df. Isolated single-token rows raise every + // probed term's df without creating a single new gram pair (no + // adjacency), so the pair estimates stay put and the <=100% cost gate + // picks the gram plan deterministically. Single tokens also never match + // any probed phrase, prefix, or MATCH_ALL golden. + def costGateBallast = (0..<500).collect { + ",\n (${1000 + it}, '${["the", "of", "and", "foo", "bar"][it % 5]}')" + }.join("") + sql """ + INSERT INTO test_common_grams_snii VALUES + (21, 'alpha gamma beta'), + (22, 'alpha beta delta'), + (23, 'foo bar and'), + (24, 'bar foo the'), + (25, 'the bar foo'), + (26, 'the and of'), + (27, 'the the of'), + (28, 'foo of thinker'), + (29, 'alpha beta gamma delta epsilon zeta eta theta iota lambda'), + (30, 'the world wide web'), + (31, 'alpha the beta of gamma and omega'), + (32, NULL)${costGateBallast} + """ + sql "INSERT INTO test_common_grams_v3_plain SELECT id, body FROM test_common_grams_snii" + sql "INSERT INTO test_common_grams_no_index SELECT id, body FROM test_common_grams_snii" + sql "SYNC" + sql "SET enable_profile = true" + sql "SET profile_level = 2" + sql "SET parallel_fragment_exec_instance_num = 1" + sql "SET enable_sql_cache = false" + sql "SET enable_inverted_index_query_cache = false" + sql "SET enable_segment_limit_pushdown = true" + + def metamorphicQueries = [ + exact_1: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha' ORDER BY id", + exact_2: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta' ORDER BY id", + exact_2_common: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of' ORDER BY id", + exact_3: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta gamma' ORDER BY id", + exact_6: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha the beta of gamma and' ORDER BY id", + exact_10: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha the beta of gamma and delta in epsilon to' ORDER BY id", + shape_nnn: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar baz' ORDER BY id", + shape_nns: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar the' ORDER BY id", + shape_nsn: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the bar' ORDER BY id", + shape_nss: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the of' ORDER BY id", + shape_snn: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo bar' ORDER BY id", + shape_sns: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo of' ORDER BY id", + shape_ssn: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of foo' ORDER BY id", + shape_sss: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + repeated_common: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the the the' ORDER BY id", + prefix_1: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alph' ORDER BY id", + prefix_2: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha be' ORDER BY id", + prefix_3: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha beta ga' ORDER BY id", + prefix_6: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma an' ORDER BY id", + prefix_10: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma and delta in epsilon t' ORDER BY id", + prefix_the_wo: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id", + prefix_foo_the: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo the' ORDER BY id", + prefix_foo_of_th: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo of th' ORDER BY id", + prefix_the_bar_ba: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the bar ba' ORDER BY id", + authoritative_miss_exact: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id", + authoritative_miss_prefix: "SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id" + ] + + order_qt_cg_any_one """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ANY 'alpha' ORDER BY id + """ + order_qt_cg_any_many """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ANY 'alpha beta' ORDER BY id + """ + order_qt_cg_any_three """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ANY 'alpha beta gamma' ORDER BY id + """ + // id < 1000 keeps the cost-gate ballast rows (which necessarily contain + // the probed stopwords) out of the any/all/regexp/match-all goldens. + order_qt_cg_any_common_many """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_ANY 'alpha the beta of gamma and' AND id < 1000 ORDER BY id + """ + order_qt_cg_all_one """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ALL 'alpha' ORDER BY id + """ + order_qt_cg_all_many """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ALL 'alpha beta' ORDER BY id + """ + order_qt_cg_all_three """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_ALL 'alpha beta gamma' ORDER BY id + """ + order_qt_cg_all_common_many """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_ALL 'alpha the beta of gamma and' ORDER BY id + """ + order_qt_cg_all_common_non_adjacent """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_ALL 'the of and' AND id < 1000 ORDER BY id + """ + + order_qt_cg_exact_1 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha' ORDER BY id + """ + order_qt_cg_exact_2 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta' ORDER BY id + """ + order_qt_cg_exact_2_common """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of' ORDER BY id + """ + order_qt_cg_exact_3 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta gamma' ORDER BY id + """ + order_qt_cg_exact_6 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and' ORDER BY id + """ + order_qt_cg_exact_10 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and delta in epsilon to' ORDER BY id + """ + + order_qt_cg_shape_nnn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar baz' ORDER BY id + """ + order_qt_cg_shape_nns """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar the' ORDER BY id + """ + order_qt_cg_shape_nsn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the bar' ORDER BY id + """ + order_qt_cg_shape_nss """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the of' ORDER BY id + """ + order_qt_cg_shape_snn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo bar' ORDER BY id + """ + order_qt_cg_shape_sns """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo of' ORDER BY id + """ + order_qt_cg_shape_ssn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of foo' ORDER BY id + """ + order_qt_cg_shape_sss """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of and' ORDER BY id + """ + order_qt_cg_repeated_common """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the the the' ORDER BY id + """ + assertProfileCounterPositive("shape_nnn", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'foo bar baz' ORDER BY id", + "SniiCommonGramsFallbackNoGram") + [ + ["exact_2_common", "the of"], + ["shape_nns", "foo bar the"], + ["shape_nss", "foo the of"], + ["shape_snn", "the foo bar"], + ["shape_sns", "the foo of"], + ["shape_ssn", "the of foo"], + ["shape_sss", "the of and"] + ].each { label, terms -> + assertProfileCounterPositive(label, + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE '${terms}' ORDER BY id", + "SniiCommonGramsGramPlans") + } + // A lone interior stopword (n-s-n) keeps every plain clause in the + // HybridV1 query plan -- the pair clauses are purely additive -- so its + // gram plan estimate is plain + pairs and can never pass the <=100% + // cost gate. Pin the gate rejecting it instead of asserting a gram + // plan that is unreachable by construction. + assertProfileCounterPositive("shape_nsn", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'foo the bar' ORDER BY id", + "SniiCommonGramsFallbackCost") + + sql "SET inverted_index_max_expansions = 50" + order_qt_cg_prefix_1 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alph' ORDER BY id + """ + order_qt_cg_prefix_2 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha be' ORDER BY id + """ + order_qt_cg_prefix_3 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha beta ga' ORDER BY id + """ + order_qt_cg_prefix_6 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma an' ORDER BY id + """ + order_qt_cg_prefix_10 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma and delta in epsilon t' ORDER BY id + """ + order_qt_cg_prefix_the_wo """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id + """ + order_qt_cg_prefix_foo_the """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo the' ORDER BY id + """ + order_qt_cg_prefix_foo_of_th """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo of th' ORDER BY id + """ + order_qt_cg_prefix_the_bar_ba """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the bar ba' ORDER BY id + """ + // Like shape_nsn: the phrase-prefix hybrid plan keeps the plain + // stopword clause alongside the pair expansions, so its estimate is + // plain + pairs and the <=100% cost gate must reject it. + assertProfileCounterPositive("prefix_common", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id", + "SniiCommonGramsFallbackCost") + def assertSloppyMatchesPlainOracle = { String label, String phrase -> + def v3SloppyRows = sql """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE '${phrase} ~1' ORDER BY id + """ + def sniiPlainRows = sql """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE '${phrase}' USING ANALYZER cg_plain_analyzer + ORDER BY id + """ + assertFalse(v3SloppyRows.isEmpty(), "sloppy phrase ${label} must have a positive match") + assertEquals(sniiPlainRows, v3SloppyRows, + "V3 sloppy phrase ${label} differs from the SNII plain oracle") + } + [ + ["1", "alpha"], + ["2", "the world"], + ["3", "the world wide"], + ["6", "alpha the beta of gamma and"], + ["10", "alpha the beta of gamma and delta in epsilon to"] + ].each { label, phrase -> assertSloppyMatchesPlainOracle(label, phrase) } + + order_qt_v3_sloppy_plain_1 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'alpha ~1' ORDER BY id + """ + order_qt_snii_plain_oracle_1 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha' USING ANALYZER cg_plain_analyzer ORDER BY id + """ + order_qt_v3_sloppy_plain_2 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'the world ~1' ORDER BY id + """ + order_qt_snii_plain_oracle_2 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the world' USING ANALYZER cg_plain_analyzer ORDER BY id + """ + order_qt_v3_sloppy_plain_3 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'the world wide ~1' ORDER BY id + """ + order_qt_snii_plain_oracle_3 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the world wide' USING ANALYZER cg_plain_analyzer ORDER BY id + """ + order_qt_v3_sloppy_plain_6 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'alpha the beta of gamma and ~1' ORDER BY id + """ + order_qt_snii_plain_oracle_6 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and' + USING ANALYZER cg_plain_analyzer ORDER BY id + """ + order_qt_v3_sloppy_plain_10 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'alpha the beta of gamma and delta in epsilon to ~1' + ORDER BY id + """ + order_qt_snii_plain_oracle_10 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and delta in epsilon to' + USING ANALYZER cg_plain_analyzer ORDER BY id + """ + order_qt_cg_regexp_namespace """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_REGEXP '^the.*' AND id < 1000 ORDER BY id + """ + String literalCommonGram = "\u001fDORIS_COMMON_GRAM_V1\u001f00000003:theof" + order_qt_cg_literal_marker_namespace """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_ANY '${literalCommonGram}' ORDER BY id + """ + order_qt_cg_regexp_leading_namespace """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_REGEXP '.*lpha' ORDER BY id + """ + String wildcardSearchOptions = + '{"default_operator":"and","default_field":"body","minimum_should_match":0,"mode":"lucene"}' + order_qt_cg_search_leading_wildcard """ + SELECT id FROM test_common_grams_snii + WHERE search('body:*lpha', '${wildcardSearchOptions}') ORDER BY id + """ + order_qt_v3_search_leading_wildcard """ + SELECT id FROM test_common_grams_v3_plain + WHERE search('body:*lpha', '${wildcardSearchOptions}') ORDER BY id + """ + order_qt_cg_search_internal_gram_namespace """ + SELECT id FROM test_common_grams_snii + WHERE search('body:*theof*', '${wildcardSearchOptions}') ORDER BY id + """ + order_qt_v3_search_internal_gram_namespace """ + SELECT id FROM test_common_grams_v3_plain + WHERE search('body:*theof*', '${wildcardSearchOptions}') ORDER BY id + """ + String commonGramMarker = "\u001fDORIS_COMMON_GRAM_V1\u001f" + order_qt_cg_search_literal_marker_namespace """ + SELECT id FROM test_common_grams_snii + WHERE search('body:*${commonGramMarker}*', '${wildcardSearchOptions}') ORDER BY id + """ + order_qt_cg_search_field_exists """ + SELECT id FROM test_common_grams_snii + WHERE search('body:**', '${wildcardSearchOptions}') AND id < 1000 ORDER BY id + """ + order_qt_cg_search_field_exists_oracle """ + SELECT id FROM test_common_grams_snii + WHERE body IS NOT NULL AND id < 1000 ORDER BY id + """ + order_qt_explicit_plain_index_routing """ + SELECT id + FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha beta' USING ANALYZER cg_plain_analyzer + ORDER BY id + """ + + // The BE-local switch is the only gate, and flipping it bumps the query-plan config + // generation that feeds the inverted index query cache key, so the very next query + // observes the new mode instead of a stale cached bitmap. + setBeConfig("enable_common_grams_query_plan", "false") + order_qt_safety_off_plain """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of and' ORDER BY id + """ + def safetyOffProfile = assertProfileCounterPositive("safety_off", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + "SniiCommonGramsFallbackKillSwitch") + assertFalse(safetyOffProfile.contains("SniiCommonGramsGramPlans:")) + setBeConfig("enable_common_grams_query_plan", "true") + + order_qt_no_index_exact """ + SELECT /*+ SET_VAR(enable_match_without_inverted_index=true) */ id + FROM test_common_grams_no_index WHERE body MATCH_PHRASE 'alpha beta gamma' ORDER BY id + """ + + sql "SET enable_inverted_index_query_cache = true" + def gramMetamorphicRows = metamorphicQueries.collectEntries { label, query -> + [(label): sql(query)] + } + setBeConfig("enable_common_grams_query_plan", "false") + def plainCacheSeedProfile = assertProfileCounterPositive("plain_cache_seed", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + "SniiCommonGramsFallbackKillSwitch") + assertFalse(plainCacheSeedProfile.contains("SniiCommonGramsGramPlans:")) + order_qt_plain_exact_1 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha' ORDER BY id + """ + order_qt_plain_exact_2 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta' ORDER BY id + """ + order_qt_plain_exact_2_common """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of' ORDER BY id + """ + order_qt_plain_exact_3 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'alpha beta gamma' ORDER BY id + """ + order_qt_plain_exact_6 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and' ORDER BY id + """ + order_qt_plain_exact_10 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and delta in epsilon to' ORDER BY id + """ + order_qt_plain_shape_nnn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar baz' ORDER BY id + """ + order_qt_plain_shape_nns """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo bar the' ORDER BY id + """ + order_qt_plain_shape_nsn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the bar' ORDER BY id + """ + order_qt_plain_shape_nss """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'foo the of' ORDER BY id + """ + order_qt_plain_shape_snn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo bar' ORDER BY id + """ + order_qt_plain_shape_sns """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the foo of' ORDER BY id + """ + order_qt_plain_shape_ssn """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of foo' ORDER BY id + """ + order_qt_plain_shape_sss """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of and' ORDER BY id + """ + order_qt_plain_repeated_common """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the the the' ORDER BY id + """ + order_qt_plain_prefix_1 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alph' ORDER BY id + """ + order_qt_plain_prefix_2 """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'alpha be' ORDER BY id + """ + order_qt_plain_prefix_3 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha beta ga' ORDER BY id + """ + order_qt_plain_prefix_6 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma an' ORDER BY id + """ + order_qt_plain_prefix_10 """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha the beta of gamma and delta in epsilon t' ORDER BY id + """ + order_qt_plain_prefix_the_wo """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id + """ + order_qt_plain_prefix_foo_the """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo the' ORDER BY id + """ + order_qt_plain_prefix_foo_of_th """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'foo of th' ORDER BY id + """ + order_qt_plain_prefix_the_bar_ba """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE_PREFIX 'the bar ba' ORDER BY id + """ + metamorphicQueries.each { label, query -> + assertEquals(gramMetamorphicRows[label], sql(query), + "CommonGrams and forced-plain results differ for ${label}") + } + + setBeConfig("enable_common_grams_query_plan", "true") + assertProfileCounterPositive("first_query_after_plan_toggle", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + "SniiCommonGramsGramPlans") + order_qt_cg_cache_reuse_across_plan_toggle """ + SELECT id FROM test_common_grams_snii WHERE body MATCH_PHRASE 'the of and' ORDER BY id + """ + assertProfileCounterPositive("cache_reuse_across_plan_toggle", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + "InvertedIndexQueryCacheHit") + sql "SET enable_inverted_index_query_cache = false" + assertProfileCounterPositive("plan_reenabled", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the of and' ORDER BY id", + "SniiCommonGramsGramPlans") + order_qt_cg_authoritative_empty """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id + """ + assertProfileCounterPositive("authoritative_empty_exact", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id", + "SniiCommonGramsAuthoritativeEmpty") + def gramAuthoritativeEmpty = sql """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id + """ + order_qt_cg_authoritative_empty_prefix """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id + """ + assertProfileCounterPositive("authoritative_empty_prefix", + "SELECT id FROM test_common_grams_snii " + + "WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id", + "SniiCommonGramsAuthoritativeEmpty") + def gramAuthoritativeEmptyPrefix = sql """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id + """ + setBeConfig("enable_common_grams_query_plan", "false") + order_qt_plain_authoritative_empty """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id + """ + def plainAuthoritativeEmpty = sql """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'the nonexistent phrase' ORDER BY id + """ + order_qt_plain_authoritative_empty_prefix """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id + """ + def plainAuthoritativeEmptyPrefix = sql """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'the nonex' ORDER BY id + """ + assertEquals(gramAuthoritativeEmpty, plainAuthoritativeEmpty) + assertEquals(gramAuthoritativeEmptyPrefix, plainAuthoritativeEmptyPrefix) + setBeConfig("enable_common_grams_query_plan", "true") + + def rankedIds = { String table, String predicate, int limit -> + sql """ + SELECT id FROM ( + SELECT id, score() AS s FROM ${table} + WHERE ${predicate} + ORDER BY s DESC LIMIT ${limit} + ) ranked ORDER BY s DESC, id + """ + } + def orderedIds = { String table, String predicate, int limit -> + sql """ + SELECT id FROM ${table} + WHERE ${predicate} + ORDER BY id LIMIT ${limit} + """ + } + [10, 100].each { limit -> + assertEquals( + rankedIds("test_common_grams_v3_plain", "body MATCH_PHRASE 'alpha beta'", limit), + rankedIds("test_common_grams_snii", "body MATCH_PHRASE 'alpha beta'", limit)) + assertEquals( + orderedIds("test_common_grams_v3_plain", "body MATCH_PHRASE_PREFIX 'alpha be'", limit), + orderedIds("test_common_grams_snii", "body MATCH_PHRASE_PREFIX 'alpha be'", limit)) + } + + qt_score_snii_exact_limit_10 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 10 + ) ranked ORDER BY s DESC, id + """ + qt_score_v3_exact_limit_10 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 10 + ) ranked ORDER BY s DESC, id + """ + qt_score_snii_exact_limit_100 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + qt_score_v3_exact_limit_100 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + qt_score_snii_prefix_limit_10 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha be' + ORDER BY s DESC LIMIT 10 + ) ranked ORDER BY s DESC, id + """ + qt_score_snii_prefix_limit_100 """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_snii + WHERE body MATCH_PHRASE_PREFIX 'alpha be' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + order_qt_v3_prefix_limit_10 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE_PREFIX 'alpha be' + ORDER BY id LIMIT 10 + """ + order_qt_v3_prefix_limit_100 """ + SELECT id FROM test_common_grams_v3_plain + WHERE body MATCH_PHRASE_PREFIX 'alpha be' + ORDER BY id LIMIT 100 + """ + setBeConfig("enable_common_grams_index_build", "false") + sql """ + CREATE TABLE test_common_grams_missing ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "SNII" + ) + """ + sql "INSERT INTO test_common_grams_missing VALUES (1, 'alpha beta the world')" + sql "INSERT INTO test_common_grams_mixed_v3_plain VALUES (1, 'alpha beta the world')" + qt_score_v3_missing_control """ + SELECT id FROM ( + SELECT id, score() AS s FROM test_common_grams_mixed_v3_plain + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + order_qt_missing_plain_fallback """ + SELECT id FROM test_common_grams_missing + WHERE body MATCH_PHRASE 'alpha beta the' ORDER BY id + """ + order_qt_missing_prefix_plain_fallback """ + SELECT id FROM test_common_grams_missing + WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id + """ + assertProfileCounterPositive("missing_prefix_fallback", + "SELECT id FROM test_common_grams_missing " + + "WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id", + "SniiCommonGramsFallbackIncompatible") + test { + sql """ + SELECT id, score() AS s FROM test_common_grams_missing + WHERE body MATCH_PHRASE 'alpha beta' ORDER BY s DESC LIMIT 100 + """ + exception "SNII semantic scoring metadata is missing" + } + + setBeConfig("enable_common_grams_index_build", "true") + sql """ + CREATE TABLE test_common_grams_mixed ( + id INT NOT NULL, + body STRING NULL, + INDEX idx_body (body) USING INVERTED PROPERTIES ( + "analyzer" = "cg_default_analyzer", + "support_phrase" = "true" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ( + "replication_num" = "1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "SNII" + ) + """ + sql "INSERT INTO test_common_grams_mixed VALUES (1, 'alpha beta the world')" + setBeConfig("enable_common_grams_index_build", "false") + sql "INSERT INTO test_common_grams_mixed VALUES (2, 'alpha beta the worker')" + sql "INSERT INTO test_common_grams_mixed_v3_plain VALUES (2, 'alpha beta the worker')" + qt_score_v3_mixed_control """ + SELECT id FROM ( + SELECT id, score() AS s FROM test_common_grams_mixed_v3_plain + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + order_qt_mixed_plain_fallback """ + SELECT id FROM test_common_grams_mixed WHERE body MATCH_PHRASE 'alpha beta' ORDER BY id + """ + order_qt_mixed_prefix_plain_fallback """ + SELECT id FROM test_common_grams_mixed + WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id + """ + assertProfileCounterPositive("mixed_prefix_fallback", + "SELECT id FROM test_common_grams_mixed " + + "WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id", + "SniiCommonGramsFallbackIncompatible") + test { + sql """ + SELECT id, score() AS s FROM test_common_grams_mixed + WHERE body MATCH_PHRASE 'alpha beta' ORDER BY s DESC LIMIT 100 + """ + exception "SNII semantic scoring metadata is missing" + } + + setBeConfig("enable_common_grams_index_build", "true") + trigger_and_wait_compaction("test_common_grams_mixed", "full", 300) + order_qt_mixed_recovered_phrase """ + SELECT id FROM test_common_grams_mixed WHERE body MATCH_PHRASE 'alpha beta' ORDER BY id + """ + order_qt_mixed_recovered_prefix """ + SELECT id FROM test_common_grams_mixed + WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id + """ + // Same structural rejection as prefix_common: the phrase-prefix hybrid + // keeps the plain stopword clause, so the gram plan cannot be cheaper. + assertProfileCounterPositive("mixed_recovered_prefix", + "SELECT id FROM test_common_grams_mixed " + + "WHERE body MATCH_PHRASE_PREFIX 'the wo' ORDER BY id", + "SniiCommonGramsFallbackCost") + qt_mixed_recovered_score """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_mixed + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + assertEquals( + rankedIds("test_common_grams_mixed_v3_plain", "body MATCH_PHRASE 'alpha beta'", 100), + rankedIds("test_common_grams_mixed", "body MATCH_PHRASE 'alpha beta'", 100)) + + trigger_and_wait_compaction("test_common_grams_snii", "full", 300) + order_qt_after_full_compaction_phrase """ + SELECT id FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha the beta of gamma and' ORDER BY id + """ + qt_after_full_compaction_score """ + SELECT id, ROUND(s, 6) FROM ( + SELECT id, score() AS s FROM test_common_grams_snii + WHERE body MATCH_PHRASE 'alpha beta' + ORDER BY s DESC LIMIT 100 + ) ranked ORDER BY s DESC, id + """ + assertEquals( + rankedIds("test_common_grams_v3_plain", "body MATCH_PHRASE 'alpha beta'", 100), + rankedIds("test_common_grams_snii", "body MATCH_PHRASE 'alpha beta'", 100)) + } finally { + originalBeConfigs.each { backendId, values -> + values.each { key, value -> + def (code, out, err) = update_be_config( + backendIdToIp.get(backendId), backendIdToHttpPort.get(backendId), + key, value) + assertEquals(0, code, + "restore ${key}=${value} failed on backend ${backendId}: ${out} ${err}") + } + } + } +} diff --git a/regression-test/suites/inverted_index_p0/storage_format/test_storage_format_snii.groovy b/regression-test/suites/inverted_index_p0/storage_format/test_storage_format_snii.groovy new file mode 100644 index 00000000000000..7800350fb6b753 --- /dev/null +++ b/regression-test/suites/inverted_index_p0/storage_format/test_storage_format_snii.groovy @@ -0,0 +1,212 @@ +// 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_storage_format_snii", "p0, nonConcurrent") { + sql "DROP TABLE IF EXISTS test_storage_format_snii" + sql "DROP TABLE IF EXISTS test_storage_format_snii_array" + sql "DROP TABLE IF EXISTS test_storage_format_snii_add_index" + sql "DROP TABLE IF EXISTS test_storage_format_snii_bkd" + sql "DROP TABLE IF EXISTS test_storage_format_snii_array_bkd" + sql "DROP TABLE IF EXISTS test_storage_format_snii_ann" + + sql """ + CREATE TABLE test_storage_format_snii ( + id INT NULL, + body TEXT NULL, + INDEX idx_body (`body`) USING INVERTED PROPERTIES( + "parser" = "english", + "support_phrase" = "true", + "lower_case" = "true" + ) COMMENT '' + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "disable_auto_compaction" = "true", + "inverted_index_storage_format" = "SNII" + ); + """ + + sql """ + INSERT INTO test_storage_format_snii VALUES + (1, 'alpha beta gamma'), + (2, 'alpha delta'), + (3, 'beta epsilon'), + (4, NULL), + (5, 'quick brown fox'), + (6, 'quick fox'); + """ + sql "sync" + + order_qt_match_any """ + SELECT id FROM test_storage_format_snii + WHERE body MATCH_ANY 'alpha' + ORDER BY id + """ + order_qt_match_all """ + SELECT id FROM test_storage_format_snii + WHERE body MATCH_ALL 'alpha beta' + ORDER BY id + """ + order_qt_match_phrase """ + SELECT id FROM test_storage_format_snii + WHERE body MATCH_PHRASE 'quick brown' + ORDER BY id + """ + order_qt_null_bitmap """ + SELECT id FROM test_storage_format_snii + WHERE body IS NULL + ORDER BY id + """ + + sql """ + CREATE TABLE test_storage_format_snii_array ( + id INT NULL, + tags ARRAY NULL, + INDEX idx_tags (`tags`) USING INVERTED COMMENT '' + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "inverted_index_storage_format" = "SNII" + ); + """ + + sql """ + INSERT INTO test_storage_format_snii_array VALUES + (1, '["alpha", "beta"]'), + (2, '["gamma"]'), + (3, NULL); + """ + sql "sync" + + order_qt_array_contains """ + SELECT id FROM test_storage_format_snii_array + WHERE array_contains(tags, 'alpha') + ORDER BY id + """ + + test { + if (isCloudMode()) { + sql "BUILD INDEX ON test_storage_format_snii" + } else { + sql "BUILD INDEX idx_body ON test_storage_format_snii" + } + exception "BUILD INDEX is not supported for SNII inverted index storage format yet" + } + + sql """ + CREATE TABLE test_storage_format_snii_add_index ( + id INT NULL, + body TEXT NULL, + score INT NULL, + scores ARRAY NULL, + embedding ARRAY NOT NULL, + INDEX idx_body_added_table (`body`) USING INVERTED COMMENT '' + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "inverted_index_storage_format" = "SNII" + ); + """ + + test { + sql """ + ALTER TABLE test_storage_format_snii_add_index + ADD INDEX idx_score_added (`score`) USING INVERTED COMMENT '' + """ + exception "SNII inverted index storage format" + } + + test { + sql """ + ALTER TABLE test_storage_format_snii_add_index + ADD INDEX idx_scores_added (`scores`) USING INVERTED COMMENT '' + """ + exception "SNII inverted index storage format" + } + + test { + sql """ + CREATE INDEX idx_ann_added ON test_storage_format_snii_add_index (`embedding`) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "1" + ) + """ + exception "ANN index is not supported in index format SNII" + } + + test { + sql """ + CREATE TABLE test_storage_format_snii_bkd ( + id INT NULL, + score INT NULL, + INDEX idx_score (`score`) USING INVERTED COMMENT '' + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "inverted_index_storage_format" = "SNII" + ); + """ + exception "SNII inverted index storage format" + } + + test { + sql """ + CREATE TABLE test_storage_format_snii_array_bkd ( + id INT NULL, + scores ARRAY NULL, + INDEX idx_scores (`scores`) USING INVERTED COMMENT '' + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "inverted_index_storage_format" = "SNII" + ); + """ + exception "SNII inverted index storage format" + } + + test { + sql """ + CREATE TABLE test_storage_format_snii_ann ( + id INT NULL, + embedding ARRAY NOT NULL, + INDEX idx_ann (`embedding`) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY RANDOM BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "inverted_index_storage_format" = "SNII" + ); + """ + exception "ANN index is not supported in index format SNII" + } +}