1use std::collections::{HashMap, HashSet};
7use std::sync::Arc;
8
9use chrono::{DateTime, Utc};
10use futures::StreamExt;
11use lance_core::deepsize::DeepSizeOf;
12use lance_io::object_store::ObjectStore;
13use object_store::path::Path;
14use roaring::RoaringBitmap;
15use uuid::Uuid;
16
17use super::pb;
18use lance_core::cache::{CacheEntryReader, CacheEntryWriter};
19use lance_core::{Error, Result};
20
21#[derive(Debug, Clone, PartialEq, DeepSizeOf)]
23pub struct IndexFile {
24 pub path: String,
26 pub size_bytes: u64,
28}
29
30#[derive(Debug, Clone, PartialEq)]
32pub struct IndexMetadata {
33 pub uuid: Uuid,
35
36 pub fields: Vec<i32>,
41
42 pub covering_fields: Vec<i32>,
52
53 pub name: String,
55
56 pub dataset_version: u64,
61
62 pub fragment_bitmap: Option<RoaringBitmap>,
68
69 pub index_details: Option<Arc<prost_types::Any>>,
74
75 pub index_version: i32,
77
78 pub created_at: Option<DateTime<Utc>>,
83
84 pub base_id: Option<u32>,
87
88 pub files: Option<Vec<IndexFile>>,
94}
95
96impl IndexMetadata {
97 pub fn effective_fragment_bitmap(
98 &self,
99 existing_fragments: &RoaringBitmap,
100 ) -> Option<RoaringBitmap> {
101 let fragment_bitmap = self.fragment_bitmap.as_ref()?;
102 Some(fragment_bitmap & existing_fragments)
103 }
104
105 pub fn file_size_map(&self) -> HashMap<String, u64> {
108 self.files
109 .as_ref()
110 .map(|files| {
111 files
112 .iter()
113 .map(|f| (f.path.clone(), f.size_bytes))
114 .collect()
115 })
116 .unwrap_or_default()
117 }
118
119 pub fn total_size_bytes(&self) -> Option<u64> {
122 self.files
123 .as_ref()
124 .map(|files| files.iter().map(|f| f.size_bytes).sum())
125 }
126
127 pub fn deleted_fragment_bitmap(
130 &self,
131 existing_fragments: &RoaringBitmap,
132 ) -> Option<RoaringBitmap> {
133 let fragment_bitmap = self.fragment_bitmap.as_ref()?;
134 Some(fragment_bitmap - existing_fragments)
135 }
136
137 pub fn results_are_row_addrs(&self) -> bool {
143 self.index_details.as_ref().is_some_and(|details| {
144 let is_fm = details
145 .type_url
146 .rsplit_once('/')
147 .is_some_and(|(_, details_type_name)| {
148 details_type_name.eq_ignore_ascii_case("lance.index.pb.FMIndexDetails")
149 });
150 details.type_url.ends_with("ZoneMapIndexDetails")
151 || details.type_url.ends_with("BloomFilterIndexDetails")
152 || is_fm
153 })
154 }
155
156 pub fn keyed_fields(&self) -> &[i32] {
165 let keyed = self.fields.len().saturating_sub(self.covering_fields.len());
166 &self.fields[..keyed]
167 }
168
169 pub fn keyed_field(&self) -> Option<i32> {
175 match self.keyed_fields() {
176 [only] => Some(*only),
177 _ => None,
178 }
179 }
180
181 pub fn validate_covering_fields(&self) -> Result<()> {
194 if self.covering_fields.is_empty() {
195 return Ok(());
196 }
197
198 let missing: Vec<i32> = self
199 .covering_fields
200 .iter()
201 .copied()
202 .filter(|f| !self.fields.contains(f))
203 .collect();
204 if !missing.is_empty() {
205 return Err(Error::invalid_input(format!(
206 "index '{}' declares covering fields {:?} but {:?} are not \
207 among its fields {:?}",
208 self.name, self.covering_fields, missing, self.fields,
209 )));
210 }
211
212 let mut seen = HashSet::with_capacity(self.covering_fields.len());
217 if let Some(duplicate) = self.covering_fields.iter().find(|f| !seen.insert(**f)) {
218 return Err(Error::invalid_input(format!(
219 "index '{}' declares covering field {} more than once in {:?}",
220 self.name, duplicate, self.covering_fields,
221 )));
222 }
223
224 if self.covering_fields.len() >= self.fields.len() {
225 return Err(Error::invalid_input(format!(
226 "index '{}' declares covering fields {:?} but its fields are {:?}; \
227 at least one field must remain indexed",
228 self.name, self.covering_fields, self.fields,
229 )));
230 }
231
232 let suffix_start = self.fields.len() - self.covering_fields.len();
233 if self.fields[suffix_start..] != self.covering_fields[..] {
234 return Err(Error::invalid_input(format!(
235 "index '{}' declares covering fields {:?} which are not the trailing \
236 entries of {:?}; covering fields must come last",
237 self.name, self.covering_fields, self.fields,
238 )));
239 }
240
241 Ok(())
242 }
243}
244
245impl DeepSizeOf for IndexMetadata {
246 fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
247 self.uuid.as_bytes().deep_size_of_children(context)
248 + self.fields.deep_size_of_children(context)
249 + self.covering_fields.deep_size_of_children(context)
250 + self.name.deep_size_of_children(context)
251 + self.dataset_version.deep_size_of_children(context)
252 + self
253 .fragment_bitmap
254 .as_ref()
255 .map(|fragment_bitmap| fragment_bitmap.serialized_size())
256 .unwrap_or(0)
257 + self.files.deep_size_of_children(context)
258 }
259}
260
261impl TryFrom<pb::IndexMetadata> for IndexMetadata {
262 type Error = Error;
263
264 fn try_from(proto: pb::IndexMetadata) -> Result<Self> {
265 let fragment_bitmap = if proto.fragment_bitmap.is_empty() {
266 None
267 } else {
268 Some(RoaringBitmap::deserialize_from(
269 &mut proto.fragment_bitmap.as_slice(),
270 )?)
271 };
272
273 let files = if proto.files.is_empty() {
274 None
275 } else {
276 Some(
277 proto
278 .files
279 .into_iter()
280 .map(|f| IndexFile {
281 path: f.path,
282 size_bytes: f.size_bytes,
283 })
284 .collect(),
285 )
286 };
287
288 let metadata = Self {
289 uuid: proto.uuid.as_ref().map(Uuid::try_from).ok_or_else(|| {
290 Error::invalid_input("uuid field does not exist in Index metadata".to_string())
291 })??,
292 name: proto.name,
293 fields: proto.fields,
294 covering_fields: proto.covering_fields,
295 dataset_version: proto.dataset_version,
296 fragment_bitmap,
297 index_details: proto.index_details.map(Arc::new),
298 index_version: proto.index_version.unwrap_or_default(),
299 created_at: proto.created_at.map(|ts| {
300 DateTime::from_timestamp_millis(ts as i64)
301 .expect("Invalid timestamp in index metadata")
302 }),
303 base_id: proto.base_id,
304 files,
305 };
306
307 metadata.validate_covering_fields()?;
314
315 Ok(metadata)
316 }
317}
318
319impl From<&IndexMetadata> for pb::IndexMetadata {
320 fn from(idx: &IndexMetadata) -> Self {
321 let mut fragment_bitmap = Vec::new();
322 if let Some(bitmap) = &idx.fragment_bitmap {
323 let mut bitmap = bitmap.clone();
329 bitmap.optimize();
330 if let Err(e) = bitmap.serialize_into(&mut fragment_bitmap) {
331 log::error!("Failed to serialize fragment bitmap: {}", e);
334 fragment_bitmap.clear();
335 }
336 }
337
338 let files = idx
339 .files
340 .as_ref()
341 .map(|files| {
342 files
343 .iter()
344 .map(|f| pb::IndexFile {
345 path: f.path.clone(),
346 size_bytes: f.size_bytes,
347 })
348 .collect()
349 })
350 .unwrap_or_default();
351
352 Self {
353 uuid: Some((&idx.uuid).into()),
354 name: idx.name.clone(),
355 fields: idx.fields.clone(),
356 covering_fields: idx.covering_fields.clone(),
357 dataset_version: idx.dataset_version,
358 fragment_bitmap,
359 index_details: idx
360 .index_details
361 .as_ref()
362 .map(|details| details.as_ref().clone()),
363 index_version: Some(idx.index_version),
364 created_at: idx.created_at.map(|dt| dt.timestamp_millis() as u64),
365 base_id: idx.base_id,
366 files,
367 }
368 }
369}
370
371type ArcAny = Arc<dyn std::any::Any + Send + Sync>;
379
380const INDEX_METADATA_TYPE_ID: &str = "lance.table.IndexMetadataList";
382const INDEX_METADATA_VERSION: u32 = 1;
384
385fn serialize_index_metadata(
386 any: &ArcAny,
387 writer: &mut CacheEntryWriter<'_>,
388) -> lance_core::Result<()> {
389 let vec = any
390 .downcast_ref::<Vec<IndexMetadata>>()
391 .expect("index_metadata_codec: wrong type (this is a bug in the cache layer)");
392 let section = pb::IndexSection {
393 indices: vec.iter().map(pb::IndexMetadata::from).collect(),
394 };
395 writer.write_header(§ion)
396}
397
398fn deserialize_index_metadata(reader: &mut CacheEntryReader<'_>) -> lance_core::Result<ArcAny> {
399 let section: pb::IndexSection = reader.read_header()?;
400 let indices: Vec<IndexMetadata> = section
401 .indices
402 .into_iter()
403 .map(IndexMetadata::try_from)
404 .collect::<lance_core::Result<_>>()?;
405 Ok(Arc::new(indices))
406}
407
408pub fn index_metadata_codec() -> lance_core::cache::CacheCodec {
409 lance_core::cache::CacheCodec::new(
410 INDEX_METADATA_TYPE_ID,
411 INDEX_METADATA_VERSION,
412 serialize_index_metadata,
413 deserialize_index_metadata,
414 )
415}
416
417pub async fn list_index_files_with_sizes(
422 object_store: &ObjectStore,
423 index_dir: &Path,
424) -> Result<Vec<IndexFile>> {
425 let mut files = Vec::new();
426 let mut stream = object_store.read_dir_all(index_dir, None);
427 while let Some(meta) = stream.next().await {
428 let meta = meta?;
429 let relative_path = meta
431 .location
432 .as_ref()
433 .strip_prefix(index_dir.as_ref())
434 .map(|s| s.trim_start_matches('/').to_string())
435 .unwrap_or_else(|| meta.location.filename().unwrap_or("").to_string());
436 files.push(IndexFile {
437 path: relative_path,
438 size_bytes: meta.size,
439 });
440 }
441 Ok(files)
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447 use rstest::rstest;
448 use std::collections::HashMap;
449
450 #[test]
451 fn test_fragment_bitmap_serialized_run_optimized() {
452 let bitmap = RoaringBitmap::from_sorted_iter(0..1_000_000).unwrap();
453 let unoptimized_size = bitmap.serialized_size();
454
455 let metadata = IndexMetadata {
456 uuid: Uuid::new_v4(),
457 name: "my_index".to_string(),
458 fields: vec![0],
459 covering_fields: vec![],
460 dataset_version: 1,
461 fragment_bitmap: Some(bitmap.clone()),
462 index_details: None,
463 index_version: 1,
464 created_at: None,
465 base_id: None,
466 files: None,
467 };
468
469 let proto = pb::IndexMetadata::from(&metadata);
470 assert!(
471 proto.fragment_bitmap.len() < unoptimized_size / 100,
472 "expected run-optimized bitmap ({} bytes) to be <1% of the \
473 unoptimized serialization ({} bytes)",
474 proto.fragment_bitmap.len(),
475 unoptimized_size
476 );
477
478 let recovered = IndexMetadata::try_from(proto).unwrap();
479 assert_eq!(recovered.fragment_bitmap, Some(bitmap));
480 }
481
482 #[test]
486 fn test_index_metadata_codec_roundtrip() {
487 let codec = index_metadata_codec();
488
489 let original = vec![
490 IndexMetadata {
491 uuid: Uuid::new_v4(),
492 name: "my_index".to_string(),
493 fields: vec![0, 1],
494 covering_fields: vec![],
495 dataset_version: 42,
496 fragment_bitmap: Some(RoaringBitmap::from_iter([1, 2, 3])),
497 index_details: None,
498 index_version: 1,
499 created_at: None,
500 base_id: None,
501 files: Some(vec![IndexFile {
502 path: "index.idx".to_string(),
503 size_bytes: 1024,
504 }]),
505 },
506 IndexMetadata {
507 uuid: Uuid::new_v4(),
508 name: "second_index".to_string(),
509 fields: vec![2],
510 covering_fields: vec![],
511 dataset_version: 43,
512 fragment_bitmap: None,
513 index_details: None,
514 index_version: 2,
515 created_at: None,
516 base_id: Some(7),
517 files: None,
518 },
519 ];
520
521 let mut store: HashMap<String, Vec<u8>> = HashMap::new();
523
524 let key = "dataset/v42/Vec<IndexMetadata>".to_string();
526 let mut buf = Vec::new();
527 let entry: Arc<dyn std::any::Any + Send + Sync> = Arc::new(original.clone());
528 codec.serialize(&entry, &mut buf).unwrap();
529 store.insert(key.clone(), buf);
530
531 let bytes = store.get(&key).unwrap();
533 let recovered = codec
534 .deserialize(&bytes::Bytes::copy_from_slice(bytes))
535 .hit()
536 .expect("entry should decode as a hit");
537 let recovered = recovered
538 .downcast::<Vec<IndexMetadata>>()
539 .expect("downcast should succeed");
540
541 assert_eq!(original.len(), recovered.len());
542 for (orig, rec) in original.iter().zip(recovered.iter()) {
543 assert_eq!(orig.uuid, rec.uuid);
544 assert_eq!(orig.name, rec.name);
545 assert_eq!(orig.fields, rec.fields);
546 assert_eq!(orig.dataset_version, rec.dataset_version);
547 assert_eq!(orig.fragment_bitmap, rec.fragment_bitmap);
548 assert_eq!(orig.index_version, rec.index_version);
549 assert_eq!(orig.base_id, rec.base_id);
550 assert_eq!(orig.files, rec.files);
551 }
552 }
553
554 #[test]
558 fn test_covering_fields_survives_proto_roundtrip() {
559 let original = IndexMetadata {
560 uuid: Uuid::new_v4(),
561 name: "covered".to_string(),
562 fields: vec![7, 11, 13],
563 covering_fields: vec![11, 13],
564 dataset_version: 1,
565 fragment_bitmap: None,
566 index_details: None,
567 index_version: 0,
568 created_at: None,
569 base_id: None,
570 files: None,
571 };
572
573 let proto = pb::IndexMetadata::from(&original);
574 assert_eq!(proto.covering_fields, vec![11, 13]);
575
576 let recovered = IndexMetadata::try_from(proto).unwrap();
577 assert_eq!(recovered, original);
578 }
579
580 #[test]
586 fn test_try_from_proto_rejects_a_malformed_covering_declaration() {
587 let mut proto = pb::IndexMetadata::from(&index_metadata_with(vec![7, 11], vec![11]));
588 proto.covering_fields = vec![7];
591
592 let err = IndexMetadata::try_from(proto)
593 .expect_err("a malformed covering declaration must not decode");
594 assert!(
595 matches!(err, Error::InvalidInput { .. }),
596 "expected InvalidInput, got {:?}",
597 err
598 );
599 assert!(
600 err.to_string().contains("must come last"),
601 "unexpected message: {err}"
602 );
603 }
604
605 #[rstest]
610 #[case::not_covered(vec![7], vec![], vec![7], Some(7))]
611 #[case::covered(vec![7, 11], vec![11], vec![7], Some(7))]
612 #[case::covered_multi(vec![7, 11, 13], vec![11, 13], vec![7], Some(7))]
613 #[case::covered_composite(vec![7, 11, 13], vec![13], vec![7, 11], None)]
614 #[case::composite(vec![7, 11], vec![], vec![7, 11], None)]
615 #[case::system_index_no_fields(vec![], vec![], vec![], None)]
616 #[case::malformed_longer_than_fields(vec![7], vec![11, 13], vec![], None)]
617 fn test_keyed_fields(
618 #[case] fields: Vec<i32>,
619 #[case] covering_fields: Vec<i32>,
620 #[case] expected_keyed: Vec<i32>,
621 #[case] expected_single: Option<i32>,
622 ) {
623 let metadata = index_metadata_with(fields, covering_fields);
624
625 assert_eq!(metadata.keyed_fields(), expected_keyed.as_slice());
626 assert_eq!(metadata.keyed_field(), expected_single);
627 }
628
629 fn index_metadata_with(fields: Vec<i32>, covering_fields: Vec<i32>) -> IndexMetadata {
630 IndexMetadata {
631 uuid: Uuid::new_v4(),
632 name: "idx".to_string(),
633 fields,
634 covering_fields,
635 dataset_version: 1,
636 fragment_bitmap: None,
637 index_details: None,
638 index_version: 0,
639 created_at: None,
640 base_id: None,
641 files: None,
642 }
643 }
644
645 #[rstest]
646 #[case::zone_map("type.googleapis.com/lance.table.ZoneMapIndexDetails", true)]
647 #[case::bloom_filter("type.googleapis.com/lance.index.pb.BloomFilterIndexDetails", true)]
648 #[case::fm("type.googleapis.com/lance.index.pb.FMIndexDetails", true)]
649 #[case::fm_case_insensitive("type.googleapis.com/LANCE.INDEX.PB.FMINDEXDETAILS", true)]
650 #[case::foreign_fm_terminal_name("type.googleapis.com/example.FMIndexDetails", false)]
651 #[case::btree("type.googleapis.com/lance.table.BTreeIndexDetails", false)]
652 fn test_results_are_row_addrs(#[case] type_url: &str, #[case] expected: bool) {
653 let mut metadata = index_metadata_with(vec![0], vec![]);
654 metadata.index_details = Some(Arc::new(prost_types::Any {
655 type_url: type_url.to_string(),
656 value: Vec::new(),
657 }));
658
659 assert_eq!(metadata.results_are_row_addrs(), expected);
660 }
661
662 #[rstest]
663 #[case::empty_is_valid(vec![7], vec![], None)]
664 #[case::system_index_no_fields(vec![], vec![], None)]
667 #[case::valid_single_covered(vec![7, 11], vec![11], None)]
668 #[case::valid_suffix(vec![7, 11, 13], vec![11, 13], None)]
669 #[case::not_a_suffix(vec![7, 11, 13], vec![11], Some("must come last"))]
670 #[case::not_a_subset(vec![7, 11], vec![99], Some("are not among its fields"))]
671 #[case::wrong_order(vec![7, 11, 13], vec![13, 11], Some("must come last"))]
672 #[case::all_fields_covered(vec![7], vec![7], Some("at least one field must remain indexed"))]
673 #[case::covers_the_search_key(vec![7, 11], vec![7, 11], Some("at least one field must remain indexed"))]
674 #[case::duplicate_covered(vec![7, 11, 11], vec![11, 11], Some("more than once"))]
677 #[case::unknown_id_reported_before_length(vec![7, 11], vec![99, 11], Some("are not among its fields"))]
680 fn test_validate_covering_fields(
681 #[case] fields: Vec<i32>,
682 #[case] covering_fields: Vec<i32>,
683 #[case] expected_error: Option<&str>,
684 ) {
685 let metadata = index_metadata_with(fields, covering_fields);
686 let result = metadata.validate_covering_fields();
687
688 match expected_error {
689 None => assert!(result.is_ok(), "expected valid, got {:?}", result),
690 Some(fragment) => {
691 let err = result.expect_err("expected a validation error");
692 assert!(
693 matches!(err, Error::InvalidInput { .. }),
694 "expected InvalidInput, got {:?}",
695 err
696 );
697 let message = err.to_string();
698 assert!(
699 message.contains(fragment),
700 "expected message to contain {:?}, got {:?}",
701 fragment,
702 message
703 );
704 }
705 }
706 }
707}