Merge the two <facetttype>_faceted_documents_ids methods into one

This commit is contained in:
Loïc Lecrenier 2022-09-05 12:51:40 +02:00 committed by Loïc Lecrenier
parent 982efab88f
commit 3d145d7f48
5 changed files with 59 additions and 62 deletions

View file

@ -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::<FieldId>()];
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::<FieldId>()];
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<RoaringBitmap> {
let mut buffer =
[0u8; main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
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::<FieldId>()];
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<RoaringBitmap> {
let mut buffer =
[0u8; main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
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::<FieldId>()];
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()),