use roaring::RoaringBitmap;
use uuid::Uuid;
use super::*;
use crate::dataset::Dataset;
use crate::{Error, Result};
use lance_core::format::pb;
use snafu::{location, Location};
#[derive(Debug, Clone)]
pub struct Index {
pub uuid: Uuid,
pub fields: Vec<i32>,
pub name: String,
pub dataset_version: u64,
pub fragment_bitmap: Option<RoaringBitmap>,
}
impl Index {
pub async fn unindexed_fragments(&self, dataset: &Dataset) -> Result<Vec<Fragment>> {
if self.dataset_version == dataset.version().version {
return Ok(vec![]);
}
if let Some(bitmap) = self.fragment_bitmap.as_ref() {
Ok(dataset
.fragments()
.iter()
.filter(|f| !bitmap.contains(f.id as u32))
.cloned()
.collect::<Vec<_>>())
} else {
let ds = dataset.checkout_version(self.dataset_version).await?;
let max_fragment_id_idx = ds.manifest.max_fragment_id().ok_or_else(|| Error::IO {
message: "No fragments in index version".to_string(),
location: location!(),
})?;
let max_fragment_id_ds =
dataset
.manifest
.max_fragment_id()
.ok_or_else(|| Error::IO {
message: "No fragments in dataset version".to_string(),
location: location!(),
})?;
if max_fragment_id_idx < max_fragment_id_ds {
dataset.manifest.fragments_since(&ds.manifest)
} else {
Ok(vec![])
}
}
}
}
impl TryFrom<&pb::IndexMetadata> for Index {
type Error = Error;
fn try_from(proto: &pb::IndexMetadata) -> Result<Self> {
let fragment_bitmap = if proto.fragment_bitmap.is_empty() {
None
} else {
Some(RoaringBitmap::deserialize_from(
&mut proto.fragment_bitmap.as_slice(),
)?)
};
Ok(Self {
uuid: proto
.uuid
.as_ref()
.map(Uuid::try_from)
.ok_or_else(|| Error::IO {
message: "uuid field does not exist in Index metadata".to_string(),
location: location!(),
})??,
name: proto.name.clone(),
fields: proto.fields.clone(),
dataset_version: proto.dataset_version,
fragment_bitmap,
})
}
}
impl From<&Index> for pb::IndexMetadata {
fn from(idx: &Index) -> Self {
let mut fragment_bitmap = Vec::new();
if let Some(bitmap) = &idx.fragment_bitmap {
if let Err(e) = bitmap.serialize_into(&mut fragment_bitmap) {
log::error!("Failed to serialize fragment bitmap: {}", e);
fragment_bitmap.clear();
}
}
Self {
uuid: Some((&idx.uuid).into()),
name: idx.name.clone(),
fields: idx.fields.clone(),
dataset_version: idx.dataset_version,
fragment_bitmap,
}
}
}