From 3d145d7f48b35739c02e3cf3a44c624cf94ce8d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Lecrenier?= Date: Mon, 5 Sep 2022 12:51:40 +0200 Subject: [PATCH] Merge the two _faceted_documents_ids methods into one --- milli/src/index.rs | 71 ++++++++------------------- milli/src/search/criteria/asc_desc.rs | 7 ++- milli/src/snapshot_tests.rs | 5 +- milli/src/update/clear_documents.rs | 16 ++++-- milli/src/update/delete_documents.rs | 22 +++++++-- 5 files changed, 59 insertions(+), 62 deletions(-) diff --git a/milli/src/index.rs b/milli/src/index.rs index 0561a77ac..40e78bf10 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -12,6 +12,7 @@ use rstar::RTree; use time::OffsetDateTime; use crate::error::{InternalError, UserError}; +use crate::facet::FacetType; use crate::fields_ids_map::FieldsIdsMap; use crate::heed_codec::facet::new::ordered_f64_codec::OrderedF64Codec; use crate::heed_codec::facet::new::str_ref::StrRefCodec; @@ -780,68 +781,38 @@ impl Index { /* faceted documents ids */ - /// Writes the documents ids that are faceted with numbers under this field id. - pub(crate) fn put_number_faceted_documents_ids( + /// Writes the documents ids that are faceted under this field id for the given facet type. + pub fn put_faceted_documents_ids( &self, wtxn: &mut RwTxn, field_id: FieldId, + facet_type: FacetType, docids: &RoaringBitmap, ) -> heed::Result<()> { - let mut buffer = - [0u8; main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::()]; - buffer[..main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()] - .copy_from_slice(main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes()); - buffer[main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()..] - .copy_from_slice(&field_id.to_be_bytes()); + let key = match facet_type { + FacetType::String => main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX, + FacetType::Number => main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX, + }; + let mut buffer = vec![0u8; key.len() + size_of::()]; + buffer[..key.len()].copy_from_slice(key.as_bytes()); + buffer[key.len()..].copy_from_slice(&field_id.to_be_bytes()); self.main.put::<_, ByteSlice, RoaringBitmapCodec>(wtxn, &buffer, docids) } - /// Retrieve all the documents ids that faceted with numbers under this field id. - pub fn number_faceted_documents_ids( + /// Retrieve all the documents ids that are faceted under this field id for the given facet type. + pub fn faceted_documents_ids( &self, rtxn: &RoTxn, field_id: FieldId, + facet_type: FacetType, ) -> heed::Result { - let mut buffer = - [0u8; main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::()]; - buffer[..main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()] - .copy_from_slice(main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes()); - buffer[main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()..] - .copy_from_slice(&field_id.to_be_bytes()); - match self.main.get::<_, ByteSlice, RoaringBitmapCodec>(rtxn, &buffer)? { - Some(docids) => Ok(docids), - None => Ok(RoaringBitmap::new()), - } - } - - /// Writes the documents ids that are faceted with strings under this field id. - pub(crate) fn put_string_faceted_documents_ids( - &self, - wtxn: &mut RwTxn, - field_id: FieldId, - docids: &RoaringBitmap, - ) -> heed::Result<()> { - let mut buffer = - [0u8; main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::()]; - buffer[..main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()] - .copy_from_slice(main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes()); - buffer[main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()..] - .copy_from_slice(&field_id.to_be_bytes()); - self.main.put::<_, ByteSlice, RoaringBitmapCodec>(wtxn, &buffer, docids) - } - - /// Retrieve all the documents ids that faceted with strings under this field id. - pub fn string_faceted_documents_ids( - &self, - rtxn: &RoTxn, - field_id: FieldId, - ) -> heed::Result { - let mut buffer = - [0u8; main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::()]; - buffer[..main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()] - .copy_from_slice(main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes()); - buffer[main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()..] - .copy_from_slice(&field_id.to_be_bytes()); + let key = match facet_type { + FacetType::String => main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX, + FacetType::Number => main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX, + }; + let mut buffer = vec![0u8; key.len() + size_of::()]; + buffer[..key.len()].copy_from_slice(key.as_bytes()); + buffer[key.len()..].copy_from_slice(&field_id.to_be_bytes()); match self.main.get::<_, ByteSlice, RoaringBitmapCodec>(rtxn, &buffer)? { Some(docids) => Ok(docids), None => Ok(RoaringBitmap::new()), diff --git a/milli/src/search/criteria/asc_desc.rs b/milli/src/search/criteria/asc_desc.rs index 23dd860e1..ccf66889e 100644 --- a/milli/src/search/criteria/asc_desc.rs +++ b/milli/src/search/criteria/asc_desc.rs @@ -6,6 +6,7 @@ use ordered_float::OrderedFloat; use roaring::RoaringBitmap; use super::{Criterion, CriterionParameters, CriterionResult}; +use crate::facet::FacetType; use crate::heed_codec::facet::new::{FacetKeyCodec, MyByteSlice}; use crate::search::criteria::{resolve_query_tree, CriteriaBuilder}; use crate::search::facet::facet_sort_ascending::ascending_facet_sort; @@ -62,8 +63,10 @@ impl<'t> AscDesc<'t> { let field_id = fields_ids_map.id(&field_name); let faceted_candidates = match field_id { Some(field_id) => { - let number_faceted = index.number_faceted_documents_ids(rtxn, field_id)?; - let string_faceted = index.string_faceted_documents_ids(rtxn, field_id)?; + let number_faceted = + index.faceted_documents_ids(rtxn, field_id, FacetType::Number)?; + let string_faceted = + index.faceted_documents_ids(rtxn, field_id, FacetType::String)?; number_faceted | string_faceted } None => RoaringBitmap::default(), diff --git a/milli/src/snapshot_tests.rs b/milli/src/snapshot_tests.rs index d054e63b5..57fd2e5fe 100644 --- a/milli/src/snapshot_tests.rs +++ b/milli/src/snapshot_tests.rs @@ -4,6 +4,7 @@ use std::path::Path; use roaring::RoaringBitmap; +use crate::facet::FacetType; use crate::heed_codec::facet::new::{FacetGroupValue, FacetKey}; use crate::{make_db_snap_from_iter, ExternalDocumentsIds, Index}; @@ -370,7 +371,7 @@ pub fn snap_number_faceted_documents_ids(index: &Index) -> String { let mut snap = String::new(); for field_id in fields_ids_map.ids() { let number_faceted_documents_ids = - index.number_faceted_documents_ids(&rtxn, field_id).unwrap(); + index.faceted_documents_ids(&rtxn, field_id, FacetType::Number).unwrap(); writeln!(&mut snap, "{field_id:<3} {}", display_bitmap(&number_faceted_documents_ids)) .unwrap(); } @@ -383,7 +384,7 @@ pub fn snap_string_faceted_documents_ids(index: &Index) -> String { let mut snap = String::new(); for field_id in fields_ids_map.ids() { let string_faceted_documents_ids = - index.string_faceted_documents_ids(&rtxn, field_id).unwrap(); + index.faceted_documents_ids(&rtxn, field_id, FacetType::String).unwrap(); writeln!(&mut snap, "{field_id:<3} {}", display_bitmap(&string_faceted_documents_ids)) .unwrap(); } diff --git a/milli/src/update/clear_documents.rs b/milli/src/update/clear_documents.rs index ba59c14cf..7d89ca89a 100644 --- a/milli/src/update/clear_documents.rs +++ b/milli/src/update/clear_documents.rs @@ -1,7 +1,7 @@ use roaring::RoaringBitmap; use time::OffsetDateTime; -use crate::{ExternalDocumentsIds, FieldDistribution, Index, Result}; +use crate::{facet::FacetType, ExternalDocumentsIds, FieldDistribution, Index, Result}; pub struct ClearDocuments<'t, 'u, 'i> { wtxn: &'t mut heed::RwTxn<'i, 'u>, @@ -55,8 +55,18 @@ impl<'t, 'u, 'i> ClearDocuments<'t, 'u, 'i> { // We clean all the faceted documents ids. for field_id in faceted_fields { - self.index.put_number_faceted_documents_ids(self.wtxn, field_id, &empty_roaring)?; - self.index.put_string_faceted_documents_ids(self.wtxn, field_id, &empty_roaring)?; + self.index.put_faceted_documents_ids( + self.wtxn, + field_id, + FacetType::Number, + &empty_roaring, + )?; + self.index.put_faceted_documents_ids( + self.wtxn, + field_id, + FacetType::String, + &empty_roaring, + )?; } // Clear the other databases. diff --git a/milli/src/update/delete_documents.rs b/milli/src/update/delete_documents.rs index 531fd2b74..ffa63f0a7 100644 --- a/milli/src/update/delete_documents.rs +++ b/milli/src/update/delete_documents.rs @@ -461,9 +461,15 @@ impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> { // Remove the documents ids from the faceted documents ids. for field_id in self.index.faceted_fields_ids(self.wtxn)? { // Remove docids from the number faceted documents ids - let mut docids = self.index.number_faceted_documents_ids(self.wtxn, field_id)?; + let mut docids = + self.index.faceted_documents_ids(self.wtxn, field_id, FacetType::Number)?; docids -= &self.to_delete_docids; - self.index.put_number_faceted_documents_ids(self.wtxn, field_id, &docids)?; + self.index.put_faceted_documents_ids( + self.wtxn, + field_id, + FacetType::Number, + &docids, + )?; remove_docids_from_field_id_docid_facet_value( self.wtxn, @@ -474,9 +480,15 @@ impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> { )?; // Remove docids from the string faceted documents ids - let mut docids = self.index.string_faceted_documents_ids(self.wtxn, field_id)?; + let mut docids = + self.index.faceted_documents_ids(self.wtxn, field_id, FacetType::String)?; docids -= &self.to_delete_docids; - self.index.put_string_faceted_documents_ids(self.wtxn, field_id, &docids)?; + self.index.put_faceted_documents_ids( + self.wtxn, + field_id, + FacetType::String, + &docids, + )?; remove_docids_from_field_id_docid_facet_value( self.wtxn, @@ -648,7 +660,7 @@ fn remove_docids_from_facet_id_docids<'a>( if !modified { return Ok(()); } - let builder = FacetsUpdateBulk::new(index, facet_type); + let builder = FacetsUpdateBulk::new_not_updating_level_0(index, facet_type); builder.execute(wtxn)?; Ok(())