use std::borrow::Cow;
use std::sync::Arc;
use ahash::{AHashMap, AHashSet};
use itertools::Itertools as _;
use crate::segment::data_types::vector_name_config::{
DenseVectorConfig, SparseVectorConfig, VectorNameConfig,
};
use crate::segment::index::field_index::CardinalityEstimation;
use crate::segment::types::{
Condition, CustomIdCheckerCondition, ExtendedPointId, Filter, SegmentConfig, SeqNumberType,
SparseVectorDataConfig, VectorDataConfig, VectorName, VectorNameBuf, WithVector,
};
#[derive(Debug)]
struct AlwaysFalseChecker;
impl CustomIdCheckerCondition for AlwaysFalseChecker {
fn estimate_cardinality(&self, _points: usize) -> CardinalityEstimation {
CardinalityEstimation::exact(0)
}
fn check(&self, _point_id: ExtendedPointId) -> bool {
false
}
}
#[derive(Debug, Clone)]
pub enum IntendedVector {
Present {
config: VectorNameConfig,
version: SeqNumberType,
supersedes_wrapped: bool,
},
Absent { version: SeqNumberType },
}
impl IntendedVector {
pub fn version(&self) -> SeqNumberType {
match self {
IntendedVector::Present { version, .. } => *version,
IntendedVector::Absent { version } => *version,
}
}
fn taints_wrapped(&self) -> bool {
matches!(
self,
IntendedVector::Absent { .. }
| IntendedVector::Present {
supersedes_wrapped: true,
..
}
)
}
}
#[derive(Debug, Default)]
pub struct ProxyVectorNameChanges {
intent: AHashMap<VectorNameBuf, IntendedVector>,
}
impl ProxyVectorNameChanges {
pub fn record_create(
&mut self,
vector_name: VectorNameBuf,
config: VectorNameConfig,
version: SeqNumberType,
wrapped_config: &SegmentConfig,
) {
let previous_taints = self
.intent
.get(&vector_name)
.is_some_and(IntendedVector::taints_wrapped);
let supersedes_wrapped =
previous_taints || wrapped_carries_stale_schema(wrapped_config, &vector_name, &config);
self.intent.insert(
vector_name,
IntendedVector::Present {
config,
version,
supersedes_wrapped,
},
);
}
pub fn record_delete(&mut self, vector_name: VectorNameBuf, version: SeqNumberType) {
self.intent
.insert(vector_name, IntendedVector::Absent { version });
}
pub fn is_empty(&self) -> bool {
self.intent.is_empty()
}
pub fn clear(&mut self) {
self.intent.clear();
}
pub fn iter_ordered(&self) -> impl Iterator<Item = (&VectorNameBuf, &IntendedVector)> {
self.intent
.iter()
.sorted_by_key(|(_, intent)| intent.version())
}
pub fn is_wrapped_data_stale(&self, vector_name: &VectorName) -> bool {
self.intent
.get(vector_name)
.is_some_and(IntendedVector::taints_wrapped)
}
pub fn redact_with_vector<'a>(
&self,
with_vector: &'a WithVector,
wrapped_config: &SegmentConfig,
) -> Cow<'a, WithVector> {
let tainted: AHashSet<&str> = self
.intent
.iter()
.filter(|(_, intent)| intent.taints_wrapped())
.map(|(name, _)| name.as_str())
.collect();
if tainted.is_empty() {
return Cow::Borrowed(with_vector);
}
match with_vector {
WithVector::Bool(false) => Cow::Borrowed(with_vector),
WithVector::Bool(true) => {
let kept: Vec<VectorNameBuf> = wrapped_config
.vector_data
.keys()
.chain(wrapped_config.sparse_vector_data.keys())
.filter(|name| !tainted.contains(name.as_str()))
.cloned()
.collect();
Cow::Owned(WithVector::Selector(kept))
}
WithVector::Selector(requested) => {
let needs_redact = requested.iter().any(|name| tainted.contains(name.as_str()));
if !needs_redact {
return Cow::Borrowed(with_vector);
}
let kept: Vec<VectorNameBuf> = requested
.iter()
.filter(|name| !tainted.contains(name.as_str()))
.cloned()
.collect();
Cow::Owned(WithVector::Selector(kept))
}
}
}
pub fn redact_filter<'a>(&self, filter: &'a Filter) -> Cow<'a, Filter> {
if !self.filter_has_stale_has_vector(filter) {
return Cow::Borrowed(filter);
}
let mut owned = filter.clone();
self.redact_filter_inplace(&mut owned);
Cow::Owned(owned)
}
fn filter_has_stale_has_vector(&self, filter: &Filter) -> bool {
let Filter {
should,
min_should,
must,
must_not,
} = filter;
let conditions = should
.iter()
.flatten()
.chain(must.iter().flatten())
.chain(must_not.iter().flatten())
.chain(min_should.iter().flat_map(|ms| ms.conditions.iter()));
for cond in conditions {
match cond {
Condition::HasVector(hv) => {
if self.is_wrapped_data_stale(&hv.has_vector) {
return true;
}
}
Condition::Nested(nested) => {
if self.filter_has_stale_has_vector(&nested.nested.filter) {
return true;
}
}
Condition::Filter(inner) => {
if self.filter_has_stale_has_vector(inner) {
return true;
}
}
Condition::Field(_) => {}
Condition::IsEmpty(_) => {}
Condition::IsNull(_) => {}
Condition::HasId(_) => {}
Condition::CustomIdChecker(_) => {}
}
}
false
}
fn redact_filter_inplace(&self, filter: &mut Filter) {
let Filter {
should,
min_should,
must,
must_not,
} = filter;
if let Some(conds) = should {
self.redact_conditions_inplace(conds);
}
if let Some(conds) = must {
self.redact_conditions_inplace(conds);
}
if let Some(conds) = must_not {
self.redact_conditions_inplace(conds);
}
if let Some(ms) = min_should {
self.redact_conditions_inplace(&mut ms.conditions);
}
}
fn redact_conditions_inplace(&self, conditions: &mut [Condition]) {
for cond in conditions.iter_mut() {
match cond {
Condition::HasVector(hv) => {
if self.is_wrapped_data_stale(&hv.has_vector) {
*cond = Condition::new_custom(Arc::new(AlwaysFalseChecker));
}
}
Condition::Nested(nested) => {
self.redact_filter_inplace(&mut nested.nested.filter);
}
Condition::Filter(inner) => {
self.redact_filter_inplace(inner);
}
Condition::Field(_) => {}
Condition::IsEmpty(_) => {}
Condition::IsNull(_) => {}
Condition::HasId(_) => {}
Condition::CustomIdChecker(_) => {}
}
}
}
pub fn merge(&mut self, other: &Self) {
for (name, other_intent) in &other.intent {
let other_taints = other_intent.taints_wrapped();
match self.intent.get_mut(name) {
None => {
self.intent.insert(name.clone(), other_intent.clone());
}
Some(self_intent) => {
if other_intent.version() > self_intent.version() {
let self_taints = self_intent.taints_wrapped();
let mut winner = other_intent.clone();
if self_taints
&& let IntendedVector::Present {
supersedes_wrapped, ..
} = &mut winner
{
*supersedes_wrapped = true;
}
*self_intent = winner;
} else if other_taints
&& let IntendedVector::Present {
supersedes_wrapped, ..
} = self_intent
{
*supersedes_wrapped = true;
}
}
}
}
}
}
fn wrapped_carries_stale_schema(
wrapped_config: &SegmentConfig,
vector_name: &VectorName,
new_config: &VectorNameConfig,
) -> bool {
match new_config {
VectorNameConfig::Dense(wrapper) => {
if wrapped_config.sparse_vector_data.contains_key(vector_name) {
return true;
}
let Some(existing) = wrapped_config.vector_data.get(vector_name) else {
return false;
};
!wrapped_dense_schema_matches(existing, &wrapper.dense)
}
VectorNameConfig::Sparse(wrapper) => {
if wrapped_config.vector_data.contains_key(vector_name) {
return true;
}
let Some(existing) = wrapped_config.sparse_vector_data.get(vector_name) else {
return false;
};
!wrapped_sparse_schema_matches(existing, &wrapper.sparse)
}
}
}
fn wrapped_dense_schema_matches(
existing: &VectorDataConfig,
new_config: &DenseVectorConfig,
) -> bool {
let DenseVectorConfig {
size,
distance,
multivector_config,
datatype,
} = new_config;
existing.size == *size
&& existing.distance == *distance
&& existing.multivector_config == *multivector_config
&& existing.datatype == *datatype
}
fn wrapped_sparse_schema_matches(
existing: &SparseVectorDataConfig,
new_config: &SparseVectorConfig,
) -> bool {
let SparseVectorConfig { modifier, datatype } = new_config;
existing.modifier == *modifier && existing.index.datatype == *datatype
}