diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index 66c18a7f12..660a9b0cc3 100644 --- a/docs/redoc/master/openapi.json +++ b/docs/redoc/master/openapi.json @@ -8349,6 +8349,11 @@ "description": "Enable HNSW graph building for this payload field. If true, builds additional HNSW links (Need payload_m > 0). Default: true.", "type": "boolean", "nullable": true + }, + "prefix": { + "description": "If true, enable prefix matching (`match: { \"prefix\": ... }`) on this field. Default: false.", + "type": "boolean", + "nullable": true } } }, @@ -9503,6 +9508,9 @@ { "$ref": "#/components/schemas/MatchPhrase" }, + { + "$ref": "#/components/schemas/MatchPrefix" + }, { "$ref": "#/components/schemas/MatchAny" }, @@ -9573,6 +9581,18 @@ } } }, + "MatchPrefix": { + "description": "Match keyword values that start with the given string.\n\nByte-wise (hence, for valid UTF-8, character-wise) and case-sensitive, consistent with exact keyword matching. Served efficiently by a keyword index created with the `prefix` option.", + "type": "object", + "required": [ + "prefix" + ], + "properties": { + "prefix": { + "type": "string" + } + } + }, "MatchAny": { "description": "Exact match on any of the given values", "type": "object", diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index df23dbb963..ceb35d8b6a 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -32,12 +32,12 @@ use super::qdrant::{ BinaryQuantization, BoolIndexParams, CompressionRatio, DatetimeIndexParams, DatetimeRange, Direction, FacetHit, FacetHitInternal, FacetValue, FacetValueInternal, FieldType, FloatIndexParams, GeoIndexParams, GeoLineString, GroupId, HardwareUsage, HasVectorCondition, - KeywordIndexParams, LookupLocation, MaxOptimizationThreads, MultiVectorComparator, - MultiVectorConfig, OrderBy, OrderValue, Range, RawVector, RecommendStrategy, RetrievedPoint, - SearchMatrixPair, SearchPointGroups, SearchPoints, ShardKeySelector, StartFrom, - StrictModeMultivector, StrictModeMultivectorConfig, StrictModeSparse, StrictModeSparseConfig, - TurboQuantBitSize, TurboQuantization, UuidIndexParams, VectorsOutput, WithLookup, raw_query, - start_from, + KeywordIndexParams, KeywordPrefixParams, LookupLocation, MaxOptimizationThreads, + MultiVectorComparator, MultiVectorConfig, OrderBy, OrderValue, Range, RawVector, + RecommendStrategy, RetrievedPoint, SearchMatrixPair, SearchPointGroups, SearchPoints, + ShardKeySelector, StartFrom, StrictModeMultivector, StrictModeMultivectorConfig, + StrictModeSparse, StrictModeSparseConfig, TurboQuantBitSize, TurboQuantization, + UuidIndexParams, VectorsOutput, WithLookup, raw_query, start_from, }; use super::stemming_algorithm::StemmingParams; use super::{ @@ -200,12 +200,14 @@ impl From for PayloadIndexParams is_tenant, on_disk, enable_hnsw, + prefix, } = params; PayloadIndexParams { index_params: Some(IndexParams::KeywordIndexParams(KeywordIndexParams { is_tenant, on_disk, enable_hnsw, + prefix: prefix.unwrap_or_default().then_some(KeywordPrefixParams {}), })), } } @@ -492,12 +494,16 @@ impl TryFrom for segment::data_types::index::KeywordIndexPar is_tenant, on_disk, enable_hnsw, + prefix, } = params; Ok(segment::data_types::index::KeywordIndexParams { r#type: KeywordIndexType::Keyword, is_tenant, on_disk, enable_hnsw, + // Presence of the (currently empty) message enables prefix + // matching. + prefix: prefix.map(|KeywordPrefixParams {}| true), }) } } @@ -2107,6 +2113,7 @@ impl TryFrom for segment::types::Match { MatchValue::TextAny(text_any) => { segment::types::Match::TextAny(segment::types::MatchTextAny { text_any }) } + MatchValue::Prefix(prefix) => segment::types::Match::Prefix(prefix.into()), }), _ => Err(Status::invalid_argument("Malformed Match condition")), } @@ -2150,6 +2157,9 @@ impl From for Match { segment::types::Match::TextAny(segment::types::MatchTextAny { text_any }) => { MatchValue::TextAny(text_any) } + segment::types::Match::Prefix(segment::types::MatchPrefix { prefix }) => { + MatchValue::Prefix(prefix) + } }; Self { match_value: Some(match_value), diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index ab33e7d977..09926207a1 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -628,6 +628,13 @@ message KeywordIndexParams { // If true, builds additional HNSW links (Need payload_m > 0). // Default: true. optional bool enable_hnsw = 3; + // If set, enable prefix matching (`match: { "prefix": ... }`) on this field. + optional KeywordPrefixParams prefix = 4; +} + +// Prefix matching options for the keyword index. Has no options yet: +// presence of this message enables prefix matching. +message KeywordPrefixParams { } message IntegerIndexParams { diff --git a/lib/api/src/grpc/proto/qdrant_common.proto b/lib/api/src/grpc/proto/qdrant_common.proto index 770fec247a..03f2b14a03 100644 --- a/lib/api/src/grpc/proto/qdrant_common.proto +++ b/lib/api/src/grpc/proto/qdrant_common.proto @@ -115,6 +115,8 @@ message Match { string phrase = 9; // Match any word in the text string text_any = 10; + // Match keywords starting with the given prefix + string prefix = 11; } } diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 5de955326f..7ddeb602b1 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -255,7 +255,7 @@ pub struct FieldCondition { #[derive(serde::Serialize)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Match { - #[prost(oneof = "r#match::MatchValue", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10")] + #[prost(oneof = "r#match::MatchValue", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11")] pub match_value: ::core::option::Option, } /// Nested message and enum types in `Match`. @@ -293,6 +293,9 @@ pub mod r#match { /// Match any word in the text #[prost(string, tag = "10")] TextAny(::prost::alloc::string::String), + /// Match keywords starting with the given prefix + #[prost(string, tag = "11")] + Prefix(::prost::alloc::string::String), } } #[derive(serde::Serialize)] @@ -1343,7 +1346,15 @@ pub struct KeywordIndexParams { /// Default: true. #[prost(bool, optional, tag = "3")] pub enable_hnsw: ::core::option::Option, + /// If set, enable prefix matching (`match: { "prefix": ... }`) on this field. + #[prost(message, optional, tag = "4")] + pub prefix: ::core::option::Option, } +/// Prefix matching options for the keyword index. Has no options yet: +/// presence of this message enables prefix matching. +#[derive(serde::Serialize)] +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct KeywordPrefixParams {} #[derive(serde::Serialize)] #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct IntegerIndexParams { diff --git a/lib/collection/src/problems/unindexed_field.rs b/lib/collection/src/problems/unindexed_field.rs index 8c71de6dc7..d92d0264b7 100644 --- a/lib/collection/src/problems/unindexed_field.rs +++ b/lib/collection/src/problems/unindexed_field.rs @@ -8,7 +8,9 @@ use http::{HeaderMap, HeaderValue, Method, Uri}; use issues::{Action, Code, ImmediateSolution, Issue, Solution}; use itertools::Itertools; use segment::common::operation_error::OperationError; -use segment::data_types::index::{TextIndexParams, TextIndexType}; +use segment::data_types::index::{ + KeywordIndexParams, KeywordIndexType, TextIndexParams, TextIndexType, +}; use segment::index::query_optimization::rescore_formula::parsed_formula::VariableId; use segment::json_path::JsonPath; use segment::types::{ @@ -240,6 +242,7 @@ fn infer_index_from_field_condition(field_condition: &FieldCondition) -> Vec infer_index_from_match_value(match_value), Match::Text(_match_text) => vec![FieldIndexType::Text], Match::Phrase(_match_text) => vec![FieldIndexType::TextPhrase], + Match::Prefix(_match_prefix) => vec![FieldIndexType::KeywordPrefix], Match::Any(match_any) => infer_index_from_any_variants(&match_any.any), Match::Except(match_except) => infer_index_from_any_variants(&match_except.except), Match::TextAny(_match_text_any) => vec![FieldIndexType::Text], @@ -548,6 +551,8 @@ enum FieldIndexType { IntMatch, IntRange, KeywordMatch, + /// Keyword index with the `prefix` option enabled. + KeywordPrefix, FloatRange, Text, TextPhrase, @@ -578,7 +583,12 @@ fn schema_capabilities(value: &PayloadFieldSchema) -> HashSet { PayloadSchemaType::Datetime => index_types.insert(FieldIndexType::DatetimeRange), }, PayloadFieldSchema::FieldParams(payload_schema_params) => match payload_schema_params { - PayloadSchemaParams::Keyword(_) => index_types.insert(FieldIndexType::KeywordMatch), + PayloadSchemaParams::Keyword(keyword_index_params) => { + if keyword_index_params.prefix.unwrap_or_default() { + index_types.insert(FieldIndexType::KeywordPrefix); + } + index_types.insert(FieldIndexType::KeywordMatch) + } PayloadSchemaParams::Integer(integer_index_params) => { if integer_index_params.lookup.unwrap_or(true) { index_types.insert(FieldIndexType::IntMatch); @@ -623,6 +633,13 @@ impl From for PayloadFieldSchema { FieldIndexType::KeywordMatch => { PayloadFieldSchema::FieldType(PayloadSchemaType::Keyword) } + FieldIndexType::KeywordPrefix => { + PayloadFieldSchema::FieldParams(PayloadSchemaParams::Keyword(KeywordIndexParams { + r#type: KeywordIndexType::Keyword, + prefix: Some(true), + ..Default::default() + })) + } FieldIndexType::FloatRange => PayloadFieldSchema::FieldType(PayloadSchemaType::Float), FieldIndexType::Text => PayloadFieldSchema::FieldType(PayloadSchemaType::Text), FieldIndexType::TextPhrase => { @@ -649,6 +666,36 @@ mod tests { use super::*; + #[test] + fn keyword_prefix_capabilities() { + // Plain keyword index (by type or by params) doesn't serve prefix. + let plain = PayloadFieldSchema::FieldType(PayloadSchemaType::Keyword); + let index_types = schema_capabilities(&plain); + assert!(index_types.contains(&FieldIndexType::KeywordMatch)); + assert!(!index_types.contains(&FieldIndexType::KeywordPrefix)); + + // Keyword index with `prefix` serves both exact and prefix match. + let with_prefix = + PayloadFieldSchema::FieldParams(PayloadSchemaParams::Keyword(KeywordIndexParams { + r#type: KeywordIndexType::Keyword, + prefix: Some(true), + ..Default::default() + })); + let index_types = schema_capabilities(&with_prefix); + assert!(index_types.contains(&FieldIndexType::KeywordMatch)); + assert!(index_types.contains(&FieldIndexType::KeywordPrefix)); + + // A prefix condition requires the prefix capability. + let condition = FieldCondition::new_match( + segment::json_path::JsonPath::new("url"), + segment::types::Match::new_prefix("https://"), + ); + assert_eq!( + infer_index_from_field_condition(&condition), + vec![FieldIndexType::KeywordPrefix], + ); + } + #[test] fn integer_index_capacities() { let params = PayloadSchemaParams::Integer(IntegerIndexParams { diff --git a/lib/edge/python/qdrant_edge.pyi b/lib/edge/python/qdrant_edge.pyi index ed008b29fa..82f9641aa5 100644 --- a/lib/edge/python/qdrant_edge.pyi +++ b/lib/edge/python/qdrant_edge.pyi @@ -29,7 +29,13 @@ ConditionType = Union[ "Filter", ] MatchType = Union[ - "MatchValue", "MatchText", "MatchTextAny", "MatchPhrase", "MatchAny", "MatchExcept" + "MatchValue", + "MatchText", + "MatchTextAny", + "MatchPhrase", + "MatchPrefix", + "MatchAny", + "MatchExcept", ] RangeType = Union["RangeFloat", "RangeDateTime"] QuantizationConfigType = Union[ @@ -1100,6 +1106,7 @@ class KeywordIndexParams: is_tenant: Optional[bool] = None, on_disk: Optional[bool] = None, enable_hnsw: Optional[bool] = None, + prefix: Optional[bool] = None, ) -> None: """ Create KeywordIndexParams. @@ -1108,6 +1115,7 @@ class KeywordIndexParams: is_tenant: Whether this field is used for tenant separation. on_disk: Whether to store index on disk. enable_hnsw: Whether to enable HNSW index for this field. + prefix: Whether to enable prefix matching for this field. """ ... @@ -1126,6 +1134,11 @@ class KeywordIndexParams: """Whether to enable HNSW index.""" ... + @property + def prefix(self) -> Optional[bool]: + """Whether prefix matching is enabled.""" + ... + class IntegerIndexParams: """Index parameters for integer fields.""" @@ -2768,6 +2781,23 @@ class MatchPhrase: """Phrase.""" ... +class MatchPrefix: + """Match keyword values starting with the given prefix.""" + + def __init__(self, prefix: str) -> None: + """ + Create a MatchPrefix. + + Args: + prefix: Prefix to match. + """ + ... + + @property + def prefix(self) -> str: + """Prefix.""" + ... + class MatchAny: """Match any of the values.""" diff --git a/lib/edge/python/src/lib.rs b/lib/edge/python/src/lib.rs index 9b3ec2b5a3..6d899a44e6 100644 --- a/lib/edge/python/src/lib.rs +++ b/lib/edge/python/src/lib.rs @@ -71,8 +71,8 @@ mod qdrant_edge { use super::types::filter::{ PyFieldCondition, PyFilter, PyGeoBoundingBox, PyGeoPoint, PyGeoPolygon, PyGeoRadius, PyHasIdCondition, PyHasVectorCondition, PyIsEmptyCondition, PyIsNullCondition, PyMatchAny, - PyMatchExcept, PyMatchPhrase, PyMatchText, PyMatchTextAny, PyMatchValue, PyMinShould, - PyNestedCondition, PyRangeDateTime, PyRangeFloat, PyValuesCount, + PyMatchExcept, PyMatchPhrase, PyMatchPrefix, PyMatchText, PyMatchTextAny, PyMatchValue, + PyMinShould, PyNestedCondition, PyRangeDateTime, PyRangeFloat, PyValuesCount, }; #[pymodule_export] use super::types::formula::{PyDecayKind, PyExpressionInterface, PyFormula}; diff --git a/lib/edge/python/src/types/filter/match.rs b/lib/edge/python/src/types/filter/match.rs index 9b7eda54ae..d0fdd9df6c 100644 --- a/lib/edge/python/src/types/filter/match.rs +++ b/lib/edge/python/src/types/filter/match.rs @@ -24,6 +24,7 @@ impl FromPyObject<'_, '_> for PyMatch { Text(PyMatchText), TextAny(PyMatchTextAny), Phrase(PyMatchPhrase), + Prefix(PyMatchPrefix), Any(PyMatchAny), Except(PyMatchExcept), } @@ -34,6 +35,7 @@ impl FromPyObject<'_, '_> for PyMatch { Match::Text(_) => {} Match::TextAny(_) => {} Match::Phrase(_) => {} + Match::Prefix(_) => {} Match::Any(_) => {} Match::Except(_) => {} } @@ -44,6 +46,7 @@ impl FromPyObject<'_, '_> for PyMatch { Helper::Text(text) => Match::Text(MatchText::from(text)), Helper::TextAny(text_any) => Match::TextAny(MatchTextAny::from(text_any)), Helper::Phrase(phrase) => Match::Phrase(MatchPhrase::from(phrase)), + Helper::Prefix(prefix) => Match::Prefix(MatchPrefix::from(prefix)), Helper::Any(any) => Match::Any(MatchAny::from(any)), Helper::Except(except) => Match::Except(MatchExcept::from(except)), }; @@ -63,6 +66,7 @@ impl<'py> IntoPyObject<'py> for PyMatch { Match::Text(text) => PyMatchText(text).into_bound_py_any(py), Match::TextAny(text_any) => PyMatchTextAny(text_any).into_bound_py_any(py), Match::Phrase(phrase) => PyMatchPhrase(phrase).into_bound_py_any(py), + Match::Prefix(prefix) => PyMatchPrefix(prefix).into_bound_py_any(py), Match::Any(any) => PyMatchAny(any).into_bound_py_any(py), Match::Except(except) => PyMatchExcept(except).into_bound_py_any(py), } @@ -76,6 +80,7 @@ impl Repr for PyMatch { Match::Text(text) => PyMatchText::wrap_ref(text).fmt(f), Match::TextAny(text_any) => PyMatchTextAny::wrap_ref(text_any).fmt(f), Match::Phrase(phrase) => PyMatchPhrase::wrap_ref(phrase).fmt(f), + Match::Prefix(prefix) => PyMatchPrefix::wrap_ref(prefix).fmt(f), Match::Any(any) => PyMatchAny::wrap_ref(any).fmt(f), Match::Except(except) => PyMatchExcept::wrap_ref(except).fmt(f), } @@ -267,6 +272,36 @@ impl PyMatchPhrase { } } +#[pyclass(name = "MatchPrefix", from_py_object)] +#[derive(Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PyMatchPrefix(pub MatchPrefix); + +#[pyclass_repr] +#[pymethods] +impl PyMatchPrefix { + #[new] + pub fn new(prefix: String) -> Self { + Self(MatchPrefix { prefix }) + } + + #[getter] + pub fn prefix(&self) -> &str { + &self.0.prefix + } + + pub fn __repr__(&self) -> String { + self.repr() + } +} + +impl PyMatchPrefix { + fn _getters(self) { + // Every field should have a getter method + let MatchPrefix { prefix: _ } = self.0; + } +} + #[pyclass(name = "MatchAny", from_py_object)] #[derive(Clone, Debug, Into, TransparentWrapper)] #[repr(transparent)] diff --git a/lib/edge/python/src/types/payload_schema/mod.rs b/lib/edge/python/src/types/payload_schema/mod.rs index 3a97c5e272..9f96a2e0bd 100644 --- a/lib/edge/python/src/types/payload_schema/mod.rs +++ b/lib/edge/python/src/types/payload_schema/mod.rs @@ -208,13 +208,19 @@ pub struct PyKeywordIndexParams(KeywordIndexParams); #[pymethods] impl PyKeywordIndexParams { #[new] - #[pyo3(signature = (is_tenant = None, on_disk = None, enable_hnsw = None))] - pub fn new(is_tenant: Option, on_disk: Option, enable_hnsw: Option) -> Self { + #[pyo3(signature = (is_tenant = None, on_disk = None, enable_hnsw = None, prefix = None))] + pub fn new( + is_tenant: Option, + on_disk: Option, + enable_hnsw: Option, + prefix: Option, + ) -> Self { Self(KeywordIndexParams { r#type: Default::default(), is_tenant, on_disk, enable_hnsw, + prefix, }) } @@ -232,6 +238,11 @@ impl PyKeywordIndexParams { pub fn enable_hnsw(&self) -> Option { self.0.enable_hnsw } + + #[getter] + pub fn prefix(&self) -> Option { + self.0.prefix + } } impl PyKeywordIndexParams { @@ -242,6 +253,7 @@ impl PyKeywordIndexParams { is_tenant: _, on_disk: _, enable_hnsw: _, + prefix: _, } = self.0; } } diff --git a/lib/edge/src/reexports.rs b/lib/edge/src/reexports.rs index 5fad9e57e6..3bc2375291 100644 --- a/lib/edge/src/reexports.rs +++ b/lib/edge/src/reexports.rs @@ -20,13 +20,14 @@ mod reexports_from_qdrant_crates { BinaryQuantizationQueryEncoding, CompressionRatio, Condition, DateTimeWrapper, Distance, ExtendedPointId as PointId, FieldCondition, Filter, GeoBoundingBox, GeoPoint, GeoPolygon, GeoRadius, HasIdCondition, HasVectorCondition, HnswConfig as HnswIndexConfig, - IsEmptyCondition, IsNullCondition, Match, MatchAny, MatchExcept, MatchPhrase, MatchText, - MatchTextAny, MatchValue, MinShould, MultiVectorComparator, MultiVectorConfig, Nested, - NestedCondition, Payload, PayloadFieldSchema, PayloadIndexInfo, PayloadSchemaParams, - PayloadSchemaType, PayloadSelector, PayloadSelectorExclude, PayloadSelectorInclude, - ProductQuantizationConfig, QuantizationConfig, QuantizationSearchParams, Range, - RangeInterface, ScalarQuantizationConfig, ScalarType, ScoredPoint, SearchParams, - ValueVariants, ValuesCount, VectorStorageDatatype, WithPayloadInterface, WithVector, + IsEmptyCondition, IsNullCondition, Match, MatchAny, MatchExcept, MatchPhrase, MatchPrefix, + MatchText, MatchTextAny, MatchValue, MinShould, MultiVectorComparator, MultiVectorConfig, + Nested, NestedCondition, Payload, PayloadFieldSchema, PayloadIndexInfo, + PayloadSchemaParams, PayloadSchemaType, PayloadSelector, PayloadSelectorExclude, + PayloadSelectorInclude, ProductQuantizationConfig, QuantizationConfig, + QuantizationSearchParams, Range, RangeInterface, ScalarQuantizationConfig, ScalarType, + ScoredPoint, SearchParams, ValueVariants, ValuesCount, VectorStorageDatatype, + WithPayloadInterface, WithVector, }; pub use segment::vector_storage::query::{ ContextPair, ContextQuery, DiscoverQuery, FeedbackItem, diff --git a/lib/segment/src/data_types/index.rs b/lib/segment/src/data_types/index.rs index ec7468df26..25c74a1b16 100644 --- a/lib/segment/src/data_types/index.rs +++ b/lib/segment/src/data_types/index.rs @@ -34,6 +34,11 @@ pub struct KeywordIndexParams { /// Default: true. #[serde(default, skip_serializing_if = "Option::is_none")] pub enable_hnsw: Option, + + /// If true, enable prefix matching (`match: { "prefix": ... }`) on this + /// field. Default: false. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prefix: Option, } // Integer diff --git a/lib/segment/src/index/field_index/bool_index/read_ops.rs b/lib/segment/src/index/field_index/bool_index/read_ops.rs index 6ecbf477fa..7d1f4cfa90 100644 --- a/lib/segment/src/index/field_index/bool_index/read_ops.rs +++ b/lib/segment/src/index/field_index/bool_index/read_ops.rs @@ -327,7 +327,8 @@ pub(super) fn condition_checker<'a, N: BoolIndexRead>( }) | Match::Text(_) | Match::TextAny(_) - | Match::Phrase(_) => None, + | Match::Phrase(_) + | Match::Prefix(_) => None, } } diff --git a/lib/segment/src/index/field_index/full_text_index/read_ops.rs b/lib/segment/src/index/field_index/full_text_index/read_ops.rs index 29ee7ac45c..f389de773b 100644 --- a/lib/segment/src/index/field_index/full_text_index/read_ops.rs +++ b/lib/segment/src/index/field_index/full_text_index/read_ops.rs @@ -16,8 +16,8 @@ use crate::index::field_index::{ }; use crate::index::payload_config::StorageType; use crate::types::{ - FieldCondition, Match, MatchAny, MatchExcept, MatchPhrase, MatchText, MatchTextAny, MatchValue, - PayloadKeyType, + FieldCondition, Match, MatchAny, MatchExcept, MatchPhrase, MatchPrefix, MatchText, + MatchTextAny, MatchValue, PayloadKeyType, }; impl FullTextIndexRead for FullTextIndex { @@ -231,7 +231,9 @@ pub fn filter<'a, T: FullTextIndexRead>( Match::TextAny(MatchTextAny { text_any }) => { index.parse_text_any_query(text_any, hw_counter) } - Match::Value(_) | Match::Any(_) | Match::Except(_) => return Ok(None), + Match::Value(_) | Match::Any(_) | Match::Except(_) | Match::Prefix(_) => { + return Ok(None); + } }?; let Some(parsed_query) = parsed_query_opt else { @@ -257,7 +259,9 @@ pub fn estimate_cardinality( Match::TextAny(MatchTextAny { text_any }) => { index.parse_text_any_query(text_any, hw_counter) } - Match::Value(_) | Match::Any(_) | Match::Except(_) => return Ok(None), + Match::Value(_) | Match::Any(_) | Match::Except(_) | Match::Prefix(_) => { + return Ok(None); + } }?; let Some(parsed_query) = parsed_query_opt else { @@ -316,7 +320,8 @@ pub fn condition_checker<'a, T: FullTextIndexRead>( Match::Phrase(MatchPhrase { phrase }) => (phrase, PayloadMatchQueryType::Phrase), Match::Value(MatchValue { value: _ }) | Match::Any(MatchAny { any: _ }) - | Match::Except(MatchExcept { except: _ }) => return Ok(None), + | Match::Except(MatchExcept { except: _ }) + | Match::Prefix(MatchPrefix { prefix: _ }) => return Ok(None), }; let query_opt = match query_type { @@ -376,6 +381,6 @@ pub fn special_check_condition( PayloadMatchQueryType::TextAny, hw_counter, )?), - Some(Match::Value(_) | Match::Any(_) | Match::Except(_)) | None => None, + Some(Match::Value(_) | Match::Any(_) | Match::Except(_) | Match::Prefix(_)) | None => None, }) } diff --git a/lib/segment/src/index/field_index/index_selector.rs b/lib/segment/src/index/field_index/index_selector.rs index ec439fc77e..f6daf0d885 100644 --- a/lib/segment/src/index/field_index/index_selector.rs +++ b/lib/segment/src/index/field_index/index_selector.rs @@ -71,15 +71,20 @@ impl IndexSelector<'_> { ); } - self.map_new(field, create_if_missing, deleted_points)? + self.map_new(field, create_if_missing, deleted_points, false)? .map(FieldIndex::IntMapIndex) } (PayloadIndexType::DatetimeIndex, PayloadSchemaParams::Datetime(_)) => self .numeric_new(field, create_if_missing, deleted_points)? .map(FieldIndex::DatetimeIndex), - (PayloadIndexType::KeywordIndex, PayloadSchemaParams::Keyword(_)) => self - .map_new(field, create_if_missing, deleted_points)? + (PayloadIndexType::KeywordIndex, PayloadSchemaParams::Keyword(params)) => self + .map_new( + field, + create_if_missing, + deleted_points, + params.prefix.unwrap_or_default(), + )? .map(FieldIndex::KeywordIndex), (PayloadIndexType::FloatIndex, PayloadSchemaParams::Float(_)) => self @@ -104,11 +109,11 @@ impl IndexSelector<'_> { .map(FieldIndex::BoolIndex), (PayloadIndexType::UuidIndex, PayloadSchemaParams::Uuid(_)) => self - .map_new(field, create_if_missing, deleted_points)? + .map_new(field, create_if_missing, deleted_points, false)? .map(FieldIndex::UuidMapIndex), (PayloadIndexType::UuidMapIndex, PayloadSchemaParams::Uuid(_)) => self - .map_new(field, create_if_missing, deleted_points)? + .map_new(field, create_if_missing, deleted_points, false)? .map(FieldIndex::UuidMapIndex), (PayloadIndexType::NullIndex, _) => { @@ -135,15 +140,20 @@ impl IndexSelector<'_> { deleted_points: &BitSlice, ) -> OperationResult>> { let indexes = match payload_schema.expand().as_ref() { - PayloadSchemaParams::Keyword(_) => self - .map_new(field, create_if_missing, deleted_points)? + PayloadSchemaParams::Keyword(params) => self + .map_new( + field, + create_if_missing, + deleted_points, + params.prefix.unwrap_or_default(), + )? .map(|index| vec![FieldIndex::KeywordIndex(index)]), PayloadSchemaParams::Integer(integer_params) => { let use_lookup = integer_params.lookup.unwrap_or(true); let use_range = integer_params.range.unwrap_or(true); let lookup = if use_lookup { - match self.map_new(field, create_if_missing, deleted_points)? { + match self.map_new(field, create_if_missing, deleted_points, false)? { Some(index) => Some(FieldIndex::IntMapIndex(index)), None => return Ok(None), } @@ -187,7 +197,7 @@ impl IndexSelector<'_> { .numeric_new(field, create_if_missing, deleted_points)? .map(|index| vec![FieldIndex::DatetimeIndex(index)]), PayloadSchemaParams::Uuid(_) => self - .map_new(field, create_if_missing, deleted_points)? + .map_new(field, create_if_missing, deleted_points, false)? .map(|index| vec![FieldIndex::UuidMapIndex(index)]), }; @@ -202,12 +212,13 @@ impl IndexSelector<'_> { deleted_points: &BitSlice, ) -> OperationResult> { let builders = match payload_schema.expand().as_ref() { - PayloadSchemaParams::Keyword(_) => { + PayloadSchemaParams::Keyword(params) => { vec![self.map_builder( field, FieldIndexBuilder::KeywordMmapIndex, FieldIndexBuilder::KeywordGridstoreIndex, deleted_points, + params.prefix.unwrap_or_default(), )] } PayloadSchemaParams::Integer(integer_params) => { @@ -220,6 +231,7 @@ impl IndexSelector<'_> { FieldIndexBuilder::IntMapMmapIndex, FieldIndexBuilder::IntMapGridstoreIndex, deleted_points, + false, )) } else { None @@ -274,6 +286,7 @@ impl IndexSelector<'_> { FieldIndexBuilder::UuidMmapIndex, FieldIndexBuilder::UuidGridstoreIndex, deleted_points, + false, )] } }; @@ -281,21 +294,27 @@ impl IndexSelector<'_> { Ok(builders) } + /// `prefix_index` enables the sorted key dictionary for prefix matching; + /// only meaningful for the keyword (string-keyed) index, other callers + /// pass `false`. fn map_new( &self, field: &JsonPath, create_if_missing: bool, deleted_points: &BitSlice, + prefix_index: bool, ) -> OperationResult>> where Vec<::Owned>: Blob + Send + Sync, { Ok(match self { + // The immutable variants detect prefix support from the presence + // of the prefix index file, written at build time. IndexSelector::NonAppendable { dir, is_on_disk } => { MapIndex::new_immutable(&map_dir(dir, field), *is_on_disk, deleted_points)? } IndexSelector::Appendable { dir } => { - MapIndex::new_mutable(map_dir(dir, field), create_if_missing)? + MapIndex::new_mutable(map_dir(dir, field), create_if_missing, prefix_index)? } }) } @@ -306,16 +325,22 @@ impl IndexSelector<'_> { make_mmap: fn(MapIndexMmapBuilder) -> FieldIndexBuilder, make_gridstore: fn(MapIndexGridstoreBuilder) -> FieldIndexBuilder, deleted_points: &BitSlice, + prefix_index: bool, ) -> FieldIndexBuilder where Vec<::Owned>: Blob + Send + Sync, { match self { - IndexSelector::NonAppendable { dir, is_on_disk } => make_mmap( - MapIndex::builder_immutable(&map_dir(dir, field), *is_on_disk, deleted_points), - ), + IndexSelector::NonAppendable { dir, is_on_disk } => { + make_mmap(MapIndex::builder_immutable( + &map_dir(dir, field), + *is_on_disk, + deleted_points, + prefix_index, + )) + } IndexSelector::Appendable { dir } => { - make_gridstore(MapIndex::builder_mutable(map_dir(dir, field))) + make_gridstore(MapIndex::builder_mutable(map_dir(dir, field), prefix_index)) } } } diff --git a/lib/segment/src/index/field_index/map_index/builders.rs b/lib/segment/src/index/field_index/map_index/builders.rs index e88e2cb8da..dae451f628 100644 --- a/lib/segment/src/index/field_index/map_index/builders.rs +++ b/lib/segment/src/index/field_index/map_index/builders.rs @@ -58,6 +58,7 @@ pub struct MapIndexMmapBuilder { pub(super) values_to_points: HashMap<::Owned, Vec>, pub(super) is_on_disk: bool, pub(super) deleted_points: BitVec, + pub(super) prefix_index: bool, } impl FieldIndexBuilderTrait for MapIndexMmapBuilder @@ -120,6 +121,7 @@ where self.values_to_points, populate, &self.deleted_points, + self.prefix_index, )?; let index = if self.is_on_disk { @@ -138,14 +140,19 @@ where { dir: PathBuf, index: Option>, + prefix_index: bool, } impl MapIndexGridstoreBuilder where Vec<::Owned>: Blob + Send + Sync, { - pub(super) fn new(dir: PathBuf) -> Self { - Self { dir, index: None } + pub(super) fn new(dir: PathBuf, prefix_index: bool) -> Self { + Self { + dir, + index: None, + prefix_index, + } } } @@ -163,7 +170,7 @@ where "index must be initialized exactly once", ); self.index.replace( - MapIndex::new_mutable(self.dir.clone(), true)?.ok_or_else(|| { + MapIndex::new_mutable(self.dir.clone(), true, self.prefix_index)?.ok_or_else(|| { OperationError::service_error("Failed to create mutable map index") })?, ); diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs index bc886c7b7f..77cf4e5055 100644 --- a/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/lifecycle.rs @@ -94,11 +94,21 @@ where log::warn!("Failed to clear mmap cache of immutable map index: {err}"); } + // In-RAM counterpart of the storage's prefix index: `Owned` ordering + // is required to match the byte order of the on-disk dictionary. + let sorted_keys = index.has_prefix_index().then(|| { + let mut keys: Vec<::Owned> = + value_to_points.keys().cloned().collect(); + keys.sort_unstable(); + keys + }); + let mut result = Self { value_to_points, value_to_points_container, deleted_value_to_points_container: BitVec::new(), point_to_values, + sorted_keys, indexed_points, values_count, storage: index, diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs index 6e34fad19b..3e413f3ba5 100644 --- a/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/mod.rs @@ -25,6 +25,11 @@ where pub(super) value_to_points_container: Vec, pub(super) deleted_value_to_points_container: BitVec, pub(super) point_to_values: ImmutablePointToValues<::Owned>, + /// Keys of `value_to_points` in ascending order, built at load time only + /// when the backing storage carries a prefix index; enables prefix range + /// scans without touching the storage. Keys whose postings become empty + /// through deletions are kept (lookups then yield empty/zero results). + pub(super) sorted_keys: Option::Owned>>, /// Amount of point which have at least one indexed payload value pub(super) indexed_points: usize, pub(super) values_count: usize, diff --git a/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs index d922e1e2a4..48b094dd27 100644 --- a/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs +++ b/lib/segment/src/index/field_index/map_index/immutable_map_index/read_ops.rs @@ -157,12 +157,18 @@ where value_to_points_container, deleted_value_to_points_container, point_to_values, + sorted_keys, indexed_points: _, values_count: _, storage: _, cached_ram_usage_bytes: _, } = self; + let sorted_keys_bytes = sorted_keys.as_ref().map_or(0, |keys| { + keys.capacity() * size_of::<::Owned>() + + keys.iter().map(|k| N::owned_heap_bytes(k)).sum::() + }); + let hashmap_entry_overhead = size_of::() + size_of::(); let vtp_base_bytes: usize = value_to_points.capacity() * (size_of::<::Owned>() @@ -178,6 +184,7 @@ where + vtp_heap_bytes + container_bytes + deleted_bytes + + sorted_keys_bytes + point_to_values.ram_usage_bytes() } } diff --git a/lib/segment/src/index/field_index/map_index/key.rs b/lib/segment/src/index/field_index/map_index/key.rs index 9c8fdc8772..df72272044 100644 --- a/lib/segment/src/index/field_index/map_index/key.rs +++ b/lib/segment/src/index/field_index/map_index/key.rs @@ -12,10 +12,31 @@ use crate::index::field_index::on_disk_point_to_values::StoredValue; use crate::types::{IntPayloadType, UuidIntType}; pub trait MapIndexKey: Key + StoredValue + Eq + Display + Debug { - type Owned: Borrow + Hash + Eq + Clone + FromStr + Default + 'static; + type Owned: Borrow + Hash + Eq + Ord + Clone + FromStr + Default + 'static; fn to_owned(&self) -> ::Owned; + /// Whether this key type supports the prefix index; gates writing its + /// file at build time. Only string keys support it. + const SUPPORTS_PREFIX_INDEX: bool = false; + + /// Natural byte representation used by the prefix index + /// ([`super::prefix_index`]). `None` for key types without prefix + /// matching semantics; only string keys support it. + /// + /// Byte-wise ordering of these representations must match `Ord` on + /// [`Self::Owned`] so that sorted in-memory structures and the on-disk + /// prefix index agree on key order. + fn prefix_index_bytes(&self) -> Option<&[u8]> { + None + } + + /// Whether a key of this type starts with the given prefix key. `false` + /// for key types without prefix matching semantics. + fn starts_with(&self, _prefix: &Self) -> bool { + false + } + /// Convert a [`FacetValue`] into this key's owned type, or `None` if the /// variant doesn't match (e.g. a keyword value against an integer index). fn from_facet_value(value: FacetValue) -> Option<::Owned>; @@ -38,6 +59,16 @@ impl MapIndexKey for str { EcoString::from(self) } + const SUPPORTS_PREFIX_INDEX: bool = true; + + fn prefix_index_bytes(&self) -> Option<&[u8]> { + Some(self.as_bytes()) + } + + fn starts_with(&self, prefix: &Self) -> bool { + str::starts_with(self, prefix) + } + fn from_facet_value(value: FacetValue) -> Option<::Owned> { match value { FacetValue::Keyword(keyword) => Some(EcoString::from(keyword)), diff --git a/lib/segment/src/index/field_index/map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/lifecycle.rs index 79d8c41572..e396ad6507 100644 --- a/lib/segment/src/index/field_index/map_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/map_index/lifecycle.rs @@ -46,8 +46,12 @@ where Ok(Some(index)) } - pub fn new_mutable(dir: PathBuf, create_if_missing: bool) -> OperationResult> { - let index = MutableMapIndex::open_gridstore(dir, create_if_missing)?; + pub fn new_mutable( + dir: PathBuf, + create_if_missing: bool, + prefix_index: bool, + ) -> OperationResult> { + let index = MutableMapIndex::open_gridstore(dir, create_if_missing, prefix_index)?; Ok(index.map(MapIndex::Mutable)) } @@ -55,6 +59,7 @@ where path: &Path, is_on_disk: bool, deleted_points: &BitSlice, + prefix_index: bool, ) -> MapIndexMmapBuilder { MapIndexMmapBuilder { path: path.to_owned(), @@ -62,11 +67,15 @@ where values_to_points: Default::default(), is_on_disk, deleted_points: deleted_points.to_owned(), + prefix_index, } } - pub fn builder_mutable(dir: PathBuf) -> super::builders::MapIndexGridstoreBuilder { - super::builders::MapIndexGridstoreBuilder::new(dir) + pub fn builder_mutable( + dir: PathBuf, + prefix_index: bool, + ) -> super::builders::MapIndexGridstoreBuilder { + super::builders::MapIndexGridstoreBuilder::new(dir, prefix_index) } pub(crate) fn flusher(&self) -> Flusher { diff --git a/lib/segment/src/index/field_index/map_index/mod.rs b/lib/segment/src/index/field_index/map_index/mod.rs index 96ee97dbcb..7f7e93e83b 100644 --- a/lib/segment/src/index/field_index/map_index/mod.rs +++ b/lib/segment/src/index/field_index/map_index/mod.rs @@ -16,6 +16,7 @@ mod lifecycle; pub mod mutable_map_index; pub mod on_disk_map_index; mod payload_index_impl; +pub mod prefix_index; pub mod read_ops; #[cfg(test)] mod tests; diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/in_memory.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/in_memory.rs index 8ae64e4558..49a1cc2d4b 100644 --- a/lib/segment/src/index/field_index/map_index/mutable_map_index/in_memory.rs +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/in_memory.rs @@ -1,5 +1,6 @@ use std::borrow::{Borrow, Cow}; -use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::collections::{BTreeSet, HashMap}; use std::iter; use common::counter::hardware_counter::HardwareCounterCell; @@ -28,18 +29,24 @@ where /// Amount of point which have at least one indexed payload value pub(in crate::index::field_index::map_index) indexed_points: usize, pub(in crate::index::field_index::map_index) values_count: usize, + /// Ordered mirror of `map`'s keys, maintained only when the index is + /// built with the `prefix` option; enables prefix range scans. Like + /// `map`, keys are never removed when their postings become empty. + pub(in crate::index::field_index::map_index) sorted_keys: + Option::Owned>>, } impl InMemoryMapIndex where Vec<::Owned>: Blob + Send + Sync, { - pub(in crate::index::field_index::map_index) fn empty() -> Self { + pub(in crate::index::field_index::map_index) fn empty(prefix_index: bool) -> Self { Self { map: HashMap::new(), point_to_values: Vec::new(), indexed_points: 0, values_count: 0, + sorted_keys: prefix_index.then(BTreeSet::new), } } @@ -57,6 +64,9 @@ where for value in values { let entry = self.map.entry(value.clone()); + if let (Entry::Vacant(_), Some(sorted_keys)) = (&entry, &mut self.sorted_keys) { + sorted_keys.insert(value.clone()); + } let inserted = entry.or_default().insert(idx); // only insert into forward index if it is not a duplicate @@ -214,8 +224,14 @@ where point_to_values, indexed_points: _, values_count: _, + sorted_keys, } = self; + let sorted_keys_bytes = sorted_keys.as_ref().map_or(0, |keys| { + keys.len() * size_of::<::Owned>() + + keys.iter().map(|k| N::owned_heap_bytes(k)).sum::() + }); + let hashmap_entry_overhead = std::mem::size_of::() + std::mem::size_of::(); let map_base_bytes = map.capacity() * (std::mem::size_of::<::Owned>() @@ -231,6 +247,6 @@ where .iter() .map(|v| v.capacity() * std::mem::size_of::<::Owned>()) .sum::(); - map_bytes + ptv_bytes + map_bytes + ptv_bytes + sorted_keys_bytes } } diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs index 84b3f81eb3..14e504c2a9 100644 --- a/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/lifecycle.rs @@ -33,7 +33,15 @@ where /// The `create_if_missing` parameter indicates whether to create a new Gridstore if it does /// not exist. If false and files don't exist, the load function will indicate nothing could be /// loaded. - pub fn open_gridstore(path: PathBuf, create_if_missing: bool) -> OperationResult> { + /// + /// `prefix_index` enables in-memory prefix range scans; it is not + /// persisted, so it must be re-supplied (from the payload schema) on every + /// open. + pub fn open_gridstore( + path: PathBuf, + create_if_missing: bool, + prefix_index: bool, + ) -> OperationResult> { let store = if create_if_missing { let options = default_gridstore_options(N::gridstore_block_size()); Gridstore::open_or_create(MmapFs, path, options, Populate::Blocking).map_err(|err| { @@ -53,7 +61,7 @@ where }; // Load in-memory index from Gridstore - let mut in_memory_index = InMemoryMapIndex::::empty(); + let mut in_memory_index = InMemoryMapIndex::::empty(prefix_index); let hw_counter = HardwareCounterCell::disposable(); let hw_counter_ref = hw_counter.ref_payload_index_io_write_counter(); diff --git a/lib/segment/src/index/field_index/map_index/mutable_map_index/read_only/lifecycle.rs b/lib/segment/src/index/field_index/map_index/mutable_map_index/read_only/lifecycle.rs index 4a073d87d2..0b965f37c3 100644 --- a/lib/segment/src/index/field_index/map_index/mutable_map_index/read_only/lifecycle.rs +++ b/lib/segment/src/index/field_index/map_index/mutable_map_index/read_only/lifecycle.rs @@ -41,7 +41,10 @@ where return Ok(None); }; - let mut in_memory_index = InMemoryMapIndex::::empty(); + // Prefix support is not wired for the read-only appendable variant: + // the gridstore carries no prefix marker, and this open path has no + // schema access. Prefix conditions fall back to slower checks. + let mut in_memory_index = InMemoryMapIndex::::empty(false); let hw_counter = HardwareCounterCell::disposable(); storage.iter::<_, GridstoreError>( storage.max_point_offset(), diff --git a/lib/segment/src/index/field_index/map_index/on_disk_map_index/lifecycle.rs b/lib/segment/src/index/field_index/map_index/on_disk_map_index/lifecycle.rs index 3fed9dea6f..eda6e0c437 100644 --- a/lib/segment/src/index/field_index/map_index/on_disk_map_index/lifecycle.rs +++ b/lib/segment/src/index/field_index/map_index/on_disk_map_index/lifecycle.rs @@ -15,6 +15,7 @@ use common::universal_io::{ use fs_err as fs; use super::super::MapIndexKey; +use super::super::prefix_index::{PREFIX_INDEX_PATH, PrefixIndex, build_prefix_index}; use super::{ CONFIG_PATH, DELETED_PATH, HASHMAP_PATH, OnDiskMapIndex, Storage, UniversalMapIndexConfig, }; @@ -57,6 +58,7 @@ where Default::default(), )?; let point_to_values = OnDiskPointToValues::open(fs, path, populate)?; + let prefix_index = PrefixIndex::open(fs, path, populate)?; let mut deleted = deleted_points.to_owned(); @@ -90,6 +92,7 @@ where value_to_points, point_to_values, deleted, + prefix_index, }, deleted_count, total_key_value_pairs: config.total_key_value_pairs, @@ -114,6 +117,9 @@ where self.path.join(DELETED_PATH), self.path.join(CONFIG_PATH), ]; + if self.storage.prefix_index.is_some() { + files.push(self.path.join(PREFIX_INDEX_PATH)); + } files.extend(self.storage.point_to_values.files()); files } @@ -124,6 +130,9 @@ where self.path.join(DELETED_PATH), self.path.join(CONFIG_PATH), ]; + if self.storage.prefix_index.is_some() { + files.push(self.path.join(PREFIX_INDEX_PATH)); + } files.extend(self.storage.point_to_values.immutable_files()); files } @@ -133,6 +142,9 @@ where pub fn populate(&self) -> OperationResult<()> { self.storage.value_to_points.populate()?; self.storage.point_to_values.populate()?; + if let Some(prefix_index) = &self.storage.prefix_index { + prefix_index.populate()?; + } Ok(()) } @@ -148,13 +160,23 @@ where value_to_points, point_to_values, deleted: _, + prefix_index, } = storage; value_to_points.clear_ram_cache()?; clear_disk_cache(&path.join(DELETED_PATH))?; point_to_values.clear_cache()?; + if let Some(prefix_index) = prefix_index { + prefix_index.clear_cache()?; + } Ok(()) } + /// Whether this index was built with the prefix option (its sorted key + /// dictionary is present on disk). + pub(in super::super) fn has_prefix_index(&self) -> bool { + self.storage.prefix_index.is_some() + } + pub(crate) fn ram_usage_bytes(&self) -> usize { self.storage.ram_usage_bytes() } @@ -173,6 +195,7 @@ where values_to_points: HashMap<::Owned, Vec>, populate: Populate, deleted_points: &BitSlice, + with_prefix_index: bool, ) -> OperationResult { fs::create_dir_all(path)?; @@ -194,6 +217,20 @@ where .map(|(value, ids)| (value.borrow(), ids.iter().copied())), )?; + if with_prefix_index && N::SUPPORTS_PREFIX_INDEX { + // The file is written (even when empty) so that its presence + // signals prefix support at load time. + let mut entries: Vec<(&[u8], usize)> = values_to_points + .iter() + .filter_map(|(value, ids)| { + let value: &N = value.borrow(); + value.prefix_index_bytes().map(|bytes| (bytes, ids.len())) + }) + .collect(); + entries.sort_unstable_by_key(|&(key, _count)| key); + build_prefix_index(path, entries.into_iter())?; + } + OnDiskPointToValues::::build_from_iter( path, point_to_values.iter().enumerate().map(|(idx, values)| { diff --git a/lib/segment/src/index/field_index/map_index/on_disk_map_index/mod.rs b/lib/segment/src/index/field_index/map_index/on_disk_map_index/mod.rs index de7f18c03e..81e604ac6c 100644 --- a/lib/segment/src/index/field_index/map_index/on_disk_map_index/mod.rs +++ b/lib/segment/src/index/field_index/map_index/on_disk_map_index/mod.rs @@ -7,6 +7,7 @@ use common::universal_io::{MmapFile, UniversalRead}; use serde::{Deserialize, Serialize}; use super::MapIndexKey; +use super::prefix_index::PrefixIndex; use crate::index::field_index::on_disk_point_to_values::OnDiskPointToValues; mod lifecycle; @@ -43,6 +44,10 @@ pub(super) struct Storage>, } impl Storage { @@ -51,10 +56,15 @@ impl Storage { value_to_points: _, point_to_values, deleted, + prefix_index, } = self; // `value_to_points` is a storage-backed hashmap with no in-memory state. - point_to_values.ram_usage_bytes() + deleted.capacity().div_ceil(u8::BITS as usize) + point_to_values.ram_usage_bytes() + + deleted.capacity().div_ceil(u8::BITS as usize) + + prefix_index + .as_ref() + .map_or(0, |prefix_index| prefix_index.ram_usage_bytes()) } } diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs index 4c3124ceb8..183b4d3219 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/int.rs @@ -279,6 +279,7 @@ fn condition_checker_impl<'a, T: MapIndexRead<'a, IntPayloadType> + 'a>( }) | Match::Text(_) | Match::TextAny(_) - | Match::Phrase(_) => None, + | Match::Phrase(_) + | Match::Prefix(_) => None, } } diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs index 596f37a475..a3a89c725f 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/str.rs @@ -1,13 +1,16 @@ +use std::collections::HashMap; use std::iter; use std::path::PathBuf; use common::counter::hardware_accumulator::HwMeasurementAcc; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; +use ecow::EcoString; use gridstore::Blob; use super::super::MapIndex; use super::super::key::MapIndexKey; +use super::super::prefix_index::{PrefixIndexStats, StrMapIndexPrefixRead}; use super::super::read_only::ReadOnlyMapIndex; use super::super::read_ops::{MapConditionChecker, MapIndexRead}; use crate::common::Flusher; @@ -15,14 +18,15 @@ use crate::common::operation_error::OperationResult; use crate::index::UniversalReadExt; use crate::index::condition_checker::ConditionCheckerEnum; use crate::index::field_index::map_index::IdIter; +use crate::index::field_index::stat_tools::number_of_selected_points; use crate::index::field_index::{ CardinalityEstimation, PayloadBlockCondition, PayloadFieldIndex, PayloadFieldIndexRead, PrimaryCondition, }; use crate::index::query_estimator::combine_should_estimations; use crate::types::{ - AnyVariants, FieldCondition, Match, MatchAny, MatchExcept, MatchValue, PayloadKeyType, - ValueVariants, + AnyVariants, FieldCondition, Match, MatchAny, MatchExcept, MatchPrefix, MatchValue, + PayloadKeyType, ValueVariants, }; impl PayloadFieldIndex for MapIndex { @@ -61,7 +65,7 @@ impl PayloadFieldIndexRead for MapIndex { condition: &FieldCondition, hw_counter: &HardwareCounterCell, ) -> OperationResult> { - Ok(estimate_cardinality_impl(self, condition, hw_counter)) + estimate_cardinality_impl(self, condition, hw_counter) } fn for_each_payload_block( @@ -104,7 +108,7 @@ where condition: &FieldCondition, hw_counter: &HardwareCounterCell, ) -> OperationResult> { - Ok(estimate_cardinality_impl(self, condition, hw_counter)) + estimate_cardinality_impl(self, condition, hw_counter) } fn for_each_payload_block( @@ -129,7 +133,7 @@ where // over `T: MapIndexRead` so a single body serves both `PayloadFieldIndexRead` // impls above. -fn filter_impl<'a, T: MapIndexRead<'a, str>>( +fn filter_impl<'a, T: MapIndexRead<'a, str> + StrMapIndexPrefixRead>( index: &'a T, condition: &'a FieldCondition, hw_counter: &'a HardwareCounterCell, @@ -158,18 +162,28 @@ fn filter_impl<'a, T: MapIndexRead<'a, str>>( AnyVariants::Strings(keywords) => Some(index.except_set(keywords, hw_counter)?), AnyVariants::Integers(_) => None, }, + Some(Match::Prefix(MatchPrefix { prefix })) => { + // `None` when this index instance has no prefix structure — the + // caller then falls back to the generic (slow) condition check. + match index.prefix_keys_with_counts(prefix, hw_counter)? { + Some(keys) => Some( + index.iter_for_values(keys.into_iter().map(|(key, _count)| key), hw_counter)?, + ), + None => None, + } + } _ => None, }; Ok(result) } -fn estimate_cardinality_impl<'a, T: MapIndexRead<'a, str>>( +fn estimate_cardinality_impl<'a, T: MapIndexRead<'a, str> + StrMapIndexPrefixRead>( index: &'a T, condition: &FieldCondition, hw_counter: &HardwareCounterCell, -) -> Option { - match &condition.r#match { +) -> OperationResult> { + let estimation = match &condition.r#match { Some(Match::Value(MatchValue { value })) => match value { ValueVariants::String(keyword) => { let mut estimation = index.match_cardinality(keyword.as_str(), hw_counter); @@ -214,16 +228,60 @@ fn estimate_cardinality_impl<'a, T: MapIndexRead<'a, str>>( } AnyVariants::Integers(_) => None, }, + Some(Match::Prefix(MatchPrefix { prefix })) => { + index.prefix_stats(prefix, hw_counter)?.map(|stats| { + prefix_cardinality(index, stats) + .with_primary_clause(PrimaryCondition::Condition(Box::new(condition.clone()))) + }) + } _ => None, + }; + Ok(estimation) +} + +/// Cardinality of a prefix match from aggregate `(keys, postings)` stats. +/// +/// A prefix selects the union of its keys' postings. The sum of counts is an +/// upper bound (a point with several values sharing the prefix is counted +/// once per value); the union cannot be smaller than the largest single +/// posting, which is at least the average. The on-disk stats are build-time +/// counts, so with deletions this is an estimate, like other on-disk +/// count-based estimations. +fn prefix_cardinality<'a, T: MapIndexRead<'a, str>>( + index: &'a T, + stats: PrefixIndexStats, +) -> CardinalityEstimation { + let PrefixIndexStats { keys, postings } = stats; + let indexed_points = index.get_indexed_points(); + + let sum = postings.min(index.get_values_count()); + let max = sum.min(indexed_points); + let min = if keys == 0 { + 0 + } else { + sum.div_ceil(keys).min(max) + }; + let exp = number_of_selected_points(indexed_points, sum).clamp(min, max); + + CardinalityEstimation { + primary_clauses: vec![], + min, + exp, + max, } } -fn for_each_payload_block_impl<'a, T: MapIndexRead<'a, str>>( +fn for_each_payload_block_impl<'a, T: MapIndexRead<'a, str> + StrMapIndexPrefixRead>( index: &'a T, threshold: usize, key: PayloadKeyType, f: &mut dyn FnMut(PayloadBlockCondition) -> OperationResult<()>, ) -> OperationResult<()> { + // Prefix blocks are disjoint from each other and from the exact-value + // blocks (see `heavy_prefix_blocks`), so emission order doesn't matter; + // blocks come largest-first for determinism. + for_each_prefix_payload_block(index, threshold, &key, f)?; + index.for_each_value(|value| { let count = index // `for_each_payload_block` is only used while building HNSW, which @@ -240,6 +298,187 @@ fn for_each_payload_block_impl<'a, T: MapIndexRead<'a, str>>( }) } +/// Emit a payload block per "heavy" prefix, so HNSW builds additional links +/// guaranteeing a navigable subgraph for prefix-filtered searches. +/// +/// No-op for indexes built without the prefix option. +fn for_each_prefix_payload_block<'a, T: MapIndexRead<'a, str> + StrMapIndexPrefixRead>( + index: &'a T, + threshold: usize, + key: &PayloadKeyType, + f: &mut dyn FnMut(PayloadBlockCondition) -> OperationResult<()>, +) -> OperationResult<()> { + // HNSW build; hardware measurement intentionally bypassed (see above). + let hw_counter = HardwareCounterCell::disposable(); + let Some(entries) = index.prefix_keys_with_counts("", &hw_counter)? else { + return Ok(()); + }; + + for (prefix, cardinality) in heavy_prefix_blocks(&entries, threshold) { + f(PayloadBlockCondition { + condition: FieldCondition::new_match(key.clone(), Match::new_prefix(&prefix)), + cardinality, + })?; + } + Ok(()) +} + +/// Enumerate the *smallest* heavy prefixes over sorted `(key, count)` +/// entries — the same principle as the geo index's `large_hashes`, which +/// emits only the deepest geohash regions above the threshold. +/// +/// A prefix qualifies when it covers at least two keys and more postings +/// than the threshold, and it is emitted only if nothing heavy is nested +/// inside it: neither a longer qualifying prefix nor a single heavy value +/// (which already gets its own exact-match block). The emitted blocks are +/// therefore mutually disjoint and disjoint from the exact-value blocks — +/// no subgraph is built twice for nested subsets. Broader prefix queries +/// span several blocks, each internally navigable, exactly like a large +/// bounding box spans several geohash tile blocks. +/// +/// All prefixes along a single-child trie chain select the same keys; each +/// chain collapses to its longest prefix — the longest common prefix (LCP) +/// of its key range. +/// +/// Single streaming pass with a stack of open LCP intervals: an interval +/// opens when consecutive keys share a longer prefix and closes when the +/// shared length drops; totals and coverage propagate from closed intervals +/// into their parents. O(total key bytes). +fn heavy_prefix_blocks(entries: &[(EcoString, usize)], threshold: usize) -> Vec<(String, usize)> { + struct OpenInterval { + /// Prefix length (in bytes) shared by every key in the interval. + lcp: usize, + keys: usize, + postings: usize, + /// Whether something heavy nested in this interval already produced + /// a block (a deeper interval or a single heavy value). + covered: bool, + } + + let lcp_len = |a: &str, b: &str| { + a.as_bytes() + .iter() + .zip(b.as_bytes()) + .take_while(|(x, y)| x == y) + .count() + }; + + // The byte-level LCP may end mid-codepoint; round down to a char boundary + // to emit a valid prefix string. Rounding can collapse a node onto its + // ancestor, so keep the larger count per emitted prefix. + let mut heavy: HashMap = HashMap::new(); + // Returns whether the interval's subtree now contains an emitted block, + // i.e. whether its ancestors must be suppressed. + let mut close = |last_key: &str, node: &OpenInterval| -> bool { + if node.keys < 2 || node.postings <= threshold { + return false; + } + if node.covered { + return true; + } + let mut boundary = node.lcp; + while boundary > 0 && !last_key.is_char_boundary(boundary) { + boundary -= 1; + } + if boundary == 0 { + // Heavy but unrepresentable at this depth; let an ancestor with + // a valid boundary produce the block instead. + return false; + } + let entry = heavy.entry(last_key[..boundary].to_string()).or_default(); + *entry = (*entry).max(node.postings); + true + }; + + let n = entries.len(); + let mut stack = vec![OpenInterval { + lcp: 0, + keys: 0, + postings: 0, + covered: false, + }]; + for i in 0..n { + let (key, count) = &entries[i]; + let lcp_left = if i == 0 { + 0 + } else { + lcp_len(&entries[i - 1].0, key) + }; + let lcp_right = if i + 1 == n { + 0 + } else { + lcp_len(key, &entries[i + 1].0) + }; + + // Close intervals deeper than what the current key shares with its + // predecessor: they ended at the previous key. + let mut carry_keys = 0; + let mut carry_postings = 0; + let mut carry_covered = false; + while stack.last().is_some_and(|top| top.lcp > lcp_left) { + let mut node = stack.pop().unwrap(); + node.keys += carry_keys; + node.postings += carry_postings; + node.covered |= carry_covered; + carry_covered = close(&entries[i - 1].0, &node); + carry_keys = node.keys; + carry_postings = node.postings; + } + // Merge the closed subtree's totals into its parent at `lcp_left`. + let top = stack.last_mut().unwrap(); + if top.lcp == lcp_left { + top.keys += carry_keys; + top.postings += carry_postings; + top.covered |= carry_covered; + } else { + stack.push(OpenInterval { + lcp: lcp_left, + keys: carry_keys, + postings: carry_postings, + covered: carry_covered, + }); + } + // A deeper interval opens if the next key shares more than the + // previous one did. + if lcp_right > lcp_left { + stack.push(OpenInterval { + lcp: lcp_right, + keys: 0, + postings: 0, + covered: false, + }); + } + // The key itself belongs to the deepest open interval. A single + // heavy value gets its own exact-match block, which counts as + // covering every prefix above it. + let top = stack.last_mut().unwrap(); + top.keys += 1; + top.postings += count; + top.covered |= *count > threshold; + } + // Close everything still open; those intervals end at the last key. + let mut carry_keys = 0; + let mut carry_postings = 0; + let mut carry_covered = false; + while let Some(mut node) = stack.pop() { + node.keys += carry_keys; + node.postings += carry_postings; + node.covered |= carry_covered; + carry_covered = match entries.last() { + Some((last_key, _)) => close(last_key, &node), + None => false, + }; + carry_keys = node.keys; + carry_postings = node.postings; + } + + let mut blocks: Vec<(String, usize)> = heavy.into_iter().collect(); + blocks.sort_unstable_by(|(prefix_a, count_a), (prefix_b, count_b)| { + count_b.cmp(count_a).then_with(|| prefix_a.cmp(prefix_b)) + }); + blocks +} + fn condition_checker_impl<'a, T: MapIndexRead<'a, str> + 'a>( index: &'a T, condition: &FieldCondition, @@ -271,6 +510,11 @@ fn condition_checker_impl<'a, T: MapIndexRead<'a, str> + 'a>( Match::Except(MatchExcept { except: AnyVariants::Strings(list), }) => Some(index.match_any_checker(hw_counter, list.clone(), true)), + // Served through the forward index; works with or without the prefix + // structures, which only accelerate `filter`/`estimate_cardinality`. + Match::Prefix(MatchPrefix { prefix }) => { + Some(index.match_prefix_checker(hw_counter, prefix.as_str())) + } // Conditions this index can't serve: Match::Text/TextAny/Phrase // (handled by FullTextIndex) and value-type mismatches (e.g. // Match::Value(Integer) against a string-keyed map). diff --git a/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs b/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs index 0af03e75a0..4fc89a6008 100644 --- a/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs +++ b/lib/segment/src/index/field_index/map_index/payload_index_impl/uuid.rs @@ -351,6 +351,7 @@ fn condition_checker_impl<'a, T: MapIndexRead<'a, UuidIntType> + 'a>( }) | Match::Text(_) | Match::TextAny(_) - | Match::Phrase(_) => None, + | Match::Phrase(_) + | Match::Prefix(_) => None, } } diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/format.rs b/lib/segment/src/index/field_index/map_index/prefix_index/format.rs new file mode 100644 index 0000000000..2ab0f554e6 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/format.rs @@ -0,0 +1,84 @@ +//! On-disk layout primitives shared by the [`writer`](super::writer) and the +//! [`reader`](super::reader). See the [module docs](super) for the file +//! format diagram. +//! +//! All records are fixed-size, little-endian [`Pod`] structs, written with +//! [`bytemuck::bytes_of`] and read back with [`read_record`] (a copying +//! [`bytemuck::pod_read_unaligned`]), so they carry no alignment requirement +//! on their position in the file. + +use bytemuck::{Pod, Zeroable}; + +pub(super) const MAGIC: [u8; 8] = *b"QdrPrfx\0"; +pub(super) const VERSION: u32 = 1; + +/// Target size of one front-coded key block. +pub(super) const BLOCK_SIZE_TARGET: usize = 4096; + +#[derive(Debug, Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub(super) struct Header { + pub(super) magic: [u8; 8], + pub(super) version: u32, + pub(super) _reserved: u32, + pub(super) key_count: u64, + pub(super) block_count: u64, + /// Size in bytes of the block index section (which starts right after the + /// header). + pub(super) block_index_size: u64, +} + +/// Per-block record of the block index section, followed in the file by the +/// block's first key (`first_key_len` bytes). +#[derive(Debug, Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub(super) struct BlockEntry { + /// Size in bytes of the block in the key blocks section. + pub(super) block_size: u32, + /// Number of keys in the block. + pub(super) key_count: u32, + pub(super) first_key_len: u32, + pub(super) _reserved: u32, + /// Sum of postings counts over the block. + pub(super) postings_count: u64, +} + +/// Per-key record within a key block, followed by the key's suffix +/// (`suffix_len` bytes). Keys are front-coded against their predecessor; +/// the first key of a block has `shared_prefix_len == 0`. +#[derive(Debug, Clone, Copy, Pod, Zeroable)] +#[repr(C)] +pub(super) struct KeyEntry { + pub(super) shared_prefix_len: u32, + pub(super) suffix_len: u32, + pub(super) postings_count: u32, +} + +/// Read one fixed-size record from the front of `bytes`; returns the record +/// and the remaining bytes, or `None` if `bytes` is too short. +pub(super) fn read_record(bytes: &[u8]) -> Option<(T, &[u8])> { + let (record, rest) = bytes.split_at_checked(size_of::())?; + Some((bytemuck::pod_read_unaligned(record), rest)) +} + +/// The smallest byte string greater than every string starting with `prefix`, +/// or `None` if no such string exists (empty prefix or all `0xFF`). +pub(super) fn prefix_successor(prefix: &[u8]) -> Option> { + let last_incrementable = prefix.iter().rposition(|&byte| byte != u8::MAX)?; + let mut successor = prefix[..=last_incrementable].to_vec(); + successor[last_incrementable] += 1; + Some(successor) +} + +/// Where `key` lies relative to the range of keys starting with `prefix`. +pub(super) fn key_vs_prefix_range(key: &[u8], prefix: &[u8]) -> std::cmp::Ordering { + if key.starts_with(prefix) { + std::cmp::Ordering::Equal + } else { + key.cmp(prefix) + } +} + +pub(super) fn common_prefix_len(a: &[u8], b: &[u8]) -> usize { + a.iter().zip(b).take_while(|(x, y)| x == y).count() +} diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/map_read.rs b/lib/segment/src/index/field_index/map_index/prefix_index/map_read.rs new file mode 100644 index 0000000000..c6a5a3d871 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/map_read.rs @@ -0,0 +1,269 @@ +//! Prefix enumeration surface of the string-keyed map index. +//! +//! Counterpart of [`MapIndexRead`][1] for the [`Match::Prefix`][2] condition: +//! every variant that carries a prefix structure (the mutable `BTreeSet`, the +//! immutable sorted key vector, or the on-disk [`PrefixIndex`][3]) enumerates +//! keys by byte-prefix through this trait. Variants without the structure +//! (built without the `prefix` option, or loaded from legacy files) return +//! `None`, which makes the caller fall back to the generic slow paths. +//! +//! [1]: super::super::read_ops::MapIndexRead +//! [2]: crate::types::Match::Prefix +//! [3]: super::reader::PrefixIndex + +use std::ops::Bound; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::universal_io::UniversalRead; +use ecow::EcoString; +use gridstore::Blob; + +use super::super::MapIndex; +use super::super::immutable_map_index::ImmutableMapIndex; +use super::super::key::MapIndexKey; +use super::super::mutable_map_index::MutableMapIndex; +use super::super::mutable_map_index::in_memory::InMemoryMapIndex; +use super::super::on_disk_map_index::OnDiskMapIndex; +use super::super::read_only::ReadOnlyMapIndex; +use super::super::read_ops::MapIndexRead; +use super::reader::PrefixIndexStats; +use crate::common::operation_error::{OperationError, OperationResult}; + +/// Prefix range scans over the keys of a string-keyed map index. +/// +/// All methods return `Ok(None)` when the index instance has no prefix +/// structure; policy (whether that should have been possible) is decided +/// upstream, per the payload schema. +pub trait StrMapIndexPrefixRead { + /// Keys starting with `prefix` in ascending byte order, each with its + /// postings count. + /// + /// Counts are exact for in-RAM variants; the on-disk variant reports + /// build-time counts, which ignore points deleted since the index was + /// built. The actual postings iteration (via + /// [`MapIndexRead::get_iterator`]) always filters deleted points, so + /// this only affects estimates. + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>>; + + /// Aggregate `(distinct keys, postings sum)` over keys starting with + /// `prefix`. Same count semantics as + /// [`Self::prefix_keys_with_counts`]; the on-disk variant computes this + /// from per-block aggregates, decoding only boundary blocks. + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>; +} + +impl StrMapIndexPrefixRead for InMemoryMapIndex { + fn prefix_keys_with_counts( + &self, + prefix: &str, + _hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + let Some(sorted_keys) = &self.sorted_keys else { + return Ok(None); + }; + let keys = sorted_keys + .range::((Bound::Included(prefix), Bound::Unbounded)) + .take_while(|key| key.starts_with(prefix)) + .map(|key| { + let count = self + .map + .get(key.as_str()) + .map_or(0, |ids| ids.len() as usize); + (key.clone(), count) + }) + .collect(); + Ok(Some(keys)) + } + + fn prefix_stats( + &self, + prefix: &str, + _hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + let Some(sorted_keys) = &self.sorted_keys else { + return Ok(None); + }; + let mut stats = PrefixIndexStats::default(); + for key in sorted_keys + .range::((Bound::Included(prefix), Bound::Unbounded)) + .take_while(|key| key.starts_with(prefix)) + { + stats.keys += 1; + stats.postings += self + .map + .get(key.as_str()) + .map_or(0, |ids| ids.len() as usize); + } + Ok(Some(stats)) + } +} + +impl StrMapIndexPrefixRead for MutableMapIndex { + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + self.in_memory_index + .prefix_keys_with_counts(prefix, hw_counter) + } + + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + self.in_memory_index.prefix_stats(prefix, hw_counter) + } +} + +impl StrMapIndexPrefixRead for ImmutableMapIndex { + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + let Some(sorted_keys) = &self.sorted_keys else { + return Ok(None); + }; + let start = sorted_keys.partition_point(|key| key.as_str() < prefix); + let keys = sorted_keys[start..] + .iter() + .take_while(|key| key.starts_with(prefix)) + .map(|key| { + // Live count; `None` means every posting of the key has been + // deleted since load. + let count = self + .get_count_for_value(key.as_str(), hw_counter) + .unwrap_or(0); + (key.clone(), count) + }) + .collect(); + Ok(Some(keys)) + } + + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + let Some(sorted_keys) = &self.sorted_keys else { + return Ok(None); + }; + let start = sorted_keys.partition_point(|key| key.as_str() < prefix); + let mut stats = PrefixIndexStats::default(); + for key in sorted_keys[start..] + .iter() + .take_while(|key| key.starts_with(prefix)) + { + stats.keys += 1; + stats.postings += self + .get_count_for_value(key.as_str(), hw_counter) + .unwrap_or(0); + } + Ok(Some(stats)) + } +} + +impl StrMapIndexPrefixRead for OnDiskMapIndex { + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + let Some(prefix_index) = &self.storage.prefix_index else { + return Ok(None); + }; + let mut keys = Vec::new(); + prefix_index.for_each_key_with_prefix( + prefix.as_bytes(), + hw_counter, + &mut |key, count| { + let key = std::str::from_utf8(key).map_err(|_| { + OperationError::service_error("Prefix index contains non-UTF-8 key") + })?; + keys.push((EcoString::from(key), count)); + Ok(()) + }, + )?; + Ok(Some(keys)) + } + + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + let Some(prefix_index) = &self.storage.prefix_index else { + return Ok(None); + }; + Ok(Some( + prefix_index.prefix_stats(prefix.as_bytes(), hw_counter)?, + )) + } +} + +impl StrMapIndexPrefixRead for MapIndex { + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + match self { + MapIndex::Mutable(index) => index.prefix_keys_with_counts(prefix, hw_counter), + MapIndex::Immutable(index) => index.prefix_keys_with_counts(prefix, hw_counter), + MapIndex::OnDisk(index) => index.prefix_keys_with_counts(prefix, hw_counter), + } + } + + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + match self { + MapIndex::Mutable(index) => index.prefix_stats(prefix, hw_counter), + MapIndex::Immutable(index) => index.prefix_stats(prefix, hw_counter), + MapIndex::OnDisk(index) => index.prefix_stats(prefix, hw_counter), + } + } +} + +impl StrMapIndexPrefixRead for ReadOnlyMapIndex +where + Vec<::Owned>: Blob + Send + Sync, +{ + fn prefix_keys_with_counts( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult>> { + match self { + // Prefix support is not wired for the read-only appendable + // variant; see `ReadOnlyAppendableMapIndex::open`. + ReadOnlyMapIndex::Appendable(_) => Ok(None), + ReadOnlyMapIndex::Immutable(index) => index.prefix_keys_with_counts(prefix, hw_counter), + ReadOnlyMapIndex::OnDisk(index) => index.prefix_keys_with_counts(prefix, hw_counter), + } + } + + fn prefix_stats( + &self, + prefix: &str, + hw_counter: &HardwareCounterCell, + ) -> OperationResult> { + match self { + ReadOnlyMapIndex::Appendable(_) => Ok(None), + ReadOnlyMapIndex::Immutable(index) => index.prefix_stats(prefix, hw_counter), + ReadOnlyMapIndex::OnDisk(index) => index.prefix_stats(prefix, hw_counter), + } + } +} diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/mod.rs b/lib/segment/src/index/field_index/map_index/prefix_index/mod.rs new file mode 100644 index 0000000000..70543e53da --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/mod.rs @@ -0,0 +1,95 @@ +//! Sorted key dictionary enabling prefix queries over the keyword map index. +//! +//! Stored in a single [`PREFIX_INDEX_PATH`] file next to the other on-disk map +//! index files. The file is an *ordered view over the keys* of +//! `values_to_points.bin`: it stores no postings — only the keys themselves +//! (front-coded, in byte-lexicographic order) and their postings counts. +//! Presence of the file is what signals "prefix matching supported" at load +//! time; absence means the index was built without the `prefix` option (or by +//! an older version) and prefix queries fall back to slower paths. +//! +//! # File format +//! +//! Fixed-size little-endian [`bytemuck::Pod`] records interleaved with raw +//! key bytes. Records are written with [`bytemuck::bytes_of`] and read back +//! by copy ([`bytemuck::pod_read_unaligned`]), so they carry no alignment +//! requirement on their position in the file. +//! +//! ```text +//! prefix_index.bin +//! ┌──────────────────────────────────────────────────────────────────┐ +//! │ Header (Pod, 40 bytes) │ +//! │ magic [u8; 8] = "QdrPrfx\0" │ +//! │ version u32 = 1 │ +//! │ _reserved u32 │ +//! │ key_count u64 total keys in the dictionary │ +//! │ block_count u64 │ +//! │ block_index_size u64 bytes in the block index section │ +//! ├──────────────────────────────────────────────────────────────────┤ +//! │ Block index — read into RAM at open time; per block: │ +//! │ BlockEntry (Pod, 24 bytes) │ +//! │ block_size u32 bytes of the block in the next │ +//! │ section │ +//! │ key_count u32 keys in the block │ +//! │ first_key_len u32 │ +//! │ _reserved u32 │ +//! │ postings_count u64 Σ postings counts over the block │ +//! │ first_key u8[first_key_len] stored in full │ +//! ├──────────────────────────────────────────────────────────────────┤ +//! │ Key blocks (~4 KiB each, ≥ 1 key) — fetched and decoded lazily; │ +//! │ per key, front-coded against the previous key in the block: │ +//! │ KeyEntry (Pod, 12 bytes) │ +//! │ shared_prefix_len u32 0 for the block's first key │ +//! │ suffix_len u32 │ +//! │ postings_count u32 │ +//! │ suffix u8[suffix_len] │ +//! └──────────────────────────────────────────────────────────────────┘ +//! ``` +//! +//! # Reading a prefix +//! +//! Keys are sorted, so all keys starting with `prefix` form the contiguous +//! range `[prefix, successor(prefix))`, where `successor` increments the last +//! non-`0xFF` byte. A lookup: +//! +//! 1. binary-searches the resident block index for the candidate block range +//! (0 storage reads: the first keys bound each block's contents); +//! 2. fetches candidate blocks and decodes them sequentially, reconstructing +//! keys from shared-prefix/suffix pairs (1 read per block); +//! 3. for aggregate statistics only the two *boundary* blocks are decoded — +//! interior blocks lie fully inside the range and contribute through the +//! per-block counts of the resident block index. +//! +//! A block's first key doubles as its separator in the binary search: any +//! string greater than the previous block's last key and not greater than the +//! block's first key would do; the full first key is the simplest correct +//! choice. +//! +//! Keys are opaque byte strings; ordering and prefix semantics are byte-wise. +//! For UTF-8 keys (the keyword index) byte-wise prefix coincides with +//! character-wise prefix. +//! +//! # Module layout +//! +//! - [`format`](self::format) — on-disk layout primitives shared by reader +//! and writer; +//! - [`writer`](self::writer) — one-pass file construction from sorted +//! entries; +//! - [`reader`](self::reader) — [`PrefixIndex`], the query surface over a +//! [`UniversalRead`](common::universal_io::UniversalRead) storage; +//! - [`map_read`](self::map_read) — [`StrMapIndexPrefixRead`], the prefix +//! enumeration trait implemented by every map index variant. + +mod format; +mod map_read; +mod reader; +mod writer; + +#[cfg(test)] +mod tests; + +pub use self::map_read::StrMapIndexPrefixRead; +pub use self::reader::{PrefixIndex, PrefixIndexStats}; +pub use self::writer::build_prefix_index; + +pub const PREFIX_INDEX_PATH: &str = "prefix_index.bin"; diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/reader.rs b/lib/segment/src/index/field_index/map_index/prefix_index/reader.rs new file mode 100644 index 0000000000..e54364076e --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/reader.rs @@ -0,0 +1,395 @@ +//! Query surface over the prefix index file: [`PrefixIndex`]. + +use std::ops::Range; +use std::path::Path; + +use common::counter::conditioned_counter::ConditionedCounter; +use common::counter::hardware_counter::HardwareCounterCell; +use common::generic_consts::Random; +use common::mmap::AdviceSetting; +use common::universal_io::{ + MmapFile, OpenOptions, Populate, ReadRange, UniversalRead, UniversalReadFileOps as _, + UniversalReadFs as _, +}; + +use super::PREFIX_INDEX_PATH; +use super::format::{ + BlockEntry, Header, KeyEntry, MAGIC, VERSION, key_vs_prefix_range, prefix_successor, + read_record, +}; +use crate::common::operation_error::{OperationError, OperationResult}; + +/// Per-block metadata, parsed from the block index section and kept resident. +pub(super) struct BlockMeta { + /// First key of the block, stored in full. Any valid separator (a string + /// greater than the previous block's last key and not greater than this + /// block's first key) would do for the binary search; the full first key + /// is the simplest correct choice. + first_key: Box<[u8]>, + /// Byte range of the block within the file. + bytes: Range, + /// Number of keys in this block. + key_count: u32, + /// Sum of postings counts over all *preceding* blocks. + postings_before: u64, + /// Sum of postings counts within this block. + postings_count: u64, +} + +/// Aggregate statistics over the keys matching a prefix. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct PrefixIndexStats { + /// Number of distinct keys with the prefix. + pub keys: usize, + /// Sum of postings counts over those keys. Counts are recorded at build + /// time, so points deleted afterwards are still included. + pub postings: usize, +} + +pub struct PrefixIndex { + storage: S, + pub(super) blocks: Vec, + key_count: usize, +} + +impl PrefixIndex { + /// Open the prefix index if its file exists; `Ok(None)` when the backing + /// map index was built without prefix support. + pub fn open(fs: &S::Fs, path: &Path, populate: Populate) -> OperationResult> { + let file_path = path.join(PREFIX_INDEX_PATH); + if !fs.exists(&file_path)? { + return Ok(None); + } + + let storage = fs.open( + &file_path, + OpenOptions { + writeable: false, + need_sequential: false, + populate, + advice: AdviceSetting::Global, + }, + Default::default(), + )?; + + let header_size = size_of::
() as u64; + let header_bytes = storage.read_bytes::(0..header_size, align_of::
())?; + let header: Header = bytemuck::try_pod_read_unaligned(header_bytes.as_ref()) + .map_err(|_| OperationError::service_error("Failed to read prefix index header"))?; + + if header.magic != MAGIC { + return Err(OperationError::service_error( + "Prefix index file has invalid magic", + )); + } + if header.version != VERSION { + return Err(OperationError::service_error(format!( + "Unsupported prefix index version {}", + header.version, + ))); + } + + let index_bytes = + storage.read_bytes::(header_size..header_size + header.block_index_size, 1)?; + let blocks = Self::parse_block_index( + index_bytes.as_ref(), + header.block_count, + header_size + header.block_index_size, + )?; + + Ok(Some(Self { + storage, + blocks, + key_count: header.key_count as usize, + })) + } + + fn parse_block_index( + mut bytes: &[u8], + block_count: u64, + blocks_section_offset: u64, + ) -> OperationResult> { + let corrupt = || OperationError::service_error("Prefix index block index is corrupt"); + + let mut blocks = Vec::with_capacity(block_count as usize); + let mut block_offset = blocks_section_offset; + let mut postings_before = 0u64; + for _ in 0..block_count { + let (entry, rest) = read_record::(bytes).ok_or_else(corrupt)?; + let first_key = rest + .get(..entry.first_key_len as usize) + .ok_or_else(corrupt)?; + bytes = &rest[entry.first_key_len as usize..]; + + blocks.push(BlockMeta { + first_key: first_key.into(), + bytes: block_offset..block_offset + u64::from(entry.block_size), + key_count: entry.key_count, + postings_before, + postings_count: entry.postings_count, + }); + block_offset += u64::from(entry.block_size); + postings_before += entry.postings_count; + } + Ok(blocks) + } + + /// Number of distinct keys in the dictionary. + pub fn key_count(&self) -> usize { + self.key_count + } + + /// RAM used by the resident block index. + pub fn ram_usage_bytes(&self) -> usize { + let Self { + storage: _, + blocks, + key_count: _, + } = self; + blocks.capacity() * size_of::() + + blocks + .iter() + .map(|block| block.first_key.len()) + .sum::() + } + + /// Range of blocks that may contain keys starting with `prefix`. + fn block_range_for_prefix(&self, prefix: &[u8]) -> Range { + // The candidate range starts in the last block whose first key is not + // greater than the prefix (an earlier block's keys are all smaller + // than this block's first key, hence smaller than any `prefix*`). + let lo = self + .blocks + .partition_point(|block| block.first_key.as_ref() <= prefix) + .saturating_sub(1); + // ...and ends before the first block whose first key already lies at + // or beyond the exclusive upper bound of the prefix range. + let hi = match prefix_successor(prefix) { + Some(succ) => self + .blocks + .partition_point(|block| block.first_key.as_ref() < succ.as_slice()), + None => self.blocks.len(), + }; + lo..hi.max(lo) + } + + /// Invoke `f(key, postings_count)` for every key starting with `prefix`, + /// in ascending byte order. + /// + /// Candidate blocks are contiguous in the file (by construction), so the + /// whole candidate range is fetched with a single storage read; the + /// over-read relative to the exact key range is bounded by the two + /// partially-matching boundary blocks. + pub fn for_each_key_with_prefix( + &self, + prefix: &[u8], + hw_counter: &HardwareCounterCell, + f: &mut dyn FnMut(&[u8], usize) -> OperationResult<()>, + ) -> OperationResult<()> { + let hw_counter = ConditionedCounter::always(hw_counter); + + let range = self.block_range_for_prefix(prefix); + let Some((first_block, last_block)) = self.blocks[range.clone()] + .first() + .zip(self.blocks[range.clone()].last()) + else { + return Ok(()); + }; + + let bytes_start = first_block.bytes.start; + hw_counter + .payload_index_io_read_counter() + .incr_delta((last_block.bytes.end - bytes_start) as usize); + + let bytes = self.storage.read::(ReadRange::new( + bytes_start, + last_block.bytes.end - bytes_start, + ))?; + + for block in &self.blocks[range] { + let block_bytes = bytes + .as_ref() + .get((block.bytes.start - bytes_start) as usize..) + .ok_or_else(block_corrupt)?; + let mut past_range = false; + decode_block(block_bytes, block.key_count, &mut |key, count| { + if key_vs_prefix_range(key, prefix).is_gt() { + past_range = true; + return Ok(()); + } + if key.starts_with(prefix) { + f(key, count)?; + } + Ok(()) + })?; + if past_range { + break; + } + } + Ok(()) + } + + /// Aggregate statistics over the keys matching `prefix`. + /// + /// Interior blocks of the candidate range are guaranteed to lie fully + /// within the prefix range and contribute through the precomputed + /// per-block counts; only the two boundary blocks are decoded. + pub fn prefix_stats( + &self, + prefix: &[u8], + hw_counter: &HardwareCounterCell, + ) -> OperationResult { + let hw_counter = ConditionedCounter::always(hw_counter); + + let range = self.block_range_for_prefix(prefix); + let mut stats = PrefixIndexStats::default(); + + // Interior blocks: first key > `prefix` (they follow the `lo` block) + // and all keys smaller than the next block's first key, which is + // still below the prefix range's upper bound. Byte-lexicographic + // containment in `[prefix, successor(prefix))` is equivalent to + // starting with `prefix`, so every key counts. + let interior = (range.start + 1)..range.end.saturating_sub(1).max(range.start + 1); + if interior.start < interior.end { + let first = &self.blocks[interior.start]; + let last = &self.blocks[interior.end - 1]; + stats.postings += + (last.postings_before + last.postings_count - first.postings_before) as usize; + stats.keys += self.blocks[interior.clone()] + .iter() + .map(|block| block.key_count as usize) + .sum::(); + } + + // Boundary blocks are decoded and filtered per key. + let mut boundary = |block_index: usize| -> OperationResult<()> { + let mut keys = 0; + let mut postings = 0; + self.read_and_decode_block(block_index, &hw_counter, &mut |key, count| { + if key.starts_with(prefix) { + keys += 1; + postings += count; + } + Ok(()) + })?; + stats.keys += keys; + stats.postings += postings; + Ok(()) + }; + + if !range.is_empty() { + boundary(range.start)?; + if range.end - range.start > 1 { + boundary(range.end - 1)?; + } + } + + Ok(stats) + } + + /// Invoke `f(key, postings_count)` for every key in the dictionary, in + /// ascending byte order. + pub fn for_each_key( + &self, + hw_counter: &HardwareCounterCell, + f: &mut dyn FnMut(&[u8], usize) -> OperationResult<()>, + ) -> OperationResult<()> { + self.for_each_key_with_prefix(b"", hw_counter, f) + } + + /// Fetch a single block from storage and decode it. + fn read_and_decode_block( + &self, + block_index: usize, + hw_counter: &ConditionedCounter<'_>, + f: &mut dyn FnMut(&[u8], usize) -> OperationResult<()>, + ) -> OperationResult<()> { + let block = &self.blocks[block_index]; + + hw_counter + .payload_index_io_read_counter() + .incr_delta((block.bytes.end - block.bytes.start) as usize); + + let bytes = self.storage.read::(ReadRange::new( + block.bytes.start, + block.bytes.end - block.bytes.start, + ))?; + decode_block(bytes.as_ref(), block.key_count, f) + } + + /// Populate all pages of the backing storage. + pub fn populate(&self) -> OperationResult<()> { + self.storage.populate()?; + Ok(()) + } + + /// Hint that pages backing this index can be reclaimed. + pub fn clear_cache(&self) -> OperationResult<()> { + self.storage.clear_ram_cache()?; + Ok(()) + } +} + +fn block_corrupt() -> OperationError { + OperationError::service_error("Prefix index key block is corrupt") +} + +/// Decode the front-coded keys of one block, reconstructing each key and +/// invoking `f(key, postings_count)` in order. +/// +/// `bytes` must start at the block's first [`KeyEntry`] and contain +/// `key_count` records of the following layout (trailing bytes beyond the +/// last record are ignored): +/// +/// ```text +/// ┌──────────────────────────────────────┬─────────────────┐ +/// │ KeyEntry (Pod, 12 bytes) │ suffix │ +/// │ shared_prefix_len u32 │ u8[suffix_len] │ × key_count +/// │ suffix_len u32 │ │ +/// │ postings_count u32 │ │ +/// └──────────────────────────────────────┴─────────────────┘ +/// ``` +/// +/// Each key is reconstructed from its predecessor: keep its first +/// `shared_prefix_len` bytes and append the suffix. The first record of a +/// block has `shared_prefix_len == 0`, so its suffix is the full key: +/// +/// ```text +/// (0, 19, _) "https://qdrant.tech" → https://qdrant.tech +/// (19, 5, _) "/docs" → https://qdrant.tech/docs +/// (13, 3, _) "com" → https://qdrant.com +/// ``` +/// +/// Records are read via [`read_record`] (`bytemuck::pod_read_unaligned`), +/// which *copies* the 12 record bytes into an aligned local instead of +/// casting a reference into the buffer. An aligned view (`cast_slice` / +/// `from_bytes`) is not an option here: the variable-length suffixes +/// interleaved between records put every record after the first at an +/// arbitrary, data-dependent offset, and the buffer itself is a slice of a +/// larger storage read starting at an arbitrary file offset — so no +/// alignment can be guaranteed by construction, and a reference cast would +/// be undefined behavior whenever the offset isn't a multiple of 4. +fn decode_block( + bytes: &[u8], + key_count: u32, + f: &mut dyn FnMut(&[u8], usize) -> OperationResult<()>, +) -> OperationResult<()> { + let mut rolling_bytes = bytes; + let mut key = Vec::new(); + for _ in 0..key_count { + let (entry, rest) = read_record::(rolling_bytes).ok_or_else(block_corrupt)?; + let suffix = rest + .get(..entry.suffix_len as usize) + .ok_or_else(block_corrupt)?; + rolling_bytes = &rest[entry.suffix_len as usize..]; + + if entry.shared_prefix_len as usize > key.len() { + return Err(block_corrupt()); + } + key.truncate(entry.shared_prefix_len as usize); + key.extend_from_slice(suffix); + + f(&key, entry.postings_count as usize)?; + } + Ok(()) +} diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/tests.rs b/lib/segment/src/index/field_index/map_index/prefix_index/tests.rs new file mode 100644 index 0000000000..ee57abcd55 --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/tests.rs @@ -0,0 +1,203 @@ +use std::collections::BTreeMap; + +use common::counter::hardware_counter::HardwareCounterCell; +use common::universal_io::{MmapFile, MmapFs, Populate}; +use itertools::Itertools as _; +use rand::rngs::StdRng; +use rand::{RngExt as _, SeedableRng as _}; +use tempfile::TempDir; + +use super::format::prefix_successor; +use super::{PrefixIndex, build_prefix_index}; + +fn build_and_open(entries: &BTreeMap, usize>) -> (TempDir, PrefixIndex) { + let dir = TempDir::with_prefix("prefix_index").unwrap(); + build_prefix_index( + dir.path(), + entries.iter().map(|(key, &count)| (key.as_slice(), count)), + ) + .unwrap(); + let index = PrefixIndex::open(&MmapFs, dir.path(), Populate::Blocking) + .unwrap() + .unwrap(); + (dir, index) +} + +fn collect_prefix(index: &PrefixIndex, prefix: &[u8]) -> Vec<(Vec, usize)> { + let hw_counter = HardwareCounterCell::disposable(); + let mut result = Vec::new(); + index + .for_each_key_with_prefix(prefix, &hw_counter, &mut |key, count| { + result.push((key.to_vec(), count)); + Ok(()) + }) + .unwrap(); + result +} + +fn naive_prefix(entries: &BTreeMap, usize>, prefix: &[u8]) -> Vec<(Vec, usize)> { + entries + .iter() + .filter(|(key, _)| key.starts_with(prefix)) + .map(|(key, &count)| (key.clone(), count)) + .collect() +} + +fn check_prefix(index: &PrefixIndex, entries: &BTreeMap, usize>, prefix: &[u8]) { + let expected = naive_prefix(entries, prefix); + assert_eq!(collect_prefix(index, prefix), expected, "prefix {prefix:?}",); + + let hw_counter = HardwareCounterCell::disposable(); + let stats = index.prefix_stats(prefix, &hw_counter).unwrap(); + assert_eq!(stats.keys, expected.len(), "prefix {prefix:?}"); + assert_eq!( + stats.postings, + expected.iter().map(|(_, count)| count).sum::(), + "prefix {prefix:?}", + ); +} + +#[test] +fn missing_file_opens_as_none() { + let dir = TempDir::with_prefix("prefix_index").unwrap(); + let index = PrefixIndex::::open(&MmapFs, dir.path(), Populate::Blocking).unwrap(); + assert!(index.is_none()); +} + +#[test] +fn empty_dictionary() { + let entries = BTreeMap::new(); + let (_dir, index) = build_and_open(&entries); + assert_eq!(index.key_count(), 0); + check_prefix(&index, &entries, b""); + check_prefix(&index, &entries, b"anything"); +} + +#[test] +fn small_dictionary() { + let entries: BTreeMap, usize> = [ + (&b"https://example.com"[..], 3), + (b"https://qdrant.tech", 7), + (b"https://qdrant.tech/docs", 2), + (b"tag", 1), + (b"tags", 5), + ] + .into_iter() + .map(|(key, count)| (key.to_vec(), count)) + .collect(); + let (_dir, index) = build_and_open(&entries); + + assert_eq!(index.key_count(), 5); + for prefix in [ + &b""[..], + b"h", + b"https://", + b"https://qdrant.", + b"https://qdrant.tech", + b"https://qdrant.tech/docs/more", + b"tag", + b"tags", + b"tagz", + b"z", + b"\xff", + ] { + check_prefix(&index, &entries, prefix); + } +} + +#[test] +fn multibyte_and_edge_keys() { + let entries: BTreeMap, usize> = [ + "".as_bytes().to_vec(), + "α".as_bytes().to_vec(), + "αβ".as_bytes().to_vec(), + "яблоко".as_bytes().to_vec(), + vec![0xFF], + vec![0xFF, 0xFF], + vec![0xFF, 0xFF, 0x01], + ] + .into_iter() + .enumerate() + .map(|(i, key)| (key, i + 1)) + .collect(); + let (_dir, index) = build_and_open(&entries); + + for prefix in [ + &b""[..], + "α".as_bytes(), + "я".as_bytes(), + &[0xCE], + &[0xFF], + &[0xFF, 0xFF], + &[0xFF, 0xFF, 0xFF], + ] { + check_prefix(&index, &entries, prefix); + } +} + +#[test] +fn multi_block_random() { + let mut rng = StdRng::seed_from_u64(42); + let mut entries = BTreeMap::new(); + // Enough long keys to span many blocks; skewed shared prefixes. + for _ in 0..5_000 { + let base = ["https://", "http://", "ftp://", ""][rng.random_range(0..4)]; + let len = rng.random_range(1..40); + let tail: String = (0..len) + .map(|_| char::from(rng.random_range(b'a'..=b'e'))) + .collect(); + entries.insert( + format!("{base}{tail}").into_bytes(), + rng.random_range(1..100), + ); + } + let (_dir, index) = build_and_open(&entries); + assert!(index.blocks.len() > 3, "test should span multiple blocks"); + assert_eq!(index.key_count(), entries.len()); + + // All keys, in order. + assert_eq!( + collect_prefix(&index, b""), + entries + .iter() + .map(|(key, &count)| (key.clone(), count)) + .collect_vec(), + ); + + for prefix in [ + &b""[..], + b"h", + b"http", + b"https://", + b"https://a", + b"https://ab", + b"https://abc", + b"ftp://e", + b"a", + b"ab", + b"nonexistent", + ] { + check_prefix(&index, &entries, prefix); + } + + // Random probes, including prefixes of existing keys. + let keys = entries.keys().cloned().collect_vec(); + for _ in 0..200 { + let key = &keys[rng.random_range(0..keys.len())]; + let len = rng.random_range(0..=key.len()); + check_prefix(&index, &entries, &key[..len]); + } +} + +#[test] +fn prefix_successor_edge_cases() { + assert_eq!(prefix_successor(b""), None); + assert_eq!(prefix_successor(&[0xFF]), None); + assert_eq!(prefix_successor(&[0xFF, 0xFF]), None); + assert_eq!(prefix_successor(b"a"), Some(b"b".to_vec())); + assert_eq!(prefix_successor(&[b'a', 0xFF]), Some(b"b".to_vec())); + assert_eq!( + prefix_successor(&[b'a', 0xFF, b'c']), + Some(vec![b'a', 0xFF, b'd']), + ); +} diff --git a/lib/segment/src/index/field_index/map_index/prefix_index/writer.rs b/lib/segment/src/index/field_index/map_index/prefix_index/writer.rs new file mode 100644 index 0000000000..a511e8b50d --- /dev/null +++ b/lib/segment/src/index/field_index/map_index/prefix_index/writer.rs @@ -0,0 +1,121 @@ +//! One-pass construction of the prefix index file. + +use std::io::{BufWriter, Write as _}; +use std::path::Path; + +use fs_err as fs; + +use super::PREFIX_INDEX_PATH; +use super::format::{ + BLOCK_SIZE_TARGET, BlockEntry, Header, KeyEntry, MAGIC, VERSION, common_prefix_len, +}; +use crate::common::operation_error::OperationResult; + +/// Build the prefix index file from `(key, postings_count)` entries sorted in +/// ascending byte order without duplicates. +pub fn build_prefix_index<'a>( + path: &Path, + entries: impl Iterator, +) -> OperationResult<()> { + let mut block_index = Vec::new(); + let mut blocks = Vec::new(); + + let mut key_count = 0u64; + let mut block_count = 0u64; + + // Current block state. + let mut block_start = 0usize; + let mut block_first_key: Vec = Vec::new(); + let mut block_key_count = 0u32; + let mut block_postings = 0u64; + let mut prev_key: Vec = Vec::new(); + + let mut flush_block = |blocks: &mut Vec, + block_start: usize, + first_key: &[u8], + key_count: u32, + postings: u64| { + block_index.extend_from_slice(bytemuck::bytes_of(&BlockEntry { + block_size: (blocks.len() - block_start) as u32, + key_count, + first_key_len: first_key.len() as u32, + _reserved: 0, + postings_count: postings, + })); + block_index.extend_from_slice(first_key); + }; + + for (key, count) in entries { + debug_assert!( + block_key_count == 0 && key_count == 0 || prev_key.as_slice() < key, + "prefix index entries must be sorted and unique", + ); + + if block_key_count > 0 && blocks.len() - block_start >= BLOCK_SIZE_TARGET { + flush_block( + &mut blocks, + block_start, + &block_first_key, + block_key_count, + block_postings, + ); + block_count += 1; + block_start = blocks.len(); + block_key_count = 0; + block_postings = 0; + } + + let shared_len = if block_key_count == 0 { + block_first_key.clear(); + block_first_key.extend_from_slice(key); + 0 + } else { + common_prefix_len(&prev_key, key) + }; + + blocks.extend_from_slice(bytemuck::bytes_of(&KeyEntry { + shared_prefix_len: shared_len as u32, + suffix_len: (key.len() - shared_len) as u32, + // Bounded by the number of points, which is a `u32` offset. + postings_count: count as u32, + })); + blocks.extend_from_slice(&key[shared_len..]); + + prev_key.clear(); + prev_key.extend_from_slice(key); + block_key_count += 1; + block_postings += count as u64; + key_count += 1; + } + + if block_key_count > 0 { + flush_block( + &mut blocks, + block_start, + &block_first_key, + block_key_count, + block_postings, + ); + block_count += 1; + } + + let header = Header { + magic: MAGIC, + version: VERSION, + _reserved: 0, + key_count, + block_count, + block_index_size: block_index.len() as u64, + }; + + let file = fs::File::create(path.join(PREFIX_INDEX_PATH))?; + let mut writer = BufWriter::new(file); + writer.write_all(bytemuck::bytes_of(&header))?; + writer.write_all(&block_index)?; + writer.write_all(&blocks)?; + writer + .into_inner() + .map_err(|err| err.into_error())? + .sync_all()?; + Ok(()) +} diff --git a/lib/segment/src/index/field_index/map_index/read_only/mod.rs b/lib/segment/src/index/field_index/map_index/read_only/mod.rs index 2272e96065..609b1e8816 100644 --- a/lib/segment/src/index/field_index/map_index/read_only/mod.rs +++ b/lib/segment/src/index/field_index/map_index/read_only/mod.rs @@ -67,7 +67,7 @@ mod tests { // Build via the writable gridstore builder (matches the existing map // tests' `IndexType::MutableGridstore` path). { - let mut builder = MapIndex::::builder_mutable(dir.path().to_path_buf()); + let mut builder = MapIndex::::builder_mutable(dir.path().to_path_buf(), false); builder.init().unwrap(); let entries: &[(PointOffsetType, &[&str])] = &[ (0, &["red", "green"]), diff --git a/lib/segment/src/index/field_index/map_index/read_ops.rs b/lib/segment/src/index/field_index/map_index/read_ops.rs index 581c41bf62..e1df04547c 100644 --- a/lib/segment/src/index/field_index/map_index/read_ops.rs +++ b/lib/segment/src/index/field_index/map_index/read_ops.rs @@ -240,6 +240,22 @@ pub trait MapIndexRead<'a, N: MapIndexKey + ?Sized + 'a>: Sized { } } + /// Condition checker for [`crate::types::Match::Prefix`]. + /// + /// Checks a point's values through the forward index, so it works in + /// every variant regardless of whether prefix structures were built. + fn match_prefix_checker( + &'a self, + hw_counter: HardwareCounterCell, + prefix: impl Borrow, + ) -> MapConditionChecker<'a, N, Self> { + MapConditionChecker { + index: self, + hw_counter, + predicate: MapPredicate::Prefix(::to_owned(prefix.borrow())), + } + } + /// Condition checker for /// - [`crate::types::Match::Any`] (when `negate` is `false`), /// - [`crate::types::Match::Except`] (when `negate` is `true`). @@ -521,6 +537,9 @@ pub struct MapConditionChecker<'a, N: MapIndexKey + ?Sized, T> { enum MapPredicate { /// For [`crate::types::Match::Value`]. Value(::Owned), + /// For [`crate::types::Match::Prefix`]; meaningful for string keys only + /// ([`MapIndexKey::starts_with`] is constant `false` elsewhere). + Prefix(::Owned), /// For [`crate::types::Match::Any`] and [`crate::types::Match::Except`], /// Linear scan version. AnyScan { @@ -546,6 +565,9 @@ where self.index .check_values_any(point_id, &self.hw_counter, |value| match &self.predicate { MapPredicate::Value(expected) => value == expected.borrow(), + MapPredicate::Prefix(prefix) => { + ::starts_with(value, prefix.borrow()) + } MapPredicate::AnyScan { list, negate } => { list.iter().any(|key| key.borrow() == value) != *negate } diff --git a/lib/segment/src/index/field_index/map_index/tests.rs b/lib/segment/src/index/field_index/map_index/tests.rs index d5460a583d..2324c73486 100644 --- a/lib/segment/src/index/field_index/map_index/tests.rs +++ b/lib/segment/src/index/field_index/map_index/tests.rs @@ -19,7 +19,7 @@ use crate::index::field_index::{ CardinalityEstimation, FieldIndexBuilderTrait, PayloadFieldIndex, PayloadFieldIndexRead, ValueIndexer, }; -use crate::types::{IntPayloadType, PayloadKeyType, UuidIntType}; +use crate::types::{FieldCondition, IntPayloadType, PayloadKeyType, UuidIntType}; /// Generous default size for the deleted-points bitslice used in tests. /// @@ -64,7 +64,7 @@ fn save_map_index( match index_type { IndexType::MutableGridstore => { - let mut builder = MapIndex::::builder_mutable(path.to_path_buf()); + let mut builder = MapIndex::::builder_mutable(path.to_path_buf(), true); builder.init().unwrap(); for (idx, values) in data.iter().enumerate() { let values: Vec = values.iter().map(&into_value).collect(); @@ -76,7 +76,7 @@ fn save_map_index( builder.finalize().unwrap(); } IndexType::Mmap | IndexType::RamMmap => { - let mut builder = MapIndex::::builder_immutable(path, false, &empty_deleted()); + let mut builder = MapIndex::::builder_immutable(path, false, &empty_deleted(), true); builder.init().unwrap(); for (idx, values) in data.iter().enumerate() { let values: Vec = values.iter().map(&into_value).collect(); @@ -99,7 +99,7 @@ where Vec<::Owned>: Blob + Send + Sync, { let index = match index_type { - IndexType::MutableGridstore => MapIndex::::new_mutable(path.to_path_buf(), true) + IndexType::MutableGridstore => MapIndex::::new_mutable(path.to_path_buf(), true, true) .unwrap() .unwrap(), IndexType::Mmap => MapIndex::::new_immutable(path, true, &empty_deleted()) @@ -128,7 +128,7 @@ where fn test_uuid_payload_index() { let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); let mut builder = - MapIndex::::builder_immutable(temp_dir.path(), false, &empty_deleted()); + MapIndex::::builder_immutable(temp_dir.path(), false, &empty_deleted(), false); builder.init().unwrap(); @@ -157,8 +157,12 @@ fn test_uuid_payload_index() { #[case(true)] fn test_index_non_ascending_insertion(#[case] on_disk: bool) { let temp_dir = Builder::new().prefix("store_dir").tempdir().unwrap(); - let mut builder = - MapIndex::::builder_immutable(temp_dir.path(), on_disk, &empty_deleted()); + let mut builder = MapIndex::::builder_immutable( + temp_dir.path(), + on_disk, + &empty_deleted(), + false, + ); builder.init().unwrap(); let data = [vec![1, 2, 3, 4, 5, 6], vec![25], vec![10, 11]]; @@ -352,7 +356,7 @@ fn test_map_index_reload(#[case] index_type: IndexType) { let deleted = deleted_with(&[1, 2, 5]); let new_index = match index_type { IndexType::MutableGridstore => { - MapIndex::::new_mutable(temp_dir.path().to_path_buf(), true) + MapIndex::::new_mutable(temp_dir.path().to_path_buf(), true, false) .unwrap() .unwrap() } @@ -553,3 +557,291 @@ fn test_for_values_map_on_disk_deleted() { .collect(); assert_eq!(result, expected); } + +// --- Prefix matching --- + +/// Ground truth for a prefix match: points with at least one value starting +/// with the prefix. +fn naive_prefix_points(data: &[Vec], prefix: &str) -> Vec { + data.iter() + .enumerate() + .filter(|(_, values)| values.iter().any(|value| value.starts_with(prefix))) + .map(|(idx, _)| idx as PointOffsetType) + .collect() +} + +fn prefix_test_data() -> Vec> { + [ + vec!["https://qdrant.tech", "https://qdrant.tech/docs"], + vec!["https://qdrant.tech"], + vec!["https://example.com"], + vec!["http://example.com", "tag"], + vec!["tags"], + vec!["αβγ", "αβδ"], + vec!["tag"], + ] + .into_iter() + .map(|values| values.into_iter().map(EcoString::from).collect()) + .collect() +} + +const PREFIX_PROBES: &[&str] = &[ + "", + "h", + "http", + "http://", + "https://", + "https://qdrant.", + "https://qdrant.tech", + "https://qdrant.tech/docs/deep", + "tag", + "tags", + "tagz", + "α", + "αβ", + "αβγ", + "nonexistent", +]; + +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_str_prefix_match(#[case] index_type: IndexType) { + use common::condition_checker::ConditionChecker as _; + use common::counter::hardware_accumulator::HwMeasurementAcc; + + use crate::json_path::JsonPath; + use crate::types::Match; + + let temp_dir = Builder::new().prefix("prefix_index_dir").tempdir().unwrap(); + let data = prefix_test_data(); + let hw_counter = HardwareCounterCell::new(); + + save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); + let index: MapIndex = load_map_index(&data, temp_dir.path(), index_type); + + for prefix in PREFIX_PROBES { + let expected = naive_prefix_points(&data, prefix); + let condition = FieldCondition::new_match(JsonPath::new("test"), Match::new_prefix(prefix)); + + // Filter iterator parity with the naive scan. + let mut result: Vec = index + .filter(&condition, &hw_counter) + .unwrap() + .unwrap_or_else(|| panic!("prefix {prefix:?} must be served by the index")) + .collect(); + result.sort_unstable(); + assert_eq!(result, expected, "prefix {prefix:?}"); + + // Cardinality bounds contain the true count (no deletions here, so + // even on-disk counts are accurate). + let estimation = index + .estimate_cardinality(&condition, &hw_counter) + .unwrap() + .unwrap_or_else(|| panic!("prefix {prefix:?} must be estimated by the index")); + assert!( + estimation.min <= expected.len() && expected.len() <= estimation.max, + "prefix {prefix:?}: {} not in [{}, {}]", + expected.len(), + estimation.min, + estimation.max, + ); + + // Condition checker parity (forward index path). + let checker = index + .condition_checker(&condition, HwMeasurementAcc::new()) + .unwrap() + .unwrap(); + for idx in 0..data.len() as PointOffsetType { + assert_eq!( + checker.check(idx).unwrap(), + expected.contains(&idx), + "prefix {prefix:?}, point {idx}", + ); + } + } +} + +/// An index built *without* the prefix option must decline prefix +/// filtering/estimation (fallback path) while still serving the per-point +/// condition checker through the forward index. +#[test] +fn test_str_prefix_match_disabled() { + use common::condition_checker::ConditionChecker as _; + use common::counter::hardware_accumulator::HwMeasurementAcc; + + use crate::json_path::JsonPath; + use crate::types::Match; + + let temp_dir = Builder::new().prefix("prefix_index_dir").tempdir().unwrap(); + let data = prefix_test_data(); + let hw_counter = HardwareCounterCell::new(); + + let mut builder = MapIndex::::builder_immutable( + temp_dir.path(), + false, + &empty_deleted(), + false, // no prefix index + ); + builder.init().unwrap(); + for (idx, values) in data.iter().enumerate() { + let values: Vec = values.iter().map(|v| v.to_string().into()).collect(); + let values: Vec<_> = values.iter().collect(); + builder + .add_point(idx as PointOffsetType, &values, &hw_counter) + .unwrap(); + } + let index = builder.finalize().unwrap(); + + let condition = FieldCondition::new_match(JsonPath::new("test"), Match::new_prefix("https://")); + assert!(index.filter(&condition, &hw_counter).unwrap().is_none()); + assert!( + index + .estimate_cardinality(&condition, &hw_counter) + .unwrap() + .is_none() + ); + + let checker = index + .condition_checker(&condition, HwMeasurementAcc::new()) + .unwrap() + .unwrap(); + let expected = naive_prefix_points(&data, "https://"); + for idx in 0..data.len() as PointOffsetType { + assert_eq!(checker.check(idx).unwrap(), expected.contains(&idx)); + } +} + +#[rstest] +#[case(IndexType::MutableGridstore)] +#[case(IndexType::Mmap)] +#[case(IndexType::RamMmap)] +fn test_str_prefix_match_after_deletion(#[case] index_type: IndexType) { + use crate::json_path::JsonPath; + use crate::types::Match; + + let temp_dir = Builder::new().prefix("prefix_index_dir").tempdir().unwrap(); + let mut data = prefix_test_data(); + let hw_counter = HardwareCounterCell::new(); + + save_map_index::(&data, temp_dir.path(), index_type, |v| v.to_string().into()); + let mut index: MapIndex = load_map_index(&data, temp_dir.path(), index_type); + + // Delete two points, one of which held the only "http://" value. + for deleted in [1, 3] { + index.remove_point(deleted).unwrap(); + data[deleted as usize].clear(); + } + + for prefix in PREFIX_PROBES { + let expected = naive_prefix_points(&data, prefix); + let condition = FieldCondition::new_match(JsonPath::new("test"), Match::new_prefix(prefix)); + let mut result: Vec = index + .filter(&condition, &hw_counter) + .unwrap() + .unwrap_or_else(|| panic!("prefix {prefix:?} must be served by the index")) + .collect(); + result.sort_unstable(); + assert_eq!(result, expected, "prefix {prefix:?}"); + } +} + +/// Prefix payload blocks follow the geo index principle (`large_hashes`): +/// only the *smallest* heavy subsets produce blocks, so prefix blocks are +/// disjoint from each other and from the exact-value blocks. +#[test] +fn test_str_prefix_payload_blocks() { + use crate::json_path::JsonPath; + use crate::types::Match; + + let temp_dir = Builder::new().prefix("prefix_index_dir").tempdir().unwrap(); + // 4 points per qdrant url (8 under "https://qdrant.tech" total, sharing + // "https://" with 4 more), 4 under "tag*". + let data: Vec> = [ + ["https://qdrant.tech/docs"; 4].as_slice(), + ["https://qdrant.tech/blog"; 4].as_slice(), + ["https://example.com"; 4].as_slice(), + ["tag"; 2].as_slice(), + ["tags"; 2].as_slice(), + ] + .into_iter() + .flatten() + .map(|value| vec![EcoString::from(*value)]) + .collect(); + + save_map_index::(&data, temp_dir.path(), IndexType::Mmap, |v| { + v.to_string().into() + }); + let index: MapIndex = load_map_index(&data, temp_dir.path(), IndexType::Mmap); + + let mut prefix_blocks = Vec::new(); + let mut value_blocks = Vec::new(); + index + .for_each_payload_block(3, JsonPath::new("test"), &mut |block| { + match block.condition.r#match.as_ref().unwrap() { + Match::Prefix(prefix) => { + prefix_blocks.push((prefix.prefix.clone(), block.cardinality)); + } + Match::Value(_) + | Match::Text(_) + | Match::TextAny(_) + | Match::Phrase(_) + | Match::Any(_) + | Match::Except(_) => value_blocks.push(block.cardinality), + } + Ok(()) + }) + .unwrap(); + + // Heavy branching nodes are "https://" (12), "https://qdrant.tech/" (8) + // and "tag" (4), but the first two contain heavy exact values (each url + // has 4 > 3 postings and gets its own exact-value block), so they are + // suppressed — only the smallest heavy subset produces a block. "tag" + // covers only light values (2 postings each) and is emitted. + assert_eq!(prefix_blocks, vec![("tag".to_string(), 4)]); + // Exact-value blocks are unaffected: the three urls with 4 postings each. + assert_eq!(value_blocks, vec![4, 4, 4]); +} + +/// `prefix_index.bin` must be tracked by `files()` / `immutable_files()` +/// (snapshot manifest) exactly when the index is built with the prefix +/// option. +#[rstest] +#[case(false)] +#[case(true)] +fn test_prefix_index_file_tracking(#[case] with_prefix: bool) { + use std::ffi::OsStr; + + use super::prefix_index::PREFIX_INDEX_PATH; + + let temp_dir = Builder::new().prefix("prefix_index_dir").tempdir().unwrap(); + let data = prefix_test_data(); + let hw_counter = HardwareCounterCell::new(); + + let mut builder = + MapIndex::::builder_immutable(temp_dir.path(), false, &empty_deleted(), with_prefix); + builder.init().unwrap(); + for (idx, values) in data.iter().enumerate() { + let values: Vec = values.iter().map(|v| v.to_string().into()).collect(); + let values: Vec<_> = values.iter().collect(); + builder + .add_point(idx as PointOffsetType, &values, &hw_counter) + .unwrap(); + } + let index = builder.finalize().unwrap(); + + let tracks_file = |files: &[std::path::PathBuf]| { + files + .iter() + .any(|file| file.file_name() == Some(OsStr::new(PREFIX_INDEX_PATH))) + }; + + assert_eq!(tracks_file(&index.files()), with_prefix); + assert_eq!(tracks_file(&index.immutable_files()), with_prefix); + // The file must actually exist on disk to be snapshottable. + assert_eq!( + temp_dir.path().join(PREFIX_INDEX_PATH).exists(), + with_prefix + ); +} diff --git a/lib/segment/src/index/field_index/schema_transition.rs b/lib/segment/src/index/field_index/schema_transition.rs index f754946d14..d9ff62d48a 100644 --- a/lib/segment/src/index/field_index/schema_transition.rs +++ b/lib/segment/src/index/field_index/schema_transition.rs @@ -107,6 +107,7 @@ mod tests { is_tenant, on_disk, enable_hnsw: None, + prefix: None, }) } @@ -209,6 +210,28 @@ mod tests { ); } + #[test] + fn keyword_prefix_change_is_incompatible() { + // Enabling or disabling prefix matching requires building or dropping + // the sorted key dictionary — a full rebuild, never an in-place swap. + let plain = wrap(keyword(Some(false), None)); + let with_prefix = wrap(PayloadSchemaParams::Keyword(KeywordIndexParams { + r#type: KeywordIndexType::Keyword, + is_tenant: None, + on_disk: Some(false), + enable_hnsw: None, + prefix: Some(true), + })); + assert_eq!( + classify(&plain, &with_prefix), + SchemaTransition::Incompatible + ); + assert_eq!( + classify(&with_prefix, &plain), + SchemaTransition::Incompatible + ); + } + #[test] fn keyword_other_field_differs_is_incompatible() { // Same on_disk, but is_tenant differs. diff --git a/lib/segment/src/payload_storage/condition_checker.rs b/lib/segment/src/payload_storage/condition_checker.rs index 152dd08ddc..15faedc3aa 100644 --- a/lib/segment/src/payload_storage/condition_checker.rs +++ b/lib/segment/src/payload_storage/condition_checker.rs @@ -8,7 +8,8 @@ use serde_json::Value; use crate::types::{ AnyVariants, CheckGeoPoint, DateTimePayloadType, FieldCondition, FloatPayloadType, GeoBoundingBox, GeoPoint, GeoPolygon, GeoRadius, Match, MatchAny, MatchExcept, MatchPhrase, - MatchText, MatchTextAny, MatchValue, Range, RangeInterface, ValueVariants, ValuesCount, + MatchPrefix, MatchText, MatchTextAny, MatchValue, Range, RangeInterface, ValueVariants, + ValuesCount, }; /// Threshold representing the point to which iterating through an IndexSet is more efficient than using hashing. @@ -190,6 +191,14 @@ impl ValueChecker for Match { | Value::Array(_) | Value::Object(_) => false, }, + Match::Prefix(MatchPrefix { prefix }) => match payload { + Value::String(stored) => stored.starts_with(prefix), + Value::Null + | Value::Bool(_) + | Value::Number(_) + | Value::Array(_) + | Value::Object(_) => false, + }, Match::Any(MatchAny { any }) => match (payload, any) { (Value::String(stored), AnyVariants::Strings(list)) => { if list.len() < INDEXSET_ITER_THRESHOLD { diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index b14d5c210b..97c7abb83d 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -2354,12 +2354,18 @@ impl Display for PayloadFieldSchema { match self { PayloadFieldSchema::FieldType(t) => write!(f, "{}", t.name()), PayloadFieldSchema::FieldParams(params) => match params { - PayloadSchemaParams::Keyword(_) - | PayloadSchemaParams::Float(_) + PayloadSchemaParams::Float(_) | PayloadSchemaParams::Geo(_) | PayloadSchemaParams::Bool(_) | PayloadSchemaParams::Datetime(_) | PayloadSchemaParams::Uuid(_) => write!(f, "{}", params.name()), + PayloadSchemaParams::Keyword(keyword_params) => { + if keyword_params.prefix.unwrap_or_default() { + write!(f, "keyword (with prefix: true)") + } else { + write!(f, "keyword") + } + } PayloadSchemaParams::Integer(integer_params) => { let range = integer_params.range.unwrap_or(true); let lookup = integer_params.lookup.unwrap_or(true); @@ -2604,6 +2610,25 @@ impl> From for MatchPhrase { } } +/// Match keyword values that start with the given string. +/// +/// Byte-wise (hence, for valid UTF-8, character-wise) and case-sensitive, +/// consistent with exact keyword matching. Served efficiently by a keyword +/// index created with the `prefix` option. +#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub struct MatchPrefix { + pub prefix: String, +} + +impl> From for MatchPrefix { + fn from(prefix: S) -> Self { + MatchPrefix { + prefix: prefix.into(), + } + } +} + /// Exact match on any of the given values #[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] @@ -2626,6 +2651,7 @@ pub enum MatchInterface { Text(MatchText), TextAny(MatchTextAny), Phrase(MatchPhrase), + Prefix(MatchPrefix), Any(MatchAny), Except(MatchExcept), } @@ -2638,6 +2664,7 @@ pub enum Match { Text(MatchText), TextAny(MatchTextAny), Phrase(MatchPhrase), + Prefix(MatchPrefix), Any(MatchAny), Except(MatchExcept), } @@ -2651,6 +2678,12 @@ impl Match { Self::Text(MatchText { text: text.into() }) } + pub fn new_prefix(prefix: &str) -> Self { + Self::Prefix(MatchPrefix { + prefix: prefix.into(), + }) + } + pub fn new_any(any: AnyVariants) -> Self { Self::Any(MatchAny { any }) } @@ -2679,6 +2712,7 @@ impl From for Match { except: except.except, }), MatchInterface::Phrase(MatchPhrase { phrase }) => Self::Phrase(MatchPhrase { phrase }), + MatchInterface::Prefix(MatchPrefix { prefix }) => Self::Prefix(MatchPrefix { prefix }), } } } @@ -3286,6 +3320,7 @@ impl FieldCondition { Match::Text(_) => 0, Match::Phrase(_) => 0, Match::TextAny(_) => 0, + Match::Prefix(_) => 0, } } } diff --git a/lib/segment/tests/integration/payload_index_test.rs b/lib/segment/tests/integration/payload_index_test.rs index f8d5842c92..e6ead67e14 100644 --- a/lib/segment/tests/integration/payload_index_test.rs +++ b/lib/segment/tests/integration/payload_index_test.rs @@ -134,7 +134,15 @@ impl TestSegments { .create_field_index( opnum, &JsonPath::new(STR_KEY), - Some(&Keyword.into()), + Some(&FieldParams(PayloadSchemaParams::Keyword( + KeywordIndexParams { + r#type: KeywordIndexType::Keyword, + is_tenant: None, + on_disk: None, + enable_hnsw: None, + prefix: Some(true), + }, + ))), &hw_counter, ) .unwrap(); @@ -312,6 +320,7 @@ impl TestSegments { is_tenant: None, on_disk: Some(true), enable_hnsw: None, + prefix: Some(true), }, ))), &hw_counter, @@ -597,6 +606,7 @@ fn test_read_operations() -> Result<()> { test_struct_payload_geo_radius_index, test_struct_payload_geo_polygon_index, test_any_matcher_cardinality_estimation, + test_prefix_match, test_struct_keyword_facet, test_mmap_keyword_facet, test_struct_keyword_facet_filtered, @@ -1527,6 +1537,52 @@ fn test_any_matcher_cardinality_estimation(test_segments: &TestSegments) -> Resu Ok(()) } +/// Prefix match must return identical results on the plain segment (payload +/// fallback), the appendable struct segment (mutable prefix structure) and +/// the mmap segment (on-disk prefix index). +fn test_prefix_match(test_segments: &TestSegments) -> Result<()> { + let hw_counter = HardwareCounterCell::new(); + + let read_with_prefix = |segment: &Segment, prefix: &str| { + let filter = Filter::new_must(Condition::Field(FieldCondition::new_match( + JsonPath::new(STR_KEY), + Match::new_prefix(prefix), + ))); + let mut points = segment + .read_filtered( + None, + None, + Some(&filter), + &Default::default(), + &hw_counter, + DeferredBehavior::VisibleOnly, + ) + .unwrap(); + points.sort_unstable(); + points + }; + + let mut matched_something = false; + for prefix in ["", "b", "bl", "re", "sol", "solid", "nonexistent-prefix"] { + let plain_result = read_with_prefix(&test_segments.plain_segment, prefix); + let struct_result = read_with_prefix(&test_segments.struct_segment, prefix); + let mmap_result = read_with_prefix(&test_segments.mmap_segment, prefix); + + ensure!( + plain_result == struct_result, + "prefix {prefix:?}: plain vs struct mismatch", + ); + ensure!( + plain_result == mmap_result, + "prefix {prefix:?}: plain vs mmap mismatch", + ); + matched_something |= !plain_result.is_empty(); + } + ensure!(matched_something, "test probes never matched anything"); + + Ok(()) +} + /// FacetParams fixture without a filter fn keyword_facet_request() -> FacetParams { let limit = 1000; diff --git a/lib/segment/tests/integration/segment_on_disk_snapshot.rs b/lib/segment/tests/integration/segment_on_disk_snapshot.rs index 7e94385ee3..fd159fb5b3 100644 --- a/lib/segment/tests/integration/segment_on_disk_snapshot.rs +++ b/lib/segment/tests/integration/segment_on_disk_snapshot.rs @@ -82,6 +82,7 @@ fn test_on_disk_segment_snapshot(#[case] format: SnapshotFormat) { is_tenant: None, on_disk: Some(true), enable_hnsw: None, + prefix: None, }), )), &hw_counter, diff --git a/tests/openapi/test_prefix_match.py b/tests/openapi/test_prefix_match.py new file mode 100644 index 0000000000..8c9f8f75fe --- /dev/null +++ b/tests/openapi/test_prefix_match.py @@ -0,0 +1,231 @@ +import pytest + +from .helpers.collection_setup import drop_collection +from .helpers.helpers import request_with_validation + +COLLECTION_NAME = "test_prefix_match" + +# id -> url payload; 4 holds an array value, 6 has no url at all. +POINT_URLS = { + 1: "https://qdrant.tech", + 2: "https://qdrant.tech/docs", + 3: "https://example.com", + 4: ["http://example.com", "https://qdrant.tech/blog"], + 5: "ftp://files.example.com", + 6: None, +} + + +def expected_ids(prefix): + result = [] + for point_id, urls in POINT_URLS.items(): + if urls is None: + continue + values = urls if isinstance(urls, list) else [urls] + if any(value.startswith(prefix) for value in values): + result.append(point_id) + return result + + +PREFIX_PROBES = [ + "https://qdrant.", + "https://", + "http", + "ftp://", + "https://example.com", + "nonexistent", + "", # matches every point with a url value +] + + +@pytest.fixture(autouse=True, scope="module") +def setup(): + create_collection(COLLECTION_NAME) + yield + drop_collection(collection_name=COLLECTION_NAME) + + +def create_collection(collection_name): + drop_collection(collection_name) + + response = request_with_validation( + api='/collections/{collection_name}', + method="PUT", + path_params={'collection_name': collection_name}, + body={ + "vectors": { + "size": 2, + "distance": "Dot", + }, + } + ) + assert response.ok + + points = [] + for point_id, urls in POINT_URLS.items(): + payload = {"tag": "even" if point_id % 2 == 0 else "odd"} + if urls is not None: + payload["url"] = urls + points.append({ + "id": point_id, + "vector": [1.0, 0.0], + "payload": payload, + }) + + response = request_with_validation( + api='/collections/{collection_name}/points', + method="PUT", + path_params={'collection_name': collection_name}, + query_params={'wait': 'true'}, + body={"points": points}, + ) + assert response.ok + + +def _prefix_filter(key, prefix): + return {"must": [{"key": key, "match": {"prefix": prefix}}]} + + +def _scroll(filter_body): + return request_with_validation( + api='/collections/{collection_name}/points/scroll', + method="POST", + path_params={'collection_name': COLLECTION_NAME}, + body={"filter": filter_body, "limit": 100}, + ) + + +def _scroll_ids(filter_body): + response = _scroll(filter_body) + assert response.ok, response.json() + return sorted(point['id'] for point in response.json()['result']['points']) + + +def _set_strict_mode(strict_mode_config): + response = request_with_validation( + api="/collections/{collection_name}", + method="PATCH", + path_params={"collection_name": COLLECTION_NAME}, + body={"strict_mode_config": strict_mode_config}, + ) + response.raise_for_status() + + +# --------------------------------------------------------------------------- +# 1. Without any index, prefix match executes via the payload fallback. +# --------------------------------------------------------------------------- + +def test_prefix_match_without_index(): + for prefix in PREFIX_PROBES: + assert _scroll_ids(_prefix_filter("url", prefix)) == expected_ids(prefix), prefix + + +# --------------------------------------------------------------------------- +# 2. Create a prefix-enabled keyword index; schema is echoed in collection +# info, and results are unchanged. +# --------------------------------------------------------------------------- + +def test_create_prefix_index(): + response = request_with_validation( + api='/collections/{collection_name}/index', + method="PUT", + path_params={'collection_name': COLLECTION_NAME}, + query_params={'wait': 'true'}, + body={ + "field_name": "url", + "field_schema": { + "type": "keyword", + "prefix": True, + }, + } + ) + assert response.ok + + response = request_with_validation( + api='/collections/{collection_name}', + method="GET", + path_params={'collection_name': COLLECTION_NAME}, + ) + assert response.ok + url_schema = response.json()['result']['payload_schema']['url'] + assert url_schema['data_type'] == "keyword" + assert url_schema['params']['prefix'] is True + + +def test_prefix_match_with_index(): + for prefix in PREFIX_PROBES: + assert _scroll_ids(_prefix_filter("url", prefix)) == expected_ids(prefix), prefix + + +def test_prefix_match_count(): + response = request_with_validation( + api='/collections/{collection_name}/points/count', + method="POST", + path_params={'collection_name': COLLECTION_NAME}, + body={"filter": _prefix_filter("url", "https://qdrant."), "exact": True}, + ) + assert response.ok + assert response.json()['result']['count'] == len(expected_ids("https://qdrant.")) + + +# --------------------------------------------------------------------------- +# 3. Facet with a prefix filter on the same key — the autocompletion flow. +# Facet counts all values of the matching points. +# --------------------------------------------------------------------------- + +def test_facet_with_prefix_filter(): + response = request_with_validation( + api='/collections/{collection_name}/facet', + method="POST", + path_params={'collection_name': COLLECTION_NAME}, + body={ + "key": "url", + "filter": _prefix_filter("url", "https://q"), + }, + ) + assert response.ok + hits = {hit['value']: hit['count'] for hit in response.json()['result']['hits']} + # Points 1, 2 and 4 match; facet counts every value they hold. + assert hits == { + "https://qdrant.tech": 1, + "https://qdrant.tech/docs": 1, + "https://qdrant.tech/blog": 1, + "http://example.com": 1, + } + + +# --------------------------------------------------------------------------- +# 4. Strict mode: prefix filtering requires a keyword index with the +# `prefix` option; a plain keyword index is not enough. +# --------------------------------------------------------------------------- + +def test_strict_mode_requires_prefix_capability(): + # Plain keyword index (no `prefix`) on another field. + response = request_with_validation( + api='/collections/{collection_name}/index', + method="PUT", + path_params={'collection_name': COLLECTION_NAME}, + query_params={'wait': 'true'}, + body={"field_name": "tag", "field_schema": "keyword"}, + ) + assert response.ok + + _set_strict_mode({ + "enabled": True, + "unindexed_filtering_retrieve": False, + }) + try: + # Exact match on the plain keyword index is allowed... + response = _scroll({"must": [{"key": "tag", "match": {"value": "odd"}}]}) + assert response.ok, response.json() + + # ...but prefix match is rejected: the index lacks the prefix option. + response = _scroll(_prefix_filter("tag", "od")) + assert response.status_code == 400, response.json() + assert "Index required but not found" in response.json()['status']['error'] + + # On the prefix-enabled index the query is allowed. + response = _scroll(_prefix_filter("url", "https://qdrant.")) + assert response.ok, response.json() + finally: + _set_strict_mode({"enabled": False})