From 41b5c4e0cdd62b72f45904c88eef5eff081f431b Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 14:06:20 -0700 Subject: [PATCH 1/6] feat(format): add file metadata size hint Co-authored-by: Cursor --- docs/src/format/table/index.md | 14 ++++++++ protos/table.proto | 17 ++++++++++ rust/lance-index/src/scalar.rs | 9 ++--- .../src/vector/distributed/index_merger.rs | 1 + rust/lance-table/benches/manifest_intern.rs | 2 ++ rust/lance-table/src/format/fragment.rs | 33 +++++++++++++++++-- rust/lance-table/src/format/index.rs | 21 +++++++++++- .../src/transaction/manifest_build.rs | 4 +++ rust/lance-table/src/transaction/proto.rs | 5 +++ rust/lance/src/dataset/files.rs | 1 + rust/lance/src/dataset/fragment.rs | 1 + .../src/dataset/tests/dataset_merge_update.rs | 4 +++ rust/lance/src/dataset/write/commit.rs | 1 + rust/lance/src/index.rs | 1 + rust/lance/src/index/vector/builder.rs | 2 ++ rust/lance/src/index/vector/ivf.rs | 6 ++++ 16 files changed, 115 insertions(+), 7 deletions(-) diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index f9da132cf3b..162ac58d2ab 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -135,6 +135,20 @@ or independently of column indices due to variable encoding widths (for Lance fi +### File Metadata Size Hints + +`DataFile.file_metadata_size_bytes` records the suffix length from the Lance +schema descriptor through the end of the data file. `IndexFile` records the +same value for Lance-format files in an index segment; it is zero for other +file formats. Combined with the corresponding total file size, a reader can +request this suffix before decoding the footer and remove a dependent metadata +request. + +The value is advisory. Zero means unavailable, values larger than the file are +invalid, and the offsets and counts decoded from the file footer remain +authoritative. A reader must preserve the ordinary footer-directed fallback +for missing, invalid, underestimated, or otherwise inconsistent hints. + !!! note "Field-to-column mapping differs between data storage versions" In **2.0**, all fields (including non-leaf fields like struct and list containers) are assigned diff --git a/protos/table.proto b/protos/table.proto index f110b471c82..aac37deb220 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -361,6 +361,14 @@ message IndexFile { string path = 1; // Size of the file in bytes uint64 size_bytes = 2; + /* Number of bytes from the Lance schema descriptor through the end of the file. + * + * A reader that knows `size_bytes` may use this advisory value to issue the + * complete metadata-suffix read without first decoding the footer. The footer + * remains authoritative, and readers must fall back when this value is zero or + * invalid. Zero also covers files that do not use the Lance file format. + */ + uint64 file_metadata_size_bytes = 3; } // Index Section, containing a list of index metadata for one dataset version. @@ -520,6 +528,15 @@ message DataFile { * Lance use it as key of the base_paths field in Manifest to determine the actual base path of the data file. */ optional uint32 base_id = 7; + + /* Number of bytes from the Lance schema descriptor through the end of the file. + * + * A reader that knows `file_size_bytes` may use this advisory value to issue + * the complete metadata-suffix read without first decoding the footer. The + * footer remains authoritative, and readers must fall back when this value is + * zero or invalid. + */ + uint64 file_metadata_size_bytes = 8; } // DataFile /* An overlay file supplies new values for a subset of (row offset, field) cells diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 6dcfc9018e0..696c1123843 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -52,8 +52,8 @@ pub use inverted::tokenizer::InvertedIndexParams; /// Convert a `Vec<`[`lance_index_core::scalar::IndexFile`]`>` to a /// `Vec<`[`lance_table::format::IndexFile`]`>`. /// -/// These two structs have identical fields; this helper bridges the crate -/// boundary without relying on orphan-rule–violating `From` impls. +/// These structs share their path and total size fields; this helper bridges the +/// crate boundary without relying on orphan-rule–violating `From` impls. pub fn index_files_to_table( files: Vec, ) -> Vec { @@ -62,6 +62,7 @@ pub fn index_files_to_table( .map(|f| lance_table::format::IndexFile { path: f.path, size_bytes: f.size_bytes, + file_metadata_size_bytes: None, }) .collect() } @@ -69,8 +70,8 @@ pub fn index_files_to_table( /// Convert a `Vec<`[`lance_table::format::IndexFile`]`>` to a /// `Vec<`[`lance_index_core::scalar::IndexFile`]`>`. /// -/// These two structs have identical fields; this helper bridges the crate -/// boundary without relying on orphan-rule–violating `From` impls. +/// These structs share their path and total size fields; this helper bridges the +/// crate boundary without relying on orphan-rule–violating `From` impls. pub fn table_files_to_index( files: Vec, ) -> Vec { diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 370011eb8bd..4081ee10cb4 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -1724,6 +1724,7 @@ async fn merge_partial_vector_auxiliary_files_inner( Ok(lance_table::format::IndexFile { path: INDEX_AUXILIARY_FILE_NAME.to_string(), size_bytes: summary.size_bytes, + file_metadata_size_bytes: None, }) } else { Err(Error::index( diff --git a/rust/lance-table/benches/manifest_intern.rs b/rust/lance-table/benches/manifest_intern.rs index 81bd57c1a22..2036d3bf325 100644 --- a/rust/lance-table/benches/manifest_intern.rs +++ b/rust/lance-table/benches/manifest_intern.rs @@ -57,6 +57,7 @@ fn make_uniform_pb_fragments(n: u64, num_fields: usize) -> Vec file_major_version: 2, file_minor_version: 0, file_size_bytes: 0, + file_metadata_size_bytes: 0, base_id: None, }], overlays: vec![], @@ -134,6 +135,7 @@ fn make_diverse_pb_fragments( file_major_version: 2, file_minor_version: 0, file_size_bytes: 0, + file_metadata_size_bytes: 0, base_id: None, }], overlays: vec![], diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 82afee22c1d..014438f12ae 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -25,7 +25,7 @@ use lance_core::error::Result; /// Lance Data File /// /// A data file is one piece of file storing data. -#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct DataFile { /// Relative path of the data file to dataset root. pub path: String, @@ -53,21 +53,40 @@ pub struct DataFile { /// The size of the file in bytes, if known. pub file_size_bytes: CachedFileSize, + /// The metadata suffix size in bytes, if known. + /// + /// The suffix starts at the Lance schema descriptor and ends at EOF. Readers + /// treat this as an advisory hint and use the file footer as authoritative. + /// `None` preserves the footer-directed open path used for older manifests. + pub file_metadata_size_bytes: Option>, + /// The base path of the datafile, when the datafile is outside the dataset. pub base_id: Option, } +impl DeepSizeOf for DataFile { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.path.deep_size_of_children(context) + + self.fields.deep_size_of_children(context) + + self.column_indices.deep_size_of_children(context) + } +} + // Custom Serialize: convert Arc<[i32]> to slice for transparent JSON output impl Serialize for DataFile { fn serialize(&self, serializer: S) -> std::result::Result { use serde::ser::SerializeStruct; - let mut s = serializer.serialize_struct("DataFile", 7)?; + let field_count = 7 + usize::from(self.file_metadata_size_bytes.is_some()); + let mut s = serializer.serialize_struct("DataFile", field_count)?; s.serialize_field("path", &self.path)?; s.serialize_field("fields", self.fields.as_ref())?; s.serialize_field("column_indices", self.column_indices.as_ref())?; s.serialize_field("file_major_version", &self.file_major_version)?; s.serialize_field("file_minor_version", &self.file_minor_version)?; s.serialize_field("file_size_bytes", &self.file_size_bytes)?; + if let Some(file_metadata_size_bytes) = self.file_metadata_size_bytes { + s.serialize_field("file_metadata_size_bytes", &file_metadata_size_bytes)?; + } s.serialize_field("base_id", &self.base_id)?; s.end() } @@ -87,6 +106,8 @@ impl<'de> Deserialize<'de> for DataFile { #[serde(default)] file_minor_version: u32, file_size_bytes: CachedFileSize, + #[serde(default)] + file_metadata_size_bytes: Option>, base_id: Option, } @@ -98,6 +119,7 @@ impl<'de> Deserialize<'de> for DataFile { file_major_version: helper.file_major_version, file_minor_version: helper.file_minor_version, file_size_bytes: helper.file_size_bytes, + file_metadata_size_bytes: helper.file_metadata_size_bytes, base_id: helper.base_id, }) } @@ -121,6 +143,7 @@ impl DataFile { file_major_version, file_minor_version, file_size_bytes: file_size_bytes.into(), + file_metadata_size_bytes: None, base_id, } } @@ -135,6 +158,7 @@ impl DataFile { file_major_version, file_minor_version, file_size_bytes: Default::default(), + file_metadata_size_bytes: None, base_id: None, } } @@ -220,6 +244,7 @@ impl From<&DataFile> for pb::DataFile { file_minor_version: df.file_minor_version, file_size_bytes: df.file_size_bytes.get().map_or(0, |v| v.get()), base_id: df.base_id, + file_metadata_size_bytes: df.file_metadata_size_bytes.map_or(0, NonZero::get), } } } @@ -236,6 +261,7 @@ impl TryFrom for DataFile { file_minor_version: proto.file_minor_version, file_size_bytes: CachedFileSize::new(proto.file_size_bytes), base_id: proto.base_id, + file_metadata_size_bytes: NonZero::new(proto.file_metadata_size_bytes), }) } } @@ -357,6 +383,7 @@ impl DataFileFieldInterner { file_minor_version: proto.file_minor_version, file_size_bytes: CachedFileSize::new(proto.file_size_bytes), base_id: proto.base_id, + file_metadata_size_bytes: NonZero::new(proto.file_metadata_size_bytes), }) } @@ -899,6 +926,7 @@ mod tests { let mut fragment = Fragment::new(123); let schema = ArrowSchema::new(vec![ArrowField::new("x", DataType::Float16, true)]); fragment.add_file_legacy("foobar.lance", &Schema::try_from(&schema).unwrap()); + fragment.files[0].file_metadata_size_bytes = NonZero::new(1_024); fragment.deletion_file = Some(DeletionFile { read_version: 123, id: 456, @@ -1011,6 +1039,7 @@ mod tests { file_major_version: MAJOR_VERSION as u32, file_minor_version: MINOR_VERSION as u32, file_size_bytes: Default::default(), + file_metadata_size_bytes: None, base_id: None, }; diff --git a/rust/lance-table/src/format/index.rs b/rust/lance-table/src/format/index.rs index bfd195a2a7f..5e67aecd818 100644 --- a/rust/lance-table/src/format/index.rs +++ b/rust/lance-table/src/format/index.rs @@ -4,6 +4,7 @@ //! Metadata for index use std::collections::{HashMap, HashSet}; +use std::num::NonZero; use std::sync::Arc; use chrono::{DateTime, Utc}; @@ -19,12 +20,24 @@ use lance_core::cache::{CacheEntryReader, CacheEntryWriter}; use lance_core::{Error, Result}; /// Metadata about a single file within an index segment. -#[derive(Debug, Clone, PartialEq, DeepSizeOf)] +#[derive(Debug, Clone, PartialEq)] pub struct IndexFile { /// Path relative to the index directory (e.g., "index.idx", "auxiliary.idx") pub path: String, /// Size of the file in bytes pub size_bytes: u64, + /// The metadata suffix size in bytes, if this is a Lance-format file. + /// + /// The suffix starts at the Lance schema descriptor and ends at EOF. Readers + /// treat this as an advisory hint and use the file footer as authoritative. + /// `None` also represents index files that use another format. + pub file_metadata_size_bytes: Option>, +} + +impl DeepSizeOf for IndexFile { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.path.deep_size_of_children(context) + } } /// Index metadata @@ -280,6 +293,7 @@ impl TryFrom for IndexMetadata { .map(|f| IndexFile { path: f.path, size_bytes: f.size_bytes, + file_metadata_size_bytes: NonZero::new(f.file_metadata_size_bytes), }) .collect(), ) @@ -344,6 +358,9 @@ impl From<&IndexMetadata> for pb::IndexMetadata { .map(|f| pb::IndexFile { path: f.path.clone(), size_bytes: f.size_bytes, + file_metadata_size_bytes: f + .file_metadata_size_bytes + .map_or(0, NonZero::get), }) .collect() }) @@ -436,6 +453,7 @@ pub async fn list_index_files_with_sizes( files.push(IndexFile { path: relative_path, size_bytes: meta.size, + file_metadata_size_bytes: None, }); } Ok(files) @@ -501,6 +519,7 @@ mod tests { files: Some(vec![IndexFile { path: "index.idx".to_string(), size_bytes: 1024, + file_metadata_size_bytes: NonZero::new(256), }]), }, IndexMetadata { diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 0c48ff83dc6..07857f5f5c2 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1764,6 +1764,7 @@ mod tests { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::new(1000), + file_metadata_size_bytes: None, base_id: None, }); @@ -1775,6 +1776,7 @@ mod tests { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::new(500), + file_metadata_size_bytes: None, base_id: None, }); @@ -1786,6 +1788,7 @@ mod tests { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::new(750), + file_metadata_size_bytes: None, base_id: None, }); @@ -1797,6 +1800,7 @@ mod tests { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::new(250), + file_metadata_size_bytes: None, base_id: None, }); diff --git a/rust/lance-table/src/transaction/proto.rs b/rust/lance-table/src/transaction/proto.rs index c51c8f16719..1c20c03b92c 100644 --- a/rust/lance-table/src/transaction/proto.rs +++ b/rust/lance-table/src/transaction/proto.rs @@ -22,6 +22,7 @@ use lance_core::{Error, Result}; use lance_file::datatypes::Fields; use roaring::RoaringBitmap; use std::collections::HashMap; +use std::num::NonZero; use std::sync::Arc; use uuid::Uuid; @@ -462,6 +463,7 @@ impl TryFrom<&pb::transaction::rewrite::RewrittenIndex> for RewrittenIndex { .map(|f| IndexFile { path: f.path.clone(), size_bytes: f.size_bytes, + file_metadata_size_bytes: NonZero::new(f.file_metadata_size_bytes), }) .collect(), ) @@ -739,6 +741,9 @@ impl From<&RewrittenIndex> for pb::transaction::rewrite::RewrittenIndex { .map(|f| pb::IndexFile { path: f.path.clone(), size_bytes: f.size_bytes, + file_metadata_size_bytes: f + .file_metadata_size_bytes + .map_or(0, NonZero::get), }) .collect() }) diff --git a/rust/lance/src/dataset/files.rs b/rust/lance/src/dataset/files.rs index cd9c05d46c7..2e4ea70f434 100644 --- a/rust/lance/src/dataset/files.rs +++ b/rust/lance/src/dataset/files.rs @@ -903,6 +903,7 @@ mod tests { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id, }; diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index a353eb148c4..824f96474b7 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -6989,6 +6989,7 @@ mod tests { file_major_version: 2, file_minor_version: 1, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id: None, }; diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index 2a5a148fe79..730655438e2 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -883,6 +883,7 @@ async fn test_datafile_partial_replacement() { file_major_version: major, file_minor_version: minor, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id: None, }; @@ -944,6 +945,7 @@ async fn test_datafile_partial_replacement() { file_major_version: major, file_minor_version: minor, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id: None, }; @@ -1044,6 +1046,7 @@ async fn test_datafile_replacement_error() { file_major_version: 2, file_minor_version: 0, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id: None, }; @@ -3266,6 +3269,7 @@ async fn test_cross_column_fast_search_blocks_column_local_stale_postings() { file_major_version, file_minor_version, file_size_bytes: CachedFileSize::unknown(), + file_metadata_size_bytes: None, base_id: None, }; let read_version = dataset.version().version; diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index c653ee7d8a3..c0c2f71764f 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -637,6 +637,7 @@ mod tests { file_major_version: major_version, file_minor_version: minor_version, file_size_bytes: CachedFileSize::new(100), + file_metadata_size_bytes: None, base_id: None, }], overlays: vec![], diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 74b63cb5010..9b807fbbd79 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -3825,6 +3825,7 @@ mod tests { files: Some(vec![lance_table::format::IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: payload.len() as u64, + file_metadata_size_bytes: None, }]), } } diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index a7525529db4..5231990dd43 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -1857,10 +1857,12 @@ impl IvfIndexBuilder IndexFile { path: INDEX_AUXILIARY_FILE_NAME.to_string(), size_bytes: storage_summary.size_bytes, + file_metadata_size_bytes: None, }, IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: index_summary.size_bytes, + file_metadata_size_bytes: None, }, ]) } diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index ab8ee6fbf52..58cfd1884ce 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -1109,6 +1109,7 @@ async fn optimize_ivf_pq_indices( IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes, + file_metadata_size_bytes: None, }, )) } @@ -1268,10 +1269,12 @@ async fn optimize_ivf_hnsw_indices( IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: index_size, + file_metadata_size_bytes: None, }, IndexFile { path: INDEX_AUXILIARY_FILE_NAME.to_string(), size_bytes: aux_size, + file_metadata_size_bytes: None, }, ], )) @@ -2237,6 +2240,7 @@ pub(crate) async fn remap_index_file( Ok(IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes, + file_metadata_size_bytes: None, }) } @@ -2305,6 +2309,7 @@ async fn write_ivf_pq_file( Ok(IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes, + file_metadata_size_bytes: None, }) } @@ -2853,6 +2858,7 @@ async fn write_root_vector_index_from_auxiliary( Ok(IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: summary.size_bytes, + file_metadata_size_bytes: None, }) } From 46a21fcfe41532d00505714aa507d9c404aac42d Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 15:45:31 -0700 Subject: [PATCH 2/6] perf(file): coalesce metadata discovery reads Co-authored-by: Cursor --- rust/lance-file/src/concat.rs | 57 ++- rust/lance-file/src/reader.rs | 500 ++++++++++++++++---- rust/lance-file/src/testing.rs | 9 +- rust/lance-file/src/versions/mod.rs | 32 +- rust/lance-file/src/versions/v2_0/writer.rs | 16 +- rust/lance-file/src/versions/v2_1/writer.rs | 18 +- rust/lance-file/src/versions/v2_2/writer.rs | 18 +- rust/lance-file/src/versions/v2_3/writer.rs | 18 +- rust/lance-file/src/writer.rs | 55 ++- 9 files changed, 599 insertions(+), 124 deletions(-) diff --git a/rust/lance-file/src/concat.rs b/rust/lance-file/src/concat.rs index c67892bf412..a298f21acce 100644 --- a/rust/lance-file/src/concat.rs +++ b/rust/lance-file/src/concat.rs @@ -11,6 +11,7 @@ use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, fmt, future::Future, + num::NonZeroU64, ops::Range, sync::Arc, }; @@ -30,10 +31,10 @@ use prost::Message; use prost_types::Any; use crate::{ - reader::{CachedFileMetadata, FileReader, RawFileMetadataOpen}, + reader::{CachedFileMetadata, FileReader, FullMetadataReadOptions, RawFileMetadataOpen}, version::ConcreteFileVersion, versions, - writer::{FileWriteSummary, FileWriterOptions}, + writer::{FileWriteResult, FileWriterOptions}, }; /// Caller-defined runtime identity of the final target for Blob-bearing parts. @@ -62,6 +63,7 @@ impl BlobTargetId { pub struct EncodedFileInput { scheduler: FileScheduler, expected_num_rows: Option, + metadata_options: FullMetadataReadOptions, } impl EncodedFileInput { @@ -70,6 +72,7 @@ impl EncodedFileInput { Self { scheduler, expected_num_rows: None, + metadata_options: FullMetadataReadOptions::default(), } } @@ -81,6 +84,14 @@ impl EncodedFileInput { self } + /// Use the exact metadata suffix size reported when this file was written. + pub fn with_metadata_size_bytes(mut self, metadata_size_bytes: NonZeroU64) -> Self { + self.metadata_options = self + .metadata_options + .with_metadata_size_bytes(metadata_size_bytes); + self + } + /// The path used to read this input. pub fn path(&self) -> &object_store::path::Path { self.scheduler.reader().path() @@ -155,7 +166,10 @@ impl DataFilePart { blob_target_id: Option, ) -> Result { validate_blob_id_range(blob_ids.as_ref())?; - let metadata = Arc::new(FileReader::read_all_metadata(&input.scheduler()).await?); + let metadata = Arc::new( + FileReader::read_all_metadata_with_options(&input.scheduler(), input.metadata_options) + .await?, + ); let schema = Arc::new(normalize_blob_footer_schema(metadata.file_schema.as_ref())); if let Some(expected_num_rows) = input.expected_num_rows && metadata.num_rows != expected_num_rows @@ -351,12 +365,13 @@ async fn validate_blob_descriptors( &blob_schema, &field_id_to_column_index, )?; - let reader = FileReader::try_open( + let reader = FileReader::try_open_with_metadata_options( input.scheduler(), Some(projection), Arc::::default(), &LanceCache::no_cache(), Default::default(), + input.metadata_options, ) .await?; let mut batches = reader @@ -551,6 +566,8 @@ pub struct FileConcatOutput { pub num_rows: u64, /// Size of the completed or reused object. pub size_bytes: u64, + /// Exact number of bytes from the schema descriptor through EOF. + pub metadata_size_bytes: NonZeroU64, } /// A compatibility reason that requires a caller-controlled decode/re-encode fallback. @@ -940,12 +957,15 @@ where }) })?; if prepared.len() == 1 && reuse_single_input { + let metadata_size_bytes = NonZeroU64::new(prepared[0].metadata.metadata_size_bytes()) + .ok_or_else(|| Error::internal("concat input has an empty metadata suffix"))?; return Ok(FileConcatResult::Reused( 0, FileConcatOutput { version: target.version, num_rows: total_rows, size_bytes: prepared[0].metadata.file_size_bytes, + metadata_size_bytes, }, )); } @@ -953,7 +973,7 @@ where let object_writer = output_factory().await?; let mut writer = versions::create_lazy_writer(target.version, object_writer, options.writer_options)?; - let write_result: Result = async { + let write_result: Result = async { let column_count = prepared[0].metadata.column_infos.len(); let mut output_pages = std::iter::repeat_with(Vec::new) .take(column_count) @@ -1018,16 +1038,20 @@ where &columns, total_rows, )?; - writer.finish().await + writer.finish_with_metadata_size().await } .await; match write_result { - Ok(summary) => Ok(FileConcatResult::Written(FileConcatOutput { - version: target.version, - num_rows: summary.num_rows, - size_bytes: summary.size_bytes, - })), + Ok(result) => { + let summary = result.summary(); + Ok(FileConcatResult::Written(FileConcatOutput { + version: target.version, + num_rows: summary.num_rows, + size_bytes: summary.size_bytes, + metadata_size_bytes: result.metadata_size_bytes(), + })) + } Err(error) => { writer.abort().await; Err(error) @@ -1080,11 +1104,12 @@ where "concat_files requires at least one complete input file", )); } - let raw_metadata = futures::future::try_join_all( - ordered_inputs - .iter() - .map(|input| FileReader::read_raw_metadata_for_dispatch(&input.scheduler)), - ) + let raw_metadata = futures::future::try_join_all(ordered_inputs.iter().map(|input| { + FileReader::read_raw_metadata_for_dispatch_with_options( + &input.scheduler, + input.metadata_options, + ) + })) .await?; if target.version == ConcreteFileVersion::V1 || raw_metadata diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index a8e4619e4d0..2b81c266234 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -6,6 +6,7 @@ use std::{ collections::{BTreeMap, BTreeSet}, fmt::Debug, io::Cursor, + num::{NonZeroU32, NonZeroU64}, ops::Range, pin::Pin, sync::Arc, @@ -57,6 +58,9 @@ pub(crate) mod structural; /// Pages larger than this will be split into multiple chunks during read pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024; +/// One column-metadata offset table entry: a u64 position and a u64 length. +const COLUMN_METADATA_OFFSET_ENTRY_SIZE: u64 = 16; + // For now, we don't use global buffers for anything other than schema. If we // use these later we should make them lazily loaded and then cached once loaded. // @@ -136,7 +140,10 @@ pub struct CachedFileMetadata { pub num_column_metadata_bytes: u64, /// The number of bytes contained in global buffers pub num_global_buffer_bytes: u64, - /// The number of bytes contained in the CMO and GBO tables + /// The number of bytes from the schema descriptor through EOF. + /// + /// The historical field name is retained for API compatibility. Prefer + /// [`Self::metadata_size_bytes`] when reading this value. pub num_footer_bytes: u64, /// The major version number stored in the file footer. pub major_version: u16, @@ -148,8 +155,8 @@ pub struct CachedFileMetadata { /// User global buffers (index >= 1) whose bytes were already captured by the /// tail read that `read_all_metadata` performs at open, keyed by buffer index. /// - /// All global buffers are laid out contiguously starting at the schema, so on - /// small/medium files they land inside the captured tail window. Retaining + /// User global buffers are written immediately before the schema, so on + /// small/medium files they can land inside the captured tail window. Retaining /// those bytes lets `read_global_buffer` serve them with zero additional I/O. /// The bytes are copied out of the tail (rather than sliced) so the much /// larger tail allocation can be dropped — we only hold what we will serve. @@ -165,6 +172,11 @@ impl CachedFileMetadata { pub fn file_size(&self) -> u64 { self.file_size_bytes } + + /// Number of bytes from the schema descriptor through EOF. + pub fn metadata_size_bytes(&self) -> u64 { + self.num_footer_bytes + } } fn column_metadata_deep_size(column_metadatas: &[pbfile::ColumnMetadata]) -> usize { @@ -389,6 +401,49 @@ impl Default for FileReaderOptions { } } +/// Advisory input for opening fully decoded current-format metadata. +#[derive(Clone, Copy, Debug, Default)] +pub struct FullMetadataReadOptions { + metadata_size_bytes: Option, +} + +impl FullMetadataReadOptions { + /// Use an exact metadata suffix size produced by the file writer. + pub fn with_metadata_size_bytes(mut self, metadata_size_bytes: NonZeroU64) -> Self { + self.metadata_size_bytes = Some(metadata_size_bytes); + self + } +} + +/// Advisory input for opening a current-format metadata index. +#[derive(Clone, Copy, Debug, Default)] +pub struct MetadataIndexReadOptions { + estimated_num_columns: Option, +} + +impl MetadataIndexReadOptions { + /// Size the initial tail read using an estimated physical-column count. + pub fn with_estimated_num_columns(mut self, estimated_num_columns: NonZeroU32) -> Self { + self.estimated_num_columns = Some(estimated_num_columns); + self + } +} + +#[derive(Clone, Copy, Debug, Default)] +enum TailReadPlan { + #[default] + Baseline, + ExactMetadata(NonZeroU64), + EstimatedMetadataIndex(NonZeroU32), +} + +struct InitialTailRead { + bytes: Bytes, + file_len: u64, + offset: u64, + retention_offset: u64, +} + #[derive(Debug, Clone)] pub(crate) struct PreparedProjection { pub column_infos: Vec>, @@ -688,15 +743,37 @@ impl FileReader { self.core.read_global_buffer(index).await } - async fn read_tail(scheduler: &FileScheduler) -> Result<(Bytes, u64)> { + async fn read_tail(scheduler: &FileScheduler, plan: TailReadPlan) -> Result { let file_size = scheduler.reader().size().await? as u64; - let begin = if file_size < scheduler.reader().block_size() as u64 { - 0 - } else { - file_size - scheduler.reader().block_size() as u64 + let baseline_read_size = file_size.min(scheduler.reader().block_size() as u64); + let requested_read_size = match plan { + TailReadPlan::Baseline => baseline_read_size, + TailReadPlan::ExactMetadata(metadata_size_bytes) => { + let metadata_size_bytes = metadata_size_bytes.get(); + if metadata_size_bytes <= file_size { + baseline_read_size.max(metadata_size_bytes) + } else { + baseline_read_size + } + } + TailReadPlan::EstimatedMetadataIndex(estimated_num_columns) => { + let estimated_cmo_size = + u64::from(estimated_num_columns.get()) * COLUMN_METADATA_OFFSET_ENTRY_SIZE; + baseline_read_size + .checked_add(estimated_cmo_size) + .unwrap_or(file_size) + .min(file_size) + } }; - let tail_bytes = scheduler.submit_single(begin..file_size, 0).await?; - Ok((tail_bytes, file_size)) + let offset = file_size - requested_read_size; + // FileScheduler may split this logical range into concurrent physical requests. + let bytes = scheduler.submit_single(offset..file_size, 0).await?; + Ok(InitialTailRead { + bytes, + file_len: file_size, + offset, + retention_offset: file_size - baseline_read_size, + }) } async fn read_range_from_tail_or_scheduler( @@ -705,11 +782,26 @@ impl FileReader { scheduler: &FileScheduler, range: Range, ) -> Result { - let tail_end = tail_offset + tail_bytes.len() as u64; + if range.start > range.end { + return Err(Error::invalid_input(format!( + "invalid metadata byte range: start={}, end={}", + range.start, range.end + ))); + } + let tail_end = tail_offset + .checked_add(tail_bytes.len() as u64) + .ok_or_else(|| Error::invalid_input_source("Tail byte range overflows".into()))?; if range.start >= tail_offset && range.end <= tail_end { let rel_start = (range.start - tail_offset) as usize; let rel_end = (range.end - tail_offset) as usize; Ok(tail_bytes.slice(rel_start..rel_end)) + } else if range.start < tail_offset && range.end > tail_offset && range.end <= tail_end { + let missing_bytes = scheduler.submit_single(range.start..tail_offset, 0).await?; + let rel_end = (range.end - tail_offset) as usize; + let mut combined = BytesMut::with_capacity((range.end - range.start) as usize); + combined.extend_from_slice(&missing_bytes); + combined.extend_from_slice(&tail_bytes[..rel_end]); + Ok(combined.freeze()) } else { scheduler.submit_single(range, 0).await } @@ -719,6 +811,7 @@ impl FileReader { gbo_table: &[BufferDescriptor], tail_bytes: &Bytes, tail_offset: u64, + retention_offset: u64, file_len: u64, ) -> Result> { let tail_end = tail_offset @@ -727,7 +820,7 @@ impl FileReader { let mut retained_buffers = BTreeMap::new(); for (index, buffer) in gbo_table.iter().enumerate().skip(1) { let range = buffer.checked_range(index, file_len)?; - if range.start >= tail_offset && range.end <= tail_end { + if range.start >= retention_offset && range.end <= tail_end { let rel_start = (range.start - tail_offset) as usize; let rel_end = (range.end - tail_offset) as usize; let bytes = Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]); @@ -951,11 +1044,22 @@ impl FileReader { Ok((num_rows, schema)) } - pub(crate) async fn read_raw_metadata_for_dispatch( + pub(crate) async fn read_raw_metadata_for_dispatch_with_options( scheduler: &FileScheduler, + options: FullMetadataReadOptions, ) -> Result { - let (tail_bytes, file_len) = Self::read_tail(scheduler).await?; - let tail_offset = file_len - tail_bytes.len() as u64; + let InitialTailRead { + bytes: tail_bytes, + file_len, + offset: tail_offset, + retention_offset, + } = Self::read_tail( + scheduler, + options + .metadata_size_bytes + .map_or(TailReadPlan::Baseline, TailReadPlan::ExactMetadata), + ) + .await?; let footer = Self::decode_footer(&tail_bytes)?; let version = ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?; @@ -999,14 +1103,16 @@ impl FileReader { let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes; let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start; // The tail read above already pulled in any global buffer that lives within - // the captured window. Copy those user buffers (index >= 1; the schema at 0 - // is decoded above and never fetched via read_global_buffer) out of the tail - // so read_global_buffer can serve them without I/O. We copy rather than slice - // so the much larger tail allocation can be released once decoding is done. + // the baseline captured window. Copy those user buffers (index >= 1; the schema + // at 0 is decoded above and never fetched via read_global_buffer) out of the + // tail so read_global_buffer can serve them without I/O. We copy rather than + // slice so the much larger tail allocation can be released once decoding is + // done. let retained_global_buffers = Self::retained_global_buffers_from_tail( &gbo_table, &tail_bytes, tail_offset, + retention_offset, file_len, )?; @@ -1031,20 +1137,39 @@ impl FileReader { async fn read_raw_metadata_index_with_known_schema( scheduler: &FileScheduler, known_schema: Option<(Arc, u64)>, + options: MetadataIndexReadOptions, ) -> Result { - let (tail_bytes, file_len) = Self::read_tail(scheduler).await?; - let tail_offset = file_len - tail_bytes.len() as u64; + let InitialTailRead { + bytes: tail_bytes, + file_len, + offset: tail_offset, + retention_offset, + } = Self::read_tail( + scheduler, + options + .estimated_num_columns + .map_or(TailReadPlan::Baseline, TailReadPlan::EstimatedMetadataIndex), + ) + .await?; let footer = Self::decode_footer(&tail_bytes)?; let file_version = Self::current_file_version(&footer)?; let gbo_table = - Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?; + Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version); + let cmo_table = Self::read_range_from_tail_or_scheduler( + &tail_bytes, + tail_offset, + scheduler, + footer.column_meta_offsets_start..footer.global_buff_offsets_start, + ); + let (gbo_table, cmo_table) = futures::try_join!(gbo_table, cmo_table)?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), )); } + let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?; let (file_schema, num_rows) = match known_schema { Some((file_schema, num_rows)) => (file_schema, num_rows), None => { @@ -1062,19 +1187,11 @@ impl FileReader { } }; - let cmo_table = Self::read_range_from_tail_or_scheduler( - &tail_bytes, - tail_offset, - scheduler, - footer.column_meta_offsets_start..footer.global_buff_offsets_start, - ) - .await?; - let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?; - let retained_global_buffers = Self::retained_global_buffers_from_tail( &gbo_table, &tail_bytes, tail_offset, + retention_offset, file_len, )?; @@ -1090,30 +1207,6 @@ impl FileReader { }) } - /// Reads the lightweight metadata index from a file. - /// - /// This reads the file schema from the schema global buffer. Use - /// [`Self::read_metadata_index_with_schema`] when the caller already has - /// the schema and row count from a higher-level metadata source. - pub(crate) async fn read_raw_metadata_index( - scheduler: &FileScheduler, - ) -> Result { - Self::read_raw_metadata_index_with_known_schema(scheduler, None).await - } - - /// Reads the metadata index without fetching the schema global buffer. - /// - /// Use this when the caller already has the file schema and physical row - /// count from an enclosing metadata layer, such as a dataset manifest. - pub(crate) async fn read_raw_metadata_index_with_schema( - scheduler: &FileScheduler, - file_schema: Arc, - num_rows: u64, - ) -> Result { - Self::read_raw_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))) - .await - } - pub(crate) fn validate_projection( projection: &ReaderProjection, metadata: &CachedFileMetadata, @@ -1952,12 +2045,33 @@ impl FileReader { cache: &LanceCache, options: FileReaderOptions, ) -> Result { - match Self::try_open_for_dispatch( + Self::try_open_with_metadata_options( + scheduler, + base_projection, + decoder_plugins, + cache, + options, + FullMetadataReadOptions::default(), + ) + .await + } + + /// Open a current-format file using advisory full-metadata options. + pub async fn try_open_with_metadata_options( + scheduler: FileScheduler, + base_projection: Option, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, + metadata_options: FullMetadataReadOptions, + ) -> Result { + match Self::try_open_for_dispatch_with_metadata_options( scheduler, base_projection, decoder_plugins, cache, options, + metadata_options, ) .await? { @@ -1973,27 +2087,31 @@ impl FileReader { } } - pub(crate) async fn try_open_for_dispatch( + pub(crate) async fn try_open_for_dispatch_with_metadata_options( scheduler: FileScheduler, base_projection: Option, decoder_plugins: Arc, cache: &LanceCache, options: FileReaderOptions, + metadata_options: FullMetadataReadOptions, ) -> Result { - let metadata = match Self::read_raw_metadata_for_dispatch(&scheduler).await? { - RawFileMetadataOpen::Legacy { - major_version, - minor_version, - } => { - return Ok(versions::OpenedFileReader::V1 { + let metadata = + match Self::read_raw_metadata_for_dispatch_with_options(&scheduler, metadata_options) + .await? + { + RawFileMetadataOpen::Legacy { major_version, minor_version, - }); - } - RawFileMetadataOpen::Current { version, metadata } => { - Arc::new(versions::finish_metadata(version, metadata)?) - } - }; + } => { + return Ok(versions::OpenedFileReader::V1 { + major_version, + minor_version, + }); + } + RawFileMetadataOpen::Current { version, metadata } => { + Arc::new(versions::finish_metadata(version, metadata)?) + } + }; let path = scheduler.reader().path().clone(); let io = Arc::new( LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size), @@ -2049,7 +2167,15 @@ impl FileReader { } pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result { - match Self::read_raw_metadata_for_dispatch(scheduler).await? { + Self::read_all_metadata_with_options(scheduler, FullMetadataReadOptions::default()).await + } + + /// Read all current-format metadata using advisory open options. + pub async fn read_all_metadata_with_options( + scheduler: &FileScheduler, + options: FullMetadataReadOptions, + ) -> Result { + match Self::read_raw_metadata_for_dispatch_with_options(scheduler, options).await? { RawFileMetadataOpen::Legacy { major_version, minor_version, @@ -2065,7 +2191,16 @@ impl FileReader { } pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result { - let index = Self::read_raw_metadata_index(scheduler).await?; + Self::read_metadata_index_with_options(scheduler, MetadataIndexReadOptions::default()).await + } + + /// Read the current-format metadata index using advisory open options. + pub async fn read_metadata_index_with_options( + scheduler: &FileScheduler, + options: MetadataIndexReadOptions, + ) -> Result { + let index = + Self::read_raw_metadata_index_with_known_schema(scheduler, None, options).await?; versions::finish_metadata_index(index) } @@ -2074,8 +2209,28 @@ impl FileReader { file_schema: Arc, num_rows: u64, ) -> Result { - let index = - Self::read_raw_metadata_index_with_schema(scheduler, file_schema, num_rows).await?; + Self::read_metadata_index_with_schema_and_options( + scheduler, + file_schema, + num_rows, + MetadataIndexReadOptions::default(), + ) + .await + } + + /// Read the metadata index using a known schema and advisory open options. + pub async fn read_metadata_index_with_schema_and_options( + scheduler: &FileScheduler, + file_schema: Arc, + num_rows: u64, + options: MetadataIndexReadOptions, + ) -> Result { + let index = Self::read_raw_metadata_index_with_known_schema( + scheduler, + Some((file_schema, num_rows)), + options, + ) + .await?; versions::finish_metadata_index(index) } @@ -2162,11 +2317,33 @@ impl ProjectedFileReader { decoder_plugins: Arc, cache: &LanceCache, options: FileReaderOptions, + ) -> Result { + Self::try_open_with_metadata_options( + scheduler, + base_projection, + decoder_plugins, + cache, + options, + MetadataIndexReadOptions::default(), + ) + .await + } + + /// Open a projected reader using advisory metadata-index options. + pub async fn try_open_with_metadata_options( + scheduler: FileScheduler, + base_projection: Option, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, + metadata_options: MetadataIndexReadOptions, ) -> Result { let base_projection = base_projection.ok_or_else(|| { Error::invalid_input("ProjectedReader requires an explicit base projection") })?; - let metadata_index = Arc::new(FileReader::read_metadata_index(&scheduler).await?); + let metadata_index = Arc::new( + FileReader::read_metadata_index_with_options(&scheduler, metadata_options).await?, + ); let path = scheduler.reader().path().clone(); let io = Arc::new( LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size), @@ -2425,13 +2602,14 @@ impl EncodedBatchReaderExt for EncodedBatch { mod tests { use std::{ collections::{BTreeMap, HashMap}, + num::{NonZeroU32, NonZeroU64}, pin::Pin, sync::Arc, }; use arrow_array::{ - DictionaryArray, Int8Array, Int32Array, ListArray, RecordBatch, RecordBatchIterator, - StringArray, UInt32Array, + ArrayRef, DictionaryArray, Int8Array, Int32Array, ListArray, RecordBatch, + RecordBatchIterator, StringArray, UInt32Array, types::{Float64Type, Int8Type, Int32Type}, }; use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; @@ -2456,12 +2634,13 @@ mod tests { use tokio::sync::mpsc; use crate::reader::{ - EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, ReaderProjection, + EncodedBatchReaderExt, FileReader, FileReaderOptions, FullMetadataReadOptions, + MetadataIndexReadOptions, ProjectedFileReader, ReaderProjection, }; use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file}; use crate::version::{ConcreteFileVersion, LanceFileVersion}; use crate::versions; - use crate::writer::{FileWriterOptions, PAGE_BUFFER_ALIGNMENT}; + use crate::writer::{FileWriteResult, FileWriterOptions, PAGE_BUFFER_ALIGNMENT}; use lance_encoding::decoder::DecoderConfig; fn footer_version(bytes: &[u8]) -> (u16, u16) { @@ -2690,6 +2869,26 @@ mod tests { .await } + async fn create_one_row_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile { + let arrow_schema = Arc::new(ArrowSchema::new( + (0..num_columns) + .map(|column_idx| Field::new(format!("c{column_idx}"), DataType::Int32, true)) + .collect::>(), + )); + let columns = (0..num_columns) + .map(|column_idx| Arc::new(Int32Array::from(vec![column_idx as i32])) as ArrayRef) + .collect(); + let batch = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], arrow_schema); + write_lance_file( + reader, + fs, + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .await + } + async fn create_wide_fixed_size_list_file(fs: &FsFixture, num_columns: usize) -> WrittenFile { let data_type = DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4); @@ -3469,6 +3668,135 @@ mod tests { assert_eq!(expected, actual); } + #[tokio::test] + async fn full_metadata_hint_removes_footer_dependency() { + let fs = FsFixture::default(); + let num_columns = + fs.object_store.block_size() / super::COLUMN_METADATA_OFFSET_ENTRY_SIZE as usize + 8; + let written = create_one_row_wide_direct_file(&fs, num_columns).await; + let file_size = written.write_result.summary().size_bytes; + let metadata_size = written.write_result.metadata_size_bytes(); + assert!(metadata_size.get() > fs.object_store.block_size() as u64); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(file_size)) + .await + .unwrap(); + fs.object_store.io_stats_incremental(); + let without_hint = FileReader::read_all_metadata(&file_scheduler) + .await + .unwrap(); + let without_hint_stats = fs.object_store.io_stats_incremental(); + assert_eq!(without_hint_stats.read_iops, 2); + assert_eq!(without_hint_stats.read_bytes, metadata_size.get()); + assert_eq!(without_hint.metadata_size_bytes(), metadata_size.get()); + + let options = FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size); + fs.object_store.io_stats_incremental(); + let with_hint = FileReader::read_all_metadata_with_options(&file_scheduler, options) + .await + .unwrap(); + let with_hint_stats = fs.object_store.io_stats_incremental(); + assert_eq!(with_hint_stats.read_iops, 1); + assert_eq!(with_hint_stats.read_bytes, metadata_size.get()); + assert_eq!(with_hint.file_schema, without_hint.file_schema); + assert_eq!(with_hint.column_metadatas, without_hint.column_metadatas); + + let stale_hints = [ + (metadata_size.get() - 1, 2), + (file_size, 1), + (file_size + 1, 2), + ]; + for (stale_hint, expected_read_iops) in stale_hints { + fs.object_store.io_stats_incremental(); + let metadata = FileReader::read_all_metadata_with_options( + &file_scheduler, + FullMetadataReadOptions::default() + .with_metadata_size_bytes(NonZeroU64::new(stale_hint).unwrap()), + ) + .await + .unwrap(); + let stats = fs.object_store.io_stats_incremental(); + assert_eq!(metadata.metadata_size_bytes(), metadata_size.get()); + assert_eq!( + stats.read_iops, expected_read_iops, + "unexpected read count for stale metadata hint {stale_hint}" + ); + } + } + + #[tokio::test] + async fn column_count_estimate_removes_cmo_dependency() { + let fs = FsFixture::default(); + let num_columns = + fs.object_store.block_size() / super::COLUMN_METADATA_OFFSET_ENTRY_SIZE as usize + 8; + let written = create_one_row_wide_direct_file(&fs, num_columns).await; + let file_size = written.write_result.summary().size_bytes; + let estimated_num_columns = NonZeroU32::new(u32::try_from(num_columns).unwrap()).unwrap(); + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(file_size)) + .await + .unwrap(); + + fs.object_store.io_stats_incremental(); + let without_estimate = + FileReader::read_metadata_index_with_schema(&file_scheduler, written.schema.clone(), 1) + .await + .unwrap(); + let without_estimate_stats = fs.object_store.io_stats_incremental(); + assert_eq!(without_estimate_stats.read_iops, 2); + + let options = + MetadataIndexReadOptions::default().with_estimated_num_columns(estimated_num_columns); + fs.object_store.io_stats_incremental(); + let with_estimate = FileReader::read_metadata_index_with_schema_and_options( + &file_scheduler, + written.schema.clone(), + 1, + options, + ) + .await + .unwrap(); + let with_estimate_stats = fs.object_store.io_stats_incremental(); + assert_eq!(with_estimate_stats.read_iops, 1); + assert_eq!( + with_estimate.column_metadata_offsets, + without_estimate.column_metadata_offsets + ); + assert_eq!(with_estimate.num_columns(), num_columns as u32); + + for (estimated_num_columns, expected_read_iops) in [ + (NonZeroU32::new(1).unwrap(), 2), + ( + NonZeroU32::new(u32::try_from(num_columns * 2).unwrap()).unwrap(), + 1, + ), + ] { + fs.object_store.io_stats_incremental(); + let metadata = FileReader::read_metadata_index_with_schema_and_options( + &file_scheduler, + written.schema.clone(), + 1, + MetadataIndexReadOptions::default() + .with_estimated_num_columns(estimated_num_columns), + ) + .await + .unwrap(); + let stats = fs.object_store.io_stats_incremental(); + assert_eq!( + metadata.column_metadata_offsets, + without_estimate.column_metadata_offsets + ); + assert_eq!( + stats.read_iops, expected_read_iops, + "unexpected read count for {} estimated columns", + estimated_num_columns + ); + } + } + #[tokio::test] async fn test_lazy_reader_loads_only_requested_column_metadata() { let fs = FsFixture::default(); @@ -4242,7 +4570,7 @@ mod tests { assert_eq!(batches.len(), 1); } - async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) { + async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) -> FileWriteResult { let lance_schema = lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new( "foo", @@ -4261,7 +4589,7 @@ mod tests { let buf_index = file_writer.add_global_buffer(buffer).await.unwrap(); assert_eq!(buf_index, 1); - file_writer.finish().await.unwrap(); + file_writer.finish_with_metadata_size().await.unwrap() } #[derive(Clone, Copy, Debug)] @@ -4372,19 +4700,27 @@ mod tests { }; let expected_read_iops = if within_window { 0 } else { 1 }; - write_file_with_global_buffer(&fs, buffer.clone()).await; + let write_result = write_file_with_global_buffer(&fs, buffer.clone()).await; let file_scheduler = fs .scheduler .open_file(&fs.tmp_path, &CachedFileSize::unknown()) .await .unwrap(); - let file_reader = FileReader::try_open( + let metadata_options = if within_window { + FullMetadataReadOptions::default() + } else { + FullMetadataReadOptions::default().with_metadata_size_bytes( + NonZeroU64::new(write_result.summary().size_bytes).unwrap(), + ) + }; + let file_reader = FileReader::try_open_with_metadata_options( file_scheduler, None, Arc::::default(), &test_cache(), FileReaderOptions::default(), + metadata_options, ) .await .unwrap(); diff --git a/rust/lance-file/src/testing.rs b/rust/lance-file/src/testing.rs index 65b340e7641..5127f6a647a 100644 --- a/rust/lance-file/src/testing.rs +++ b/rust/lance-file/src/testing.rs @@ -17,7 +17,10 @@ use lance_io::{ use crate::reader::{FileReader, FileReaderOptions}; use crate::version::ConcreteFileVersion; -use crate::{versions, writer::FileWriterOptions}; +use crate::{ + versions, + writer::{FileWriteResult, FileWriterOptions}, +}; pub struct FsFixture { pub tmp_path: TempObjFile, @@ -43,6 +46,7 @@ pub struct WrittenFile { pub schema: Arc, pub data: Vec, pub field_id_mapping: Vec<(u32, u32)>, + pub write_result: FileWriteResult, } pub async fn write_lance_file( @@ -67,11 +71,12 @@ pub async fn write_lance_file( } let field_id_mapping = file_writer.field_id_to_column_indices().to_vec(); file_writer.add_schema_metadata("foo", "bar"); - file_writer.finish().await.unwrap(); + let write_result = file_writer.finish_with_metadata_size().await.unwrap(); WrittenFile { schema: Arc::new(lance_schema), data, field_id_mapping, + write_result, } } diff --git a/rust/lance-file/src/versions/mod.rs b/rust/lance-file/src/versions/mod.rs index 580e29bf28b..855fc4c0771 100644 --- a/rust/lance-file/src/versions/mod.rs +++ b/rust/lance-file/src/versions/mod.rs @@ -21,8 +21,8 @@ use crate::{ format::pbfile, reader::{ BufferDescriptor, CachedFileMetadata, FileMetadataIndex, FileMetadataProvider, FileReader, - FileReaderOptions, PreparedProjection, ProjectedFileReader, RawFileMetadata, - ReadProjection, ReaderProjection, + FileReaderOptions, FullMetadataReadOptions, PreparedProjection, ProjectedFileReader, + RawFileMetadata, ReadProjection, ReaderProjection, }, version::ConcreteFileVersion, writer::{FileWriter, FileWriterOptions}, @@ -322,7 +322,33 @@ pub async fn open_self_described_reader( cache: &LanceCache, options: FileReaderOptions, ) -> Result { - FileReader::try_open_for_dispatch(scheduler, None, decoder_plugins, cache, options).await + open_self_described_reader_with_metadata_options( + scheduler, + decoder_plugins, + cache, + options, + FullMetadataReadOptions::default(), + ) + .await +} + +/// Open a self-described file using advisory full-metadata options. +pub async fn open_self_described_reader_with_metadata_options( + scheduler: FileScheduler, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, + metadata_options: FullMetadataReadOptions, +) -> Result { + FileReader::try_open_for_dispatch_with_metadata_options( + scheduler, + None, + decoder_plugins, + cache, + options, + metadata_options, + ) + .await } /// Create a current-format writer for an exact file version. diff --git a/rust/lance-file/src/versions/v2_0/writer.rs b/rust/lance-file/src/versions/v2_0/writer.rs index a627027e421..5baba5dddb0 100644 --- a/rust/lance-file/src/versions/v2_0/writer.rs +++ b/rust/lance-file/src/versions/v2_0/writer.rs @@ -34,7 +34,7 @@ use crate::format::pb; use crate::format::pbfile; use crate::format::pbfile::DirectEncoding; use crate::writer::{ - ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, FileWriteSummary, FileWriterOptions, + ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, FileWriteResult, FileWriteSummary, FileWriterOptions, PAGE_BUFFER_ALIGNMENT, }; @@ -753,6 +753,11 @@ impl Writer { /// /// Returns a summary of the completed file write. pub async fn finish(&mut self) -> Result { + Ok(self.finish_with_metadata_size().await?.summary()) + } + + /// Finishes writing the file and returns the exact metadata suffix size. + pub async fn finish_with_metadata_size(&mut self) -> Result { // 1. flush any remaining data and write out those pages let mut external_buffers = OutOfLineBuffers::new(self.tell().await?, PAGE_BUFFER_ALIGNMENT as u64); @@ -776,6 +781,10 @@ impl Writer { // 3. write global buffers (we write the schema here) let global_buffer_offsets = self.write_global_buffers().await?; + let metadata_start = global_buffer_offsets + .first() + .map(|(position, _)| *position) + .ok_or_else(|| Error::internal("schema descriptor was not written"))?; let num_global_buffers = global_buffer_offsets.len() as u32; // 4. write the column metadatas @@ -809,10 +818,11 @@ impl Writer { // 7. close the writer let write_result = ObjectWriter::shutdown(self.writer.as_mut()).await?; - Ok(FileWriteSummary { + let summary = FileWriteSummary { num_rows: self.rows_written, size_bytes: write_result.size as u64, - }) + }; + FileWriteResult::try_new(summary, metadata_start) } pub async fn abort(&mut self) { diff --git a/rust/lance-file/src/versions/v2_1/writer.rs b/rust/lance-file/src/versions/v2_1/writer.rs index 1599a9fd6e1..297cd73340e 100644 --- a/rust/lance-file/src/versions/v2_1/writer.rs +++ b/rust/lance-file/src/versions/v2_1/writer.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use arrow_array::{ArrayRef, RecordBatch}; use bytes::{BufMut, Bytes}; -use lance_core::{Result, datatypes::Schema}; +use lance_core::{Error, Result, datatypes::Schema}; use lance_encoding::{ compression_config::CompressionParams, encoder::{BatchEncoder, EncodedBatch}, @@ -17,7 +17,7 @@ use tokio::io::AsyncWriteExt; use crate::{ format::{MAGIC, pbfile}, writer::{ - FileWriteSummary, FileWriterOptions, + FileWriteResult, FileWriteSummary, FileWriterOptions, structural::{EncodedBatchBody, EncodingPipeline, StructuralFileSink, encode_batch_body}, }, }; @@ -150,12 +150,21 @@ impl Writer { /// Finish the v2.1 file and close its object writer. pub async fn finish(&mut self) -> Result { + Ok(self.finish_with_metadata_size().await?.summary()) + } + + /// Finish the v2.1 file and return the exact metadata suffix size. + pub async fn finish_with_metadata_size(&mut self) -> Result { // The order below is the v2.1 wire contract. self.encoding.flush(&mut self.sink).await?; self.encoding.finish_encoders(&mut self.sink).await?; let descriptor = self.encoding.make_file_descriptor()?; let global_buffer_offsets = self.sink.write_global_buffers(descriptor).await?; + let metadata_start = global_buffer_offsets + .first() + .map(|(position, _)| *position) + .ok_or_else(|| Error::internal("schema descriptor was not written"))?; let num_global_buffers = global_buffer_offsets.len() as u32; let column_metadata_start = self.sink.tell().await?; @@ -178,10 +187,11 @@ impl Writer { output.write_u16_le(1).await?; output.write_all(MAGIC).await?; - Ok(FileWriteSummary { + let summary = FileWriteSummary { num_rows: self.encoding.rows_written(), size_bytes: self.sink.shutdown().await?, - }) + }; + FileWriteResult::try_new(summary, metadata_start) } /// Abandon this write. diff --git a/rust/lance-file/src/versions/v2_2/writer.rs b/rust/lance-file/src/versions/v2_2/writer.rs index bfe5494cc1f..59f16589297 100644 --- a/rust/lance-file/src/versions/v2_2/writer.rs +++ b/rust/lance-file/src/versions/v2_2/writer.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use arrow_array::{ArrayRef, RecordBatch}; use bytes::{BufMut, Bytes}; -use lance_core::{Result, datatypes::Schema}; +use lance_core::{Error, Result, datatypes::Schema}; use lance_encoding::{ compression_config::CompressionParams, encoder::{BatchEncoder, EncodedBatch}, @@ -17,7 +17,7 @@ use tokio::io::AsyncWriteExt; use crate::{ format::{MAGIC, pbfile}, writer::{ - FileWriteSummary, FileWriterOptions, + FileWriteResult, FileWriteSummary, FileWriterOptions, structural::{EncodedBatchBody, EncodingPipeline, StructuralFileSink, encode_batch_body}, }, }; @@ -150,12 +150,21 @@ impl Writer { /// Finish the v2.2 file and close its object writer. pub async fn finish(&mut self) -> Result { + Ok(self.finish_with_metadata_size().await?.summary()) + } + + /// Finish the v2.2 file and return the exact metadata suffix size. + pub async fn finish_with_metadata_size(&mut self) -> Result { // The order below is the v2.2 wire contract. self.encoding.flush(&mut self.sink).await?; self.encoding.finish_encoders(&mut self.sink).await?; let descriptor = self.encoding.make_file_descriptor()?; let global_buffer_offsets = self.sink.write_global_buffers(descriptor).await?; + let metadata_start = global_buffer_offsets + .first() + .map(|(position, _)| *position) + .ok_or_else(|| Error::internal("schema descriptor was not written"))?; let num_global_buffers = global_buffer_offsets.len() as u32; let column_metadata_start = self.sink.tell().await?; @@ -178,10 +187,11 @@ impl Writer { output.write_u16_le(2).await?; output.write_all(MAGIC).await?; - Ok(FileWriteSummary { + let summary = FileWriteSummary { num_rows: self.encoding.rows_written(), size_bytes: self.sink.shutdown().await?, - }) + }; + FileWriteResult::try_new(summary, metadata_start) } /// Abandon this write. diff --git a/rust/lance-file/src/versions/v2_3/writer.rs b/rust/lance-file/src/versions/v2_3/writer.rs index c18c7aa6e18..fe612fb3d78 100644 --- a/rust/lance-file/src/versions/v2_3/writer.rs +++ b/rust/lance-file/src/versions/v2_3/writer.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use arrow_array::{ArrayRef, RecordBatch}; use bytes::{BufMut, Bytes}; -use lance_core::{Result, datatypes::Schema}; +use lance_core::{Error, Result, datatypes::Schema}; use lance_encoding::{ compression_config::CompressionParams, encoder::{BatchEncoder, EncodedBatch}, @@ -17,7 +17,7 @@ use tokio::io::AsyncWriteExt; use crate::{ format::{MAGIC, pbfile}, writer::{ - FileWriteSummary, FileWriterOptions, + FileWriteResult, FileWriteSummary, FileWriterOptions, structural::{EncodedBatchBody, EncodingPipeline, StructuralFileSink, encode_batch_body}, }, }; @@ -150,12 +150,21 @@ impl Writer { /// Finish the v2.3 file and close its object writer. pub async fn finish(&mut self) -> Result { + Ok(self.finish_with_metadata_size().await?.summary()) + } + + /// Finish the v2.3 file and return the exact metadata suffix size. + pub async fn finish_with_metadata_size(&mut self) -> Result { // The order below is the v2.3 wire contract. self.encoding.flush(&mut self.sink).await?; self.encoding.finish_encoders(&mut self.sink).await?; let descriptor = self.encoding.make_file_descriptor()?; let global_buffer_offsets = self.sink.write_global_buffers(descriptor).await?; + let metadata_start = global_buffer_offsets + .first() + .map(|(position, _)| *position) + .ok_or_else(|| Error::internal("schema descriptor was not written"))?; let num_global_buffers = global_buffer_offsets.len() as u32; let column_metadata_start = self.sink.tell().await?; @@ -178,10 +187,11 @@ impl Writer { output.write_u16_le(3).await?; output.write_all(MAGIC).await?; - Ok(FileWriteSummary { + let summary = FileWriteSummary { num_rows: self.encoding.rows_written(), size_bytes: self.sink.shutdown().await?, - }) + }; + FileWriteResult::try_new(summary, metadata_start) } /// Abandon this write. diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index dcba2658225..1ef6ad57ed2 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::sync::Arc; +use std::{num::NonZero, sync::Arc}; use arrow_array::{ArrayRef, RecordBatch}; use bytes::Bytes; -use lance_core::{Result, datatypes::Schema}; +use lance_core::{Error, Result, datatypes::Schema}; use lance_encoding::decoder::{ColumnInfo, PageEncoding}; use lance_io::object_store::ObjectStore; use object_store::path::Path; @@ -29,6 +29,44 @@ pub struct FileWriteSummary { pub size_bytes: u64, } +/// Result of a completed current-format Lance file write. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FileWriteResult { + /// The existing summary of the completed write. + summary: FileWriteSummary, + /// The exact number of bytes from the schema descriptor through EOF. + metadata_size_bytes: NonZero, +} + +impl FileWriteResult { + /// Return the row count and total file size. + pub fn summary(&self) -> FileWriteSummary { + self.summary + } + + /// Return the exact number of bytes from the schema descriptor through EOF. + pub fn metadata_size_bytes(&self) -> NonZero { + self.metadata_size_bytes + } + + pub(crate) fn try_new(summary: FileWriteSummary, metadata_start: u64) -> Result { + let metadata_size_bytes = summary + .size_bytes + .checked_sub(metadata_start) + .and_then(NonZero::new) + .ok_or_else(|| { + Error::internal(format!( + "invalid metadata suffix: metadata_start={}, file_size={}", + metadata_start, summary.size_bytes + )) + })?; + Ok(Self { + summary, + metadata_size_bytes, + }) + } +} + /// Runtime options shared by all current-format writers. /// /// These options control buffering and execution only. Select the file grammar @@ -235,11 +273,16 @@ impl FileWriter { /// Finish the file and close its object writer. pub async fn finish(&mut self) -> Result { + Ok(self.finish_with_metadata_size().await?.summary()) + } + + /// Finish the file and return the exact metadata suffix size. + pub async fn finish_with_metadata_size(&mut self) -> Result { match self { - Self::V2_0(writer) => writer.finish().await, - Self::V2_1(writer) => writer.finish().await, - Self::V2_2(writer) => writer.finish().await, - Self::V2_3(writer) => writer.finish().await, + Self::V2_0(writer) => writer.finish_with_metadata_size().await, + Self::V2_1(writer) => writer.finish_with_metadata_size().await, + Self::V2_2(writer) => writer.finish_with_metadata_size().await, + Self::V2_3(writer) => writer.finish_with_metadata_size().await, } } From f1573cb20e5ddfbab7889fd5440ef1b7458ab70c Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 15:45:37 -0700 Subject: [PATCH 3/6] perf: use metadata open hints across table and index Co-authored-by: Cursor --- rust/lance-index-core/src/scalar.rs | 13 +- rust/lance-index/protos-cache/cache.proto | 5 + rust/lance-index/src/scalar.rs | 3 +- .../src/scalar/inverted/builder.rs | 2 + rust/lance-index/src/scalar/lance_format.rs | 85 ++++++++- .../src/vector/distributed/index_merger.rs | 5 +- rust/lance-table/src/format/fragment.rs | 29 +++ rust/lance-table/src/format/index.rs | 12 ++ rust/lance/src/dataset.rs | 5 +- rust/lance/src/dataset/data_file_part.rs | 19 +- rust/lance/src/dataset/fragment.rs | 71 ++++++-- rust/lance/src/dataset/fragment/write.rs | 4 +- rust/lance/src/dataset/index.rs | 4 +- rust/lance/src/dataset/optimize.rs | 10 +- .../lance/src/dataset/optimize/binary_copy.rs | 26 ++- .../lance/src/dataset/tests/data_file_part.rs | 9 + rust/lance/src/dataset/write.rs | 38 ++-- rust/lance/src/index.rs | 49 ++++-- rust/lance/src/index/api.rs | 37 +++- rust/lance/src/index/append.rs | 14 +- rust/lance/src/index/vector.rs | 16 +- rust/lance/src/index/vector/builder.rs | 10 +- rust/lance/src/index/vector/ivf.rs | 5 +- .../src/index/vector/ivf/partition_serde.rs | 2 + rust/lance/src/index/vector/ivf/v2.rs | 166 ++++++++++++------ 25 files changed, 503 insertions(+), 136 deletions(-) diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index 80da1633a26..324f10f455a 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -19,6 +19,7 @@ use roaring::{RoaringBitmap, RoaringTreemap}; use serde::Serialize; use std::collections::HashMap; use std::fmt::Debug; +use std::num::NonZeroU64; use std::pin::Pin; use std::{any::Any, sync::Arc}; @@ -26,12 +27,22 @@ use crate::metrics::MetricsCollector; use crate::{Index, IndexParams, IndexType}; /// Metadata about a single file within an index segment. -#[derive(Debug, Clone, PartialEq, DeepSizeOf)] +#[derive(Debug, Clone, PartialEq)] pub struct IndexFile { /// Path relative to the index directory (e.g., "index.idx", "auxiliary.idx") pub path: String, /// Size of the file in bytes pub size_bytes: u64, + /// Metadata suffix size for a Lance-format file. + /// + /// `None` means the size is unavailable or the file uses another format. + pub file_metadata_size_bytes: Option, +} + +impl DeepSizeOf for IndexFile { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.path.deep_size_of_children(context) + } } pub const LANCE_SCALAR_INDEX: &str = "__lance_scalar_index"; diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index 80298e96633..14f0866908e 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -201,6 +201,11 @@ message IvfStateHeader { string quantizer_metadata_json = 7; uint64 index_file_size = 9; uint64 aux_file_size = 10; + /* Lance metadata suffix sizes used to avoid dependent footer reads. + * Zero means the value was unavailable when the state was cached. + */ + uint64 index_file_metadata_size_bytes = 11; + uint64 aux_file_metadata_size_bytes = 12; } /* RabitQ quantizer. Sections: sub-index IPC, rotate-matrix IPC (Matrix rotation diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 696c1123843..a15d0dc97be 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -62,7 +62,7 @@ pub fn index_files_to_table( .map(|f| lance_table::format::IndexFile { path: f.path, size_bytes: f.size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: f.file_metadata_size_bytes, }) .collect() } @@ -80,6 +80,7 @@ pub fn table_files_to_index( .map(|f| lance_index_core::scalar::IndexFile { path: f.path, size_bytes: f.size_bytes, + file_metadata_size_bytes: f.file_metadata_size_bytes, }) .collect() } diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 10e613678c2..c6665080aa7 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -2951,6 +2951,7 @@ mod tests { Ok(IndexFile { path: self.path.clone(), size_bytes: 0, + file_metadata_size_bytes: None, }) } @@ -2961,6 +2962,7 @@ mod tests { Ok(IndexFile { path: self.path.clone(), size_bytes: 0, + file_metadata_size_bytes: None, }) } } diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index cca61bda539..7fac43c5f4b 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -12,7 +12,9 @@ use futures::TryStreamExt; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result, cache::LanceCache}; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; -use lance_file::reader::{FileReader as CurrentFileReader, FileReaderOptions}; +use lance_file::reader::{ + FileReader as CurrentFileReader, FileReaderOptions, FullMetadataReadOptions, +}; use lance_file::version::ConcreteFileVersion; use lance_file::versions::v1::reader::FileReader as V1FileReader; use lance_file::versions::{self, OpenedFileReader}; @@ -25,6 +27,7 @@ use lance_table::format::list_index_files_with_sizes; use object_store::path::Path; use std::cmp::min; use std::collections::HashMap; +use std::num::NonZeroU64; use std::pin::Pin; use std::{any::Any, sync::Arc}; @@ -44,6 +47,8 @@ pub struct LanceIndexStore { // Partition priority views share this immutable map. Cloning all file names // for every partition would make request rebinding quadratic in partitions. file_sizes: Arc>, + /// Cached metadata suffix sizes for Lance-format files. + file_metadata_sizes: Arc>, format_version: ConcreteFileVersion, /// Base I/O priority for all requests this store submits to `scheduler`. io_priority: u64, @@ -92,6 +97,7 @@ impl LanceIndexStore { metadata_cache, scheduler, file_sizes: Arc::default(), + file_metadata_sizes: Arc::default(), format_version, io_priority: 0, } @@ -106,6 +112,15 @@ impl LanceIndexStore { self } + /// Set metadata suffix sizes used to avoid dependent footer discovery reads. + pub fn with_file_metadata_sizes( + mut self, + file_metadata_sizes: HashMap, + ) -> Self { + self.file_metadata_sizes = Arc::new(file_metadata_sizes); + self + } + /// The base I/O priority all this store's requests are submitted at. pub fn io_priority(&self) -> u64 { self.io_priority @@ -148,10 +163,12 @@ impl IndexWriter for LanceIndexWriter { } async fn finish(&mut self) -> Result { - let summary = self.inner.finish().await?; + let result = self.inner.finish_with_metadata_size().await?; + let summary = result.summary(); Ok(IndexFile { path: self.path.clone(), size_bytes: summary.size_bytes, + file_metadata_size_bytes: Some(result.metadata_size_bytes()), }) } @@ -162,10 +179,12 @@ impl IndexWriter for LanceIndexWriter { metadata.into_iter().for_each(|(k, v)| { self.inner.add_schema_metadata(k, v); }); - let summary = self.inner.finish().await?; + let result = self.inner.finish_with_metadata_size().await?; + let summary = result.summary(); Ok(IndexFile { path: self.path.clone(), size_bytes: summary.size_bytes, + file_metadata_size_bytes: Some(result.metadata_size_bytes()), }) } } @@ -460,11 +479,18 @@ impl IndexStore for LanceIndexStore { .scheduler .open_file_with_priority(&path, self.io_priority, &cached_size) .await?; - match versions::open_self_described_reader( + let metadata_options = self.file_metadata_sizes.get(name).copied().map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ); + match versions::open_self_described_reader_with_metadata_options( file_scheduler, Arc::::default(), &self.metadata_cache, FileReaderOptions::default(), + metadata_options, ) .await? { @@ -503,6 +529,7 @@ impl IndexStore for LanceIndexStore { Ok(IndexFile { path: new_name.to_string(), size_bytes: result.size as u64, + file_metadata_size_bytes: self.file_metadata_sizes.get(name).copied(), }) } _ => { @@ -532,6 +559,7 @@ impl IndexStore for LanceIndexStore { Ok(IndexFile { path: new_name.to_string(), size_bytes: result.size as u64, + file_metadata_size_bytes: self.file_metadata_sizes.get(name).copied(), }) } @@ -547,6 +575,7 @@ impl IndexStore for LanceIndexStore { .map(|f| IndexFile { path: f.path, size_bytes: f.size_bytes, + file_metadata_size_bytes: f.file_metadata_size_bytes, }) .collect()) } @@ -711,6 +740,54 @@ mod tests { assert_eq!(actual, expected); } + #[tokio::test] + async fn metadata_hint_removes_scalar_index_footer_dependency() { + let tempdir = TempDir::default(); + let (object_store, index_dir) = ObjectStore::from_uri(tempdir.obj_path().as_ref()) + .await + .unwrap(); + let file_name = "metadata.lance"; + let mut schema = Schema::empty(); + schema.metadata.insert( + "large_metadata".to_string(), + "x".repeat(2 * object_store.block_size()), + ); + let write_store = LanceIndexStore::new( + object_store.clone(), + index_dir.clone(), + Arc::new(LanceCache::no_cache()), + ); + let mut writer = write_store + .new_index_file(file_name, Arc::new(schema)) + .await + .unwrap(); + let index_file = writer.finish().await.unwrap(); + let metadata_size = index_file.file_metadata_size_bytes.unwrap(); + assert!(metadata_size.get() > object_store.block_size() as u64); + + let file_sizes = HashMap::from([(file_name.to_string(), index_file.size_bytes)]); + let without_hint = LanceIndexStore::new( + object_store.clone(), + index_dir.clone(), + Arc::new(LanceCache::no_cache()), + ) + .with_file_sizes(file_sizes.clone()); + object_store.io_stats_incremental(); + without_hint.open_index_file(file_name).await.unwrap(); + assert_eq!(object_store.io_stats_incremental().read_iops, 2); + + let with_hint = LanceIndexStore::new( + object_store.clone(), + index_dir, + Arc::new(LanceCache::no_cache()), + ) + .with_file_sizes(file_sizes) + .with_file_metadata_sizes(HashMap::from([(file_name.to_string(), metadata_size)])); + object_store.io_stats_incremental(); + with_hint.open_index_file(file_name).await.unwrap(); + assert_eq!(object_store.io_stats_incremental().read_iops, 1); + } + #[tokio::test] async fn test_basic_btree() { let tempdir = TempDir::default(); diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 4081ee10cb4..94a8a6fa1ca 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -1718,13 +1718,14 @@ async fn merge_partial_vector_auxiliary_files_inner( } let dt2 = distance_type.ok_or_else(|| Error::index("Distance type missing".to_string()))?; write_unified_ivf_and_index_metadata(w, &ivf_model, dt2, idx_type_final).await?; - let summary = w.finish().await?; + let result = w.finish_with_metadata_size().await?; + let summary = result.summary(); progress.stage_progress("write_auxiliary_index", 1).await?; progress.stage_complete("write_auxiliary_index").await?; Ok(lance_table::format::IndexFile { path: INDEX_AUXILIARY_FILE_NAME.to_string(), size_bytes: summary.size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: Some(result.metadata_size_bytes()), }) } else { Err(Error::index( diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 014438f12ae..9e3072ebe73 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -148,6 +148,27 @@ impl DataFile { } } + /// Attach the exact metadata suffix size reported by the file writer. + pub fn with_file_metadata_size_bytes(mut self, file_metadata_size_bytes: NonZero) -> Self { + self.file_metadata_size_bytes = Some(file_metadata_size_bytes); + self + } + + /// Estimate the number of physical columns from manifest column mappings. + /// + /// The file footer remains authoritative. Counting non-negative entries + /// keeps malformed or sparse mapping values from becoming read sizes. + pub fn estimated_num_columns(&self) -> Option> { + let estimated_num_columns = self + .column_indices + .iter() + .filter(|column_index| **column_index >= 0) + .count(); + u32::try_from(estimated_num_columns) + .ok() + .and_then(NonZero::new) + } + /// Create a new `DataFile` whose fields and column indices will be set later. pub fn new_unstarted(path: impl Into, file_version: ConcreteFileVersion) -> Self { let (file_major_version, file_minor_version) = file_version.to_data_file_numbers(); @@ -1048,4 +1069,12 @@ mod tests { .validate(&base_path) .expect("validation should allow extra columns without field ids"); } + + #[test] + fn physical_column_estimate_counts_mappings_not_values() { + let mut data_file = DataFile::new_unstarted("foo.lance", ConcreteFileVersion::V2_1); + data_file.column_indices = Arc::from([0, -1, 7, i32::MAX, -2]); + + assert_eq!(data_file.estimated_num_columns(), NonZero::new(3)); + } } diff --git a/rust/lance-table/src/format/index.rs b/rust/lance-table/src/format/index.rs index 5e67aecd818..6b8a56bfe95 100644 --- a/rust/lance-table/src/format/index.rs +++ b/rust/lance-table/src/format/index.rs @@ -129,6 +129,18 @@ impl IndexMetadata { .unwrap_or_default() } + /// Returns known Lance metadata suffix sizes by relative file path. + pub fn file_metadata_size_map(&self) -> HashMap> { + self.files + .iter() + .flatten() + .filter_map(|file| { + file.file_metadata_size_bytes + .map(|size| (file.path.clone(), size)) + }) + .collect() + } + /// Returns the total size of all files in this index segment in bytes. /// Returns None if file information is not available. pub fn total_size_bytes(&self) -> Option { diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 7cdd1de1685..375a268a1ff 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -2420,6 +2420,8 @@ impl Dataset { } let file_size_nz = NonZero::new(file_size); + let metadata_size_bytes = NonZero::new(file_metadata.metadata_size_bytes()) + .ok_or_else(|| Error::internal("opened file has an empty metadata suffix"))?; Ok(DataFile::new( path, fields, @@ -2427,7 +2429,8 @@ impl Dataset { lance_file_format, file_size_nz, base_id, - )) + ) + .with_file_metadata_size_bytes(metadata_size_bytes)) } /// Resolve the data directory for a given base_id. diff --git a/rust/lance/src/dataset/data_file_part.rs b/rust/lance/src/dataset/data_file_part.rs index 926044a414f..32c4cd8bfab 100644 --- a/rust/lance/src/dataset/data_file_part.rs +++ b/rust/lance/src/dataset/data_file_part.rs @@ -38,6 +38,8 @@ pub struct DataFilePart { pub(super) blob_ids: Option>, pub(super) num_rows: u64, pub(super) size_bytes: NonZeroU64, + #[serde(default)] + pub(super) metadata_size_bytes: Option, } impl DataFilePart { @@ -51,6 +53,11 @@ impl DataFilePart { self.size_bytes.get() } + /// Expected metadata suffix size, when recorded by the part writer. + pub fn metadata_size_bytes(&self) -> Option { + self.metadata_size_bytes.map(NonZeroU64::get) + } + pub(super) async fn open_all( dataset: &Dataset, target: &DataFileTarget, @@ -117,13 +124,13 @@ impl DataFilePart { let file = scheduler .open_file(&path, &CachedFileSize::new(size)) .await?; + let mut input = EncodedFileInput::new(file).with_expected_num_rows(part.num_rows()); + if let Some(metadata_size_bytes) = part.metadata_size_bytes { + input = input.with_metadata_size_bytes(metadata_size_bytes); + } opened.push( - OpenedDataFilePart::open( - EncodedFileInput::new(file).with_expected_num_rows(part.num_rows()), - part.blob_ids.clone(), - target.blob_target_id(), - ) - .await?, + OpenedDataFilePart::open(input, part.blob_ids.clone(), target.blob_target_id()) + .await?, ); } Ok(opened) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 824f96474b7..afc1b2cff0b 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -42,7 +42,8 @@ use lance_core::{ use lance_datafusion::utils::StreamingWriteSource; use lance_encoding::decoder::DecoderPlugins; use lance_file::reader::{ - CachedFileMetadata, FileMetadataIndex, FileReaderOptions, ProjectedFileReader, + CachedFileMetadata, FileMetadataIndex, FileReaderOptions, FullMetadataReadOptions, + MetadataIndexReadOptions, ProjectedFileReader, }; use lance_file::version::ConcreteFileVersion; use lance_file::versions::v1::reader::{FileReader as V1FileReader, read_batch as v1_read_batch}; @@ -1218,17 +1219,24 @@ impl FileFragment { .clone() .or_else(|| self.dataset.file_reader_options.clone()) .unwrap_or_default(); + let estimated_num_columns = data_file.estimated_num_columns(); let prefer_indexed = metadata_mode == MetadataMode::LazyAllowed && reader_projection.column_indices.len().saturating_mul(4) - < data_file - .column_indices - .iter() - .filter(|column_index| **column_index >= 0) - .count(); + < estimated_num_columns.map_or(0, |num_columns| num_columns.get() as usize); let known_schema = self .metadata .physical_rows .map(|num_rows| (data_file_schema.clone(), num_rows as u64)); + let full_metadata_options = data_file.file_metadata_size_bytes.map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ); + let metadata_index_options = + estimated_num_columns.map_or_else(MetadataIndexReadOptions::default, |num_columns| { + MetadataIndexReadOptions::default().with_estimated_num_columns(num_columns) + }); let encodings_io = Arc::new( LanceEncodingsIo::new(file_scheduler.clone()) @@ -1240,7 +1248,11 @@ impl FileFragment { prefer_indexed, || async { let metadata_index = self - .get_file_metadata_index(&file_scheduler, known_schema.clone()) + .get_file_metadata_index_with_options( + &file_scheduler, + known_schema.clone(), + metadata_index_options, + ) .await?; if (reader_projection.column_indices.len() as u32).saturating_mul(4) >= metadata_index.num_columns() @@ -1261,7 +1273,9 @@ impl FileFragment { )) }, || async { - let file_metadata = self.get_file_metadata(&file_scheduler).await?; + let file_metadata = self + .get_file_metadata_with_options(&file_scheduler, full_metadata_options) + .await?; ProjectedFileReader::try_open_with_file_metadata( encodings_io.clone(), path.clone(), @@ -1809,6 +1823,15 @@ impl FileFragment { pub async fn get_file_metadata( &self, file_scheduler: &FileScheduler, + ) -> Result> { + self.get_file_metadata_with_options(file_scheduler, FullMetadataReadOptions::default()) + .await + } + + async fn get_file_metadata_with_options( + &self, + file_scheduler: &FileScheduler, + options: FullMetadataReadOptions, ) -> Result> { let path = file_scheduler.reader().path(); let cache = self.dataset.metadata_cache.file_metadata_cache(path); @@ -1816,7 +1839,11 @@ impl FileFragment { let file_metadata = cache .get_or_insert_with_key(FileMetadataCacheKey, || async { let file_metadata: CachedFileMetadata = - lance_file::reader::FileReader::read_all_metadata(file_scheduler).await?; + lance_file::reader::FileReader::read_all_metadata_with_options( + file_scheduler, + options, + ) + .await?; Ok(file_metadata) }) .await?; @@ -1827,6 +1854,20 @@ impl FileFragment { &self, file_scheduler: &FileScheduler, known_schema: Option<(Arc, u64)>, + ) -> Result> { + self.get_file_metadata_index_with_options( + file_scheduler, + known_schema, + MetadataIndexReadOptions::default(), + ) + .await + } + + async fn get_file_metadata_index_with_options( + &self, + file_scheduler: &FileScheduler, + known_schema: Option<(Arc, u64)>, + options: MetadataIndexReadOptions, ) -> Result> { let path = file_scheduler.reader().path(); let cache = self.dataset.metadata_cache.file_metadata_cache(path); @@ -1834,14 +1875,19 @@ impl FileFragment { let metadata_index = cache .get_or_insert_with_key(FileMetadataIndexCacheKey, || async { let metadata_index = if let Some((file_schema, num_rows)) = known_schema { - lance_file::reader::FileReader::read_metadata_index_with_schema( + lance_file::reader::FileReader::read_metadata_index_with_schema_and_options( file_scheduler, file_schema, num_rows, + options, ) .await? } else { - lance_file::reader::FileReader::read_metadata_index(file_scheduler).await? + lance_file::reader::FileReader::read_metadata_index_with_options( + file_scheduler, + options, + ) + .await? }; Ok(metadata_index) }) @@ -6911,7 +6957,7 @@ mod tests { assert!(empty_narrow_stream.next().await.is_none()); let narrow_metadata_stats = dataset.object_store.as_ref().io_stats_incremental(); assert!( - narrow_metadata_stats.read_iops <= 3, + narrow_metadata_stats.read_iops <= 2, "expected lazy metadata open to skip the schema buffer read, iops={}, bytes={}", narrow_metadata_stats.read_iops, narrow_metadata_stats.read_bytes @@ -6926,6 +6972,7 @@ mod tests { assert!(empty_full_stream.next().await.is_none()); let full_metadata_stats = dataset.object_store.as_ref().io_stats_incremental(); + assert_eq!(full_metadata_stats.read_iops, 1); assert!( full_metadata_stats.read_bytes > narrow_metadata_stats.read_bytes * 4, "expected narrow lazy metadata read to fetch much less than full metadata, narrow={} bytes, full={} bytes", diff --git a/rust/lance/src/dataset/fragment/write.rs b/rust/lance/src/dataset/fragment/write.rs index b641ee16cc5..39a193c940f 100644 --- a/rust/lance/src/dataset/fragment/write.rs +++ b/rust/lance/src/dataset/fragment/write.rs @@ -181,7 +181,8 @@ impl<'a> FragmentCreateBuilder<'a> { writer.write_batches(batch_chunk.iter()).await?; } - let write_summary = writer.finish().await?; + let write_result = writer.finish_with_metadata_size().await?; + let write_summary = write_result.summary(); fragment.physical_rows = Some(write_summary.num_rows as usize); if matches!(fragment.physical_rows, Some(0)) { @@ -204,6 +205,7 @@ impl<'a> FragmentCreateBuilder<'a> { fragment.files[0].fields = field_ids; fragment.files[0].column_indices = column_indices; fragment.files[0].file_size_bytes = CachedFileSize::new(write_summary.size_bytes); + fragment.files[0].file_metadata_size_bytes = Some(write_result.metadata_size_bytes()); progress.complete(&fragment).await?; diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index a3e0917d432..f54472b5902 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -197,7 +197,9 @@ impl LanceIndexStoreExt for LanceIndexStore { let object_store = dataset.object_store_for_index(index).await?; let store = Self::with_format_version(object_store, index_dir, Arc::new(cache), format_version); - Ok(store.with_file_sizes(index.file_size_map())) + Ok(store + .with_file_sizes(index.file_size_map()) + .with_file_metadata_sizes(index.file_metadata_size_map())) } } diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 6e3894b5381..7065f797cb9 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -664,7 +664,15 @@ pub(super) async fn can_use_binary_copy_current( let file_scheduler = scan_scheduler .open_file_with_priority(&full_path, 0, &data_file.file_size_bytes) .await?; - let file_meta = LFReader::read_all_metadata(&file_scheduler).await?; + let metadata_options = data_file.file_metadata_size_bytes.map_or_else( + lance_file::reader::FullMetadataReadOptions::default, + |metadata_size_bytes| { + lance_file::reader::FullMetadataReadOptions::default() + .with_metadata_size_bytes(metadata_size_bytes) + }, + ); + let file_meta = + LFReader::read_all_metadata_with_options(&file_scheduler, metadata_options).await?; // Binary copy only preserves page and column-buffer bytes. The output file's footer // (including global buffers) is re-generated, not copied from inputs. // diff --git a/rust/lance/src/dataset/optimize/binary_copy.rs b/rust/lance/src/dataset/optimize/binary_copy.rs index c76e0ea300f..5ae81c3e08c 100644 --- a/rust/lance/src/dataset/optimize/binary_copy.rs +++ b/rust/lance/src/dataset/optimize/binary_copy.rs @@ -9,10 +9,10 @@ use crate::dataset::fragment::write::generate_random_filename; use crate::datatypes::Schema; use lance_core::Error; use lance_encoding::decoder::{ColumnInfo, PageInfo as DecPageInfo}; -use lance_file::reader::FileReader as LFReader; +use lance_file::reader::{FileReader as LFReader, FullMetadataReadOptions}; use lance_file::version::ConcreteFileVersion; use lance_file::versions as file_versions; -use lance_file::writer::{FileWriter, FileWriterOptions}; +use lance_file::writer::{FileWriteResult, FileWriter, FileWriterOptions}; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_table::format::{DataFile, Fragment}; use prost::Message; @@ -83,7 +83,8 @@ async fn finalize_current_output_file( let mut writer = current_writer .take() .ok_or_else(|| Error::internal("binary copy output writer was not initialized"))?; - flush_footer(&mut writer, schema, &final_cols, total_rows_in_current).await?; + let write_result = + flush_footer(&mut writer, schema, &final_cols, total_rows_in_current).await?; // Register the newly closed output file as a fragment data file let mut fragment = Fragment::new(0); @@ -94,6 +95,9 @@ async fn finalize_current_output_file( let mut data_file = DataFile::new_unstarted(filename, version); data_file.fields = field_ids.into(); data_file.column_indices = field_column_indices.into(); + let write_summary = write_result.summary(); + data_file.file_size_bytes = lance_io::utils::CachedFileSize::new(write_summary.size_bytes); + data_file.file_metadata_size_bytes = Some(write_result.metadata_size_bytes()); fragment.files.push(data_file); fragment.physical_rows = Some(total_rows_in_current as usize); Ok(fragment) @@ -180,7 +184,14 @@ pub async fn rewrite_files_binary_copy( let file_scheduler = scan_scheduler .open_file_with_priority(&full_path, 0, &df.file_size_bytes) .await?; - let file_meta = LFReader::read_all_metadata(&file_scheduler).await?; + let metadata_options = df.file_metadata_size_bytes.map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ); + let file_meta = + LFReader::read_all_metadata_with_options(&file_scheduler, metadata_options).await?; let src_column_infos = file_meta.column_infos.clone(); // Initialize current_page_table if current_page_table.is_empty() { @@ -421,7 +432,7 @@ pub async fn rewrite_files_binary_copy( /// This function does not manually craft the footer. Instead it: /// - Pads the current `ObjectWriter` position to a 64‑byte boundary (required for v2_1+ readers). /// - Initializes the active `FileWriter` from the collected column metadata. -/// - Calls `FileWriter::finish()` to emit column metadata, offset tables, global buffers +/// - Calls `FileWriter::finish_with_metadata_size()` to emit column metadata, offset tables, global buffers /// (schema descriptor), version, and to close the writer. /// /// Preconditions: @@ -433,9 +444,8 @@ async fn flush_footer( schema: &Schema, final_cols: &[Arc], total_rows_in_current: u64, -) -> Result<()> { +) -> Result { writer.write_external_buffer(&[]).await?; writer.initialize_with_external_columns(schema.clone(), final_cols, total_rows_in_current)?; - writer.finish().await?; - Ok(()) + writer.finish_with_metadata_size().await } diff --git a/rust/lance/src/dataset/tests/data_file_part.rs b/rust/lance/src/dataset/tests/data_file_part.rs index 35ae388c7c4..593ba09f754 100644 --- a/rust/lance/src/dataset/tests/data_file_part.rs +++ b/rust/lance/src/dataset/tests/data_file_part.rs @@ -296,6 +296,14 @@ async fn managed_part_rejects_invalid_descriptions_and_duplicate_inputs() { .unwrap(); let part = write_part(&dataset, &target, None, batch).await; let description = serde_json::to_value(&part).unwrap(); + assert!(part.metadata_size_bytes().is_some()); + let mut older_description = description.clone(); + older_description + .as_object_mut() + .unwrap() + .remove("metadata_size_bytes"); + let older_part = serde_json::from_value::(older_description).unwrap(); + assert_eq!(older_part.metadata_size_bytes(), None); let mut zero_size = description.clone(); zero_size["size_bytes"] = serde_json::json!(0); let error = serde_json::from_value::(zero_size).unwrap_err(); @@ -401,6 +409,7 @@ async fn abandoned_target_cleanup_includes_failed_writes_and_preserves_other_tar .concat_data_file_parts(&target, std::slice::from_ref(&first)) .await .unwrap(); + assert!(assembled.file_metadata_size_bytes.is_some()); let assembled_path = dataset .data_file_dir_for_base(base_id) .unwrap() diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index ca738434481..e5bcc636a73 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -224,20 +224,24 @@ impl Dataset { if let Some(preprocessor) = preprocessor.as_mut() { preprocessor.finish().await?; } - writer.finish().await + writer.finish_with_metadata_size().await } .await; match write_result { - Ok(summary) => Ok(DataFilePart { - target_file_name: target.file_name.clone(), - base_id: target.base_id, - file_name, - blob_ids: part_blob_ids, - num_rows: summary.num_rows, - size_bytes: NonZeroU64::new(summary.size_bytes) - .ok_or_else(|| Error::internal("completed part has zero file size"))?, - }), + Ok(result) => { + let summary = result.summary(); + Ok(DataFilePart { + target_file_name: target.file_name.clone(), + base_id: target.base_id, + file_name, + blob_ids: part_blob_ids, + num_rows: summary.num_rows, + size_bytes: NonZeroU64::new(summary.size_bytes) + .ok_or_else(|| Error::internal("completed part has zero file size"))?, + metadata_size_bytes: Some(result.metadata_size_bytes()), + }) + } Err(error) => { writer.abort().await; if let Some(preprocessor) = preprocessor.as_mut() { @@ -344,7 +348,8 @@ impl Dataset { target.version, NonZeroU64::new(output.size_bytes), target.base_id, - )) + ) + .with_file_metadata_size_bytes(output.metadata_size_bytes)) } } @@ -1992,14 +1997,16 @@ impl GenericWriter for V2WriterAdapter { .iter() .map(|(_, column_index)| *column_index as i32) .collect::>(); - let write_summary = self.writer.finish().await?; + let write_result = self.writer.finish_with_metadata_size().await?; let mut data_file = self .data_file .take() .ok_or_else(|| Error::internal("current writer was already finished"))?; data_file.fields = field_ids.into(); data_file.column_indices = column_indices.into(); + let write_summary = write_result.summary(); data_file.file_size_bytes = NonZero::new(write_summary.size_bytes).into(); + data_file.file_metadata_size_bytes = Some(write_result.metadata_size_bytes()); Ok((write_summary.num_rows as u32, data_file)) } @@ -2946,6 +2953,13 @@ mod tests { "version: {}", version ); + if version.resolve() == ConcreteFileVersion::V1 { + assert_eq!(fragment.files[0].file_metadata_size_bytes, None); + } else { + let file_size = fragment.files[0].file_size_bytes.get().unwrap(); + let metadata_size = fragment.files[0].file_metadata_size_bytes.unwrap(); + assert!(metadata_size <= file_size); + } } } diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 9b807fbbd79..6f81ee6f670 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -73,7 +73,7 @@ use vector::details::{ vector_details_as_json, }; pub(crate) use vector::details::{vector_index_details, vector_index_details_default}; -use vector::ivf::v2::{IVFIndex, IvfStateEntryBox}; +use vector::ivf::v2::{IVFIndex, IvfFileOpenHints, IvfStateEntryBox}; use vector::utils::get_vector_type; mod api; @@ -396,6 +396,18 @@ pub(crate) async fn build_index_metadata_from_segments( prune_stale_segment_coverage(dataset, &mut segments, false, false).await?; let new_indices = futures::stream::iter(segments.into_iter().map(|segment| async move { + let known_file_hints = segment + .known_files() + .map(|files| { + files + .iter() + .filter_map(|file| { + file.file_metadata_size_bytes + .map(|size| (file.path.clone(), (file.size_bytes, size))) + }) + .collect::>() + }) + .unwrap_or_default(); let ( uuid, fragment_bitmap, @@ -428,6 +440,13 @@ pub(crate) async fn build_index_metadata_from_segments( if is_inverted_index { retain_committed_inverted_files(&mut files); } + for file in &mut files { + if let Some((known_size, metadata_size)) = known_file_hints.get(&file.path) + && *known_size == file.size_bytes + { + file.file_metadata_size_bytes = Some(*metadata_size); + } + } Ok::<_, Error>(IndexMetadata { uuid, name: index_name.to_string(), @@ -3042,6 +3061,7 @@ impl DatasetIndexInternalExt for Dataset { .join(uuid.to_string()) .join(INDEX_FILE_NAME); let file_sizes = index_meta.file_size_map(); + let file_metadata_sizes = index_meta.file_metadata_size_map(); let reader: Arc = vector::open_index_file( object_store.as_ref(), &index_file, @@ -3114,21 +3134,20 @@ impl DatasetIndexInternalExt for Dataset { } (0, 3) | (2, _) => { + let file_open_hints = IvfFileOpenHints::new(file_sizes, file_metadata_sizes); let scheduler = ScanScheduler::new( object_store.clone(), SchedulerConfig::max_bandwidth(&object_store), ); - let cached_size = file_sizes - .get(INDEX_FILE_NAME) - .map(|&size| CachedFileSize::new(size)) - .unwrap_or_else(CachedFileSize::unknown); + let cached_size = CachedFileSize::new(file_open_hints.file_size(INDEX_FILE_NAME)); let file = scheduler.open_file(&index_file, &cached_size).await?; - let reader = lance_file::reader::FileReader::try_open( + let reader = lance_file::reader::FileReader::try_open_with_metadata_options( file, None, Default::default(), &self.metadata_cache.file_metadata_cache(&index_file), FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_FILE_NAME), ) .await?; let index_metadata = reader @@ -3156,7 +3175,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3169,7 +3188,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3188,7 +3207,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3202,7 +3221,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3216,7 +3235,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3231,7 +3250,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3244,7 +3263,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3259,7 +3278,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) @@ -3273,7 +3292,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_sizes, + file_open_hints, ) .await?; Ok(wrap_ivf(ivf)) diff --git a/rust/lance/src/index/api.rs b/rust/lance/src/index/api.rs index c17d821b608..6a9707ea13a 100644 --- a/rust/lance/src/index/api.rs +++ b/rust/lance/src/index/api.rs @@ -8,7 +8,7 @@ use datafusion::execution::SendableRecordBatchStream; use lance_index::{ FtsPrewarmResult, IndexParams, IndexType, PrewarmOptions, optimize::OptimizeOptions, }; -use lance_table::format::IndexMetadata; +use lance_table::format::{IndexFile, IndexMetadata}; use roaring::RoaringBitmap; use uuid::Uuid; @@ -37,6 +37,8 @@ pub struct IndexSegment { index_version: i32, /// Dataset version at which this segment's physical contents were built. dataset_version: u64, + /// Previously reported file metadata, used to preserve immutable open hints. + known_files: Option>, } impl IndexSegment { @@ -63,9 +65,23 @@ impl IndexSegment { index_details, index_version, dataset_version, + known_files: None, } } + /// Attach the immutable files already produced for this segment. + /// + /// Matching path and total-size metadata is preserved when the segment is + /// committed; the directory listing remains authoritative. + pub fn with_files(mut self, files: Vec) -> Self { + self.known_files = Some(files); + self + } + + pub(crate) fn known_files(&self) -> Option<&[IndexFile]> { + self.known_files.as_deref() + } + /// Return the UUID of this segment. pub fn uuid(&self) -> Uuid { self.uuid @@ -173,7 +189,8 @@ impl IntoIndexSegment for IndexMetadata { )) })?; - Ok(IndexSegment::new( + let known_files = self.files; + let segment = IndexSegment::new( self.uuid, fragment_bitmap.iter(), self.fields, @@ -181,7 +198,11 @@ impl IntoIndexSegment for IndexMetadata { self.index_version, self.dataset_version, self.covering_fields, - )) + ); + Ok(match known_files { + Some(files) => segment.with_files(files), + None => segment, + }) } } @@ -393,12 +414,20 @@ mod tests { index_version: 5, created_at: None, base_id: None, - files: None, + files: Some(vec![IndexFile { + path: "index.idx".to_string(), + size_bytes: 1_024, + file_metadata_size_bytes: std::num::NonZeroU64::new(256), + }]), }; let segment = metadata.into_index_segment().unwrap(); assert_eq!(segment.fields(), [3, 7]); assert_eq!(segment.dataset_version(), 42); + assert_eq!( + segment.known_files().unwrap()[0].file_metadata_size_bytes, + std::num::NonZeroU64::new(256) + ); } /// Segments arrive from callers this build never validated, so the keyed diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 245c861de72..34e72739f26 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use futures::{FutureExt, TryStreamExt}; use lance_core::{Error, Result}; -use lance_file::reader::FileReaderOptions; +use lance_file::reader::{FileReaderOptions, FullMetadataReadOptions}; use lance_index::{ INDEX_FILE_NAME, IndexType, frag_reuse::CompactFragReuseIndex, @@ -21,7 +21,6 @@ use lance_index::{ }, }; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; -use lance_io::utils::CachedFileSize; use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::{Fragment, IndexMetadata}; use prost::Message; @@ -274,17 +273,24 @@ async fn try_harvest_seeds( .clone() .join(crate::dataset::DATA_DIR) .join(data_file.path.as_str()); - let Ok(file_scheduler) = scheduler.open_file(&path, &CachedFileSize::unknown()).await + let Ok(file_scheduler) = scheduler.open_file(&path, &data_file.file_size_bytes).await else { return Ok(None); }; + let metadata_options = data_file.file_metadata_size_bytes.map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ); - let Ok(reader) = lance_file::reader::FileReader::try_open( + let Ok(reader) = lance_file::reader::FileReader::try_open_with_metadata_options( file_scheduler, None, Default::default(), &dataset.metadata_cache.file_metadata_cache(&path), FileReaderOptions::default(), + metadata_options, ) .await else { diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index d9e749d44ad..314703da6c8 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -2253,10 +2253,9 @@ mod tests { ); } - /// `open_index_file` looks up sizes in `IndexMetadata::file_size_map()` by - /// bare file name. This pins that a freshly created HNSW index records both - /// the main and auxiliary files under those exact names with nonzero sizes, - /// which is what lets the open path skip the HEAD. + /// `open_index_file` looks up size hints by bare file name. This pins that a + /// freshly created HNSW index records both the main and auxiliary files under + /// those exact names with nonzero total and metadata suffix sizes. #[tokio::test] async fn test_hnsw_index_records_file_sizes() { use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME}; @@ -2303,6 +2302,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); let index = indices.iter().find(|idx| idx.name == "hnsw").unwrap(); let file_sizes = index.file_size_map(); + let metadata_sizes = index.file_metadata_size_map(); assert!( file_sizes.get(INDEX_FILE_NAME).copied().unwrap_or(0) > 0, @@ -2316,6 +2316,14 @@ mod tests { > 0, "manifest should record a nonzero {INDEX_AUXILIARY_FILE_NAME} size, got {file_sizes:?}" ); + for file_name in [INDEX_FILE_NAME, INDEX_AUXILIARY_FILE_NAME] { + assert!( + metadata_sizes + .get(file_name) + .is_some_and(|metadata_size| metadata_size.get() <= file_sizes[file_name]), + "manifest should record a valid metadata suffix for {file_name}, got {metadata_sizes:?}" + ); + } } #[tokio::test] diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 5231990dd43..11ca533280c 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -1848,8 +1848,10 @@ impl IvfIndexBuilder serde_json::to_string(&partition_index_metadata)?, ); - let storage_summary = storage_writer.finish().await?; - let index_summary = index_writer.finish().await?; + let storage_result = storage_writer.finish_with_metadata_size().await?; + let index_result = index_writer.finish_with_metadata_size().await?; + let storage_summary = storage_result.summary(); + let index_summary = index_result.summary(); log::info!("merging {} partitions done", ivf.num_partitions()); @@ -1857,12 +1859,12 @@ impl IvfIndexBuilder IndexFile { path: INDEX_AUXILIARY_FILE_NAME.to_string(), size_bytes: storage_summary.size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: Some(storage_result.metadata_size_bytes()), }, IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: index_summary.size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: Some(index_result.metadata_size_bytes()), }, ]) } diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 58cfd1884ce..a866b27b69f 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -2851,14 +2851,15 @@ async fn write_root_vector_index_from_auxiliary( v2_writer.write_batch(&empty_batch).await?; } - let summary = v2_writer.finish().await?; + let result = v2_writer.finish_with_metadata_size().await?; + let summary = result.summary(); progress.stage_progress("write_root_index", 1).await?; progress.stage_complete("write_root_index").await?; Ok(IndexFile { path: INDEX_FILE_NAME.to_string(), size_bytes: summary.size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: Some(result.metadata_size_bytes()), }) } diff --git a/rust/lance/src/index/vector/ivf/partition_serde.rs b/rust/lance/src/index/vector/ivf/partition_serde.rs index a1869357980..a333bf73a44 100644 --- a/rust/lance/src/index/vector/ivf/partition_serde.rs +++ b/rust/lance/src/index/vector/ivf/partition_serde.rs @@ -1142,6 +1142,8 @@ mod tests { quantization_type: QuantizationType::Flat, index_file_size: 1024, aux_file_size: 512, + index_file_metadata_size_bytes: std::num::NonZeroU64::new(256), + aux_file_metadata_size_bytes: std::num::NonZeroU64::new(128), rq_search_cache: empty_rabit_search_cache_cell(), }; diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index d078994fbca..25662e63642 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -9,6 +9,7 @@ use std::{ any::Any, borrow::Cow, collections::{BinaryHeap, HashMap}, + num::NonZeroU64, ops::Range, sync::{ Arc, LazyLock, Mutex, OnceLock, @@ -41,7 +42,9 @@ use lance_core::utils::tracing::{IO_TYPE_LOAD_VECTOR_PART, TRACE_IO_EVENTS}; use lance_core::{Error, ROW_ID, Result}; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; use lance_file::LanceEncodingsIo; -use lance_file::reader::{CachedFileMetadata, FileReader, FileReaderOptions, ReaderProjection}; +use lance_file::reader::{ + CachedFileMetadata, FileReader, FileReaderOptions, FullMetadataReadOptions, ReaderProjection, +}; use lance_index::cache_pb::IvfStateHeader; use lance_index::frag_reuse::{CompactFragReuseIndex, CompactFragReuseIndexHandle}; use lance_index::metrics::{LocalMetricsCollector, MetricsCollector}; @@ -124,6 +127,9 @@ pub(crate) struct IvfIndexState { /// when reconstructing from cache. pub(crate) index_file_size: u64, pub(crate) aux_file_size: u64, + /// Metadata suffix sizes used to avoid dependent footer reads. + pub(crate) index_file_metadata_size_bytes: Option, + pub(crate) aux_file_metadata_size_bytes: Option, /// Runtime-only cache, intentionally excluded from the CacheCodec wire format. pub(crate) rq_search_cache: RabitSearchCacheCell, } @@ -646,6 +652,10 @@ impl CacheCodecImpl for IvfStateEntryBox { quantization_type, index_file_size: header.index_file_size, aux_file_size: header.aux_file_size, + index_file_metadata_size_bytes: NonZeroU64::new( + header.index_file_metadata_size_bytes, + ), + aux_file_metadata_size_bytes: NonZeroU64::new(header.aux_file_metadata_size_bytes), rq_search_cache: empty_rabit_search_cache_cell(), }))) } @@ -717,6 +727,12 @@ impl IvfStateEntry for IvfIndexState { quantizer_metadata_json, index_file_size: self.index_file_size, aux_file_size: self.aux_file_size, + index_file_metadata_size_bytes: self + .index_file_metadata_size_bytes + .map_or(0, NonZeroU64::get), + aux_file_metadata_size_bytes: self + .aux_file_metadata_size_bytes + .map_or(0, NonZeroU64::get), }; let ivf_bytes = pb::Ivf::try_from(&self.ivf)?.encode_to_vec(); let aux_ivf_bytes = pb::Ivf::try_from(&self.aux_ivf)?.encode_to_vec(); @@ -780,12 +796,44 @@ impl CacheKey for FileMetadataCacheKey { fn write_key(&self, _builder: &mut KeyBuilder) {} } +#[derive(Debug, Default)] +pub(crate) struct IvfFileOpenHints { + file_sizes: HashMap, + metadata_sizes: HashMap, +} + +impl IvfFileOpenHints { + pub(crate) fn new( + file_sizes: HashMap, + metadata_sizes: HashMap, + ) -> Self { + Self { + file_sizes, + metadata_sizes, + } + } + + pub(crate) fn file_size(&self, name: &str) -> u64 { + self.file_sizes.get(name).copied().unwrap_or(0) + } + + pub(crate) fn metadata_options(&self, name: &str) -> FullMetadataReadOptions { + self.metadata_sizes.get(name).copied().map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ) + } +} + /// Open a FileReader, reusing cached file metadata if available. async fn open_reader_cached( scheduler: &Arc, path: &Path, cache: &LanceCache, known_file_size: u64, + metadata_options: FullMetadataReadOptions, ) -> Result { let file_cache = cache.with_key_prefix(path.as_ref()); // CachedFileSize::new(0) == CachedFileSize::unknown(); passing the raw @@ -807,12 +855,13 @@ async fn open_reader_cached( .await } else { let file_scheduler = scheduler.open_file(path, &cached_size).await?; - let reader = FileReader::try_open( + let reader = FileReader::try_open_with_metadata_options( file_scheduler, None, Arc::::default(), cache, FileReaderOptions::default(), + metadata_options, ) .await?; // File metadata is store-free, so it outlives the reader opened here: @@ -1365,7 +1414,7 @@ impl IVFIndex { frag_reuse_index: Option>, file_metadata_cache: &LanceCache, index_cache: LanceCache, - file_sizes: HashMap, + file_open_hints: IvfFileOpenHints, ) -> Result { let io_parallelism = object_store.io_parallelism(); let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); @@ -1376,18 +1425,35 @@ impl IVFIndex { .clone() .join(uuid_str.as_str()) .join(INDEX_FILE_NAME); - let cached_size = file_sizes - .get(INDEX_FILE_NAME) - .map(|&size| CachedFileSize::new(size)) - .unwrap_or_else(CachedFileSize::unknown); - let index_reader = FileReader::try_open( - scheduler.open_file(&uri, &cached_size).await?, - None, - Arc::::default(), - file_metadata_cache, - FileReaderOptions::default(), - ) - .await?; + let aux_path = index_dir + .clone() + .join(uuid_str.as_str()) + .join(INDEX_AUXILIARY_FILE_NAME); + let index_size = CachedFileSize::new(file_open_hints.file_size(INDEX_FILE_NAME)); + let aux_size = CachedFileSize::new(file_open_hints.file_size(INDEX_AUXILIARY_FILE_NAME)); + let open_index = async { + FileReader::try_open_with_metadata_options( + scheduler.open_file(&uri, &index_size).await?, + None, + Arc::::default(), + file_metadata_cache, + FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_FILE_NAME), + ) + .await + }; + let open_aux = async { + FileReader::try_open_with_metadata_options( + scheduler.open_file(&aux_path, &aux_size).await?, + None, + Arc::::default(), + file_metadata_cache, + FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_AUXILIARY_FILE_NAME), + ) + .await + }; + let (index_reader, storage_reader) = tokio::try_join!(open_index, open_aux)?; let index_metadata: IndexMetadata = serde_json::from_str( index_reader .schema() @@ -1415,26 +1481,6 @@ impl IVFIndex { .ok_or(Error::index(format!("{} not found", S::metadata_key())))?; let sub_index_metadata: Vec = serde_json::from_str(sub_index_metadata)?; - let aux_cached_size = file_sizes - .get(INDEX_AUXILIARY_FILE_NAME) - .map(|&size| CachedFileSize::new(size)) - .unwrap_or_else(CachedFileSize::unknown); - let storage_reader = FileReader::try_open( - scheduler - .open_file( - &index_dir - .clone() - .join(uuid_str.as_str()) - .join(INDEX_AUXILIARY_FILE_NAME), - &aux_cached_size, - ) - .await?, - None, - Arc::::default(), - file_metadata_cache, - FileReaderOptions::default(), - ) - .await?; let frag_reuse_index = frag_reuse_index .clone() .map(|index| Arc::new(CompactFragReuseIndexHandle(index)) as Arc); @@ -1447,10 +1493,6 @@ impl IVFIndex { .with_key_prefix(uri.as_ref()) .insert_with_key(&FileMetadataCacheKey, index_reader.metadata().clone()) .await; - let aux_path = index_dir - .clone() - .join(uuid_str.as_str()) - .join(INDEX_AUXILIARY_FILE_NAME); file_metadata_cache .with_key_prefix(aux_path.as_ref()) .insert_with_key(&FileMetadataCacheKey, storage.reader().metadata().clone()) @@ -1826,6 +1868,12 @@ impl IVFIndex { quantization_type, index_file_size: self.reader.metadata().file_size(), aux_file_size: self.storage.reader().metadata().file_size(), + index_file_metadata_size_bytes: NonZeroU64::new( + self.reader.metadata().metadata_size_bytes(), + ), + aux_file_metadata_size_bytes: NonZeroU64::new( + self.storage.reader().metadata().metadata_size_bytes(), + ), rq_search_cache: rabit_search_cache_cell(self.rq_search_cache.clone()), })) } @@ -2511,20 +2559,32 @@ async fn reconstruct_typed( // fresh readers to the object store supplied for this reconstruction. let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); let scheduler = ScanScheduler::new(object_store, scheduler_config); - let index_reader = open_reader_cached( - &scheduler, - &index_path, - file_metadata_cache, - state.index_file_size, - ) - .await?; - let aux_reader = open_reader_cached( - &scheduler, - &aux_path, - file_metadata_cache, - state.aux_file_size, - ) - .await?; + let (index_reader, aux_reader) = tokio::try_join!( + open_reader_cached( + &scheduler, + &index_path, + file_metadata_cache, + state.index_file_size, + state.index_file_metadata_size_bytes.map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ), + ), + open_reader_cached( + &scheduler, + &aux_path, + file_metadata_cache, + state.aux_file_size, + state.aux_file_metadata_size_bytes.map_or_else( + FullMetadataReadOptions::default, + |metadata_size_bytes| { + FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) + }, + ), + ), + )?; let frag_reuse_index = frag_reuse_index .map(|index| Arc::new(CompactFragReuseIndexHandle(index)) as Arc); From 860e5d81be841477fda78b3b2191333f1a6936b3 Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 16:38:07 -0700 Subject: [PATCH 4/6] fix(bindings): initialize metadata size hints Co-authored-by: Cursor --- java/lance-jni/src/fragment.rs | 1 + java/lance-jni/src/transaction.rs | 1 + python/src/fragment.rs | 1 + python/src/transaction.rs | 6 +++++- 4 files changed, 8 insertions(+), 1 deletion(-) diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index 9c35cf1a8fc..f59010ff9c7 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -1027,6 +1027,7 @@ impl FromJObjectWithEnv for JObject<'_> { file_major_version, file_minor_version, file_size_bytes, + file_metadata_size_bytes: None, base_id, }) } diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index eb06d599f2d..96a5de75e93 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -144,6 +144,7 @@ impl FromJObjectWithEnv for JObject<'_> { "newIndexFiles.sizeBytes", env.call_method(self, "getSizeBytes", "()J", &[])?.j()?, )?, + file_metadata_size_bytes: None, }) } } diff --git a/python/src/fragment.rs b/python/src/fragment.rs index ae7135443e0..a393f8d9d1d 100644 --- a/python/src/fragment.rs +++ b/python/src/fragment.rs @@ -966,6 +966,7 @@ impl FromPyObject<'_, '_> for PyLance { file_major_version: ob.getattr("file_major_version")?.extract()?, file_minor_version: ob.getattr("file_minor_version")?.extract()?, file_size_bytes, + file_metadata_size_bytes: None, base_id: ob.getattr("base_id")?.extract()?, })) } diff --git a/python/src/transaction.rs b/python/src/transaction.rs index b77d04e2513..d43b286574b 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -28,7 +28,11 @@ impl FromPyObject<'_, '_> for PyLance { fn extract(ob: Borrowed<'_, '_, PyAny>) -> PyResult { let path = ob.getattr("path")?.extract()?; let size_bytes = ob.getattr("size_bytes")?.extract()?; - Ok(Self(IndexFile { path, size_bytes })) + Ok(Self(IndexFile { + path, + size_bytes, + file_metadata_size_bytes: None, + })) } } From 9430aa3f24db878c9008e65b187c0bf1b10af2b1 Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 17:44:03 -0700 Subject: [PATCH 5/6] perf(file): eliminate redundant metadata reads Co-authored-by: Cursor --- java/lance-jni/src/fragment.rs | 34 +- java/lance-jni/src/transaction.rs | 25 +- .../java/org/lance/fragment/DataFile.java | 47 ++- .../main/java/org/lance/index/IndexFile.java | 29 +- .../src/test/java/org/lance/FragmentTest.java | 15 + .../test/java/org/lance/TransactionTest.java | 9 +- python/python/lance/dataset.py | 18 +- python/python/lance/fragment.py | 11 + python/python/tests/test_commit_index.py | 23 +- python/python/tests/test_fragment.py | 18 +- python/python/tests/test_table_ops.py | 11 +- python/src/fragment.rs | 6 +- python/src/transaction.rs | 11 +- rust/lance-file/src/concat.rs | 31 +- rust/lance-file/src/reader.rs | 298 ++++++++++-------- rust/lance-index/src/scalar/lance_format.rs | 81 +++-- rust/lance-table/src/format/index.rs | 12 - .../src/transaction/manifest_build.rs | 3 +- rust/lance/src/dataset/fragment.rs | 9 +- rust/lance/src/dataset/index.rs | 12 +- .../lance/src/dataset/tests/data_file_part.rs | 7 + rust/lance/src/index.rs | 205 ++++++------ rust/lance/src/index/vector.rs | 71 +++-- rust/lance/src/index/vector/ivf/v2.rs | 111 ++++--- 24 files changed, 725 insertions(+), 372 deletions(-) diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index f59010ff9c7..42d89e01064 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -17,7 +17,7 @@ use lance::table::format::{ }; use lance_io::utils::CachedFileSize; use lance_table::rowids::{RowIdSequence, write_row_ids}; -use std::iter::once; +use std::{iter::once, num::NonZeroU64}; use roaring::RoaringBitmap; @@ -685,7 +685,7 @@ fn inner_encode_row_ids(env: &mut JNIEnv, row_ids: &JLongArray) -> Result JObject::null(), }; let base_id = convert_to_java_integer(env, self.base_id)?; + let file_metadata_size_bytes = match self.file_metadata_size_bytes { + Some(size) => { + let size = i64::try_from(size.get()).map_err(|_| { + Error::runtime_error(format!( + "Cannot convert DataFile.file_metadata_size_bytes={} to Java long", + size.get() + )) + })?; + JLance(size).into_java(env)? + } + None => JObject::null(), + }; Ok(env.new_object( DATA_FILE_CLASS, DATA_FILE_CONSTRUCTOR_SIG, @@ -756,6 +768,7 @@ impl IntoJava for &DataFile { JValueGen::Int(self.file_minor_version as i32), JValueGen::Object(&file_size_bytes), JValueGen::Object(&base_id), + JValueGen::Object(&file_metadata_size_bytes), ], )?) } @@ -1019,6 +1032,21 @@ impl FromJObjectWithEnv for JObject<'_> { .extract_object(env)?; let file_size_bytes = file_size_bytes.map_or(Default::default(), |r| CachedFileSize::new(r as u64)); + let file_metadata_size_bytes: Option = env + .call_method(self, "getFileMetadataSizeBytes", "()Ljava/lang/Long;", &[])? + .l()? + .extract_object(env)?; + let file_metadata_size_bytes = match file_metadata_size_bytes { + Some(size) => { + let size = u64::try_from(size).map_err(|_| { + Error::input_error(format!( + "DataFile.fileMetadataSizeBytes must be non-negative, got {size}" + )) + })?; + NonZeroU64::new(size) + } + None => None, + }; let base_id = get_base_id(env, self)?; Ok(DataFile { path, @@ -1027,7 +1055,7 @@ impl FromJObjectWithEnv for JObject<'_> { file_major_version, file_minor_version, file_size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes, base_id, }) } diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 96a5de75e93..e07f9e437b1 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -37,6 +37,7 @@ use prost_types::Any; use roaring::RoaringBitmap; use std::cell::Cell; use std::collections::HashMap; +use std::num::NonZeroU64; use std::sync::Arc; use uuid::Uuid; @@ -125,12 +126,21 @@ fn export_unsigned_longs<'a>( impl IntoJava for &IndexFile { fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result> { let path = env.new_string(&self.path)?; + let file_metadata_size_bytes = match self.file_metadata_size_bytes { + Some(size) => JLance(u64_to_jlong( + "newIndexFiles.fileMetadataSizeBytes", + size.get(), + )?) + .into_java(env)?, + None => JObject::null(), + }; Ok(env.new_object( "org/lance/index/IndexFile", - "(Ljava/lang/String;J)V", + "(Ljava/lang/String;JLjava/lang/Long;)V", &[ JValue::Object(&path), JValue::Long(u64_to_jlong("newIndexFiles.sizeBytes", self.size_bytes)?), + JValue::Object(&file_metadata_size_bytes), ], )?) } @@ -138,13 +148,24 @@ impl IntoJava for &IndexFile { impl FromJObjectWithEnv for JObject<'_> { fn extract_object(&self, env: &mut JNIEnv<'_>) -> Result { + let file_metadata_size_bytes: Option = env + .call_method(self, "getFileMetadataSizeBytes", "()Ljava/lang/Long;", &[])? + .l()? + .extract_object(env)?; + let file_metadata_size_bytes = match file_metadata_size_bytes { + Some(size) => { + let size = nonnegative_jlong_to_u64("newIndexFiles.fileMetadataSizeBytes", size)?; + NonZeroU64::new(size) + } + None => None, + }; Ok(IndexFile { path: env.get_string_from_method(self, "getPath")?, size_bytes: nonnegative_jlong_to_u64( "newIndexFiles.sizeBytes", env.call_method(self, "getSizeBytes", "()J", &[])?.j()?, )?, - file_metadata_size_bytes: None, + file_metadata_size_bytes, }) } } diff --git a/java/src/main/java/org/lance/fragment/DataFile.java b/java/src/main/java/org/lance/fragment/DataFile.java index 1120a5286c1..2764fe3796a 100644 --- a/java/src/main/java/org/lance/fragment/DataFile.java +++ b/java/src/main/java/org/lance/fragment/DataFile.java @@ -29,6 +29,7 @@ public class DataFile implements Serializable { private final int fileMinorVersion; private final Long fileSizeBytes; private final Integer baseId; + private final Long fileMetadataSizeBytes; public DataFile( String path, @@ -38,6 +39,41 @@ public DataFile( int fileMinorVersion, Long fileSizeBytes, Integer baseId) { + this( + path, + fields, + columnIndices, + fileMajorVersion, + fileMinorVersion, + fileSizeBytes, + baseId, + null); + } + + /** + * Creates data file metadata with an optional exact metadata suffix size. + * + * @param path path to the data file + * @param fields field IDs stored in the file + * @param columnIndices physical column indices for the fields + * @param fileMajorVersion major storage format version + * @param fileMinorVersion minor storage format version + * @param fileSizeBytes total file size, or null if unknown + * @param baseId base path ID, or null for the dataset base + * @param fileMetadataSizeBytes metadata suffix size; zero or null if unknown + */ + public DataFile( + String path, + int[] fields, + int[] columnIndices, + int fileMajorVersion, + int fileMinorVersion, + Long fileSizeBytes, + Integer baseId, + Long fileMetadataSizeBytes) { + if (fileMetadataSizeBytes != null && fileMetadataSizeBytes < 0) { + throw new IllegalArgumentException("fileMetadataSizeBytes must be non-negative"); + } this.path = path; this.fields = fields; this.columnIndices = columnIndices; @@ -45,6 +81,8 @@ public DataFile( this.fileMinorVersion = fileMinorVersion; this.fileSizeBytes = fileSizeBytes; this.baseId = baseId; + this.fileMetadataSizeBytes = + fileMetadataSizeBytes != null && fileMetadataSizeBytes == 0 ? null : fileMetadataSizeBytes; } public String getPath() { @@ -75,6 +113,11 @@ public Optional getBaseId() { return Optional.ofNullable(baseId); } + /** Returns the exact metadata suffix size, or null if it is unavailable. */ + public Long getFileMetadataSizeBytes() { + return fileMetadataSizeBytes; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -85,7 +128,8 @@ public boolean equals(Object o) { && Objects.equals(path, that.path) && Arrays.equals(fields, that.fields) && Arrays.equals(columnIndices, that.columnIndices) - && Objects.equals(fileSizeBytes, that.fileSizeBytes); + && Objects.equals(fileSizeBytes, that.fileSizeBytes) + && Objects.equals(fileMetadataSizeBytes, that.fileMetadataSizeBytes); } @Override @@ -98,6 +142,7 @@ public String toString() { .add("fileMinorVersion", fileMinorVersion) .add("fileSizeBytes", fileSizeBytes) .add("baseId", baseId) + .add("fileMetadataSizeBytes", fileMetadataSizeBytes) .toString(); } } diff --git a/java/src/main/java/org/lance/index/IndexFile.java b/java/src/main/java/org/lance/index/IndexFile.java index bb6fdc9cd3f..477f571dc15 100644 --- a/java/src/main/java/org/lance/index/IndexFile.java +++ b/java/src/main/java/org/lance/index/IndexFile.java @@ -21,13 +21,30 @@ public final class IndexFile { private final String path; private final long sizeBytes; + private final Long fileMetadataSizeBytes; public IndexFile(String path, long sizeBytes) { + this(path, sizeBytes, null); + } + + /** + * Creates index file metadata with an optional exact metadata suffix size. + * + * @param path path relative to the index directory + * @param sizeBytes total file size + * @param fileMetadataSizeBytes metadata suffix size; zero or null if unavailable + */ + public IndexFile(String path, long sizeBytes, Long fileMetadataSizeBytes) { this.path = Objects.requireNonNull(path, "path must not be null"); if (sizeBytes < 0) { throw new IllegalArgumentException("sizeBytes must be non-negative"); } + if (fileMetadataSizeBytes != null && fileMetadataSizeBytes < 0) { + throw new IllegalArgumentException("fileMetadataSizeBytes must be non-negative"); + } this.sizeBytes = sizeBytes; + this.fileMetadataSizeBytes = + fileMetadataSizeBytes != null && fileMetadataSizeBytes == 0 ? null : fileMetadataSizeBytes; } public String getPath() { @@ -38,17 +55,24 @@ public long getSizeBytes() { return sizeBytes; } + /** Returns the exact metadata suffix size, or null if it is unavailable. */ + public Long getFileMetadataSizeBytes() { + return fileMetadataSizeBytes; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; IndexFile indexFile = (IndexFile) o; - return sizeBytes == indexFile.sizeBytes && Objects.equals(path, indexFile.path); + return sizeBytes == indexFile.sizeBytes + && Objects.equals(path, indexFile.path) + && Objects.equals(fileMetadataSizeBytes, indexFile.fileMetadataSizeBytes); } @Override public int hashCode() { - return Objects.hash(path, sizeBytes); + return Objects.hash(path, sizeBytes, fileMetadataSizeBytes); } @Override @@ -56,6 +80,7 @@ public String toString() { return MoreObjects.toStringHelper(this) .add("path", path) .add("sizeBytes", sizeBytes) + .add("fileMetadataSizeBytes", fileMetadataSizeBytes) .toString(); } } diff --git a/java/src/test/java/org/lance/FragmentTest.java b/java/src/test/java/org/lance/FragmentTest.java index 025aaa24479..18844f77135 100644 --- a/java/src/test/java/org/lance/FragmentTest.java +++ b/java/src/test/java/org/lance/FragmentTest.java @@ -81,6 +81,15 @@ void testFragmentCreateFfiArray(@TempDir Path tempDir) { @Test void testFragmentCreate(@TempDir Path tempDir) throws Exception { + assertNull( + new org.lance.fragment.DataFile("unknown.lance", new int[0], new int[0], 2, 2, 1L, null, 0L) + .getFileMetadataSizeBytes()); + assertThrows( + IllegalArgumentException.class, + () -> + new org.lance.fragment.DataFile( + "invalid.lance", new int[0], new int[0], 2, 2, 1L, null, -1L)); + String datasetPath = tempDir.resolve("new_fragment").toString(); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset testDataset = @@ -89,6 +98,9 @@ void testFragmentCreate(@TempDir Path tempDir) throws Exception { int rowCount = 21; FragmentMetadata fragmentMeta = testDataset.createNewFragment(rowCount); assertEquals(fragmentMeta.getFiles(), fragmentMeta.getReferencedLanceFiles()); + Long fileMetadataSizeBytes = fragmentMeta.getFiles().get(0).getFileMetadataSizeBytes(); + assertNotNull(fileMetadataSizeBytes); + assertTrue(fileMetadataSizeBytes > 0); // Commit fragment FragmentOperation.Append appendOp = new FragmentOperation.Append(Arrays.asList(fragmentMeta)); @@ -97,6 +109,9 @@ void testFragmentCreate(@TempDir Path tempDir) throws Exception { assertEquals(2, dataset.latestVersion()); assertEquals(rowCount, dataset.countRows()); Fragment fragment = dataset.getFragments().get(0); + assertEquals( + fileMetadataSizeBytes, + fragment.metadata().getFiles().get(0).getFileMetadataSizeBytes()); try (LanceScanner scanner = fragment.newScan()) { Schema schemaRes = scanner.schema(); diff --git a/java/src/test/java/org/lance/TransactionTest.java b/java/src/test/java/org/lance/TransactionTest.java index afb0eaac2b8..a363f09332a 100644 --- a/java/src/test/java/org/lance/TransactionTest.java +++ b/java/src/test/java/org/lance/TransactionTest.java @@ -58,6 +58,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -922,6 +923,11 @@ public void testRewrittenIndexFilesModelPreservesAbsentAndEmpty() { @Test public void testRewrittenIndexFilesRoundTrip(@TempDir Path tempDir) { + assertNull(new IndexFile("unknown.idx", 123, 0L).getFileMetadataSizeBytes()); + IllegalArgumentException invalidMetadataSize = + assertThrows(IllegalArgumentException.class, () -> new IndexFile("index.idx", 123, -1L)); + assertTrue(invalidMetadataSize.getMessage().contains("must be non-negative")); + String datasetPath = tempDir.resolve("rewritten_index_files").toString(); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset testDataset = @@ -935,7 +941,7 @@ public void testRewrittenIndexFilesRoundTrip(@TempDir Path tempDir) { .withIndexName("btree_id") .build()); Index oldIndex = dataset.getIndexes().get(0); - IndexFile indexFile = new IndexFile("index.idx", 123); + IndexFile indexFile = new IndexFile("index.idx", 123, 37L); RewrittenIndex rewrittenIndex = RewrittenIndex.builder() .oldId(oldIndex.uuid()) @@ -960,6 +966,7 @@ public void testRewrittenIndexFilesRoundTrip(@TempDir Path tempDir) { operation.rewrittenIndices().get(0).getNewIndexFiles().orElseThrow().get(0); assertEquals("index.idx", readFile.getPath()); assertEquals(123, readFile.getSizeBytes()); + assertEquals(37L, readFile.getFileMetadataSizeBytes()); } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 92239e66fd4..8e018509301 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -5952,10 +5952,26 @@ class ExecuteResult(TypedDict): @dataclass class IndexFile: - """Metadata about a file in an index segment.""" + """Metadata about a file in an index segment. + + ``file_metadata_size_bytes`` is the exact metadata suffix size when + available. Zero is treated as unavailable. + """ path: str size_bytes: int + file_metadata_size_bytes: Optional[int] = None + + def __post_init__(self) -> None: + if ( + self.file_metadata_size_bytes is not None + and self.file_metadata_size_bytes < 0 + ): + raise ValueError( + "file_metadata_size_bytes must be non-negative, " + f"got {self.file_metadata_size_bytes}" + ) + self.file_metadata_size_bytes = self.file_metadata_size_bytes or None @dataclass diff --git a/python/python/lance/fragment.py b/python/python/lance/fragment.py index bbffb826009..c118857b24f 100644 --- a/python/python/lance/fragment.py +++ b/python/python/lance/fragment.py @@ -239,6 +239,9 @@ class DataFile: The minor version of the data storage format. file_size_bytes : Optional[int] The size of the data file in bytes, if available. + file_metadata_size_bytes : Optional[int] + The exact metadata suffix size in bytes, if available. Zero is treated + as unavailable. """ _path: str @@ -248,6 +251,7 @@ class DataFile: file_minor_version: int = 0 file_size_bytes: Optional[int] = None base_id: Optional[int] = None + file_metadata_size_bytes: Optional[int] = None def __init__( self, @@ -258,7 +262,13 @@ def __init__( file_minor_version: int = 0, file_size_bytes: Optional[int] = None, base_id: Optional[int] = None, + file_metadata_size_bytes: Optional[int] = None, ): + if file_metadata_size_bytes is not None and file_metadata_size_bytes < 0: + raise ValueError( + "file_metadata_size_bytes must be non-negative, " + f"got {file_metadata_size_bytes}" + ) # TODO: only we eliminate the path method, we can remove this self._path = path self.fields = fields @@ -267,6 +277,7 @@ def __init__( self.file_minor_version = file_minor_version self.file_size_bytes = file_size_bytes self.base_id = base_id + self.file_metadata_size_bytes = file_metadata_size_bytes or None def __repr__(self): # pretend we have a 'path' attribute diff --git a/python/python/tests/test_commit_index.py b/python/python/tests/test_commit_index.py index f4939434957..0c43f2ce010 100644 --- a/python/python/tests/test_commit_index.py +++ b/python/python/tests/test_commit_index.py @@ -111,6 +111,15 @@ def test_commit_index_with_files(dataset_with_index, test_table, tmp_path): """Test that the files field on Index round-trips through commit.""" from lance.dataset import Index, IndexFile + assert ( + IndexFile( + path="unknown.idx", size_bytes=1, file_metadata_size_bytes=0 + ).file_metadata_size_bytes + is None + ) + with pytest.raises(ValueError, match="must be non-negative"): + IndexFile(path="invalid.idx", size_bytes=1, file_metadata_size_bytes=-1) + # Get info about the existing index created by the fixture original_desc = dataset_with_index.describe_indices()[0] index_id = original_desc.segments[0].uuid @@ -134,7 +143,11 @@ def test_commit_index_with_files(dataset_with_index, test_table, tmp_path): # Create IndexFile objects with custom sizes to verify they round-trip index_files = [ - IndexFile(path="index.idx", size_bytes=1024), + IndexFile( + path="index.idx", + size_bytes=1024, + file_metadata_size_bytes=256, + ), IndexFile(path="auxiliary.bin", size_bytes=2048), ] @@ -176,9 +189,11 @@ def test_commit_index_with_files(dataset_with_index, test_table, tmp_path): assert len(committed_index.files) == 2 # Verify the file sizes match what we set - files_by_path = {f.path: f.size_bytes for f in committed_index.files} - assert files_by_path["index.idx"] == 1024 - assert files_by_path["auxiliary.bin"] == 2048 + files_by_path = {f.path: f for f in committed_index.files} + assert files_by_path["index.idx"].size_bytes == 1024 + assert files_by_path["index.idx"].file_metadata_size_bytes == 256 + assert files_by_path["auxiliary.bin"].size_bytes == 2048 + assert files_by_path["auxiliary.bin"].file_metadata_size_bytes is None def test_commit_index_with_index_details(dataset_with_index, test_table, tmp_path): diff --git a/python/python/tests/test_fragment.py b/python/python/tests/test_fragment.py index fc765be2950..26a87f95576 100644 --- a/python/python/tests/test_fragment.py +++ b/python/python/tests/test_fragment.py @@ -254,12 +254,26 @@ def test_dataset_progress(tmp_path: Path): def test_fragment_meta(): + assert ( + lance.fragment.DataFile( + "unknown.lance", [], file_metadata_size_bytes=0 + ).file_metadata_size_bytes + is None + ) + with pytest.raises(ValueError, match="must be non-negative"): + lance.fragment.DataFile("invalid.lance", [], file_metadata_size_bytes=-1) + # Intentionally leaving off column_indices / version fields to make sure # we can handle backwards compatibility (though not clear we need to) data = { "id": 0, "files": [ - {"path": "0.lance", "fields": [0], "file_size_bytes": 100}, + { + "path": "0.lance", + "fields": [0], + "file_size_bytes": 100, + "file_metadata_size_bytes": 25, + }, {"path": "1.lance", "fields": [1]}, ], "deletion_file": None, @@ -271,7 +285,9 @@ def test_fragment_meta(): assert len(meta.files) == 2 with pytest.warns(DeprecationWarning): assert meta.files[0].path() == "0.lance" + assert meta.files[0].file_metadata_size_bytes == 25 assert meta.files[1].path == "1.lance" + assert meta.files[1].file_metadata_size_bytes is None assert repr(meta) == ( "FragmentMetadata(id=0, files=[DataFile(path='0.lance', fields=[0], " diff --git a/python/python/tests/test_table_ops.py b/python/python/tests/test_table_ops.py index 24cf3409f91..ef2a4d404f2 100644 --- a/python/python/tests/test_table_ops.py +++ b/python/python/tests/test_table_ops.py @@ -129,6 +129,8 @@ def test_data_file_create_basic(tmp_path: str): assert df.file_major_version == int(stable_version().split(".")[0]) assert df.file_minor_version == int(stable_version().split(".")[1]) assert df.file_size_bytes is not None and df.file_size_bytes > 0 + assert df.file_metadata_size_bytes is not None + assert 0 < df.file_metadata_size_bytes <= df.file_size_bytes def test_data_file_create_subset_columns(tmp_path: str): @@ -180,7 +182,14 @@ def test_data_file_create_end_to_end(tmp_path: str): read_version=ds.version, ) - result = lance.dataset(tmp_path).to_table() + committed = lance.dataset(tmp_path) + committed_file = next( + data_file + for data_file in committed.get_fragments()[0].data_files() + if data_file.path == new_file_name + ) + assert committed_file.file_metadata_size_bytes == df.file_metadata_size_bytes + result = committed.to_table() assert result.column("b").to_pylist() == list(range(200, 300)) assert result.column("a").to_pylist() == list(range(100)) diff --git a/python/src/fragment.rs b/python/src/fragment.rs index a393f8d9d1d..5542cf3d845 100644 --- a/python/src/fragment.rs +++ b/python/src/fragment.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::fmt::Write as _; +use std::num::NonZeroU64; use std::sync::Arc; use arrow::ffi_stream::ArrowArrayStreamReader; @@ -957,6 +958,8 @@ impl FromPyObject<'_, '_> for PyLance { fn extract(ob: Borrowed<'_, '_, PyAny>) -> PyResult { let file_size_bytes: Option = ob.getattr("file_size_bytes")?.extract()?; let file_size_bytes = CachedFileSize::new(file_size_bytes.unwrap_or(0)); + let file_metadata_size_bytes: Option = + ob.getattr("file_metadata_size_bytes")?.extract()?; let fields: Vec = ob.getattr("fields")?.extract()?; let column_indices: Vec = ob.getattr("column_indices")?.extract()?; Ok(Self(DataFile { @@ -966,7 +969,7 @@ impl FromPyObject<'_, '_> for PyLance { file_major_version: ob.getattr("file_major_version")?.extract()?, file_minor_version: ob.getattr("file_minor_version")?.extract()?, file_size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: file_metadata_size_bytes.and_then(NonZeroU64::new), base_id: ob.getattr("base_id")?.extract()?, })) } @@ -992,6 +995,7 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataFile> { self.0.file_minor_version, file_size_bytes, self.0.base_id, + self.0.file_metadata_size_bytes.map(NonZeroU64::get), )) } } diff --git a/python/src/transaction.rs b/python/src/transaction.rs index d43b286574b..72b0fcb87e8 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -19,6 +19,7 @@ use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; use pyo3::{intern, prelude::*}; use roaring::RoaringBitmap; use std::collections::HashMap; +use std::num::NonZeroU64; use std::sync::Arc; use uuid::Uuid; @@ -28,10 +29,12 @@ impl FromPyObject<'_, '_> for PyLance { fn extract(ob: Borrowed<'_, '_, PyAny>) -> PyResult { let path = ob.getattr("path")?.extract()?; let size_bytes = ob.getattr("size_bytes")?.extract()?; + let file_metadata_size_bytes: Option = + ob.getattr("file_metadata_size_bytes")?.extract()?; Ok(Self(IndexFile { path, size_bytes, - file_metadata_size_bytes: None, + file_metadata_size_bytes: file_metadata_size_bytes.and_then(NonZeroU64::new), })) } } @@ -49,7 +52,11 @@ impl<'py> IntoPyObject<'py> for PyLance<&IndexFile> { let cls = namespace .getattr("IndexFile") .expect("Failed to get IndexFile class"); - cls.call1((self.0.path.clone(), self.0.size_bytes)) + cls.call1(( + self.0.path.clone(), + self.0.size_bytes, + self.0.file_metadata_size_bytes.map(NonZeroU64::get), + )) } } diff --git a/rust/lance-file/src/concat.rs b/rust/lance-file/src/concat.rs index a298f21acce..079f6539096 100644 --- a/rust/lance-file/src/concat.rs +++ b/rust/lance-file/src/concat.rs @@ -31,7 +31,11 @@ use prost::Message; use prost_types::Any; use crate::{ - reader::{CachedFileMetadata, FileReader, FullMetadataReadOptions, RawFileMetadataOpen}, + io::LanceEncodingsIo, + reader::{ + CachedFileMetadata, FileReader, FileReaderOptions, FullMetadataReadOptions, + RawFileMetadataOpen, + }, version::ConcreteFileVersion, versions, writer::{FileWriteResult, FileWriterOptions}, @@ -214,13 +218,8 @@ impl DataFilePart { let has_blob_v2 = schema.fields_pre_order().any(|field| field.is_blob_v2()); if has_blob_v2 { - validate_blob_descriptors( - &input, - metadata.as_ref(), - schema.as_ref(), - blob_ids.as_ref(), - ) - .await?; + validate_blob_descriptors(&input, &metadata, schema.as_ref(), blob_ids.as_ref()) + .await?; if blob_target_id.is_none() { return Err(Error::invalid_input(format!( "part at '{}' contains Blob v2 columns but no Blob target ID was provided", @@ -329,7 +328,7 @@ fn validate_blob_id_range(blob_ids: Option<&Range>) -> Result<()> { async fn validate_blob_descriptors( input: &EncodedFileInput, - metadata: &CachedFileMetadata, + metadata: &Arc, schema: &Schema, blob_ids: Option<&Range>, ) -> Result<()> { @@ -365,15 +364,21 @@ async fn validate_blob_descriptors( &blob_schema, &field_id_to_column_index, )?; - let reader = FileReader::try_open_with_metadata_options( - input.scheduler(), + let options = FileReaderOptions::default(); + let io = Arc::new( + LanceEncodingsIo::new(input.scheduler()).with_read_chunk_size(options.read_chunk_size), + ); + let reader = FileReader::try_open_with_file_metadata( + io, + input.path().clone(), Some(projection), Arc::::default(), + metadata.clone(), &LanceCache::no_cache(), - Default::default(), - input.metadata_options, + options, ) .await?; + debug_assert!(Arc::ptr_eq(reader.metadata(), metadata)); let mut batches = reader .read_stream( ReadBatchParams::RangeFull, diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 2b81c266234..63c5f1219ea 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -437,13 +437,52 @@ enum TailReadPlan { EstimatedMetadataIndex(NonZeroU32), } -struct InitialTailRead { +struct TailRead { bytes: Bytes, file_len: u64, offset: u64, retention_offset: u64, } +impl TailRead { + async fn extend_to(&mut self, start: u64, scheduler: &FileScheduler) -> Result<()> { + if start > self.file_len { + return Err(Error::invalid_input(format!( + "metadata tail start {start} is outside file of size {}", + self.file_len + ))); + } + if start >= self.offset { + return Ok(()); + } + let missing_bytes = scheduler.submit_single(start..self.offset, 0).await?; + let mut combined = BytesMut::with_capacity(missing_bytes.len() + self.bytes.len()); + combined.extend_from_slice(&missing_bytes); + combined.extend_from_slice(&self.bytes); + self.bytes = combined.freeze(); + self.offset = start; + Ok(()) + } + + fn slice(&self, range: Range) -> Result { + if range.start > range.end || range.start < self.offset || range.end > self.file_len { + return Err(Error::invalid_input(format!( + "metadata byte range {}..{} is outside loaded tail {}..{}", + range.start, range.end, self.offset, self.file_len + ))); + } + let start = usize::try_from(range.start - self.offset) + .map_err(|_| Error::invalid_input_source("Metadata range start overflows".into()))?; + let end = usize::try_from(range.end - self.offset) + .map_err(|_| Error::invalid_input_source("Metadata range end overflows".into()))?; + Ok(self.bytes.slice(start..end)) + } + + fn suffix_from(&self, start: u64) -> Result { + self.slice(start..self.file_len) + } +} + #[derive(Debug, Clone)] pub(crate) struct PreparedProjection { pub column_infos: Vec>, @@ -743,7 +782,7 @@ impl FileReader { self.core.read_global_buffer(index).await } - async fn read_tail(scheduler: &FileScheduler, plan: TailReadPlan) -> Result { + async fn read_tail(scheduler: &FileScheduler, plan: TailReadPlan) -> Result { let file_size = scheduler.reader().size().await? as u64; let baseline_read_size = file_size.min(scheduler.reader().block_size() as u64); let requested_read_size = match plan { @@ -757,18 +796,19 @@ impl FileReader { } } TailReadPlan::EstimatedMetadataIndex(estimated_num_columns) => { - let estimated_cmo_size = - u64::from(estimated_num_columns.get()) * COLUMN_METADATA_OFFSET_ENTRY_SIZE; - baseline_read_size - .checked_add(estimated_cmo_size) - .unwrap_or(file_size) + u64::from(estimated_num_columns.get()) + .checked_mul(COLUMN_METADATA_OFFSET_ENTRY_SIZE) + .and_then(|estimated_cmo_size| { + baseline_read_size.checked_add(estimated_cmo_size) + }) + .unwrap_or(baseline_read_size) .min(file_size) } }; let offset = file_size - requested_read_size; // FileScheduler may split this logical range into concurrent physical requests. let bytes = scheduler.submit_single(offset..file_size, 0).await?; - Ok(InitialTailRead { + Ok(TailRead { bytes, file_len: file_size, offset, @@ -776,37 +816,6 @@ impl FileReader { }) } - async fn read_range_from_tail_or_scheduler( - tail_bytes: &Bytes, - tail_offset: u64, - scheduler: &FileScheduler, - range: Range, - ) -> Result { - if range.start > range.end { - return Err(Error::invalid_input(format!( - "invalid metadata byte range: start={}, end={}", - range.start, range.end - ))); - } - let tail_end = tail_offset - .checked_add(tail_bytes.len() as u64) - .ok_or_else(|| Error::invalid_input_source("Tail byte range overflows".into()))?; - if range.start >= tail_offset && range.end <= tail_end { - let rel_start = (range.start - tail_offset) as usize; - let rel_end = (range.end - tail_offset) as usize; - Ok(tail_bytes.slice(rel_start..rel_end)) - } else if range.start < tail_offset && range.end > tail_offset && range.end <= tail_end { - let missing_bytes = scheduler.submit_single(range.start..tail_offset, 0).await?; - let rel_end = (range.end - tail_offset) as usize; - let mut combined = BytesMut::with_capacity((range.end - range.start) as usize); - combined.extend_from_slice(&missing_bytes); - combined.extend_from_slice(&tail_bytes[..rel_end]); - Ok(combined.freeze()) - } else { - scheduler.submit_single(range, 0).await - } - } - fn retained_global_buffers_from_tail( gbo_table: &[BufferDescriptor], tail_bytes: &Bytes, @@ -953,36 +962,6 @@ impl FileReader { Ok(Arc::from(offsets)) } - async fn optimistic_tail_read( - data: &Bytes, - start_pos: u64, - scheduler: &FileScheduler, - file_len: u64, - ) -> Result { - let num_bytes_needed = file_len.checked_sub(start_pos).ok_or_else(|| { - Error::invalid_input_source( - format!( - "Tail read position {} is outside file of size {}", - start_pos, file_len - ) - .into(), - ) - })? as usize; - if data.len() >= num_bytes_needed { - Ok(data.slice((data.len() - num_bytes_needed)..)) - } else { - let num_bytes_missing = (num_bytes_needed - data.len()) as u64; - let start = file_len - num_bytes_needed as u64; - let missing_bytes = scheduler - .submit_single(start..start + num_bytes_missing, 0) - .await?; - let mut combined = BytesMut::with_capacity(data.len() + num_bytes_missing as usize); - combined.extend(missing_bytes); - combined.extend(data); - Ok(combined.freeze()) - } - } - fn do_decode_gbo_table(gbo_bytes: &Bytes, footer: &Footer) -> Result> { let mut global_bufs_cursor = Cursor::new(gbo_bytes); @@ -1011,22 +990,13 @@ impl FileReader { Ok(()) } - async fn decode_gbo_table( - tail_bytes: &Bytes, + fn decode_gbo_table( + tail: &TailRead, file_len: u64, - scheduler: &FileScheduler, footer: &Footer, version: ConcreteFileVersion, ) -> Result> { - // This could, in theory, trigger another IOP but the GBO table should never be large - // enough for that to happen - let gbo_bytes = Self::optimistic_tail_read( - tail_bytes, - footer.global_buff_offsets_start, - scheduler, - file_len, - ) - .await?; + let gbo_bytes = tail.suffix_from(footer.global_buff_offsets_start)?; let gbo_table = Self::do_decode_gbo_table(&gbo_bytes, footer)?; Self::validate_gbo_table(&gbo_table, file_len, version)?; Ok(gbo_table) @@ -1048,19 +1018,14 @@ impl FileReader { scheduler: &FileScheduler, options: FullMetadataReadOptions, ) -> Result { - let InitialTailRead { - bytes: tail_bytes, - file_len, - offset: tail_offset, - retention_offset, - } = Self::read_tail( + let mut tail = Self::read_tail( scheduler, options .metadata_size_bytes .map_or(TailReadPlan::Baseline, TailReadPlan::ExactMetadata), ) .await?; - let footer = Self::decode_footer(&tail_bytes)?; + let footer = Self::decode_footer(&tail.bytes)?; let version = ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?; if version == ConcreteFileVersion::V1 { @@ -1070,8 +1035,9 @@ impl FileReader { }); } - let gbo_table = - Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, version).await?; + tail.extend_to(footer.global_buff_offsets_start, scheduler) + .await?; + let gbo_table = Self::decode_gbo_table(&tail, tail.file_len, &footer, version)?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), @@ -1079,24 +1045,24 @@ impl FileReader { } let schema_start = gbo_table[0].position; let schema_size = gbo_table[0].size; - let num_footer_bytes = file_len.checked_sub(schema_start).ok_or_else(|| { + let num_footer_bytes = tail.file_len.checked_sub(schema_start).ok_or_else(|| { Error::invalid_input_source( format!( "Schema position {} is outside file of size {}", - schema_start, file_len + schema_start, tail.file_len ) .into(), ) })?; - let all_metadata_bytes = - Self::optimistic_tail_read(&tail_bytes, schema_start, scheduler, file_len).await?; - let schema_bytes = all_metadata_bytes.slice(0..schema_size as usize); + tail.extend_to(schema_start, scheduler).await?; + let schema_end = schema_start + .checked_add(schema_size) + .ok_or_else(|| Error::invalid_input_source("Schema byte range overflows".into()))?; + let schema_bytes = tail.slice(schema_start..schema_end)?; let (num_rows, schema) = Self::decode_schema(schema_bytes)?; - let column_metadata_start = (footer.column_meta_start - schema_start) as usize; - let column_metadata_end = (footer.global_buff_offsets_start - schema_start) as usize; let column_metadata_bytes = - all_metadata_bytes.slice(column_metadata_start..column_metadata_end); + tail.slice(footer.column_meta_start..footer.global_buff_offsets_start)?; let column_metadatas = Self::read_all_column_metadata(column_metadata_bytes, &footer)?; let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::(); @@ -1110,10 +1076,10 @@ impl FileReader { // done. let retained_global_buffers = Self::retained_global_buffers_from_tail( &gbo_table, - &tail_bytes, - tail_offset, - retention_offset, - file_len, + &tail.bytes, + tail.offset, + tail.retention_offset, + tail.file_len, )?; Ok(RawFileMetadataOpen::Current { @@ -1128,7 +1094,7 @@ impl FileReader { num_global_buffer_bytes, num_footer_bytes, footer, - file_size_bytes: file_len, + file_size_bytes: tail.file_len, retained_global_buffers, }, }) @@ -1139,31 +1105,24 @@ impl FileReader { known_schema: Option<(Arc, u64)>, options: MetadataIndexReadOptions, ) -> Result { - let InitialTailRead { - bytes: tail_bytes, - file_len, - offset: tail_offset, - retention_offset, - } = Self::read_tail( + let mut tail = Self::read_tail( scheduler, options .estimated_num_columns .map_or(TailReadPlan::Baseline, TailReadPlan::EstimatedMetadataIndex), ) .await?; - let footer = Self::decode_footer(&tail_bytes)?; + let footer = Self::decode_footer(&tail.bytes)?; let file_version = Self::current_file_version(&footer)?; - let gbo_table = - Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version); - let cmo_table = Self::read_range_from_tail_or_scheduler( - &tail_bytes, - tail_offset, - scheduler, - footer.column_meta_offsets_start..footer.global_buff_offsets_start, - ); - let (gbo_table, cmo_table) = futures::try_join!(gbo_table, cmo_table)?; + let metadata_index_start = footer + .column_meta_offsets_start + .min(footer.global_buff_offsets_start); + tail.extend_to(metadata_index_start, scheduler).await?; + let gbo_table = Self::decode_gbo_table(&tail, tail.file_len, &footer, file_version)?; + let cmo_table = + tail.slice(footer.column_meta_offsets_start..footer.global_buff_offsets_start)?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), @@ -1174,14 +1133,9 @@ impl FileReader { Some((file_schema, num_rows)) => (file_schema, num_rows), None => { let schema_buffer = &gbo_table[0]; - let schema_range = schema_buffer.checked_range(0, file_len)?; - let schema_bytes = Self::read_range_from_tail_or_scheduler( - &tail_bytes, - tail_offset, - scheduler, - schema_range, - ) - .await?; + let schema_range = schema_buffer.checked_range(0, tail.file_len)?; + tail.extend_to(schema_range.start, scheduler).await?; + let schema_bytes = tail.slice(schema_range)?; let (num_rows, schema) = Self::decode_schema(schema_bytes)?; (Arc::new(schema), num_rows) } @@ -1189,10 +1143,10 @@ impl FileReader { let retained_global_buffers = Self::retained_global_buffers_from_tail( &gbo_table, - &tail_bytes, - tail_offset, - retention_offset, - file_len, + &tail.bytes, + tail.offset, + tail.retention_offset, + tail.file_len, )?; Ok(FileMetadataIndex { @@ -1202,7 +1156,7 @@ impl FileReader { column_metadata_offsets, num_columns: footer.num_columns, version: file_version, - file_size_bytes: file_len, + file_size_bytes: tail.file_len, retained_global_buffers, }) } @@ -3726,6 +3680,69 @@ mod tests { } } + #[tokio::test] + async fn full_metadata_fallback_extends_one_accumulated_suffix() { + let fs = FsFixture::default(); + let num_user_global_buffers = + fs.object_store.block_size() / super::COLUMN_METADATA_OFFSET_ENTRY_SIZE as usize + 8; + let write_result = write_file_with_global_buffers( + &fs, + (0..num_user_global_buffers).map(|_| Bytes::from_static(b"x")), + ) + .await; + let summary = write_result.summary(); + let metadata_size = write_result.metadata_size_bytes(); + assert!(metadata_size.get() > fs.object_store.block_size() as u64); + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(summary.size_bytes)) + .await + .unwrap(); + + fs.object_store.io_stats_incremental(); + let metadata = FileReader::read_all_metadata(&file_scheduler) + .await + .unwrap(); + let stats = fs.object_store.io_stats_incremental(); + + assert_eq!(stats.read_iops, 3); + assert_eq!(stats.read_bytes, metadata_size.get()); + assert_eq!(metadata.file_buffers.len(), num_user_global_buffers + 1); + } + + #[tokio::test] + async fn reader_reuses_decoded_metadata_without_io() { + let fs = FsFixture::default(); + create_some_file(&fs, ConcreteFileVersion::V2_1).await; + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let path = file_scheduler.reader().path().clone(); + let metadata = Arc::new( + FileReader::read_all_metadata(&file_scheduler) + .await + .unwrap(), + ); + + fs.object_store.io_stats_incremental(); + let reader = FileReader::try_open_with_file_metadata( + Arc::new(crate::io::LanceEncodingsIo::new(file_scheduler)), + path, + None, + Arc::::default(), + metadata.clone(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + assert!(Arc::ptr_eq(reader.metadata(), &metadata)); + assert_eq!(fs.object_store.io_stats_incremental().read_iops, 0); + } + #[tokio::test] async fn column_count_estimate_removes_cmo_dependency() { let fs = FsFixture::default(); @@ -4570,7 +4587,10 @@ mod tests { assert_eq!(batches.len(), 1); } - async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) -> FileWriteResult { + async fn write_file_with_global_buffers( + fs: &FsFixture, + buffers: impl IntoIterator, + ) -> FileWriteResult { let lance_schema = lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new( "foo", @@ -4586,8 +4606,10 @@ mod tests { ) .unwrap(); - let buf_index = file_writer.add_global_buffer(buffer).await.unwrap(); - assert_eq!(buf_index, 1); + for (buffer_index, buffer) in buffers.into_iter().enumerate() { + let actual_index = file_writer.add_global_buffer(buffer).await.unwrap(); + assert_eq!(actual_index as usize, buffer_index + 1); + } file_writer.finish_with_metadata_size().await.unwrap() } @@ -4635,7 +4657,7 @@ mod tests { #[case] expected_message: &str, ) { let fs = FsFixture::default(); - write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await; + write_file_with_global_buffers(&fs, [Bytes::from_static(b"hello")]).await; let mut file_bytes = fs .object_store @@ -4700,7 +4722,7 @@ mod tests { }; let expected_read_iops = if within_window { 0 } else { 1 }; - let write_result = write_file_with_global_buffer(&fs, buffer.clone()).await; + let write_result = write_file_with_global_buffers(&fs, [buffer.clone()]).await; let file_scheduler = fs .scheduler @@ -4839,7 +4861,7 @@ mod tests { async fn test_read_global_buffer_out_of_range() { let fs = FsFixture::default(); - write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await; + write_file_with_global_buffers(&fs, [Bytes::from_static(b"hello")]).await; let file_scheduler = fs .scheduler diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index 7fac43c5f4b..2f333ac46c5 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -31,6 +31,12 @@ use std::num::NonZeroU64; use std::pin::Pin; use std::{any::Any, sync::Arc}; +#[derive(Debug, Clone, Copy)] +struct IndexFileOpenHint { + size_bytes: u64, + file_metadata_size_bytes: Option, +} + /// An index store that serializes scalar indices using the lance format /// /// Scalar indices are made up of named collections of record batches. This @@ -42,13 +48,11 @@ pub struct LanceIndexStore { index_dir: Path, metadata_cache: Arc, scheduler: Arc, - /// Cached file sizes (filename -> size in bytes) - /// When set, used to avoid HEAD calls when opening files + /// Cached file information keyed by relative path. + /// When set, sizes avoid HEAD calls and metadata suffix sizes avoid dependent reads. // Partition priority views share this immutable map. Cloning all file names // for every partition would make request rebinding quadratic in partitions. - file_sizes: Arc>, - /// Cached metadata suffix sizes for Lance-format files. - file_metadata_sizes: Arc>, + files: Arc>, format_version: ConcreteFileVersion, /// Base I/O priority for all requests this store submits to `scheduler`. io_priority: u64, @@ -96,8 +100,7 @@ impl LanceIndexStore { index_dir, metadata_cache, scheduler, - file_sizes: Arc::default(), - file_metadata_sizes: Arc::default(), + files: Arc::default(), format_version, io_priority: 0, } @@ -107,17 +110,35 @@ impl LanceIndexStore { /// /// The map should contain relative paths (e.g., "index.idx") as keys /// and file sizes in bytes as values. - pub fn with_file_sizes(mut self, file_sizes: HashMap) -> Self { - self.file_sizes = Arc::new(file_sizes); - self + pub fn with_file_sizes(self, file_sizes: HashMap) -> Self { + self.with_index_files( + file_sizes + .into_iter() + .map(|(path, size_bytes)| IndexFile { + path, + size_bytes, + file_metadata_size_bytes: None, + }) + .collect(), + ) } - /// Set metadata suffix sizes used to avoid dependent footer discovery reads. - pub fn with_file_metadata_sizes( - mut self, - file_metadata_sizes: HashMap, - ) -> Self { - self.file_metadata_sizes = Arc::new(file_metadata_sizes); + /// Set cached information for files already recorded in an index manifest. + pub fn with_index_files(mut self, files: Vec) -> Self { + self.files = Arc::new( + files + .into_iter() + .map(|file| { + ( + file.path, + IndexFileOpenHint { + size_bytes: file.size_bytes, + file_metadata_size_bytes: file.file_metadata_size_bytes, + }, + ) + }) + .collect(), + ); self } @@ -471,20 +492,21 @@ impl IndexStore for LanceIndexStore { let path = self.index_file_path(name)?; // Use cached file size if available, otherwise unknown (requires HEAD call) let cached_size = self - .file_sizes + .files .get(name) - .map(|&size| CachedFileSize::new(size)) + .map(|file| CachedFileSize::new(file.size_bytes)) .unwrap_or_else(CachedFileSize::unknown); let file_scheduler = self .scheduler .open_file_with_priority(&path, self.io_priority, &cached_size) .await?; - let metadata_options = self.file_metadata_sizes.get(name).copied().map_or_else( - FullMetadataReadOptions::default, - |metadata_size_bytes| { + let metadata_options = self + .files + .get(name) + .and_then(|file| file.file_metadata_size_bytes) + .map_or_else(FullMetadataReadOptions::default, |metadata_size_bytes| { FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) - }, - ); + }); match versions::open_self_described_reader_with_metadata_options( file_scheduler, Arc::::default(), @@ -529,7 +551,10 @@ impl IndexStore for LanceIndexStore { Ok(IndexFile { path: new_name.to_string(), size_bytes: result.size as u64, - file_metadata_size_bytes: self.file_metadata_sizes.get(name).copied(), + file_metadata_size_bytes: self + .files + .get(name) + .and_then(|file| file.file_metadata_size_bytes), }) } _ => { @@ -559,7 +584,10 @@ impl IndexStore for LanceIndexStore { Ok(IndexFile { path: new_name.to_string(), size_bytes: result.size as u64, - file_metadata_size_bytes: self.file_metadata_sizes.get(name).copied(), + file_metadata_size_bytes: self + .files + .get(name) + .and_then(|file| file.file_metadata_size_bytes), }) } @@ -781,8 +809,7 @@ mod tests { index_dir, Arc::new(LanceCache::no_cache()), ) - .with_file_sizes(file_sizes) - .with_file_metadata_sizes(HashMap::from([(file_name.to_string(), metadata_size)])); + .with_index_files(vec![index_file]); object_store.io_stats_incremental(); with_hint.open_index_file(file_name).await.unwrap(); assert_eq!(object_store.io_stats_incremental().read_iops, 1); diff --git a/rust/lance-table/src/format/index.rs b/rust/lance-table/src/format/index.rs index 6b8a56bfe95..5e67aecd818 100644 --- a/rust/lance-table/src/format/index.rs +++ b/rust/lance-table/src/format/index.rs @@ -129,18 +129,6 @@ impl IndexMetadata { .unwrap_or_default() } - /// Returns known Lance metadata suffix sizes by relative file path. - pub fn file_metadata_size_map(&self) -> HashMap> { - self.files - .iter() - .flatten() - .filter_map(|file| { - file.file_metadata_size_bytes - .map(|size| (file.path.clone(), size)) - }) - .collect() - } - /// Returns the total size of all files in this index segment in bytes. /// Returns None if file information is not available. pub fn total_size_bytes(&self) -> Option { diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 07857f5f5c2..4515ba170b4 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1014,9 +1014,10 @@ impl Transaction { && file.file_major_version == new_file.file_major_version && file.file_minor_version == new_file.file_minor_version { - // assign the new file path / size / base to the fragment + // assign the new file path / sizes / base to the fragment file.path = new_file.path.clone(); file.file_size_bytes = new_file.file_size_bytes.clone(); + file.file_metadata_size_bytes = new_file.file_metadata_size_bytes; file.base_id = new_file.base_id; replaced_in_place = true; } diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index afc1b2cff0b..400e80bee52 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -1221,8 +1221,9 @@ impl FileFragment { .unwrap_or_default(); let estimated_num_columns = data_file.estimated_num_columns(); let prefer_indexed = metadata_mode == MetadataMode::LazyAllowed - && reader_projection.column_indices.len().saturating_mul(4) - < estimated_num_columns.map_or(0, |num_columns| num_columns.get() as usize); + && estimated_num_columns.is_some_and(|num_columns| { + reader_projection.prefers_indexed_metadata(num_columns.get() as usize) + }); let known_schema = self .metadata .physical_rows @@ -1254,8 +1255,8 @@ impl FileFragment { metadata_index_options, ) .await?; - if (reader_projection.column_indices.len() as u32).saturating_mul(4) - >= metadata_index.num_columns() + if !reader_projection + .prefers_indexed_metadata(metadata_index.num_columns() as usize) { return Ok(None); } diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index f54472b5902..189663d9dec 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -19,7 +19,7 @@ use lance_core::{Error, Result}; use lance_file::version::ConcreteFileVersion; use lance_index::is_system_index; use lance_index::pb::VectorIndexDetails; -use lance_index::scalar::lance_format::LanceIndexStore; +use lance_index::scalar::{lance_format::LanceIndexStore, table_files_to_index}; use lance_table::format::IndexMetadata; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -197,9 +197,13 @@ impl LanceIndexStoreExt for LanceIndexStore { let object_store = dataset.object_store_for_index(index).await?; let store = Self::with_format_version(object_store, index_dir, Arc::new(cache), format_version); - Ok(store - .with_file_sizes(index.file_size_map()) - .with_file_metadata_sizes(index.file_metadata_size_map())) + Ok(store.with_index_files( + index + .files + .clone() + .map(table_files_to_index) + .unwrap_or_default(), + )) } } diff --git a/rust/lance/src/dataset/tests/data_file_part.rs b/rust/lance/src/dataset/tests/data_file_part.rs index 593ba09f754..0c3cdf8653d 100644 --- a/rust/lance/src/dataset/tests/data_file_part.rs +++ b/rust/lance/src/dataset/tests/data_file_part.rs @@ -141,9 +141,16 @@ async fn concatenates_parts_in_caller_order_without_reusing_staging_files() { .concat_data_file_parts(&target, &ordered_parts) .await .unwrap(); + let metadata_size_bytes = data_file.file_metadata_size_bytes; let replacement = DataReplacementGroup(only_fragment(&dataset).id() as u64, data_file); assert_eq!(replacement.1.path, target.file_name.as_str()); let dataset = commit(&dataset, replacement).await.unwrap(); + let committed_file = dataset.manifest.fragments[0] + .files + .iter() + .find(|file| file.path == target.file_name) + .unwrap(); + assert_eq!(committed_file.file_metadata_size_bytes, metadata_size_bytes); let batch = dataset.scan().try_into_batch().await.unwrap(); assert_eq!( batch["id"].as_primitive::().values(), diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 6f81ee6f670..0db37ae94a2 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -21,9 +21,9 @@ use lance_core::utils::parse::parse_env_as_bool; use lance_core::utils::tracing::{ IO_TYPE_OPEN_FRAG_REUSE, IO_TYPE_OPEN_MEM_WAL, IO_TYPE_OPEN_VECTOR, TRACE_IO_EVENTS, }; -use lance_file::reader::FileReaderOptions; +use lance_file::reader::{FileReader, FileReaderOptions}; use lance_file::versions::v1::reader::FileReader as V1FileReader; -use lance_index::INDEX_METADATA_SCHEMA_KEY; +use lance_file::versions::{self as file_versions, OpenedFileReader}; pub use lance_index::IndexParams; use lance_index::frag_reuse::{ CompactFragReuseIndex, CompactFragReuseIndexHandle, FRAG_REUSE_INDEX_NAME, @@ -48,6 +48,7 @@ use lance_index::{ FtsPrewarmDiagnostics, FtsPrewarmOptions, FtsPrewarmResult, FtsPrewarmSegmentStatus, INDEX_FILE_NAME, Index, IndexType, PrewarmOptions, pb, vector::VectorIndex, }; +use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_METADATA_SCHEMA_KEY}; use lance_index::{ IndexCriteria, is_system_index, metrics::{MetricsCollector, NoOpMetricsCollector}, @@ -58,7 +59,6 @@ use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::traits::Reader; use lance_io::utils::{ CachedFileSize, read_last_block, read_message, read_message_from_buf, read_metadata_offset, - read_version, }; use lance_table::format::{DataFile, Fragment, SelfDescribingFileReader}; use lance_table::format::{IndexFile, IndexMetadata, list_index_files_with_sizes}; @@ -73,7 +73,7 @@ use vector::details::{ vector_details_as_json, }; pub(crate) use vector::details::{vector_index_details, vector_index_details_default}; -use vector::ivf::v2::{IVFIndex, IvfFileOpenHints, IvfStateEntryBox}; +use vector::ivf::v2::{IVFIndex, IvfFileOpenContext, IvfFileOpenHints, IvfStateEntryBox}; use vector::utils::get_vector_type; mod api; @@ -3056,23 +3056,54 @@ impl DatasetIndexInternalExt for Dataset { let frag_reuse_index = self.open_frag_reuse_index(metrics).await?; let index_dir = self.indice_files_dir(&index_meta)?; - let index_file = index_dir - .clone() - .join(uuid.to_string()) - .join(INDEX_FILE_NAME); - let file_sizes = index_meta.file_size_map(); - let file_metadata_sizes = index_meta.file_metadata_size_map(); - let reader: Arc = vector::open_index_file( - object_store.as_ref(), - &index_file, - INDEX_FILE_NAME, - &file_sizes, - ) - .await? - .into(); - - let tailing_bytes = read_last_block(reader.as_ref()).await?; - let (major_version, minor_version) = read_version(&tailing_bytes)?; + let segment_dir = index_dir.clone().join(uuid.to_string()); + let index_file = segment_dir.clone().join(INDEX_FILE_NAME); + let auxiliary_file = segment_dir.join(INDEX_AUXILIARY_FILE_NAME); + let file_open_hints = IvfFileOpenHints::new(index_meta.files.clone().unwrap_or_default()); + let scheduler = ScanScheduler::new( + object_store.clone(), + SchedulerConfig::max_bandwidth(&object_store), + ); + let cached_size = CachedFileSize::new(file_open_hints.file_size(INDEX_FILE_NAME)); + let file = scheduler.open_file(&index_file, &cached_size).await?; + let object_reader = file.reader().clone(); + let index_file_cache = self.metadata_cache.file_metadata_cache(&index_file); + let open_index = file_versions::open_self_described_reader_with_metadata_options( + file, + Default::default(), + &index_file_cache, + FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_FILE_NAME), + ); + // Current writers record an exact auxiliary suffix. Requiring that hint + // avoids speculative legacy I/O while putting both current-file opens in + // the same dependency stage. + let prefetch_auxiliary = async { + if file_open_hints + .metadata_size_bytes(INDEX_AUXILIARY_FILE_NAME) + .is_none() + { + return Ok(None); + } + let cached_size = + CachedFileSize::new(file_open_hints.file_size(INDEX_AUXILIARY_FILE_NAME)); + let file = scheduler.open_file(&auxiliary_file, &cached_size).await?; + FileReader::try_open_with_metadata_options( + file, + None, + Default::default(), + &self.metadata_cache.file_metadata_cache(&auxiliary_file), + FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_AUXILIARY_FILE_NAME), + ) + .await + .map(Some) + }; + // Defer propagating the speculative result until root dispatch proves + // this is a current-format index, so advisory data cannot break legacy. + let (opened_reader, prefetched_storage_reader) = + tokio::join!(open_index, prefetch_auxiliary); + let opened_reader = opened_reader?; // Namespace the index cache by the UUID of the index. v2+ partition // entries are store-free and remain reusable across object-store @@ -3089,67 +3120,57 @@ impl DatasetIndexInternalExt for Dataset { // the index file is in lance format since version (0,2) // TODO: we need to change the legacy IVF_PQ to be in lance format - let result: Result<(Arc, Option)> = match ( - major_version, - minor_version, - ) { - (0, 1) | (0, 0) => { - info!(target: TRACE_IO_EVENTS, index_uuid=%uuid, r#type=IO_TYPE_OPEN_VECTOR, version="0.1", index_type="IVF_PQ"); - let proto = open_index_proto(reader.as_ref()).await?; - match &proto.implementation { - Some(Implementation::VectorIndex(vector_index)) => { - let dataset = Arc::new(self.clone()); - let idx = vector::open_vector_index( - dataset, - uuid, - vector_index, - reader, - frag_reuse_index, - ) - .await?; - Ok((idx, None::)) + let result: Result<(Arc, Option)> = match opened_reader { + OpenedFileReader::V1 { + major_version, + minor_version, + } => match (major_version, minor_version) { + (0, 1) | (0, 0) => { + info!(target: TRACE_IO_EVENTS, index_uuid=%uuid, r#type=IO_TYPE_OPEN_VECTOR, version="0.1", index_type="IVF_PQ"); + let proto = open_index_proto(object_reader.as_ref()).await?; + match &proto.implementation { + Some(Implementation::VectorIndex(vector_index)) => { + let dataset = Arc::new(self.clone()); + let idx = vector::open_vector_index( + dataset, + uuid, + vector_index, + object_reader, + frag_reuse_index, + ) + .await?; + Ok((idx, None::)) + } + None => Err(Error::internal( + "Index proto was missing implementation field", + )), } - None => Err(Error::internal( - "Index proto was missing implementation field", - )), } - } - (0, 2) => { - info!(target: TRACE_IO_EVENTS, index_uuid=%uuid, r#type=IO_TYPE_OPEN_VECTOR, version="0.2", index_type="IVF_PQ"); - let reader = V1FileReader::try_new_self_described_from_reader( - reader.clone(), - Some(&self.metadata_cache.file_metadata_cache(&index_file)), - ) - .await?; - let idx = vector::open_vector_index_v2( - Arc::new(self.clone()), - column, - uuid, - reader, - frag_reuse_index, - ) - .await?; - Ok((idx, None::)) - } + (0, 2) => { + info!(target: TRACE_IO_EVENTS, index_uuid=%uuid, r#type=IO_TYPE_OPEN_VECTOR, version="0.2", index_type="IVF_PQ"); + let reader = V1FileReader::try_new_self_described_from_reader( + object_reader, + Some(&self.metadata_cache.file_metadata_cache(&index_file)), + ) + .await?; + let idx = vector::open_vector_index_v2( + Arc::new(self.clone()), + column, + uuid, + reader, + frag_reuse_index, + ) + .await?; + Ok((idx, None::)) + } - (0, 3) | (2, _) => { - let file_open_hints = IvfFileOpenHints::new(file_sizes, file_metadata_sizes); - let scheduler = ScanScheduler::new( - object_store.clone(), - SchedulerConfig::max_bandwidth(&object_store), - ); - let cached_size = CachedFileSize::new(file_open_hints.file_size(INDEX_FILE_NAME)); - let file = scheduler.open_file(&index_file, &cached_size).await?; - let reader = lance_file::reader::FileReader::try_open_with_metadata_options( - file, - None, - Default::default(), - &self.metadata_cache.file_metadata_cache(&index_file), - FileReaderOptions::default(), - file_open_hints.metadata_options(INDEX_FILE_NAME), - ) - .await?; + _ => Err(Error::index( + "unsupported index version (maybe need to upgrade your lance version)" + .to_owned(), + )), + }, + OpenedFileReader::Current(reader) => { let index_metadata = reader .schema() .metadata @@ -3162,6 +3183,12 @@ impl DatasetIndexInternalExt for Dataset { let (field_path, field) = resolve_index_column(self.schema(), &index_meta, column)?; let (_, element_type) = get_vector_type(self.schema(), &field_path)?; + let file_open_context = IvfFileOpenContext::new( + scheduler, + reader, + prefetched_storage_reader?, + file_open_hints, + ); info!(target: TRACE_IO_EVENTS, index_uuid=%uuid, r#type=IO_TYPE_OPEN_VECTOR, version="0.3", index_type=index_metadata.index_type); @@ -3175,7 +3202,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3188,7 +3215,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3207,7 +3234,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3221,7 +3248,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3235,7 +3262,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3250,7 +3277,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3263,7 +3290,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3278,7 +3305,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3292,7 +3319,7 @@ impl DatasetIndexInternalExt for Dataset { frag_reuse_index, self.metadata_cache.as_ref(), index_cache, - file_open_hints, + file_open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -3304,10 +3331,6 @@ impl DatasetIndexInternalExt for Dataset { ))), } } - - _ => Err(Error::index( - "unsupported index version (maybe need to upgrade your lance version)".to_owned(), - )), }; let (index, ivf_entry) = result?; metrics.record_index_load(); diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 314703da6c8..85a59dd48a3 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -2256,6 +2256,10 @@ mod tests { /// `open_index_file` looks up size hints by bare file name. This pins that a /// freshly created HNSW index records both the main and auxiliary files under /// those exact names with nonzero total and metadata suffix sizes. + /// + /// A cold open must also reuse the main reader selected during format + /// dispatch instead of reading the same root metadata range again. + /// Its concurrently prefetched auxiliary reader must also be reused. #[tokio::test] async fn test_hnsw_index_records_file_sizes() { use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME}; @@ -2301,27 +2305,56 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); let index = indices.iter().find(|idx| idx.name == "hnsw").unwrap(); - let file_sizes = index.file_size_map(); - let metadata_sizes = index.file_metadata_size_map(); - - assert!( - file_sizes.get(INDEX_FILE_NAME).copied().unwrap_or(0) > 0, - "manifest should record a nonzero {INDEX_FILE_NAME} size, got {file_sizes:?}" - ); - assert!( - file_sizes - .get(INDEX_AUXILIARY_FILE_NAME) - .copied() - .unwrap_or(0) - > 0, - "manifest should record a nonzero {INDEX_AUXILIARY_FILE_NAME} size, got {file_sizes:?}" - ); + let files = index.files.as_ref().expect("manifest should record files"); for file_name in [INDEX_FILE_NAME, INDEX_AUXILIARY_FILE_NAME] { + let file = files + .iter() + .find(|file| file.path == file_name) + .unwrap_or_else(|| panic!("manifest should record {file_name}, got {files:?}")); + assert!(file.size_bytes > 0); assert!( - metadata_sizes - .get(file_name) - .is_some_and(|metadata_size| metadata_size.get() <= file_sizes[file_name]), - "manifest should record a valid metadata suffix for {file_name}, got {metadata_sizes:?}" + file.file_metadata_size_bytes + .is_some_and(|metadata_size| metadata_size.get() <= file.size_bytes), + "manifest should record a valid metadata suffix for {file_name}, got {files:?}" + ); + } + let index_uuid = index.uuid; + drop(indices); + drop(dataset); + + let dataset = Dataset::open(&uri).await.unwrap(); + let segment_path = dataset.indices_dir().join(index_uuid.to_string()); + let index_paths = [ + segment_path.clone().join(INDEX_FILE_NAME), + segment_path.join(INDEX_AUXILIARY_FILE_NAME), + ]; + dataset.object_store.io_stats_incremental(); + let _index = dataset + .open_vector_index("vector", &index_uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let stats = dataset.object_store.io_stats_incremental(); + for index_path in index_paths { + let requests = stats + .requests + .iter() + .filter(|request| request.path == index_path) + .map(|request| { + ( + request.method, + request.range.as_ref().map(|range| (range.start, range.end)), + ) + }) + .collect::>(); + let unique_requests = requests + .iter() + .copied() + .collect::>(); + assert!(!requests.is_empty(), "no reads for {index_path}: {stats:?}"); + assert_eq!( + requests.len(), + unique_requests.len(), + "{index_path} requests were issued more than once: {stats:?}" ); } } diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 25662e63642..18fba7c4172 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -8,7 +8,7 @@ use std::marker::PhantomData; use std::{ any::Any, borrow::Cow, - collections::{BinaryHeap, HashMap}, + collections::BinaryHeap, num::NonZeroU64, ops::Range, sync::{ @@ -86,6 +86,7 @@ use lance_io::{ }; use lance_linalg::distance::DistanceType; use lance_select::RowAddrTreeMap; +use lance_table::format::IndexFile as TableIndexFile; use object_store::path::Path; use prost::Message; use roaring::RoaringBitmap; @@ -798,27 +799,30 @@ impl CacheKey for FileMetadataCacheKey { #[derive(Debug, Default)] pub(crate) struct IvfFileOpenHints { - file_sizes: HashMap, - metadata_sizes: HashMap, + files: Vec, } impl IvfFileOpenHints { - pub(crate) fn new( - file_sizes: HashMap, - metadata_sizes: HashMap, - ) -> Self { - Self { - file_sizes, - metadata_sizes, - } + pub(crate) fn new(files: Vec) -> Self { + Self { files } } pub(crate) fn file_size(&self, name: &str) -> u64 { - self.file_sizes.get(name).copied().unwrap_or(0) + self.files + .iter() + .find(|file| file.path == name) + .map_or(0, |file| file.size_bytes) + } + + pub(crate) fn metadata_size_bytes(&self, name: &str) -> Option { + self.files + .iter() + .find(|file| file.path == name) + .and_then(|file| file.file_metadata_size_bytes) } pub(crate) fn metadata_options(&self, name: &str) -> FullMetadataReadOptions { - self.metadata_sizes.get(name).copied().map_or_else( + self.metadata_size_bytes(name).map_or_else( FullMetadataReadOptions::default, |metadata_size_bytes| { FullMetadataReadOptions::default().with_metadata_size_bytes(metadata_size_bytes) @@ -827,6 +831,30 @@ impl IvfFileOpenHints { } } +#[derive(Debug)] +pub(crate) struct IvfFileOpenContext { + scheduler: Arc, + index_reader: FileReader, + storage_reader: Option, + hints: IvfFileOpenHints, +} + +impl IvfFileOpenContext { + pub(crate) fn new( + scheduler: Arc, + index_reader: FileReader, + storage_reader: Option, + hints: IvfFileOpenHints, + ) -> Self { + Self { + scheduler, + index_reader, + storage_reader, + hints, + } + } +} + /// Open a FileReader, reusing cached file metadata if available. async fn open_reader_cached( scheduler: &Arc, @@ -1414,11 +1442,15 @@ impl IVFIndex { frag_reuse_index: Option>, file_metadata_cache: &LanceCache, index_cache: LanceCache, - file_open_hints: IvfFileOpenHints, + file_open_context: IvfFileOpenContext, ) -> Result { let io_parallelism = object_store.io_parallelism(); - let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); - let scheduler = ScanScheduler::new(object_store, scheduler_config); + let IvfFileOpenContext { + scheduler, + index_reader, + storage_reader, + hints: file_open_hints, + } = file_open_context; let uuid_str = uuid.to_string(); let uri = index_dir @@ -1429,31 +1461,22 @@ impl IVFIndex { .clone() .join(uuid_str.as_str()) .join(INDEX_AUXILIARY_FILE_NAME); - let index_size = CachedFileSize::new(file_open_hints.file_size(INDEX_FILE_NAME)); - let aux_size = CachedFileSize::new(file_open_hints.file_size(INDEX_AUXILIARY_FILE_NAME)); - let open_index = async { - FileReader::try_open_with_metadata_options( - scheduler.open_file(&uri, &index_size).await?, - None, - Arc::::default(), - file_metadata_cache, - FileReaderOptions::default(), - file_open_hints.metadata_options(INDEX_FILE_NAME), - ) - .await - }; - let open_aux = async { - FileReader::try_open_with_metadata_options( - scheduler.open_file(&aux_path, &aux_size).await?, - None, - Arc::::default(), - file_metadata_cache, - FileReaderOptions::default(), - file_open_hints.metadata_options(INDEX_AUXILIARY_FILE_NAME), - ) - .await + let storage_reader = match storage_reader { + Some(storage_reader) => storage_reader, + None => { + let aux_size = + CachedFileSize::new(file_open_hints.file_size(INDEX_AUXILIARY_FILE_NAME)); + FileReader::try_open_with_metadata_options( + scheduler.open_file(&aux_path, &aux_size).await?, + None, + Arc::::default(), + file_metadata_cache, + FileReaderOptions::default(), + file_open_hints.metadata_options(INDEX_AUXILIARY_FILE_NAME), + ) + .await? + } }; - let (index_reader, storage_reader) = tokio::try_join!(open_index, open_aux)?; let index_metadata: IndexMetadata = serde_json::from_str( index_reader .schema() @@ -1503,10 +1526,10 @@ impl IVFIndex { let use_residual_scratch = Self::use_residual_scratch(&ivf, use_query_residual); let rq_search_cache = Self::build_rq_search_cache(&ivf, &storage)?; - // The scheduler is freshly created above and, at this point, has served - // only the open-time reads (file footers, IVF centroids, quantization - // metadata) -- partition reads happen later, during queries. So its - // cumulative stats are exactly the one-time index-open I/O. + // At this point, the scheduler has served only the open-time reads + // (file footers, IVF centroids, quantization metadata) -- partition reads + // happen later, during queries. So its cumulative stats are exactly the + // one-time index-open I/O. let open_io_stats = scheduler.stats(); let read_projection = Self::read_projection(&index_reader)?; From f0c8bb1a191367a79238a9003e7045ed90041fcd Mon Sep 17 00:00:00 2001 From: Ali Arslan Date: Fri, 11 Sep 2026 18:03:48 -0700 Subject: [PATCH 6/6] fix(file): bound advisory metadata prefetch Co-authored-by: Cursor --- rust/lance-file/src/reader.rs | 69 +++++++++++++++++++++++++++++++++- rust/lance-io/src/scheduler.rs | 9 +++++ 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 63c5f1219ea..b715957ca05 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -785,7 +785,7 @@ impl FileReader { async fn read_tail(scheduler: &FileScheduler, plan: TailReadPlan) -> Result { let file_size = scheduler.reader().size().await? as u64; let baseline_read_size = file_size.min(scheduler.reader().block_size() as u64); - let requested_read_size = match plan { + let planned_read_size = match plan { TailReadPlan::Baseline => baseline_read_size, TailReadPlan::ExactMetadata(metadata_size_bytes) => { let metadata_size_bytes = metadata_size_bytes.get(); @@ -805,6 +805,17 @@ impl FileReader { .min(file_size) } }; + let io_buffer_size_bytes = scheduler.io_buffer_size_bytes(); + // Footer validation happens after this read, so keep advisory overread within + // the scheduler's configured unread-I/O budget. With no finite byte budget, + // use the authoritative footer-directed path. + let requested_read_size = if planned_read_size > baseline_read_size + && (io_buffer_size_bytes == 0 || planned_read_size > io_buffer_size_bytes) + { + baseline_read_size + } else { + planned_read_size + }; let offset = file_size - requested_read_size; // FileScheduler may split this logical range into concurrent physical requests. let bytes = scheduler.submit_single(offset..file_size, 0).await?; @@ -2582,7 +2593,11 @@ mod tests { encoder::{EncodedBatch, EncodingOptions, encode_batch}, format::pb21, }; - use lance_io::{stream::RecordBatchStream, utils::CachedFileSize}; + use lance_io::{ + scheduler::{ScanScheduler, SchedulerConfig}, + stream::RecordBatchStream, + utils::CachedFileSize, + }; use log::debug; use rstest::rstest; use tokio::sync::mpsc; @@ -3678,6 +3693,30 @@ mod tests { "unexpected read count for stale metadata hint {stale_hint}" ); } + + assert!(file_size > metadata_size.get()); + let bounded_scheduler = ScanScheduler::new( + fs.object_store.clone(), + SchedulerConfig::new(metadata_size.get()), + ); + let bounded_file_scheduler = bounded_scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(file_size)) + .await + .unwrap(); + for (hint, expected_read_iops) in [(metadata_size.get(), 1), (file_size, 2)] { + fs.object_store.io_stats_incremental(); + let metadata = FileReader::read_all_metadata_with_options( + &bounded_file_scheduler, + FullMetadataReadOptions::default() + .with_metadata_size_bytes(NonZeroU64::new(hint).unwrap()), + ) + .await + .unwrap(); + let stats = fs.object_store.io_stats_incremental(); + assert_eq!(metadata.metadata_size_bytes(), metadata_size.get()); + assert_eq!(stats.read_iops, expected_read_iops); + assert_eq!(stats.read_bytes, metadata_size.get()); + } } #[tokio::test] @@ -3812,6 +3851,32 @@ mod tests { estimated_num_columns ); } + + let bounded_scheduler = ScanScheduler::new( + fs.object_store.clone(), + SchedulerConfig::new(fs.object_store.block_size() as u64), + ); + let bounded_file_scheduler = bounded_scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(file_size)) + .await + .unwrap(); + fs.object_store.io_stats_incremental(); + let metadata = FileReader::read_metadata_index_with_schema_and_options( + &bounded_file_scheduler, + written.schema.clone(), + 1, + MetadataIndexReadOptions::default() + .with_estimated_num_columns(NonZeroU32::new(u32::MAX).unwrap()), + ) + .await + .unwrap(); + let stats = fs.object_store.io_stats_incremental(); + assert_eq!( + metadata.column_metadata_offsets, + without_estimate.column_metadata_offsets + ); + assert_eq!(stats.read_iops, 2); + assert_eq!(stats.read_bytes, without_estimate_stats.read_bytes); } #[tokio::test] diff --git a/rust/lance-io/src/scheduler.rs b/rust/lance-io/src/scheduler.rs index 1876fe3feba..3d5717e01a6 100644 --- a/rust/lance-io/src/scheduler.rs +++ b/rust/lance-io/src/scheduler.rs @@ -803,6 +803,7 @@ pub struct ScanScheduler { object_store: Arc, io_queue: IoQueueType, stats: IoStats, + io_buffer_size_bytes: u64, } impl Debug for ScanScheduler { @@ -916,6 +917,7 @@ impl ScanScheduler { object_store, io_queue, stats, + io_buffer_size_bytes: config.io_buffer_size_bytes, }) } @@ -1366,6 +1368,13 @@ impl FileScheduler { .map_ok(|vec_bytes| vec_bytes.into_iter().next().unwrap()) } + /// Returns the configured byte budget for unread scheduled I/O. + /// + /// A value of zero means byte-based backpressure is disabled. + pub fn io_buffer_size_bytes(&self) -> u64 { + self.root.io_buffer_size_bytes + } + /// Provides access to the underlying reader /// /// Do not use this for reading data as it will bypass any I/O scheduling!