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 {
151 self.index_details.as_ref().is_some_and(|details| {
152 let is_fm = details
153 .type_url
154 .rsplit_once('/')
155 .is_some_and(|(_, details_type_name)| {
156 details_type_name.eq_ignore_ascii_case("lance.index.pb.FMIndexDetails")
157 });
158 details.type_url.ends_with("ZoneMapIndexDetails")
159 || details.type_url.ends_with("BloomFilterIndexDetails")
160 || is_fm
161 || Self::is_btree_addr_domain(details, self.index_version)
162 })
163 }
164
165 fn is_btree_addr_domain(details: &prost_types::Any, index_version: i32) -> bool {
171 details.type_url.ends_with("BTreeIndexDetails") && index_version >= 1
172 }
173
174 pub fn supports_eager_remap(&self) -> bool {
187 self.index_details
188 .as_ref()
189 .is_some_and(|details| Self::is_btree_addr_domain(details, self.index_version))
190 }
191
192 pub fn keyed_fields(&self) -> &[i32] {
201 let keyed = self.fields.len().saturating_sub(self.covering_fields.len());
202 &self.fields[..keyed]
203 }
204
205 pub fn keyed_field(&self) -> Option<i32> {
211 match self.keyed_fields() {
212 [only] => Some(*only),
213 _ => None,
214 }
215 }
216
217 pub fn validate_covering_fields(&self) -> Result<()> {
230 if self.covering_fields.is_empty() {
231 return Ok(());
232 }
233
234 let missing: Vec<i32> = self
235 .covering_fields
236 .iter()
237 .copied()
238 .filter(|f| !self.fields.contains(f))
239 .collect();
240 if !missing.is_empty() {
241 return Err(Error::invalid_input(format!(
242 "index '{}' declares covering fields {:?} but {:?} are not \
243 among its fields {:?}",
244 self.name, self.covering_fields, missing, self.fields,
245 )));
246 }
247
248 let mut seen = HashSet::with_capacity(self.covering_fields.len());
253 if let Some(duplicate) = self.covering_fields.iter().find(|f| !seen.insert(**f)) {
254 return Err(Error::invalid_input(format!(
255 "index '{}' declares covering field {} more than once in {:?}",
256 self.name, duplicate, self.covering_fields,
257 )));
258 }
259
260 if self.covering_fields.len() >= self.fields.len() {
261 return Err(Error::invalid_input(format!(
262 "index '{}' declares covering fields {:?} but its fields are {:?}; \
263 at least one field must remain indexed",
264 self.name, self.covering_fields, self.fields,
265 )));
266 }
267
268 let suffix_start = self.fields.len() - self.covering_fields.len();
269 if self.fields[suffix_start..] != self.covering_fields[..] {
270 return Err(Error::invalid_input(format!(
271 "index '{}' declares covering fields {:?} which are not the trailing \
272 entries of {:?}; covering fields must come last",
273 self.name, self.covering_fields, self.fields,
274 )));
275 }
276
277 Ok(())
278 }
279}
280
281impl DeepSizeOf for IndexMetadata {
282 fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
283 self.uuid.as_bytes().deep_size_of_children(context)
284 + self.fields.deep_size_of_children(context)
285 + self.covering_fields.deep_size_of_children(context)
286 + self.name.deep_size_of_children(context)
287 + self.dataset_version.deep_size_of_children(context)
288 + self
289 .fragment_bitmap
290 .as_ref()
291 .map(|fragment_bitmap| fragment_bitmap.serialized_size())
292 .unwrap_or(0)
293 + self.files.deep_size_of_children(context)
294 + self
295 .index_details
296 .as_ref()
297 .map(|details| {
298 if context.mark_seen(Arc::as_ptr(details) as usize) {
299 std::mem::size_of::<prost_types::Any>()
300 + details.type_url.capacity()
301 + details.value.capacity()
302 } else {
303 0
304 }
305 })
306 .unwrap_or(0)
307 }
308}
309
310impl TryFrom<pb::IndexMetadata> for IndexMetadata {
311 type Error = Error;
312
313 fn try_from(proto: pb::IndexMetadata) -> Result<Self> {
314 let fragment_bitmap = if proto.fragment_bitmap.is_empty() {
315 None
316 } else {
317 Some(RoaringBitmap::deserialize_from(
318 &mut proto.fragment_bitmap.as_slice(),
319 )?)
320 };
321
322 let files = if proto.files.is_empty() {
323 None
324 } else {
325 Some(
326 proto
327 .files
328 .into_iter()
329 .map(|f| IndexFile {
330 path: f.path,
331 size_bytes: f.size_bytes,
332 })
333 .collect(),
334 )
335 };
336
337 let metadata = Self {
338 uuid: proto.uuid.as_ref().map(Uuid::try_from).ok_or_else(|| {
339 Error::invalid_input("uuid field does not exist in Index metadata".to_string())
340 })??,
341 name: proto.name,
342 fields: proto.fields,
343 covering_fields: proto.covering_fields,
344 dataset_version: proto.dataset_version,
345 fragment_bitmap,
346 index_details: proto.index_details.map(Arc::new),
347 index_version: proto.index_version.unwrap_or_default(),
348 created_at: proto.created_at.map(|ts| {
349 DateTime::from_timestamp_millis(ts as i64)
350 .expect("Invalid timestamp in index metadata")
351 }),
352 base_id: proto.base_id,
353 files,
354 };
355
356 metadata.validate_covering_fields()?;
363
364 Ok(metadata)
365 }
366}
367
368impl From<&IndexMetadata> for pb::IndexMetadata {
369 fn from(idx: &IndexMetadata) -> Self {
370 let mut fragment_bitmap = Vec::new();
371 if let Some(bitmap) = &idx.fragment_bitmap {
372 let mut bitmap = bitmap.clone();
378 bitmap.optimize();
379 if let Err(e) = bitmap.serialize_into(&mut fragment_bitmap) {
380 log::error!("Failed to serialize fragment bitmap: {}", e);
383 fragment_bitmap.clear();
384 }
385 }
386
387 let files = idx
388 .files
389 .as_ref()
390 .map(|files| {
391 files
392 .iter()
393 .map(|f| pb::IndexFile {
394 path: f.path.clone(),
395 size_bytes: f.size_bytes,
396 })
397 .collect()
398 })
399 .unwrap_or_default();
400
401 Self {
402 uuid: Some((&idx.uuid).into()),
403 name: idx.name.clone(),
404 fields: idx.fields.clone(),
405 covering_fields: idx.covering_fields.clone(),
406 dataset_version: idx.dataset_version,
407 fragment_bitmap,
408 index_details: idx
409 .index_details
410 .as_ref()
411 .map(|details| details.as_ref().clone()),
412 index_version: Some(idx.index_version),
413 clustering_id: None,
414 metadata: Default::default(),
415 created_at: idx.created_at.map(|dt| dt.timestamp_millis() as u64),
416 base_id: idx.base_id,
417 files,
418 }
419 }
420}
421
422type ArcAny = Arc<dyn std::any::Any + Send + Sync>;
430
431const INDEX_METADATA_TYPE_ID: &str = "lance.table.IndexMetadataList";
433const INDEX_METADATA_VERSION: u32 = 1;
435
436fn serialize_index_metadata(
437 any: &ArcAny,
438 writer: &mut CacheEntryWriter<'_>,
439) -> lance_core::Result<()> {
440 let vec = any
441 .downcast_ref::<Vec<IndexMetadata>>()
442 .expect("index_metadata_codec: wrong type (this is a bug in the cache layer)");
443 let section = pb::IndexSection {
444 indices: vec.iter().map(pb::IndexMetadata::from).collect(),
445 };
446 writer.write_header(§ion)
447}
448
449fn deserialize_index_metadata(reader: &mut CacheEntryReader<'_>) -> lance_core::Result<ArcAny> {
450 let section: pb::IndexSection = reader.read_header()?;
451 let indices: Vec<IndexMetadata> = section
452 .indices
453 .into_iter()
454 .map(IndexMetadata::try_from)
455 .collect::<lance_core::Result<_>>()?;
456 Ok(Arc::new(indices))
457}
458
459pub fn index_metadata_codec() -> lance_core::cache::CacheCodec {
460 lance_core::cache::CacheCodec::new(
461 INDEX_METADATA_TYPE_ID,
462 INDEX_METADATA_VERSION,
463 serialize_index_metadata,
464 deserialize_index_metadata,
465 )
466}
467
468pub async fn list_index_files_with_sizes(
473 object_store: &ObjectStore,
474 index_dir: &Path,
475) -> Result<Vec<IndexFile>> {
476 let mut files = Vec::new();
477 let mut stream = object_store.read_dir_all(index_dir, None);
478 while let Some(meta) = stream.next().await {
479 let meta = meta?;
480 let relative_path = meta
482 .location
483 .as_ref()
484 .strip_prefix(index_dir.as_ref())
485 .map(|s| s.trim_start_matches('/').to_string())
486 .unwrap_or_else(|| meta.location.filename().unwrap_or("").to_string());
487 files.push(IndexFile {
488 path: relative_path,
489 size_bytes: meta.size,
490 });
491 }
492 Ok(files)
493}
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498 use rstest::rstest;
499 use std::collections::HashMap;
500
501 #[test]
502 fn test_fragment_bitmap_serialized_run_optimized() {
503 let bitmap = RoaringBitmap::from_sorted_iter(0..1_000_000).unwrap();
504 let unoptimized_size = bitmap.serialized_size();
505
506 let metadata = IndexMetadata {
507 uuid: Uuid::new_v4(),
508 name: "my_index".to_string(),
509 fields: vec![0],
510 covering_fields: vec![],
511 dataset_version: 1,
512 fragment_bitmap: Some(bitmap.clone()),
513 index_details: None,
514 index_version: 1,
515 created_at: None,
516 base_id: None,
517 files: None,
518 };
519
520 let proto = pb::IndexMetadata::from(&metadata);
521 assert!(
522 proto.fragment_bitmap.len() < unoptimized_size / 100,
523 "expected run-optimized bitmap ({} bytes) to be <1% of the \
524 unoptimized serialization ({} bytes)",
525 proto.fragment_bitmap.len(),
526 unoptimized_size
527 );
528
529 let recovered = IndexMetadata::try_from(proto).unwrap();
530 assert_eq!(recovered.fragment_bitmap, Some(bitmap));
531 }
532
533 #[test]
534 fn test_deep_size_counts_index_details() {
535 fn details() -> Arc<prost_types::Any> {
537 let mut type_url = String::with_capacity(128);
538 type_url.push_str("/lance.table.FragmentReuseIndexDetails");
539 let mut value = Vec::with_capacity(12_000);
540 value.resize(10_000, 0u8);
541 assert_eq!((type_url.capacity(), value.capacity()), (128, 12_000));
542 Arc::new(prost_types::Any { type_url, value })
543 }
544 let details_size = std::mem::size_of::<prost_types::Any>() + 128 + 12_000;
546
547 let meta = IndexMetadata {
548 uuid: Uuid::new_v4(),
549 name: "__lance_frag_reuse".to_string(),
550 fields: vec![],
551 covering_fields: vec![],
552 dataset_version: 1,
553 fragment_bitmap: None,
554 index_details: None,
555 index_version: 0,
556 created_at: None,
557 base_id: None,
558 files: None,
559 };
560 let with = |details: Arc<prost_types::Any>| IndexMetadata {
561 index_details: Some(details),
562 ..meta.clone()
563 };
564
565 assert_eq!(
566 with(details()).deep_size_of() - meta.deep_size_of(),
567 details_size
568 );
569
570 let two_without = vec![meta.clone(), meta.clone()].deep_size_of();
571
572 let shared = details();
574 let shared = vec![with(shared.clone()), with(shared)];
575 assert_eq!(shared.deep_size_of() - two_without, details_size);
576
577 let independent = vec![with(details()), with(details())];
579 assert_eq!(independent.deep_size_of() - two_without, 2 * details_size);
580 }
581
582 #[test]
586 fn test_index_metadata_codec_roundtrip() {
587 let codec = index_metadata_codec();
588
589 let original = vec![
590 IndexMetadata {
591 uuid: Uuid::new_v4(),
592 name: "my_index".to_string(),
593 fields: vec![0, 1],
594 covering_fields: vec![],
595 dataset_version: 42,
596 fragment_bitmap: Some(RoaringBitmap::from_iter([1, 2, 3])),
597 index_details: None,
598 index_version: 1,
599 created_at: None,
600 base_id: None,
601 files: Some(vec![IndexFile {
602 path: "index.idx".to_string(),
603 size_bytes: 1024,
604 }]),
605 },
606 IndexMetadata {
607 uuid: Uuid::new_v4(),
608 name: "second_index".to_string(),
609 fields: vec![2],
610 covering_fields: vec![],
611 dataset_version: 43,
612 fragment_bitmap: None,
613 index_details: None,
614 index_version: 2,
615 created_at: None,
616 base_id: Some(7),
617 files: None,
618 },
619 ];
620
621 let mut store: HashMap<String, Vec<u8>> = HashMap::new();
623
624 let key = "dataset/v42/Vec<IndexMetadata>".to_string();
626 let mut buf = Vec::new();
627 let entry: Arc<dyn std::any::Any + Send + Sync> = Arc::new(original.clone());
628 codec.serialize(&entry, &mut buf).unwrap();
629 store.insert(key.clone(), buf);
630
631 let bytes = store.get(&key).unwrap();
633 let recovered = codec
634 .deserialize(&bytes::Bytes::copy_from_slice(bytes))
635 .hit()
636 .expect("entry should decode as a hit");
637 let recovered = recovered
638 .downcast::<Vec<IndexMetadata>>()
639 .expect("downcast should succeed");
640
641 assert_eq!(original.len(), recovered.len());
642 for (orig, rec) in original.iter().zip(recovered.iter()) {
643 assert_eq!(orig.uuid, rec.uuid);
644 assert_eq!(orig.name, rec.name);
645 assert_eq!(orig.fields, rec.fields);
646 assert_eq!(orig.dataset_version, rec.dataset_version);
647 assert_eq!(orig.fragment_bitmap, rec.fragment_bitmap);
648 assert_eq!(orig.index_version, rec.index_version);
649 assert_eq!(orig.base_id, rec.base_id);
650 assert_eq!(orig.files, rec.files);
651 }
652 }
653
654 #[test]
658 fn test_covering_fields_survives_proto_roundtrip() {
659 let original = IndexMetadata {
660 uuid: Uuid::new_v4(),
661 name: "covered".to_string(),
662 fields: vec![7, 11, 13],
663 covering_fields: vec![11, 13],
664 dataset_version: 1,
665 fragment_bitmap: None,
666 index_details: None,
667 index_version: 0,
668 created_at: None,
669 base_id: None,
670 files: None,
671 };
672
673 let proto = pb::IndexMetadata::from(&original);
674 assert_eq!(proto.covering_fields, vec![11, 13]);
675
676 let recovered = IndexMetadata::try_from(proto).unwrap();
677 assert_eq!(recovered, original);
678 }
679
680 #[test]
686 fn test_try_from_proto_rejects_a_malformed_covering_declaration() {
687 let mut proto = pb::IndexMetadata::from(&index_metadata_with(vec![7, 11], vec![11]));
688 proto.covering_fields = vec![7];
691
692 let err = IndexMetadata::try_from(proto)
693 .expect_err("a malformed covering declaration must not decode");
694 assert!(
695 matches!(err, Error::InvalidInput { .. }),
696 "expected InvalidInput, got {:?}",
697 err
698 );
699 assert!(
700 err.to_string().contains("must come last"),
701 "unexpected message: {err}"
702 );
703 }
704
705 #[rstest]
710 #[case::not_covered(vec![7], vec![], vec![7], Some(7))]
711 #[case::covered(vec![7, 11], vec![11], vec![7], Some(7))]
712 #[case::covered_multi(vec![7, 11, 13], vec![11, 13], vec![7], Some(7))]
713 #[case::covered_composite(vec![7, 11, 13], vec![13], vec![7, 11], None)]
714 #[case::composite(vec![7, 11], vec![], vec![7, 11], None)]
715 #[case::system_index_no_fields(vec![], vec![], vec![], None)]
716 #[case::malformed_longer_than_fields(vec![7], vec![11, 13], vec![], None)]
717 fn test_keyed_fields(
718 #[case] fields: Vec<i32>,
719 #[case] covering_fields: Vec<i32>,
720 #[case] expected_keyed: Vec<i32>,
721 #[case] expected_single: Option<i32>,
722 ) {
723 let metadata = index_metadata_with(fields, covering_fields);
724
725 assert_eq!(metadata.keyed_fields(), expected_keyed.as_slice());
726 assert_eq!(metadata.keyed_field(), expected_single);
727 }
728
729 fn index_metadata_with(fields: Vec<i32>, covering_fields: Vec<i32>) -> IndexMetadata {
730 IndexMetadata {
731 uuid: Uuid::new_v4(),
732 name: "idx".to_string(),
733 fields,
734 covering_fields,
735 dataset_version: 1,
736 fragment_bitmap: None,
737 index_details: None,
738 index_version: 0,
739 created_at: None,
740 base_id: None,
741 files: None,
742 }
743 }
744
745 #[rstest]
746 #[case::zone_map("type.googleapis.com/lance.table.ZoneMapIndexDetails", true)]
747 #[case::bloom_filter("type.googleapis.com/lance.index.pb.BloomFilterIndexDetails", true)]
748 #[case::fm("type.googleapis.com/lance.index.pb.FMIndexDetails", true)]
749 #[case::fm_case_insensitive("type.googleapis.com/LANCE.INDEX.PB.FMINDEXDETAILS", true)]
750 #[case::foreign_fm_terminal_name("type.googleapis.com/example.FMIndexDetails", false)]
751 #[case::btree_legacy("type.googleapis.com/lance.table.BTreeIndexDetails", false)]
752 fn test_results_are_row_addrs(#[case] type_url: &str, #[case] expected: bool) {
753 let mut metadata = index_metadata_with(vec![0], vec![]);
754 metadata.index_details = Some(Arc::new(prost_types::Any {
755 type_url: type_url.to_string(),
756 value: Vec::new(),
757 }));
758 metadata.index_version = 0;
759
760 assert_eq!(metadata.results_are_row_addrs(), expected);
761 }
762
763 #[test]
764 fn test_results_are_row_addrs_btree_is_version_gated() {
765 let mut metadata = index_metadata_with(vec![0], vec![]);
766 metadata.index_details = Some(Arc::new(prost_types::Any {
767 type_url: "type.googleapis.com/lance.table.BTreeIndexDetails".to_string(),
768 value: Vec::new(),
769 }));
770
771 metadata.index_version = 0;
772 assert!(
773 !metadata.results_are_row_addrs(),
774 "a legacy BTree segment stores row ids directly"
775 );
776
777 metadata.index_version = 1;
778 assert!(
779 metadata.results_are_row_addrs(),
780 "a BTree segment built with address-domain support stores row addresses"
781 );
782 }
783
784 #[rstest]
785 #[case::empty_is_valid(vec![7], vec![], None)]
786 #[case::system_index_no_fields(vec![], vec![], None)]
789 #[case::valid_single_covered(vec![7, 11], vec![11], None)]
790 #[case::valid_suffix(vec![7, 11, 13], vec![11, 13], None)]
791 #[case::not_a_suffix(vec![7, 11, 13], vec![11], Some("must come last"))]
792 #[case::not_a_subset(vec![7, 11], vec![99], Some("are not among its fields"))]
793 #[case::wrong_order(vec![7, 11, 13], vec![13, 11], Some("must come last"))]
794 #[case::all_fields_covered(vec![7], vec![7], Some("at least one field must remain indexed"))]
795 #[case::covers_the_search_key(vec![7, 11], vec![7, 11], Some("at least one field must remain indexed"))]
796 #[case::duplicate_covered(vec![7, 11, 11], vec![11, 11], Some("more than once"))]
799 #[case::unknown_id_reported_before_length(vec![7, 11], vec![99, 11], Some("are not among its fields"))]
802 fn test_validate_covering_fields(
803 #[case] fields: Vec<i32>,
804 #[case] covering_fields: Vec<i32>,
805 #[case] expected_error: Option<&str>,
806 ) {
807 let metadata = index_metadata_with(fields, covering_fields);
808 let result = metadata.validate_covering_fields();
809
810 match expected_error {
811 None => assert!(result.is_ok(), "expected valid, got {:?}", result),
812 Some(fragment) => {
813 let err = result.expect_err("expected a validation error");
814 assert!(
815 matches!(err, Error::InvalidInput { .. }),
816 "expected InvalidInput, got {:?}",
817 err
818 );
819 let message = err.to_string();
820 assert!(
821 message.contains(fragment),
822 "expected message to contain {:?}, got {:?}",
823 fragment,
824 message
825 );
826 }
827 }
828 }
829}