Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[ENH] Enable test_embeddings #2376

Merged
merged 1 commit into from
Jun 19, 2024
Merged

[ENH] Enable test_embeddings #2376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/_python-tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ jobs:
"chromadb/test/property/test_collections.py",
"chromadb/test/property/test_add.py",
"chromadb/test/property/test_filtering.py",
"chromadb/test/property/test_embeddings.py",
"chromadb/test/property/test_collections_with_database_tenant.py",
"chromadb/test/property/test_collections_with_database_tenant_overwrite.py",
"chromadb/test/ingest/test_producer_consumer.py",
Expand Down
2 changes: 1 addition & 1 deletion rust/worker/src/execution/operators/brute_force_knn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ impl Operator<BruteForceKnnOperatorInput, BruteForceKnnOperatorOutput> for Brute
i += 1;
}

trace!("Brute force Knn result. distances: {:?}", sorted_distances);
tracing::info!("Brute force Knn result. distances: {:?}", sorted_distances);
Ok(BruteForceKnnOperatorOutput {
user_ids: sorted_user_ids,
embeddings: sorted_embeddings,
Expand Down
66 changes: 59 additions & 7 deletions rust/worker/src/execution/operators/get_vectors_operator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,10 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO
let logs = input.log_records.clone();
let mut remaining_search_user_ids: HashSet<String> =
HashSet::from_iter(input.search_user_ids.iter().cloned());
let mut original_search_user_ids: HashSet<String> =
HashSet::from_iter(input.search_user_ids.iter().cloned());
for (log_record, _) in logs.iter() {
if remaining_search_user_ids.contains(&log_record.record.id) {
if original_search_user_ids.contains(&log_record.record.id) {
HammadB marked this conversation as resolved.
Show resolved Hide resolved
match log_record.record.operation {
crate::types::Operation::Add => {
// If there is a record segment, validate the add
Expand Down Expand Up @@ -146,6 +148,17 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO
));
}
}
} else {
// Record segment is uninitialized.
// If the user id is already present in the log set, skip it
// We use the first add log entry for a user id if a
// user has multiple log entries
if output_vectors.contains_key(&log_record.record.id) {
continue;
}
let vector = log_record.record.embedding.as_ref().expect("Invariant violation. The log record for an add does not have an embedding.");
output_vectors.insert(log_record.record.id.clone(), vector.clone());
remaining_search_user_ids.remove(&log_record.record.id);
}
}
crate::types::Operation::Update => {
Expand Down Expand Up @@ -174,9 +187,26 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO
}
Ok(false) => {
// The record does not exist in the record segment,
// the update is faulty.
// We skip the update
continue;
// If the user id is present in the output set then it means
// that it was inserted previously. Update the embedding
// if it exists.
if !output_vectors.contains_key(&log_record.record.id) {
continue;
}
match &log_record.record.embedding {
Some(vector) => {
// This will overwrite the vector if it already exists
// (e.g if it was added previously in the log)
output_vectors.insert(
log_record.record.id.clone(),
vector.clone(),
);
remaining_search_user_ids.remove(&log_record.record.id);
}
None => {
// Nothing to do with this as the vector was not updated
}
}
}
Err(e) => {
// If there is an error, we skip the update
Expand All @@ -185,6 +215,26 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO
));
}
}
} else {
// Record segment is uninitialized.
// If the user id is present in the output set then it means
// that it was inserted previously. Update the embedding
// if it exists.
if !output_vectors.contains_key(&log_record.record.id) {
continue;
}
match &log_record.record.embedding {
Some(vector) => {
// This will overwrite the vector if it already exists
// (e.g if it was added previously in the log)
output_vectors
.insert(log_record.record.id.clone(), vector.clone());
remaining_search_user_ids.remove(&log_record.record.id);
}
None => {
// Nothing to do with this as the vector was not updated
}
}
}
}
crate::types::Operation::Upsert => {
Expand Down Expand Up @@ -222,9 +272,11 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO

let mut ids = Vec::new();
let mut vectors = Vec::new();
for (id, vector) in output_vectors.drain() {
ids.push(id);
vectors.push(vector);
for id in &input.search_user_ids {
if output_vectors.contains_key(id) {
ids.push(id.clone());
vectors.push(output_vectors.remove(id).unwrap());
}
}
return Ok(GetVectorsOperatorOutput { ids, vectors });
}
Expand Down
1 change: 1 addition & 0 deletions rust/worker/src/execution/orchestration/get_vectors.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ impl GetVectorsOrchestrator {
.expect("Invariant violation. Record segment is not set.");
let blockfile_provider = self.blockfile_provider.clone();
let operator = GetVectorsOperator::new();
tracing::info!("get_vectors with search ids {:?}", self.search_user_ids);
let input = GetVectorsOperatorInput::new(
record_segment.clone(),
blockfile_provider,
Expand Down
6 changes: 6 additions & 0 deletions rust/worker/src/execution/orchestration/hnsw.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,12 @@ impl HnswQueryOrchestrator {
let finish_dependency_count = query_vectors.len() as u32;
// pre-allocate the result vectors
let results = Some(Vec::with_capacity(query_vectors.len()));
tracing::info!(
"Performing KNN for k = {}, allowed_ids = {:?}, num query vectors = {:?}",
k,
allowed_ids,
query_vectors.len()
);

HnswQueryOrchestrator {
state: ExecutionState::Pending,
Expand Down
Loading