From daced255d89cdb48debcaa7f0c5a8e84d25d21a3 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 06:25:36 +0800 Subject: [PATCH] refactor: compose exact current-format readers --- java/lance-jni/src/file_reader.rs | 19 +- python/src/file.rs | 3 +- rust/lance-file/src/compatibility_tests.rs | 370 --- rust/lance-file/src/lib.rs | 13 +- rust/lance-file/src/reader.rs | 2068 +++++++---------- rust/lance-file/src/reader/structural.rs | 445 ++++ rust/lance-file/src/version.rs | 2 +- rust/lance-file/src/versions/mod.rs | 283 ++- rust/lance-file/src/versions/v2_0/mod.rs | 127 +- rust/lance-file/src/versions/v2_0/reader.rs | 415 ++++ rust/lance-file/src/versions/v2_1/mod.rs | 40 +- rust/lance-file/src/versions/v2_1/reader.rs | 304 +++ rust/lance-file/src/versions/v2_2/mod.rs | 37 +- rust/lance-file/src/versions/v2_2/reader.rs | 308 +++ rust/lance-file/src/versions/v2_3/mod.rs | 44 +- rust/lance-file/src/versions/v2_3/reader.rs | 341 +++ rust/lance-file/src/writer.rs | 4 + rust/lance-file/src/writer_tests.rs | 1887 +++++++++++++++ .../test_data/exact_versions/README.md | 14 +- rust/lance-index/src/scalar/lance_format.rs | 23 +- .../src/vector/distributed/index_merger.rs | 2 +- rust/lance/src/dataset/fragment.rs | 193 +- rust/lance/src/dataset/tests/dataset_index.rs | 6 +- rust/lance/src/index/vector/ivf.rs | 2 +- rust/lance/src/index/vector/ivf/v2.rs | 2 +- 25 files changed, 5152 insertions(+), 1800 deletions(-) delete mode 100644 rust/lance-file/src/compatibility_tests.rs create mode 100644 rust/lance-file/src/reader/structural.rs create mode 100644 rust/lance-file/src/versions/v2_0/reader.rs create mode 100644 rust/lance-file/src/versions/v2_1/reader.rs create mode 100644 rust/lance-file/src/versions/v2_2/reader.rs create mode 100644 rust/lance-file/src/versions/v2_3/reader.rs create mode 100644 rust/lance-file/src/writer_tests.rs diff --git a/java/lance-jni/src/file_reader.rs b/java/lance-jni/src/file_reader.rs index 3df9766d066..c7ba27d9d11 100644 --- a/java/lance-jni/src/file_reader.rs +++ b/java/lance-jni/src/file_reader.rs @@ -23,8 +23,8 @@ use lance::io::ObjectStore; use lance_core::cache::LanceCache; use lance_core::datatypes::{BlobHandling, OnMissing, Projection, Schema}; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; -use lance_encoding::version::LanceFileVersion; use lance_file::reader::{FileReader, FileReaderOptions, ReaderProjection}; +use lance_file::versions as file_versions; use lance_io::object_store::{ObjectStoreParams, ObjectStoreRegistry}; use lance_io::{ ReadBatchParams, @@ -262,18 +262,17 @@ pub extern "system" fn Java_org_lance_file_LanceFileReader_readAllNative( let transformed_schema = projection.to_bare_schema(); - let field_id_to_column_index = base_schema - .fields_pre_order() - .filter(|field| { - file_version < LanceFileVersion::V2_1 - || field.is_leaf() - || field.is_packed_struct() + let (field_ids, column_indices) = + file_versions::data_file_columns(file_version, &base_schema); + let field_id_to_column_index = field_ids + .into_iter() + .zip(column_indices) + .filter_map(|(field_id, column_index)| { + (column_index >= 0).then_some((field_id as u32, column_index as u32)) }) - .enumerate() - .map(|(idx, field)| (field.id as u32, idx as u32)) .collect::>(); - Some(ReaderProjection::from_field_ids( + Some(file_versions::reader_projection_from_field_ids( file_version, &transformed_schema, &field_id_to_column_index, diff --git a/python/src/file.rs b/python/src/file.rs index 1b879044df4..b1110b16f52 100644 --- a/python/src/file.rs +++ b/python/src/file.rs @@ -26,7 +26,6 @@ use lance_core::utils::path::LancePathExt; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; use lance_file::reader::{ BufferDescriptor, CachedFileMetadata, FileReader, FileReaderOptions, FileStatistics, - ReaderProjection, }; use lance_file::writer::{FileWriter, FileWriterOptions}; use lance_file::{LanceEncodingsIo, version::LanceFileVersion, versions as file_versions}; @@ -820,7 +819,7 @@ impl LanceFileReader { let mut base_projection = None; if let Some(columns) = columns { base_projection = Some( - ReaderProjection::from_column_names( + file_versions::reader_projection_from_column_names( file_metadata.version(), &file_metadata.file_schema, &columns.iter().map(|s| s.as_str()).collect::>(), diff --git a/rust/lance-file/src/compatibility_tests.rs b/rust/lance-file/src/compatibility_tests.rs deleted file mode 100644 index 4fdff0bcb65..00000000000 --- a/rust/lance-file/src/compatibility_tests.rs +++ /dev/null @@ -1,370 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -use std::collections::HashMap; -use std::sync::Arc; - -use arrow_array::builder::StringDictionaryBuilder; -use arrow_array::cast::AsArray; -use arrow_array::types::{Int8Type, Int32Type}; -use arrow_array::{ - Array, ArrayRef, Int32Array, LargeBinaryArray, ListArray, RecordBatch, StringArray, -}; -use arrow_schema::{DataType, Field, Schema as ArrowSchema}; -use futures::TryStreamExt; -use lance_core::cache::LanceCache; -use lance_core::datatypes::Schema as LanceSchema; -use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; -use lance_io::ReadBatchParams; -use lance_io::traits::Writer; -use lance_io::utils::CachedFileSize; -use rstest::rstest; -use tokio::io::AsyncWriteExt; - -use crate::reader::{FileReader, FileReaderOptions}; -use crate::testing::FsFixture; -use crate::version::ConcreteFileVersion; -use crate::versions; -use crate::versions::v1::reader::FileReader as V1Reader; -use crate::versions::v1::writer::{ - FileWriter as V1Writer, FileWriterOptions as V1WriterOptions, NotSelfDescribing, -}; -use crate::writer::FileWriterOptions; - -fn compatibility_fixture_batch() -> RecordBatch { - let row_count = 4097; - let ids = Arc::new(Int32Array::from_iter_values(0..row_count)) as ArrayRef; - let names = Arc::new(StringArray::from_iter((0..row_count).map(|index| { - (index % 7 != 0).then(|| format!("value-{index:04}-deterministic-fixture")) - }))) as ArrayRef; - let items = Arc::new(ListArray::from_iter_primitive::( - (0..row_count).map(|index| { - (index % 11 != 0).then(|| { - vec![ - Some(index), - (index % 5 != 0).then_some(index * 2), - Some(index * 3), - ] - }) - }), - )) as ArrayRef; - let mut categories = StringDictionaryBuilder::::new(); - for index in 0..row_count { - if index % 13 == 0 { - categories.append_null(); - } else { - categories - .append(match index % 3 { - 0 => "red", - 1 => "green", - _ => "blue", - }) - .unwrap(); - } - } - let categories = Arc::new(categories.finish()) as ArrayRef; - let blobs = Arc::new(LargeBinaryArray::from_iter_values( - (0..row_count).map(|index| format!("blob-{index:04}-deterministic-payload").into_bytes()), - )) as ArrayRef; - - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("name", DataType::Utf8, true).with_metadata(HashMap::from([( - "lance-encoding:compression".to_string(), - "none".to_string(), - )])), - Field::new( - "items", - DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), - true, - ), - Field::new( - "category", - DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), - true, - ) - .with_metadata(HashMap::from([( - "lance-encoding:dict-values-compression".to_string(), - "none".to_string(), - )])), - Field::new("blob", DataType::LargeBinary, true).with_metadata(HashMap::from([( - "lance-encoding:blob".to_string(), - "true".to_string(), - )])), - ])); - RecordBatch::try_new(schema, vec![ids, names, items, categories, blobs]).unwrap() -} - -fn stable_fixture(version: ConcreteFileVersion) -> &'static [u8] { - match version { - ConcreteFileVersion::V1 => include_bytes!("../test_data/exact_versions/v1.lance"), - ConcreteFileVersion::V2_0 => { - include_bytes!("../test_data/exact_versions/v2_0.lance") - } - ConcreteFileVersion::V2_1 => { - include_bytes!("../test_data/exact_versions/v2_1.lance") - } - ConcreteFileVersion::V2_2 => { - include_bytes!("../test_data/exact_versions/v2_2.lance") - } - ConcreteFileVersion::V2_3 => { - unreachable!("v2.3 is unstable and has no compatibility fixture") - } - } -} - -fn assert_blob_column_eq(actual: &dyn Array, expected: &dyn Array) { - let actual = actual.as_binary::(); - let expected = expected.as_binary::(); - assert_eq!(actual.len(), expected.len()); - for index in 0..actual.len() { - assert_eq!( - actual.is_null(index), - expected.is_null(index), - "blob validity differs at row {index}" - ); - if actual.is_valid(index) { - assert_eq!( - actual.value(index), - expected.value(index), - "blob payload differs at row {index}" - ); - } - } -} - -fn assert_wire_bytes_equal(actual: &[u8], expected: &[u8]) { - if let Some(offset) = actual - .iter() - .zip(expected) - .position(|(actual, expected)| actual != expected) - { - panic!( - "wire fixture first differs at byte {offset}: actual={}, expected={}", - actual[offset], expected[offset] - ); - } - assert_eq!( - actual.len(), - expected.len(), - "wire fixture length changed after a common {}-byte prefix", - actual.len().min(expected.len()) - ); -} - -async fn write_current_fixture( - version: ConcreteFileVersion, - batch: &RecordBatch, - schema: &LanceSchema, -) -> Vec { - let fs = FsFixture::default(); - let object_writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); - let options = FileWriterOptions { - data_cache_bytes: Some(1), - max_page_bytes: Some(1024), - ..Default::default() - }; - let summary = match version { - ConcreteFileVersion::V1 => unreachable!("v1 uses its manifest-backed writer"), - ConcreteFileVersion::V2_0 => { - let mut writer = - versions::v2_0::create_writer(object_writer, schema.clone(), options).unwrap(); - for offset in (0..batch.num_rows()).step_by(1024) { - let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); - writer.write_batch(&slice).await.unwrap(); - } - writer.finish().await.unwrap() - } - ConcreteFileVersion::V2_1 => { - let mut writer = - versions::v2_1::create_writer(object_writer, schema.clone(), options).unwrap(); - for offset in (0..batch.num_rows()).step_by(1024) { - let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); - writer.write_batch(&slice).await.unwrap(); - } - writer.finish().await.unwrap() - } - ConcreteFileVersion::V2_2 => { - let mut writer = - versions::v2_2::create_writer(object_writer, schema.clone(), options).unwrap(); - for offset in (0..batch.num_rows()).step_by(1024) { - let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); - writer.write_batch(&slice).await.unwrap(); - } - writer.finish().await.unwrap() - } - ConcreteFileVersion::V2_3 => { - let mut writer = - versions::v2_3::create_writer(object_writer, schema.clone(), options).unwrap(); - for offset in (0..batch.num_rows()).step_by(1024) { - let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); - writer.write_batch(&slice).await.unwrap(); - } - writer.finish().await.unwrap() - } - }; - fs.object_store - .open(&fs.tmp_path) - .await - .unwrap() - .get_range(0..summary.size_bytes as usize) - .await - .unwrap() - .to_vec() -} - -async fn assert_current_reader_roundtrip( - fixture: &[u8], - version: ConcreteFileVersion, - expected: &RecordBatch, -) { - let fs = FsFixture::default(); - let mut fixture_writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); - fixture_writer.write_all(fixture).await.unwrap(); - Writer::shutdown(fixture_writer.as_mut()).await.unwrap(); - let scheduler = fs - .scheduler - .open_file(&fs.tmp_path, &CachedFileSize::new(fixture.len() as u64)) - .await - .unwrap(); - let reader = FileReader::try_open( - scheduler, - None, - Arc::::default(), - &LanceCache::no_cache(), - FileReaderOptions::default(), - ) - .await - .unwrap(); - assert_eq!( - ConcreteFileVersion::from(reader.metadata().version()), - version - ); - assert!( - reader - .metadata() - .column_metadatas - .iter() - .any(|metadata| metadata.pages.len() > 1) - ); - let batches = reader - .read_stream( - ReadBatchParams::RangeFull, - 1024, - 16, - FilterExpression::no_filter(), - ) - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected.num_rows() - ); - assert!( - batches - .iter() - .all(|actual| actual.schema_ref() == expected.schema_ref()) - ); - let mut row_offset = 0; - for actual in &batches { - let expected = expected.slice(row_offset, actual.num_rows()); - assert_blob_column_eq(actual.column(4).as_ref(), expected.column(4).as_ref()); - row_offset += actual.num_rows(); - } - assert_eq!(row_offset, expected.num_rows()); -} - -#[rstest] -#[case::v2_0(ConcreteFileVersion::V2_0)] -#[case::v2_1(ConcreteFileVersion::V2_1)] -#[case::v2_2(ConcreteFileVersion::V2_2)] -#[tokio::test] -async fn stable_current_writer_and_reader_are_wire_compatible( - #[case] version: ConcreteFileVersion, -) { - let batch = compatibility_fixture_batch(); - let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); - schema.set_dictionary(&batch).unwrap(); - - let actual = write_current_fixture(version, &batch, &schema).await; - let expected = stable_fixture(version); - assert_wire_bytes_equal(&actual, expected); - assert_current_reader_roundtrip(expected, version, &batch).await; -} - -#[tokio::test] -async fn v2_3_output_is_deterministic_within_the_current_revision() { - let batch = compatibility_fixture_batch(); - let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); - schema.set_dictionary(&batch).unwrap(); - - let first = write_current_fixture(ConcreteFileVersion::V2_3, &batch, &schema).await; - let second = write_current_fixture(ConcreteFileVersion::V2_3, &batch, &schema).await; - assert_eq!(first, second); - assert_eq!( - &first[first.len() - 8..], - &[2, 0, 3, 0, b'L', b'A', b'N', b'C'] - ); - assert_current_reader_roundtrip(&first, ConcreteFileVersion::V2_3, &batch).await; -} - -#[tokio::test] -async fn v1_writer_and_reader_are_wire_compatible() { - let expected = stable_fixture(ConcreteFileVersion::V1); - let batch = compatibility_fixture_batch(); - let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); - schema.set_dictionary(&batch).unwrap(); - let fs = FsFixture::default(); - let mut writer = V1Writer::::try_new( - fs.object_store.as_ref(), - &fs.tmp_path, - schema.clone(), - &V1WriterOptions { - collect_stats_for_fields: Some(Vec::new()), - }, - ) - .await - .unwrap(); - for offset in (0..batch.num_rows()).step_by(1024) { - let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); - writer.write(std::slice::from_ref(&slice)).await.unwrap(); - } - let summary = writer.finish().await.unwrap(); - let actual = fs - .object_store - .open(&fs.tmp_path) - .await - .unwrap() - .get_range(0..summary.size_bytes as usize) - .await - .unwrap(); - assert_wire_bytes_equal(actual.as_ref(), expected); - - let fixture_fs = FsFixture::default(); - let mut fixture_writer = fixture_fs - .object_store - .create(&fixture_fs.tmp_path) - .await - .unwrap(); - fixture_writer.write_all(expected).await.unwrap(); - Writer::shutdown(fixture_writer.as_mut()).await.unwrap(); - let reader = V1Reader::try_new( - fixture_fs.object_store.as_ref(), - &fixture_fs.tmp_path, - schema.clone(), - ) - .await - .unwrap(); - let actual_batch = reader - .read_range(0..batch.num_rows(), &schema) - .await - .unwrap(); - assert_eq!(reader.num_batches(), 5); - assert_eq!(actual_batch.num_rows(), batch.num_rows()); - assert_eq!(actual_batch.column(0).to_data(), batch.column(0).to_data()); - assert_eq!(actual_batch.column(1).to_data(), batch.column(1).to_data()); - assert_blob_column_eq(actual_batch.column(4).as_ref(), batch.column(4).as_ref()); -} diff --git a/rust/lance-file/src/lib.rs b/rust/lance-file/src/lib.rs index 1ffa697cc9c..f48563dbc4a 100644 --- a/rust/lance-file/src/lib.rs +++ b/rust/lance-file/src/lib.rs @@ -10,24 +10,23 @@ pub mod version; pub mod versions; pub mod writer; -#[cfg(test)] -mod compatibility_tests; - pub use io::LanceEncodingsIo; use format::MAGIC; use lance_core::{Error, Result}; use lance_io::object_store::ObjectStore; use object_store::path::Path; -use version::{ConcreteFileVersion, LanceFileVersion}; +use version::ConcreteFileVersion; pub async fn determine_file_version( store: &ObjectStore, path: &Path, known_size: Option, -) -> Result { +) -> Result { let size = match known_size { - None => store.size(path).await.unwrap() as usize, + None => usize::try_from(store.size(path).await?).map_err(|_| { + Error::invalid_input(format!("file {} is too large for this platform", path)) + })?, Some(size) => size, }; if size < 8 { @@ -53,5 +52,5 @@ pub async fn determine_file_version( let major_version = u16::from_le_bytes([footer[0], footer[1]]); let minor_version = u16::from_le_bytes([footer[2], footer[3]]); - ConcreteFileVersion::from_footer_numbers(major_version, minor_version).map(Into::into) + ConcreteFileVersion::from_footer_numbers(major_version, minor_version) } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 30e35915ab5..5b9558ec30b 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -4,6 +4,7 @@ use std::{ borrow::Cow, collections::{BTreeMap, BTreeSet}, + fmt::Debug, io::Cursor, ops::Range, pin::Pin, @@ -12,6 +13,7 @@ use std::{ use arrow_array::RecordBatchReader; use arrow_schema::Schema as ArrowSchema; +use async_trait::async_trait; use byteorder::{ByteOrder, LittleEndian, ReadBytesExt}; use bytes::{Bytes, BytesMut}; use futures::{Stream, StreamExt, stream::BoxStream}; @@ -19,16 +21,14 @@ use lance_core::deepsize::{Context, DeepSizeOf}; use lance_encoding::{ EncodingsIo, decoder::{ - ColumnInfo, DecoderConfig, DecoderPlugins, FilterExpression, PageEncoding, PageInfo, - ReadBatchTask, RequestedRows, SchedulerDecoderConfig, schedule_and_decode, - schedule_and_decode_blocking, + ColumnInfo, DecoderConfig, DecoderPlugins, FilterExpression, PageEncoding, ReadBatchTask, + RequestedRows, SchedulerDecoderConfig, schedule_and_decode, schedule_and_decode_blocking, }, encoder::EncodedBatch, - version::LanceFileVersion, }; use log::debug; use object_store::path::Path; -use prost::{Message, Name}; +use prost::Message; use lance_core::{ Error, Result, @@ -45,12 +45,14 @@ use lance_io::{ use crate::{ datatypes::{Fields, FieldsWithMeta}, - format::{MAGIC, MAJOR_VERSION, MINOR_VERSION, pb, pbfile}, + format::{MAGIC, pb, pbfile}, io::LanceEncodingsIo, version::ConcreteFileVersion, - writer::PAGE_BUFFER_ALIGNMENT, + versions, }; +pub(crate) mod structural; + /// Default chunk size for reading large pages (8MiB) /// Pages larger than this will be split into multiple chunks during read pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024; @@ -103,8 +105,11 @@ pub struct CachedFileMetadata { pub num_global_buffer_bytes: u64, /// The number of bytes contained in the CMO and GBO tables pub num_footer_bytes: u64, + /// The major version number stored in the file footer. pub major_version: u16, + /// The minor version number stored in the file footer. pub minor_version: u16, + pub version: ConcreteFileVersion, /// The actual total file size in bytes, as reported by the object store. pub file_size_bytes: u64, /// User global buffers (index >= 1) whose bytes were already captured by the @@ -177,14 +182,14 @@ impl DeepSizeOf for CachedFileMetadata { /// hold decoded metadata for every column. #[derive(Debug, DeepSizeOf)] pub struct FileMetadataIndex { - file_schema: Arc, - num_rows: u64, - file_buffers: Vec, - column_metadata_offsets: Arc<[(u64, u64)]>, - num_columns: u32, - version: LanceFileVersion, - file_size_bytes: u64, - retained_global_buffers: BTreeMap, + pub(crate) file_schema: Arc, + pub(crate) num_rows: u64, + pub(crate) file_buffers: Vec, + pub(crate) column_metadata_offsets: Arc<[(u64, u64)]>, + pub(crate) num_columns: u32, + pub(crate) version: ConcreteFileVersion, + pub(crate) file_size_bytes: u64, + pub(crate) retained_global_buffers: BTreeMap, } impl FileMetadataIndex { @@ -230,18 +235,8 @@ impl CacheKey for ColumnMetadataCacheKey { } impl CachedFileMetadata { - pub fn version(&self) -> LanceFileVersion { - match (self.major_version, self.minor_version) { - (0, 3) => LanceFileVersion::V2_0, - (2, 0) => LanceFileVersion::V2_0, - (2, 1) => LanceFileVersion::V2_1, - (2, 2) => LanceFileVersion::V2_2, - (2, 3) => LanceFileVersion::V2_3, - _ => panic!( - "Unsupported version: {}.{}", - self.major_version, self.minor_version - ), - } + pub fn version(&self) -> ConcreteFileVersion { + self.version } } @@ -314,123 +309,14 @@ pub struct ReaderProjection { } impl ReaderProjection { - fn from_field_ids_helper<'a>( - file_version: LanceFileVersion, - fields: impl Iterator, - field_id_to_column_index: &BTreeMap, - column_indices: &mut Vec, - ) -> Result<()> { - for field in fields { - let is_structural = file_version >= LanceFileVersion::V2_1; - let (contributes, recurse) = field_column_shape(field, is_structural); - // In the 2.0 system we needed ids for intermediate fields. In 2.1+ - // we only need ids for leaf fields. - if contributes - && let Some(column_idx) = field_id_to_column_index.get(&(field.id as u32)).copied() - { - column_indices.push(column_idx); - } - if recurse { - Self::from_field_ids_helper( - file_version, - field.children.iter(), - field_id_to_column_index, - column_indices, - )?; - } - } - Ok(()) - } - - /// Creates a projection using a mapping from field IDs to column indices - /// - /// You can obtain such a mapping when the file is written using the - /// [`crate::writer::FileWriter::field_id_to_column_indices`] method. - pub fn from_field_ids( - file_version: LanceFileVersion, - schema: &Schema, - field_id_to_column_index: &BTreeMap, - ) -> Result { - let mut column_indices = Vec::new(); - Self::from_field_ids_helper( - file_version, - schema.fields.iter(), - field_id_to_column_index, - &mut column_indices, - )?; - let projection = Self { - schema: Arc::new(schema.clone()), - column_indices, - }; - Ok(projection) - } - - /// Creates a projection that reads the entire file - /// - /// If the schema provided is not the schema of the entire file then - /// the projection will be invalid and the read will fail. - /// If the field is a `struct datatype` with `packed` set to true in the field metadata, - /// the whole struct has one column index. - /// To support nested `packed-struct encoding`, this method need to be further adjusted. - pub fn from_whole_schema(schema: &Schema, version: LanceFileVersion) -> Self { - let schema = Arc::new(schema.clone()); - let is_structural = version >= LanceFileVersion::V2_1; - let mut column_indices = vec![]; - let mut curr_column_idx = 0; - let mut packed_struct_fields_num = 0; - for field in schema.fields_pre_order() { - if packed_struct_fields_num > 0 { - packed_struct_fields_num -= 1; - continue; - } - if field.is_packed_struct() { - column_indices.push(curr_column_idx); - curr_column_idx += 1; - packed_struct_fields_num = field.children.len(); - } else if field.children.is_empty() || !is_structural { - column_indices.push(curr_column_idx); - curr_column_idx += 1; - } - } - Self { - schema, - column_indices, - } - } - - /// Creates a projection that reads the specified columns provided by name + /// Returns whether this projection is selective enough to benefit from + /// loading column metadata through the file's metadata index. /// - /// The syntax for column names is the same as [`lance_core::datatypes::Schema::project`] - /// - /// If the schema provided is not the schema of the entire file then - /// the projection will be invalid and the read will fail. - pub fn from_column_names( - file_version: LanceFileVersion, - schema: &Schema, - column_names: &[&str], - ) -> Result { - let field_id_to_column_index = schema - .fields_pre_order() - // In the 2.0 system we needed ids for intermediate fields. In 2.1+ - // we only need ids for leaf fields. - .filter(|field| { - file_version < LanceFileVersion::V2_1 || field.is_leaf() || field.is_packed_struct() - }) - .enumerate() - .map(|(idx, field)| (field.id as u32, idx as u32)) - .collect::>(); - let projected = schema.project(column_names)?; - let mut column_indices = Vec::new(); - Self::from_field_ids_helper( - file_version, - projected.fields.iter(), - &field_id_to_column_index, - &mut column_indices, - )?; - Ok(Self { - schema: Arc::new(projected), - column_indices, - }) + /// The caller must already have selected a file format that supports indexed + /// metadata. This method only evaluates the projection shape and selectivity. + pub fn prefers_indexed_metadata(&self, total_columns: usize) -> bool { + FileMetadataProvider::projection_matches_indexed_metadata(self) + && self.column_indices.len().saturating_mul(4) < total_columns } } @@ -462,28 +348,52 @@ impl Default for FileReaderOptions { } #[derive(Debug, Clone)] -struct PreparedProjection { - column_infos: Vec>, - decoder_projection: ReaderProjection, +pub(crate) struct PreparedProjection { + pub column_infos: Vec>, + pub decoder_projection: ReaderProjection, } #[derive(Debug, Clone)] -enum FileMetadataProvider { +pub(crate) enum FileMetadataProvider { Full(Arc), Indexed(Arc), } +/// Executable projection behavior selected by an exact file-version module. +/// +/// The shared reader invokes this behavior but never interprets a version or +/// accepted-grammar profile. +#[async_trait] +pub(crate) trait ReadProjection: Debug + Send + Sync { + fn validate_indexed( + &self, + projection: &ReaderProjection, + metadata_index: &FileMetadataIndex, + ) -> Result<()>; + + fn read_length(&self, prepared: &PreparedProjection) -> Result; + + async fn prepare( + &self, + metadata_provider: &FileMetadataProvider, + projection: &ReaderProjection, + io: &Arc, + cache: &Arc, + ) -> Result<(PreparedProjection, u64)>; +} + #[derive(Debug, Clone)] -struct FileReadCore { - scheduler: Arc, - base_projection: ReaderProjection, - metadata_provider: FileMetadataProvider, - decoder_plugins: Arc, - cache: Arc, - options: FileReaderOptions, +pub(crate) struct DecodeEngine { + pub scheduler: Arc, + pub base_projection: ReaderProjection, + pub metadata_provider: FileMetadataProvider, + pub read_projection: Arc, + pub decoder_plugins: Arc, + pub cache: Arc, + pub options: FileReaderOptions, } -/// A projection-scoped reader for Lance files. +/// A projection-scoped reader for a current-format Lance file. /// /// This reader fixes a base projection at construction time. All later reads /// must stay within that projection, which lets the reader load only the column @@ -492,51 +402,71 @@ struct FileReadCore { /// file metadata. #[derive(Debug, Clone)] pub struct ProjectedFileReader { - core: FileReadCore, + core: DecodeEngine, } -/// A Lance file reader backed by fully decoded file metadata. +/// A current-format Lance file reader backed by fully decoded metadata. #[derive(Debug, Clone)] pub struct FileReader { - core: FileReadCore, - metadata: Arc, + pub(crate) core: DecodeEngine, + pub(crate) metadata: Arc, +} + +pub(crate) fn tasks_to_record_batch_stream( + schema: Arc, + tasks: Pin + Send>>, + batch_readahead: u32, +) -> Pin> { + let arrow_schema = Arc::new(ArrowSchema::from(schema.as_ref())); + let batches = tasks + .map(|task| task.task) + .buffered(batch_readahead as usize) + .boxed(); + Box::pin(RecordBatchStreamAdapter::new(arrow_schema, batches)) +} + +pub(crate) enum RawFileMetadataOpen { + Legacy { + major_version: u16, + minor_version: u16, + }, + Current { + version: ConcreteFileVersion, + metadata: RawFileMetadata, + }, +} + +pub(crate) struct RawFileMetadata { + pub file_schema: Arc, + pub column_metadatas: Vec, + pub num_rows: u64, + pub file_buffers: Vec, + pub num_data_bytes: u64, + pub num_column_metadata_bytes: u64, + pub num_global_buffer_bytes: u64, + pub num_footer_bytes: u64, + pub footer: Footer, + pub file_size_bytes: u64, + pub retained_global_buffers: BTreeMap, } + #[derive(Debug)] -struct Footer { +pub(crate) struct Footer { #[allow(dead_code)] - column_meta_start: u64, + pub column_meta_start: u64, // We don't use this today because we always load metadata for every column // and don't yet support "metadata projection" #[allow(dead_code)] - column_meta_offsets_start: u64, - global_buff_offsets_start: u64, - num_global_buffers: u32, - num_columns: u32, - major_version: u16, - minor_version: u16, + pub column_meta_offsets_start: u64, + pub global_buff_offsets_start: u64, + pub num_global_buffers: u32, + pub num_columns: u32, + pub major_version: u16, + pub minor_version: u16, } const FOOTER_LEN: usize = 40; -// How a field maps onto physical columns, shared by the projection-building and -// projection-validation walks so they stay in lockstep. In the 2.0 layout every -// ordinary field (including structs and lists) has its own column; in 2.1 only -// leaves do. Blob/packed-struct fields are opaque in all versions: they are a -// single column with no descent, including unloaded blob descriptor schemas. -// Returns `(contributes, recurse)`: whether the field has its own column and -// whether to walk into its children. The DFS order is the field's own column (if -// any) followed by its children, so a field's root (first) column is always the -// first entry of its sub-slice. -fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) { - if field.is_blob() || field.is_packed_struct() { - return (true, false); - } - - let contributes = !is_structural || field.children.is_empty(); - let recurse = !field.children.is_empty(); - (contributes, recurse) -} - // Count the V2.1 physical columns required to reconstruct a projected field. // This is the same DFS shape consumed by `ColumnInfoIter`: ordinary structural // nodes are transparent and leaves contribute columns. Indexed metadata loading @@ -550,100 +480,45 @@ fn indexed_projection_column_count(field: &Field) -> Option { return None; } - let (contributes, recurse) = field_column_shape(field, true); - let initial = usize::from(contributes); - if !recurse { - return Some(initial); + if field.children.is_empty() { + return Some(1); } - field.children.iter().try_fold(initial, |count, child| { + field.children.iter().try_fold(0usize, |count, child| { count.checked_add(indexed_projection_column_count(child)?) }) } -// Whether a field's children each cover the same rows as the field itself. Struct -// children do (one value per parent row), so they must share its length. List, -// map, and fixed-size-list items have an independent cardinality (item count, not -// row count) and are validated only against themselves. -fn children_share_parent_length(field: &Field) -> bool { - field.logical_type.is_struct() -} - -// Validate one field's slice of a projection's flat `column_indices`, returning -// the field's top-level row count (the page-row sum of its root column). Walks the -// same DFS order as `from_field_ids_helper`, advancing `cursor` past every column -// the field contributes. -// -// `comparable` tracks whether the field's row count shares the read's top-level -// cardinality. A struct's children must all match that count -- the decoders -// combine them assuming equal lengths and would otherwise panic or read past a -// shorter child -- so the equality check runs only while `comparable` holds. Once -// the walk descends through a list/map/fixed-size-list its items have an -// independent cardinality (item count, not row count), so `comparable` turns off -// for that whole subtree and a nested struct's children are no longer compared. -fn validate_field_length Result>( - field: &Field, - is_structural: bool, - comparable: bool, - column_indices: &[u32], - cursor: &mut usize, - column_len: &F, -) -> Result { - let (contributes, recurse) = field_column_shape(field, is_structural); - let mut field_rows: Option = None; - if contributes { - let column = *column_indices.get(*cursor).ok_or_else(|| { - Error::invalid_input(format!( - "projection supplied fewer column indices than its fields require \ - (ran out at field '{}')", - field.name - )) - })?; - *cursor += 1; - field_rows = Some(column_len(column as usize)?); - } - if recurse { - // Only enforce equal-length children for a struct whose own count is still - // at the top-level cardinality; below a list/map/fixed-size-list the items - // have an independent cardinality, so neither this field nor its - // descendants are comparable. - let enforce_children = comparable && children_share_parent_length(field); - for child in &field.children { - let child_rows = validate_field_length( - child, - is_structural, - enforce_children, - column_indices, - cursor, - column_len, - )?; - // A struct that contributes no column of its own (the 2.1 layout) - // takes its row count from its first child. - let expected = *field_rows.get_or_insert(child_rows); - if enforce_children && child_rows != expected { - return Err(Error::invalid_input(format!( - "cannot read field '{}': its children have differing lengths \ - (child '{}' has {} rows, but the field has {}); a struct's \ - children must all have the same length", - field.name, child.name, child_rows, expected - ))); - } - } - } - field_rows.ok_or_else(|| { - Error::invalid_input(format!( - "projected field '{}' maps to no columns", - field.name - )) - }) -} - // The reader combines a projection's columns into rectangular batches, so they // must all have the same length. Returns that common length, or a descriptive // error (naming each column's length) when they differ. Ordinary files always // pass; only files written with `FileWriter::write_column` whose columns ended up // unequal can fail, and those must be read separately. -fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result { +pub(crate) fn normalized_column_num_rows(info: &ColumnInfo) -> Result { + info.page_infos.iter().try_fold(0_u64, |rows, page| { + let page_rows = match &page.encoding { + PageEncoding::Structural(layout) => match &layout.layout { + Some(pbenc21::page_layout::Layout::SparseLayout(sparse)) => sparse + .structural_layers + .first() + .and_then(|layer| layer.layer.as_ref()) + .map_or(page.num_rows, |layer| match layer { + pbenc21::sparse_structural_layer::Layer::Validity(layer) => layer.num_slots, + pbenc21::sparse_structural_layer::Layer::List(layer) => layer.num_slots, + pbenc21::sparse_structural_layer::Layer::FixedSizeList(layer) => { + layer.num_slots + } + }), + _ => page.num_rows, + }, + _ => page.num_rows, + }; + rows.checked_add(page_rows) + .ok_or_else(|| Error::invalid_input_source("Column row count overflows u64".into())) + }) +} + +pub(crate) fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result { let first = field_lengths.first().map_or(0, |&(_, len)| len); if field_lengths.iter().all(|&(_, len)| len == first) { return Ok(first); @@ -660,6 +535,30 @@ fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result { } impl FileReader { + pub(crate) fn base_projection(&self) -> &ReaderProjection { + &self.core.base_projection + } + + pub(crate) fn full_projection(&self, projection: ReaderProjection) -> PreparedProjection { + PreparedProjection { + column_infos: self.metadata.column_infos.clone(), + decoder_projection: projection, + } + } + + pub(crate) async fn read_prepared_tasks( + &self, + params: ReadBatchParams, + batch_size: u32, + prepared: PreparedProjection, + read_len: u64, + filter: FilterExpression, + ) -> Result + Send>>> { + self.core + .read_prepared_tasks(params, batch_size, prepared, read_len, filter) + .await + } + pub fn with_scheduler(&self, scheduler: Arc) -> Self { Self { core: self.core.with_scheduler(scheduler), @@ -819,14 +718,6 @@ impl FileReader { let major_version = cursor.read_u16::()?; let minor_version = cursor.read_u16::()?; - if major_version == MAJOR_VERSION as u16 && minor_version == MINOR_VERSION as u16 { - return Err(Error::version_conflict( - "Attempt to use the lance v2 reader to read a legacy file".to_string(), - major_version, - minor_version, - )); - } - let magic_bytes = footer_bytes.slice(len - 4..); if magic_bytes.as_ref() != MAGIC { return Err(Error::invalid_input(format!( @@ -845,6 +736,22 @@ impl FileReader { }) } + fn current_file_version(footer: &Footer) -> Result { + let version = + ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?; + match version { + ConcreteFileVersion::V1 => Err(Error::version_conflict( + "Attempt to use the lance v2 reader to read a legacy file".to_string(), + footer.major_version, + footer.minor_version, + )), + ConcreteFileVersion::V2_0 + | ConcreteFileVersion::V2_1 + | ConcreteFileVersion::V2_2 + | ConcreteFileVersion::V2_3 => Ok(version), + } + } + // TODO: Once we have coalesced I/O we should only read the column metadatas that we need fn read_all_column_metadata( column_metadata_bytes: Bytes, @@ -936,25 +843,12 @@ impl FileReader { } } - fn do_decode_gbo_table( - gbo_bytes: &Bytes, - footer: &Footer, - version: LanceFileVersion, - ) -> Result> { + fn do_decode_gbo_table(gbo_bytes: &Bytes, footer: &Footer) -> Result> { let mut global_bufs_cursor = Cursor::new(gbo_bytes); let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize); - for buffer_index in 0..footer.num_global_buffers { + for _ in 0..footer.num_global_buffers { let buf_pos = global_bufs_cursor.read_u64::()?; - if version >= LanceFileVersion::V2_1 && buf_pos % PAGE_BUFFER_ALIGNMENT as u64 != 0 { - return Err(Error::invalid_input_source( - format!( - "Global buffer {} position {} is not aligned to {} bytes", - buffer_index, buf_pos, PAGE_BUFFER_ALIGNMENT - ) - .into(), - )); - } let buf_size = global_bufs_cursor.read_u64::()?; global_buffers.push(BufferDescriptor { position: buf_pos, @@ -970,7 +864,6 @@ impl FileReader { file_len: u64, scheduler: &FileScheduler, footer: &Footer, - version: LanceFileVersion, ) -> Result> { // This could, in theory, trigger another IOP but the GBO table should never be large // enough for that to happen @@ -981,7 +874,7 @@ impl FileReader { file_len, ) .await?; - Self::do_decode_gbo_table(&gbo_bytes, footer, version) + Self::do_decode_gbo_table(&gbo_bytes, footer) } fn decode_schema(schema_bytes: Bytes) -> Result<(u64, lance_core::datatypes::Schema)> { @@ -996,31 +889,24 @@ impl FileReader { Ok((num_rows, schema)) } - // TODO: Support late projection. Currently, if we want to perform a - // projected read of a file, we load all of the column metadata, and then - // only read the column data that is requested. This is fine for most cases. - // - // However, if there are many columns then loading all of the column metadata - // may be expensive. We should support a mode where we only load the column - // metadata for the columns that are requested (the file format supports this). - // - // The main challenge is that we either need to ignore the column metadata cache - // or have a more sophisticated cache that can cache per-column metadata. - // - // Also, if the number of columns is fairly small, it's faster to read them as a - // single IOP, but we can fix this through coalescing. - pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result { - // 1. read the footer + pub(crate) async fn read_raw_metadata_for_dispatch( + scheduler: &FileScheduler, + ) -> Result { let (tail_bytes, file_len) = Self::read_tail(scheduler).await?; let tail_offset = file_len - tail_bytes.len() as u64; let footer = Self::decode_footer(&tail_bytes)?; + let version = + ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?; + if version == ConcreteFileVersion::V1 { + return Ok(RawFileMetadataOpen::Legacy { + major_version: footer.major_version, + minor_version: footer.minor_version, + }); + } - let file_version: LanceFileVersion = - ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)? - .into(); - - let gbo_table = - Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?; + // Exact readers validate their own alignment contract after this + // version-free table parse. + let gbo_table = Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer).await?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), @@ -1028,19 +914,12 @@ impl FileReader { } let schema_start = gbo_table[0].position; let schema_size = gbo_table[0].size; - let num_footer_bytes = file_len - schema_start; - - // By default we read all column metadatas. We do NOT read the column metadata buffers - // at this point. We only want to read the column metadata for columns we are actually loading. 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); let (num_rows, schema) = Self::decode_schema(schema_bytes)?; - // Next, read the metadata for the columns - // This is both the column metadata and the CMO table 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 = @@ -1050,35 +929,28 @@ impl FileReader { let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::(); 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; - - let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version)?; - - // 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. let retained_global_buffers = Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset); - Ok(CachedFileMetadata { - file_schema: Arc::new(schema), - column_metadatas, - column_infos, - num_rows, - num_data_bytes, - num_column_metadata_bytes, - num_global_buffer_bytes, - num_footer_bytes, - file_buffers: gbo_table, - major_version: footer.major_version, - minor_version: footer.minor_version, - file_size_bytes: file_len, - retained_global_buffers, + Ok(RawFileMetadataOpen::Current { + version, + metadata: RawFileMetadata { + file_schema: Arc::new(schema), + column_metadatas, + num_rows, + file_buffers: gbo_table, + num_data_bytes, + num_column_metadata_bytes, + num_global_buffer_bytes, + num_footer_bytes, + footer, + file_size_bytes: file_len, + retained_global_buffers, + }, }) } - async fn read_metadata_index_with_known_schema( + async fn read_raw_metadata_index_with_known_schema( scheduler: &FileScheduler, known_schema: Option<(Arc, u64)>, ) -> Result { @@ -1086,12 +958,9 @@ impl FileReader { let tail_offset = file_len - tail_bytes.len() as u64; let footer = Self::decode_footer(&tail_bytes)?; - let file_version: LanceFileVersion = - ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)? - .into(); + let file_version = Self::current_file_version(&footer)?; - let gbo_table = - Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?; + let gbo_table = Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer).await?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), @@ -1142,176 +1011,26 @@ impl FileReader { /// 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 async fn read_metadata_index(scheduler: &FileScheduler) -> Result { - Self::read_metadata_index_with_known_schema(scheduler, None).await + 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 async fn read_metadata_index_with_schema( + pub(crate) async fn read_raw_metadata_index_with_schema( scheduler: &FileScheduler, file_schema: Arc, num_rows: u64, ) -> Result { - Self::read_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))).await - } - - fn fetch_encoding(encoding: &pbfile::Encoding) -> Result { - match &encoding.location { - Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source( - "Indirect file encodings are not supported".into(), - )), - Some(pbfile::encoding::Location::Direct(encoding)) => { - let encoding_buf = Bytes::from(encoding.encoding.clone()); - let encoding_any = prost_types::Any::decode(encoding_buf).map_err(|error| { - Error::invalid_input_source( - format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(), - ) - })?; - encoding_any.to_msg::().map_err(|error| { - Error::invalid_input_source( - format!("Invalid direct {} encoding: {error}", M::NAME).into(), - ) - }) - } - Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source( - format!("Missing {} encoding description", M::NAME).into(), - )), - None => Err(Error::invalid_input_source( - format!("Missing {} encoding location", M::NAME).into(), - )), - } - } - - fn meta_to_col_infos( - column_metadatas: &[pbfile::ColumnMetadata], - file_version: LanceFileVersion, - ) -> Result>> { - column_metadatas - .iter() - .enumerate() - .map(|(col_idx, col_meta)| { - let col_idx = u32::try_from(col_idx).map_err(|_| { - Error::invalid_input_source("File has more than u32::MAX columns".into()) - })?; - Self::meta_to_col_info(col_idx, col_meta, file_version) - }) - .collect() - } - - fn meta_to_col_info( - col_idx: u32, - col_meta: &pbfile::ColumnMetadata, - file_version: LanceFileVersion, - ) -> Result> { - let page_infos = col_meta - .pages - .iter() - .enumerate() - .map(|(page_idx, page)| { - let num_rows = page.length; - let encoding = match file_version { - LanceFileVersion::V2_0 => { - PageEncoding::Array(Self::fetch_encoding::( - page.encoding.as_ref().ok_or_else(|| { - Error::invalid_input_source( - format!( - "Column {} page {} is missing its encoding", - col_idx, page_idx - ) - .into(), - ) - })?, - )?) - } - _ => { - PageEncoding::Structural(Self::fetch_encoding::( - page.encoding.as_ref().ok_or_else(|| { - Error::invalid_input_source( - format!( - "Column {} page {} is missing its encoding", - col_idx, page_idx - ) - .into(), - ) - })?, - )?) - } - }; - if page.buffer_offsets.len() != page.buffer_sizes.len() { - return Err(Error::invalid_input_source( - format!( - "Column {} page {} has {} buffer offsets but {} buffer sizes", - col_idx, - page_idx, - page.buffer_offsets.len(), - page.buffer_sizes.len() - ) - .into(), - )); - } - let buffer_offsets_and_sizes = Arc::from( - page.buffer_offsets - .iter() - .zip(page.buffer_sizes.iter()) - .map(|(offset, size)| -> Result<_> { - if file_version >= LanceFileVersion::V2_1 - && offset % PAGE_BUFFER_ALIGNMENT as u64 != 0 - { - return Err(Error::invalid_input_source( - format!( - "Column {} page {} buffer offset {} is not aligned to {} bytes", - col_idx, page_idx, offset, PAGE_BUFFER_ALIGNMENT - ) - .into(), - )); - } - Ok((*offset, *size)) - }) - .collect::>>()?, - ); - Ok(PageInfo { - buffer_offsets_and_sizes, - encoding, - num_rows, - priority: page.priority, - }) - }) - .collect::>>()?; - if col_meta.buffer_offsets.len() != col_meta.buffer_sizes.len() { - return Err(Error::invalid_input_source( - format!( - "Column {} has {} buffer offsets but {} buffer sizes", - col_idx, - col_meta.buffer_offsets.len(), - col_meta.buffer_sizes.len() - ) - .into(), - )); - } - let buffer_offsets_and_sizes = Arc::from( - col_meta - .buffer_offsets - .iter() - .zip(col_meta.buffer_sizes.iter()) - .map(|(offset, size)| (*offset, *size)) - .collect::>(), - ); - Ok(Arc::new(ColumnInfo { - index: col_idx, - page_infos: Arc::from(page_infos), - buffer_offsets_and_sizes, - encoding: Self::fetch_encoding(col_meta.encoding.as_ref().ok_or_else(|| { - Error::invalid_input_source( - format!("Column {} is missing its encoding", col_idx).into(), - ) - })?)?, - })) + Self::read_raw_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))) + .await } - fn validate_projection( + pub(crate) fn validate_projection( projection: &ReaderProjection, metadata: &CachedFileMetadata, ) -> Result<()> { @@ -1340,68 +1059,6 @@ impl FileReader { Ok(()) } - /// Opens a new file reader without any pre-existing knowledge - /// - /// This will read the file schema from the file itself and thus requires a bit more I/O - /// - /// A `base_projection` can also be provided. If provided, then the projection will apply - /// to all reads from the file that do not specify their own projection. - pub async fn try_open( - scheduler: FileScheduler, - base_projection: Option, - decoder_plugins: Arc, - cache: &LanceCache, - options: FileReaderOptions, - ) -> Result { - let file_metadata = Arc::new(Self::read_all_metadata(&scheduler).await?); - let path = scheduler.reader().path().clone(); - - // Create LanceEncodingsIo with read chunk size from options - let encodings_io = - LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size); - - Self::try_open_with_file_metadata( - Arc::new(encodings_io), - path, - base_projection, - decoder_plugins, - file_metadata, - cache, - options, - ) - .await - } - - /// Same as `try_open` but with the file metadata already loaded. - /// - /// This method also can accept any kind of `EncodingsIo` implementation allowing - /// for custom strategies to be used for I/O scheduling (e.g. for takes on fast - /// disks it may be better to avoid asynchronous overhead). - /// Opens a data reader backed by fully decoded file metadata. - pub async fn try_open_with_file_metadata( - scheduler: Arc, - path: Path, - base_projection: Option, - decoder_plugins: Arc, - file_metadata: Arc, - cache: &LanceCache, - options: FileReaderOptions, - ) -> Result { - let cache = Arc::new(cache.with_key_prefix(path.as_ref())); - let core = FileReadCore::try_new( - scheduler, - base_projection, - decoder_plugins, - FileMetadataProvider::Full(file_metadata.clone()), - cache, - options, - )?; - Ok(Self { - core, - metadata: file_metadata, - }) - } - // The actual decoder needs all the column infos that make up a type. In other words, if // the first type in the schema is Struct then the decoder will need 3 column infos. // @@ -1557,97 +1214,11 @@ impl FileReader { .await } - /// Creates a stream of "read tasks" to read the data from the file - /// - /// The arguments are similar to [`Self::read_stream_projected`] but instead of returning a stream - /// of record batches it returns a stream of "read tasks". - /// - /// The tasks should be consumed with some kind of `buffered` argument if CPU parallelism is desired. - /// - /// Note that "read task" is probably a bit imprecise. The tasks are actually "decode tasks". The - /// reading happens asynchronously in the background. In other words, a single read task may map to - /// multiple I/O operations or a single I/O operation may map to multiple read tasks. - /// - /// # Why is this async? - /// - /// Constructing the read stream requires running the decode scheduler's - /// `initialize` step, which performs the metadata I/O (chunk metadata, - /// dictionaries, repetition index, ...) needed to plan the read. We - /// drive that I/O on the awaiting task rather than smuggling it into - /// the stream's first poll. This way callers control where the - /// scheduling I/O runs (typically inside a per-fragment - /// `tokio::spawn`), planning errors surface from the await instead of - /// from the first stream item, and small reads can also complete the - /// synchronous scheduling step before returning (see - /// [`DecoderConfig::inline_scheduling`]). - pub async fn read_tasks( + fn take_rows_blocking( &self, - params: ReadBatchParams, + indices: Vec, batch_size: u32, - projection: Option, - filter: FilterExpression, - ) -> Result + Send>>> { - self.core - .read_tasks(params, batch_size, projection, filter) - .await - } - - /// Reads data from the file as a stream of record batches - /// - /// * `params` - Specifies the range (or indices) of data to read - /// * `batch_size` - The maximum size of a single batch. A batch may be smaller - /// if it is the last batch or if it is not possible to create a batch of the - /// requested size. - /// - /// For example, if the batch size is 1024 and one of the columns is a string - /// column then there may be some ranges of 1024 rows that contain more than - /// 2^31 bytes of string data (which is the maximum size of a string column - /// in Arrow). In this case smaller batches may be emitted. - /// * `batch_readahead` - The number of batches to read ahead. This controls the - /// amount of CPU parallelism of the read. In other words it controls how many - /// batches will be decoded in parallel. It has no effect on the I/O parallelism - /// of the read (how many I/O requests are in flight at once). - /// - /// This parameter also is also related to backpressure. If the consumer of the - /// stream is slow then the reader will build up RAM. - /// * `projection` - A projection to apply to the read. This controls which columns - /// are read from the file. The projection is NOT applied on top of the base - /// projection. The projection is applied directly to the file schema. - /// - /// # Why is this async? - /// - /// This delegates to [`Self::read_tasks`], which awaits the decode - /// scheduler's `initialize` step (and, for small reads, the synchronous - /// scheduling that follows) before returning. See `read_tasks` for - /// details on why this work is performed up front rather than on the - /// stream's first poll. - pub async fn read_stream_projected( - &self, - params: ReadBatchParams, - batch_size: u32, - batch_readahead: u32, - projection: ReaderProjection, - filter: FilterExpression, - ) -> Result>> { - let arrow_schema = Arc::new(ArrowSchema::from(projection.schema.as_ref())); - let tasks_stream = self - .read_tasks(params, batch_size, Some(projection), filter) - .await?; - let batch_stream = tasks_stream - .map(|task| task.task) - .buffered(batch_readahead as usize) - .boxed(); - Ok(Box::pin(RecordBatchStreamAdapter::new( - arrow_schema, - batch_stream, - ))) - } - - fn take_rows_blocking( - &self, - indices: Vec, - batch_size: u32, - projection: ReaderProjection, + projection: ReaderProjection, filter: FilterExpression, ) -> Result> { let column_infos = self.collect_columns_from_projection(&projection)?; @@ -1761,36 +1332,15 @@ impl FileReader { ) } - /// Read data from the file as an iterator of record batches - /// - /// This is a blocking variant of [`Self::read_stream_projected`] that runs entirely in the - /// calling thread. It will block on I/O if the decode is faster than the I/O. It is useful - /// for benchmarking and potentially from "take"ing small batches from fast disks. - /// - /// Large scans of in-memory data will still benefit from threading (and should therefore not - /// use this method) because we can parallelize the decode. - /// - /// Note: calling this from within a tokio runtime will panic. It is acceptable to call this - /// from a spawn_blocking context. - pub fn read_stream_projected_blocking( + pub(crate) fn read_prepared_blocking( &self, params: ReadBatchParams, batch_size: u32, - projection: Option, + prepared: PreparedProjection, + read_len: u64, filter: FilterExpression, ) -> Result> { - let projection = projection.unwrap_or_else(|| self.core.base_projection.clone()); - Self::validate_projection(&projection, &self.metadata)?; - // Apply the same projection-length validation as the async path. This - // reader is always backed by full metadata, so we can build the prepared - // projection synchronously (no column-metadata I/O) and reuse the shared - // check. `read_len` is the projection's common column length, which - // `RangeFull`/`RangeFrom` resolve against rather than `num_rows`. - let prepared = PreparedProjection { - column_infos: self.metadata.column_infos.clone(), - decoder_projection: projection.clone(), - }; - let read_len = self.core.prepared_read_length(&prepared)?; + let projection = prepared.decoder_projection; let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { if bound > read_len || (bound == read_len && inclusive) { Err(Error::invalid_input(format!( @@ -1802,17 +1352,13 @@ impl FileReader { }; match ¶ms { ReadBatchParams::Indices(indices) => { - for idx in indices { - match idx { - None => { - return Err(Error::invalid_input("Null value in indices array")); - } - Some(idx) => { - verify_bound(¶ms, idx as u64, true)?; - } + for index in indices { + match index { + None => return Err(Error::invalid_input("Null value in indices array")), + Some(index) => verify_bound(¶ms, index as u64, true)?, } } - let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect(); + let indices = indices.iter().map(|index| index.unwrap() as u64).collect(); self.take_rows_blocking(indices, batch_size, projection, filter) } ReadBatchParams::Range(range) => { @@ -1851,76 +1397,48 @@ impl FileReader { } } - /// Reads data from the file as a stream of record batches - /// - /// This is similar to [`Self::read_stream_projected`] but uses the base projection - /// provided when the file was opened (or reads all columns if the file was - /// opened without a base projection) - /// - /// # Why is this async? - /// - /// This delegates to [`Self::read_stream_projected`], which awaits the - /// decode scheduler's `initialize` step before returning the stream. - /// See [`Self::read_tasks`] for the rationale. - pub async fn read_stream( - &self, - params: ReadBatchParams, - batch_size: u32, - batch_readahead: u32, - filter: FilterExpression, - ) -> Result>> { - self.read_stream_projected( - params, - batch_size, - batch_readahead, - self.core.base_projection.clone(), - filter, - ) - .await - } - pub fn schema(&self) -> &Arc { self.core.schema() } } impl FileMetadataProvider { - fn version(&self) -> LanceFileVersion { + pub(crate) fn version(&self) -> ConcreteFileVersion { match self { - Self::Full(metadata) => metadata.version(), + Self::Full(metadata) => metadata.version, Self::Indexed(metadata_index) => metadata_index.version, } } - fn num_rows(&self) -> u64 { + pub(crate) fn num_rows(&self) -> u64 { match self { Self::Full(metadata) => metadata.num_rows, Self::Indexed(metadata_index) => metadata_index.num_rows, } } - fn schema(&self) -> &Arc { + pub(crate) fn schema(&self) -> &Arc { match self { Self::Full(metadata) => &metadata.file_schema, Self::Indexed(metadata_index) => &metadata_index.file_schema, } } - fn file_buffers(&self) -> &Vec { + pub(crate) fn file_buffers(&self) -> &Vec { match self { Self::Full(metadata) => &metadata.file_buffers, Self::Indexed(metadata_index) => &metadata_index.file_buffers, } } - fn retained_global_buffers(&self) -> &BTreeMap { + pub(crate) fn retained_global_buffers(&self) -> &BTreeMap { match self { Self::Full(metadata) => &metadata.retained_global_buffers, Self::Indexed(metadata_index) => &metadata_index.retained_global_buffers, } } - fn file_statistics(&self) -> Option { + pub(crate) fn file_statistics(&self) -> Option { let metadata = match self { Self::Full(metadata) => metadata, Self::Indexed(_) => return None, @@ -1930,11 +1448,8 @@ impl FileMetadataProvider { )) } - fn supports_indexed_projection( - projection: &ReaderProjection, - version: LanceFileVersion, - ) -> bool { - if version < LanceFileVersion::V2_1 || projection.schema.fields.is_empty() { + pub(crate) fn projection_matches_indexed_metadata(projection: &ReaderProjection) -> bool { + if projection.schema.fields.is_empty() { return false; } @@ -1948,7 +1463,7 @@ impl FileMetadataProvider { == Some(projection.column_indices.len()) } - fn validate_indexed_projection( + pub(crate) fn validate_indexed_projection_structure( projection: &ReaderProjection, metadata_index: &FileMetadataIndex, ) -> Result<()> { @@ -1973,24 +1488,19 @@ impl FileMetadataProvider { ))); } } - if !Self::supports_indexed_projection(projection, metadata_index.version) { - return Err(Error::not_supported(format!( - "lazy column metadata loading requires a V2.1+ ordinary structural projection without blob or packed-struct fields whose physical-column count matches the projection; got file version {:?}, {} schema fields, and {} column indices", - metadata_index.version, - projection.schema.fields.len(), - projection.column_indices.len() - ))); - } Ok(()) } - fn validate_projection(&self, projection: &ReaderProjection) -> Result<()> { - match self { - Self::Full(metadata) => FileReader::validate_projection(projection, metadata), - Self::Indexed(metadata_index) => { - Self::validate_indexed_projection(projection, metadata_index) - } - } + pub(crate) fn indexed_projection_error( + projection: &ReaderProjection, + metadata_index: &FileMetadataIndex, + ) -> Error { + Error::not_supported(format!( + "lazy column metadata loading requires a V2.1+ ordinary structural projection without blob or packed-struct fields whose physical-column count matches the projection; got file version {:?}, {} schema fields, and {} column indices", + metadata_index.version, + projection.schema.fields.len(), + projection.column_indices.len() + )) } fn column_metadata_range( @@ -2016,12 +1526,16 @@ impl FileMetadataProvider { Ok(position..end) } - async fn load_indexed_column_infos( + pub(crate) async fn load_indexed_column_infos( metadata_index: &FileMetadataIndex, io: &Arc, cache: &Arc, column_indices: &[u32], - ) -> Result>> { + decode_column: F, + ) -> Result>> + where + F: Fn(u32, &pbfile::ColumnMetadata) -> Result>, + { let mut column_infos = vec![None; column_indices.len()]; let mut missing_columns = Vec::new(); @@ -2046,11 +1560,7 @@ impl FileMetadataProvider { missing_columns.into_iter().zip(metadata_bytes) { let column_metadata = pbfile::ColumnMetadata::decode(bytes)?; - let column_info = FileReader::meta_to_col_info( - column_index, - &column_metadata, - metadata_index.version, - )?; + let column_info = decode_column(column_index, &column_metadata)?; let cached = Arc::new(CachedColumnMetadata { column_metadata, column_info: column_info.clone(), @@ -2074,83 +1584,41 @@ impl FileMetadataProvider { }) .collect() } - - async fn prepare_projection( - &self, - projection: &ReaderProjection, - io: &Arc, - cache: &Arc, - ) -> Result { - self.validate_projection(projection)?; - match self { - Self::Full(metadata) => Ok(PreparedProjection { - column_infos: metadata.column_infos.clone(), - decoder_projection: projection.clone(), - }), - Self::Indexed(metadata_index) => { - let column_infos = Self::load_indexed_column_infos( - metadata_index, - io, - cache, - &projection.column_indices, - ) - .await?; - let decoder_projection = ReaderProjection { - schema: projection.schema.clone(), - column_indices: (0..projection.column_indices.len()) - .map(|idx| idx as u32) - .collect(), - }; - Ok(PreparedProjection { - column_infos, - decoder_projection, - }) - } - } - } } -impl FileReadCore { - fn try_new( +impl DecodeEngine { + pub(crate) fn try_new( scheduler: Arc, - base_projection: Option, + base_projection: ReaderProjection, decoder_plugins: Arc, metadata_provider: FileMetadataProvider, + read_projection: Arc, cache: Arc, options: FileReaderOptions, ) -> Result { - if let Some(base_projection) = base_projection.as_ref() { - metadata_provider.validate_projection(base_projection)?; - } - let base_projection = base_projection.unwrap_or(ReaderProjection::from_whole_schema( - metadata_provider.schema().as_ref(), - metadata_provider.version(), - )); Ok(Self { scheduler, base_projection, metadata_provider, + read_projection, decoder_plugins, cache, options, }) } - fn with_scheduler(&self, scheduler: Arc) -> Self { + pub(crate) fn with_scheduler(&self, scheduler: Arc) -> Self { Self { scheduler, base_projection: self.base_projection.clone(), metadata_provider: self.metadata_provider.clone(), + read_projection: self.read_projection.clone(), decoder_plugins: self.decoder_plugins.clone(), cache: self.cache.clone(), options: self.options.clone(), } } - fn version(&self) -> LanceFileVersion { - self.metadata_provider.version() - } - fn num_rows(&self) -> u64 { self.metadata_provider.num_rows() } @@ -2159,7 +1627,7 @@ impl FileReadCore { self.metadata_provider.schema() } - async fn read_global_buffer(&self, index: u32) -> Result { + pub(crate) async fn read_global_buffer(&self, index: u32) -> Result { let file_buffers = self.metadata_provider.file_buffers(); let buffer_desc = file_buffers.get(index as usize).ok_or_else(|| { Error::invalid_input(format!( @@ -2188,82 +1656,6 @@ impl FileReadCore { }) } - // The common length to read across a prepared projection, after validating - // its columns can be combined into rectangular batches. Each top-level field - // is checked for internal consistency (see `validate_field_length`); the - // top-level fields must then share a length, since one read combines them. - // Ordinary files always pass (every column has `num_rows` rows); files - // written with `FileWriter::write_column` whose columns ended up unequal are - // rejected here and must be read separately. - // - // `column_infos` and `decoder_projection.column_indices` line up for both - // metadata providers: the full provider keeps absolute indices into the whole - // file, while the indexed (lazy) provider loads only the projected columns and - // renumbers them 0..N -- in either case `column_infos[column_index]` is the - // requested column. - fn prepared_read_length(&self, prepared: &PreparedProjection) -> Result { - let is_structural = self.version() >= LanceFileVersion::V2_1; - let column_infos = &prepared.column_infos; - let column_len = |column: usize| -> Result { - let info = column_infos.get(column).ok_or_else(|| { - Error::invalid_input(format!( - "projection references column index {} but only {} columns are available", - column, - column_infos.len() - )) - })?; - info.page_infos.iter().try_fold(0_u64, |rows, page| { - let page_rows = match &page.encoding { - PageEncoding::Structural(layout) => match &layout.layout { - Some(pbenc21::page_layout::Layout::SparseLayout(sparse)) => sparse - .structural_layers - .first() - .and_then(|layer| layer.layer.as_ref()) - .map_or(page.num_rows, |layer| match layer { - pbenc21::sparse_structural_layer::Layer::Validity(layer) => { - layer.num_slots - } - pbenc21::sparse_structural_layer::Layer::List(layer) => { - layer.num_slots - } - pbenc21::sparse_structural_layer::Layer::FixedSizeList(layer) => { - layer.num_slots - } - }), - _ => page.num_rows, - }, - _ => page.num_rows, - }; - rows.checked_add(page_rows).ok_or_else(|| { - Error::invalid_input_source("Column row count overflows u64".into()) - }) - }) - }; - let column_indices = &prepared.decoder_projection.column_indices; - let fields = &prepared.decoder_projection.schema.fields; - let mut cursor = 0usize; - let mut field_lengths = Vec::with_capacity(fields.len()); - for field in fields { - let rows = validate_field_length( - field, - is_structural, - true, - column_indices, - &mut cursor, - &column_len, - )?; - field_lengths.push((field.name.as_str(), rows)); - } - if cursor != column_indices.len() { - return Err(Error::invalid_input(format!( - "projection supplied {} column indices but its fields require {}", - column_indices.len(), - cursor - ))); - } - verify_uniform_lengths(&field_lengths) - } - async fn read_range( &self, range: Range, @@ -2287,134 +1679,389 @@ impl FileReadCore { .await } - async fn take_rows( + async fn take_rows( + &self, + indices: Vec, + batch_size: u32, + prepared: PreparedProjection, + ) -> Result> { + FileReader::do_take_rows( + prepared.column_infos, + self.scheduler.clone(), + self.cache.clone(), + self.decoder_plugins.clone(), + indices, + batch_size, + prepared.decoder_projection, + FilterExpression::no_filter(), + self.options.decoder_config.clone(), + self.options.batch_size_bytes, + ) + .await + } + + async fn read_ranges( + &self, + ranges: Vec>, + batch_size: u32, + prepared: PreparedProjection, + filter: FilterExpression, + ) -> Result> { + FileReader::do_read_ranges( + prepared.column_infos, + self.scheduler.clone(), + self.cache.clone(), + self.decoder_plugins.clone(), + ranges, + batch_size, + prepared.decoder_projection, + filter, + self.options.decoder_config.clone(), + self.options.batch_size_bytes, + ) + .await + } + + pub(crate) async fn read_prepared_tasks( + &self, + params: ReadBatchParams, + batch_size: u32, + prepared: PreparedProjection, + read_len: u64, + filter: FilterExpression, + ) -> Result + Send>>> { + let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { + if bound > read_len || (bound == read_len && inclusive) { + Err(Error::invalid_input(format!( + "cannot read {params:?} from columns with {read_len} rows" + ))) + } else { + Ok(()) + } + }; + match ¶ms { + ReadBatchParams::Indices(indices) => { + for idx in indices { + match idx { + None => { + return Err(Error::invalid_input("Null value in indices array")); + } + Some(idx) => { + verify_bound(¶ms, idx as u64, true)?; + } + } + } + let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect(); + self.take_rows(indices, batch_size, prepared).await + } + ReadBatchParams::Range(range) => { + verify_bound(¶ms, range.end as u64, false)?; + self.read_range( + range.start as u64..range.end as u64, + batch_size, + prepared, + filter, + ) + .await + } + ReadBatchParams::Ranges(ranges) => { + let mut ranges_u64 = Vec::with_capacity(ranges.len()); + for range in ranges.as_ref() { + verify_bound(¶ms, range.end, false)?; + ranges_u64.push(range.start..range.end); + } + self.read_ranges(ranges_u64, batch_size, prepared, filter) + .await + } + ReadBatchParams::RangeFrom(range) => { + verify_bound(¶ms, range.start as u64, true)?; + self.read_range(range.start as u64..read_len, batch_size, prepared, filter) + .await + } + ReadBatchParams::RangeTo(range) => { + verify_bound(¶ms, range.end as u64, false)?; + self.read_range(0..range.end as u64, batch_size, prepared, filter) + .await + } + ReadBatchParams::RangeFull => { + self.read_range(0..read_len, batch_size, prepared, filter) + .await + } + } + } +} + +impl ProjectedFileReader { + pub(crate) fn base_projection(&self) -> &ReaderProjection { + &self.core.base_projection + } + + pub(crate) async fn read_prepared_tasks( + &self, + params: ReadBatchParams, + batch_size: u32, + prepared: PreparedProjection, + read_len: u64, + filter: FilterExpression, + ) -> Result + Send>>> { + self.core + .read_prepared_tasks(params, batch_size, prepared, read_len, filter) + .await + } + + /// Returns a clone of this reader using a different scheduler. + pub fn with_scheduler(&self, scheduler: Arc) -> Self { + Self { + core: self.core.with_scheduler(scheduler), + } + } + + /// Returns the number of rows in the file. + pub fn num_rows(&self) -> u64 { + self.core.num_rows() + } + + /// Returns the file schema visible to this reader. + pub fn schema(&self) -> &Arc { + self.core.schema() + } + + /// Returns file statistics when this reader has full file metadata. + pub fn file_statistics(&self) -> Option { + self.core.metadata_provider.file_statistics() + } + + #[cfg(test)] + pub(crate) fn metadata_index(&self) -> Option<&Arc> { + match &self.core.metadata_provider { + FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index), + FileMetadataProvider::Full(_) => None, + } + } + + /// Reads a global buffer by index. + pub async fn read_global_buffer(&self, index: u32) -> Result { + self.core.read_global_buffer(index).await + } +} + +impl FileReader { + #[cfg(test)] + fn scheduler(&self) -> Arc { + self.core.scheduler.clone() + } + + pub async fn try_open( + scheduler: FileScheduler, + base_projection: Option, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, + ) -> Result { + match Self::try_open_for_dispatch( + scheduler, + base_projection, + decoder_plugins, + cache, + options, + ) + .await? + { + versions::OpenedFileReader::V1 { + major_version, + minor_version, + } => Err(Error::version_conflict( + "Attempt to use the Lance current-format reader to read a v1 file".to_string(), + major_version, + minor_version, + )), + versions::OpenedFileReader::Current(reader) => Ok(reader), + } + } + + pub(crate) async fn try_open_for_dispatch( + scheduler: FileScheduler, + base_projection: Option, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, + ) -> Result { + let metadata = match Self::read_raw_metadata_for_dispatch(&scheduler).await? { + RawFileMetadataOpen::Legacy { + major_version, + minor_version, + } => { + 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), + ); + Self::try_open_with_file_metadata( + io, + path, + base_projection, + decoder_plugins, + metadata, + cache, + options, + ) + .await + .map(versions::OpenedFileReader::Current) + } + + pub async fn try_open_with_file_metadata( + scheduler: Arc, + path: Path, + base_projection: Option, + decoder_plugins: Arc, + metadata: Arc, + cache: &LanceCache, + options: FileReaderOptions, + ) -> Result { + if metadata.version == ConcreteFileVersion::V1 { + return Err(Error::version_conflict( + "Attempt to use the Lance current-format reader with v1 metadata".to_string(), + metadata.major_version, + metadata.minor_version, + )); + } + let read_projection = versions::read_projection(metadata.version)?; + let has_explicit_projection = base_projection.is_some(); + let base_projection = base_projection.unwrap_or_else(|| { + versions::reader_projection_from_whole_schema(&metadata.file_schema, metadata.version) + }); + if has_explicit_projection { + Self::validate_projection(&base_projection, &metadata)?; + } + let cache = Arc::new(cache.with_key_prefix(path.as_ref())); + let core = DecodeEngine::try_new( + scheduler, + base_projection, + decoder_plugins, + FileMetadataProvider::Full(metadata.clone()), + read_projection, + cache, + options, + )?; + Ok(Self { core, metadata }) + } + + pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result { + match Self::read_raw_metadata_for_dispatch(scheduler).await? { + RawFileMetadataOpen::Legacy { + major_version, + minor_version, + } => Err(Error::version_conflict( + "Attempt to use the Lance current-format reader to read v1 metadata".to_string(), + major_version, + minor_version, + )), + RawFileMetadataOpen::Current { version, metadata } => { + versions::finish_metadata(version, metadata) + } + } + } + + pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result { + let index = Self::read_raw_metadata_index(scheduler).await?; + versions::finish_metadata_index(index) + } + + pub async fn read_metadata_index_with_schema( + scheduler: &FileScheduler, + file_schema: Arc, + num_rows: u64, + ) -> Result { + let index = + Self::read_raw_metadata_index_with_schema(scheduler, file_schema, num_rows).await?; + versions::finish_metadata_index(index) + } + + pub fn version(&self) -> ConcreteFileVersion { + self.metadata.version + } + + async fn prepare(&self, projection: ReaderProjection) -> Result<(PreparedProjection, u64)> { + self.core + .read_projection + .prepare( + &self.core.metadata_provider, + &projection, + &self.core.scheduler, + &self.core.cache, + ) + .await + } + + pub async fn read_tasks( + &self, + params: ReadBatchParams, + batch_size: u32, + projection: Option, + filter: FilterExpression, + ) -> Result + Send>>> { + let projection = projection.unwrap_or_else(|| self.base_projection().clone()); + let (prepared, read_len) = self.prepare(projection).await?; + self.read_prepared_tasks(params, batch_size, prepared, read_len, filter) + .await + } + + pub async fn read_stream_projected( + &self, + params: ReadBatchParams, + batch_size: u32, + batch_readahead: u32, + projection: ReaderProjection, + filter: FilterExpression, + ) -> Result>> { + let schema = projection.schema.clone(); + let tasks = self + .read_tasks(params, batch_size, Some(projection), filter) + .await?; + Ok(tasks_to_record_batch_stream(schema, tasks, batch_readahead)) + } + + pub fn read_stream_projected_blocking( &self, - indices: Vec, + params: ReadBatchParams, batch_size: u32, - prepared: PreparedProjection, - ) -> Result> { - FileReader::do_take_rows( - prepared.column_infos, - self.scheduler.clone(), - self.cache.clone(), - self.decoder_plugins.clone(), - indices, - batch_size, - prepared.decoder_projection, - FilterExpression::no_filter(), - self.options.decoder_config.clone(), - self.options.batch_size_bytes, - ) - .await + projection: Option, + filter: FilterExpression, + ) -> Result> { + let projection = projection.unwrap_or_else(|| self.base_projection().clone()); + Self::validate_projection(&projection, self.metadata())?; + let prepared = self.full_projection(projection); + let read_len = self.core.read_projection.read_length(&prepared)?; + self.read_prepared_blocking(params, batch_size, prepared, read_len, filter) } - async fn read_ranges( + pub async fn read_stream( &self, - ranges: Vec>, + params: ReadBatchParams, batch_size: u32, - prepared: PreparedProjection, + batch_readahead: u32, filter: FilterExpression, - ) -> Result> { - FileReader::do_read_ranges( - prepared.column_infos, - self.scheduler.clone(), - self.cache.clone(), - self.decoder_plugins.clone(), - ranges, + ) -> Result>> { + self.read_stream_projected( + params, batch_size, - prepared.decoder_projection, + batch_readahead, + self.base_projection().clone(), filter, - self.options.decoder_config.clone(), - self.options.batch_size_bytes, ) .await } - - async fn read_tasks( - &self, - params: ReadBatchParams, - batch_size: u32, - projection: Option, - filter: FilterExpression, - ) -> Result + Send>>> { - let projection = projection.unwrap_or_else(|| self.base_projection.clone()); - let prepared = self - .metadata_provider - .prepare_projection(&projection, &self.scheduler, &self.cache) - .await?; - // All projected columns must share a length: the reader combines them - // into rectangular batches. Ordinary files satisfy this (every column - // has `num_rows` rows); files written with `FileWriter::write_column` - // may not, and such columns must be read separately. `read_len` is that - // common length, which `RangeFull`/`RangeFrom` resolve against (rather - // than `num_rows`, the file's longest column). - let read_len = self.prepared_read_length(&prepared)?; - let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { - if bound > read_len || (bound == read_len && inclusive) { - Err(Error::invalid_input(format!( - "cannot read {params:?} from columns with {read_len} rows" - ))) - } else { - Ok(()) - } - }; - match ¶ms { - ReadBatchParams::Indices(indices) => { - for idx in indices { - match idx { - None => { - return Err(Error::invalid_input("Null value in indices array")); - } - Some(idx) => { - verify_bound(¶ms, idx as u64, true)?; - } - } - } - let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect(); - self.take_rows(indices, batch_size, prepared).await - } - ReadBatchParams::Range(range) => { - verify_bound(¶ms, range.end as u64, false)?; - self.read_range( - range.start as u64..range.end as u64, - batch_size, - prepared, - filter, - ) - .await - } - ReadBatchParams::Ranges(ranges) => { - let mut ranges_u64 = Vec::with_capacity(ranges.len()); - for range in ranges.as_ref() { - verify_bound(¶ms, range.end, false)?; - ranges_u64.push(range.start..range.end); - } - self.read_ranges(ranges_u64, batch_size, prepared, filter) - .await - } - ReadBatchParams::RangeFrom(range) => { - verify_bound(¶ms, range.start as u64, true)?; - self.read_range(range.start as u64..read_len, batch_size, prepared, filter) - .await - } - ReadBatchParams::RangeTo(range) => { - verify_bound(¶ms, range.end as u64, false)?; - self.read_range(0..range.end as u64, batch_size, prepared, filter) - .await - } - ReadBatchParams::RangeFull => { - self.read_range(0..read_len, batch_size, prepared, filter) - .await - } - } - } } impl ProjectedFileReader { - /// Opens a data reader backed by indexed column metadata. - /// - /// `base_projection` must be a supported indexed projection. Reads that do - /// not pass an explicit projection use this base projection. pub async fn try_open( scheduler: FileScheduler, base_projection: Option, @@ -2422,13 +2069,16 @@ impl ProjectedFileReader { cache: &LanceCache, options: FileReaderOptions, ) -> Result { - let base_projection = Self::require_indexed_base_projection(base_projection)?; + 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 path = scheduler.reader().path().clone(); - let encodings_io = - LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size); + let io = Arc::new( + LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size), + ); Self::try_open_with_metadata_index( - Arc::new(encodings_io), + io, path, Some(base_projection), decoder_plugins, @@ -2439,11 +2089,6 @@ impl ProjectedFileReader { .await } - /// Opens a data reader from a previously loaded metadata index. - /// - /// `base_projection` must be a supported indexed projection. Use - /// [`Self::try_open_with_file_metadata`] when the default read should cover - /// the whole file schema. pub async fn try_open_with_metadata_index( scheduler: Arc, path: Path, @@ -2453,94 +2098,74 @@ impl ProjectedFileReader { cache: &LanceCache, options: FileReaderOptions, ) -> Result { - let base_projection = Self::require_indexed_base_projection(base_projection)?; + if metadata_index.version == ConcreteFileVersion::V1 { + return Err(Error::version_conflict( + "Attempt to use the Lance projected current-format reader with v1 metadata" + .to_string(), + 0, + 2, + )); + } + let base_projection = base_projection.ok_or_else(|| { + Error::invalid_input("ProjectedReader requires an explicit base projection") + })?; + let read_projection = versions::read_projection(metadata_index.version)?; + read_projection.validate_indexed(&base_projection, &metadata_index)?; let cache = Arc::new(cache.with_key_prefix(path.as_ref())); - let core = FileReadCore::try_new( + let core = DecodeEngine::try_new( scheduler, - Some(base_projection), + base_projection, decoder_plugins, FileMetadataProvider::Indexed(metadata_index), + read_projection, cache, options, )?; Ok(Self { core }) } - fn require_indexed_base_projection( - base_projection: Option, - ) -> Result { - base_projection.ok_or_else(|| { - Error::invalid_input("ProjectedFileReader requires an explicit base projection") - }) - } - pub async fn try_open_with_file_metadata( scheduler: Arc, path: Path, base_projection: Option, decoder_plugins: Arc, - file_metadata: Arc, + metadata: Arc, cache: &LanceCache, options: FileReaderOptions, ) -> Result { + if metadata.version == ConcreteFileVersion::V1 { + return Err(Error::version_conflict( + "Attempt to use the Lance projected current-format reader with v1 metadata" + .to_string(), + metadata.major_version, + metadata.minor_version, + )); + } + let read_projection = versions::read_projection(metadata.version)?; + let has_explicit_projection = base_projection.is_some(); + let base_projection = base_projection.unwrap_or_else(|| { + versions::reader_projection_from_whole_schema(&metadata.file_schema, metadata.version) + }); + if has_explicit_projection { + FileReader::validate_projection(&base_projection, &metadata)?; + } let cache = Arc::new(cache.with_key_prefix(path.as_ref())); - let core = FileReadCore::try_new( + let core = DecodeEngine::try_new( scheduler, base_projection, decoder_plugins, - FileMetadataProvider::Full(file_metadata), + FileMetadataProvider::Full(metadata), + read_projection, cache, options, )?; Ok(Self { core }) } - /// Returns whether a projection can be served by indexed column metadata. - pub fn supports_projection(projection: &ReaderProjection, version: LanceFileVersion) -> bool { - FileMetadataProvider::supports_indexed_projection(projection, version) - } - - /// Returns a clone of this reader using a different scheduler. - pub fn with_scheduler(&self, scheduler: Arc) -> Self { - Self { - core: self.core.with_scheduler(scheduler), - } - } - - /// Returns the Lance file version. - pub fn version(&self) -> LanceFileVersion { - self.core.version() - } - - /// Returns the number of rows in the file. - pub fn num_rows(&self) -> u64 { - self.core.num_rows() - } - - /// Returns the file schema visible to this reader. - pub fn schema(&self) -> &Arc { - self.core.schema() - } - - /// Returns file statistics when this reader has full file metadata. - pub fn file_statistics(&self) -> Option { - self.core.metadata_provider.file_statistics() - } - - #[cfg(test)] - fn metadata_index(&self) -> Option<&Arc> { - match &self.core.metadata_provider { - FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index), - FileMetadataProvider::Full(_) => None, - } - } - - /// Reads a global buffer by index. - pub async fn read_global_buffer(&self, index: u32) -> Result { - self.core.read_global_buffer(index).await + pub fn version(&self) -> ConcreteFileVersion { + self.core.metadata_provider.version() } - /// Creates a stream of read tasks for the requested rows and projection. pub async fn read_tasks( &self, params: ReadBatchParams, @@ -2548,8 +2173,18 @@ impl ProjectedFileReader { projection: Option, filter: FilterExpression, ) -> Result + Send>>> { - self.core - .read_tasks(params, batch_size, projection, filter) + let projection = projection.unwrap_or_else(|| self.base_projection().clone()); + let (prepared, read_len) = self + .core + .read_projection + .prepare( + &self.core.metadata_provider, + &projection, + &self.core.scheduler, + &self.core.cache, + ) + .await?; + self.read_prepared_tasks(params, batch_size, prepared, read_len, filter) .await } } @@ -2604,11 +2239,7 @@ pub fn describe_encoding(page: &pbfile::column_metadata::Page) -> String { } pub trait EncodedBatchReaderExt { - fn try_from_mini_lance( - bytes: Bytes, - schema: &Schema, - version: LanceFileVersion, - ) -> Result + fn try_from_mini_lance(bytes: Bytes, schema: &Schema) -> Result where Self: Sized; fn try_from_self_described_lance(bytes: Bytes) -> Result @@ -2617,16 +2248,13 @@ pub trait EncodedBatchReaderExt { } impl EncodedBatchReaderExt for EncodedBatch { - fn try_from_mini_lance( - bytes: Bytes, - schema: &Schema, - file_version: LanceFileVersion, - ) -> Result + fn try_from_mini_lance(bytes: Bytes, schema: &Schema) -> Result where Self: Sized, { - let projection = ReaderProjection::from_whole_schema(schema, file_version); let footer = FileReader::decode_footer(&bytes)?; + let file_version = FileReader::current_file_version(&footer)?; + let projection = versions::reader_projection_from_whole_schema(schema, file_version); // Next, read the metadata for the columns // This is both the column metadata and the CMO table @@ -2636,11 +2264,7 @@ impl EncodedBatchReaderExt for EncodedBatch { let column_metadatas = FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?; - let file_version: LanceFileVersion = - ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)? - .into(); - - let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?; + let page_table = versions::decode_column_metadata(file_version, &column_metadatas)?; Ok(Self { data: bytes, @@ -2659,15 +2283,13 @@ impl EncodedBatchReaderExt for EncodedBatch { Self: Sized, { let footer = FileReader::decode_footer(&bytes)?; - let file_version: LanceFileVersion = - ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)? - .into(); + let file_version = FileReader::current_file_version(&footer)?; let gbo_table = FileReader::do_decode_gbo_table( &bytes.slice(footer.global_buff_offsets_start as usize..), &footer, - file_version, )?; + versions::validate_global_buffers(file_version, &gbo_table)?; if gbo_table.is_empty() { return Err(Error::internal( "File did not contain any global buffers, schema expected".to_string(), @@ -2678,7 +2300,7 @@ impl EncodedBatchReaderExt for EncodedBatch { let schema_bytes = bytes.slice(schema_start..(schema_start + schema_size)); let (_, schema) = FileReader::decode_schema(schema_bytes)?; - let projection = ReaderProjection::from_whole_schema(&schema, file_version); + let projection = versions::reader_projection_from_whole_schema(&schema, file_version); // Next, read the metadata for the columns // This is both the column metadata and the CMO table @@ -2688,7 +2310,7 @@ impl EncodedBatchReaderExt for EncodedBatch { let column_metadatas = FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?; - let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?; + let page_table = versions::decode_column_metadata(file_version, &column_metadatas)?; Ok(Self { data: bytes, @@ -2730,7 +2352,6 @@ mod tests { }, encoder::{EncodedBatch, EncodingOptions, encode_batch}, format::pb21, - version::LanceFileVersion, }; use lance_io::{stream::RecordBatchStream, utils::CachedFileSize}; use log::debug; @@ -2738,8 +2359,7 @@ mod tests { use tokio::sync::mpsc; use crate::reader::{ - EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, - ReaderProjection, validate_field_length, verify_uniform_lengths, + EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, ReaderProjection, }; use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file}; use crate::version::ConcreteFileVersion; @@ -2747,14 +2367,6 @@ mod tests { use crate::writer::FileWriterOptions; use lance_encoding::decoder::DecoderConfig; - fn footer_version(bytes: &[u8]) -> (u16, u16) { - let version_start = bytes.len() - 8; - ( - u16::from_le_bytes([bytes[version_start], bytes[version_start + 1]]), - u16::from_le_bytes([bytes[version_start + 2], bytes[version_start + 3]]), - ) - } - #[tokio::test] async fn sparse_file_writer_reader_scan_range_and_take_roundtrip() { let fs = FsFixture::default(); @@ -2819,10 +2431,10 @@ mod tests { ) .await .unwrap(); - assert_eq!(file_reader.metadata.column_infos.len(), 2); + assert_eq!(file_reader.metadata().column_infos.len(), 2); assert!( file_reader - .metadata + .metadata() .column_infos .iter() .flat_map(|column| column.page_infos.iter()) @@ -2882,7 +2494,15 @@ mod tests { assert_eq!(take, vec![batch.take(&indices).unwrap()]); } - async fn create_some_file(fs: &FsFixture, version: LanceFileVersion) -> WrittenFile { + fn footer_version(bytes: &[u8]) -> (u16, u16) { + let version_start = bytes.len() - 8; + ( + u16::from_le_bytes([bytes[version_start], bytes[version_start + 1]]), + u16::from_le_bytes([bytes[version_start + 2], bytes[version_start + 3]]), + ) + } + + async fn create_some_file(fs: &FsFixture, version: ConcreteFileVersion) -> WrittenFile { let location_type = DataType::Struct(Fields::from(vec![ Field::new("x", DataType::Float64, true), Field::new("y", DataType::Float64, true), @@ -2894,18 +2514,12 @@ mod tests { .col("location", array::rand_type(&location_type)) .col("categories", array::rand_type(&categories_type)) .col("binary", array::rand_type(&DataType::Binary)); - if version <= LanceFileVersion::V2_0 { + if version == ConcreteFileVersion::V2_0 { reader = reader.col("large_bin", array::rand_type(&DataType::LargeBinary)); } let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100)); - write_lance_file( - reader, - fs, - ConcreteFileVersion::from(version), - FileWriterOptions::default(), - ) - .await + write_lance_file(reader, fs, version, FileWriterOptions::default()).await } async fn create_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile { @@ -3042,7 +2656,7 @@ mod tests { async fn test_round_trip() { let fs = FsFixture::default(); - let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await; + let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await; let file_size = fs.object_store.size(&fs.tmp_path).await.unwrap() as usize; let footer = fs @@ -3058,7 +2672,7 @@ mod tests { crate::determine_file_version(&fs.object_store, &fs.tmp_path, Some(file_size)) .await .unwrap(), - LanceFileVersion::V2_0 + ConcreteFileVersion::V2_0 ); for read_size in [32, 1024, 1024 * 1024] { @@ -3077,6 +2691,13 @@ mod tests { .await .unwrap(); + assert_eq!( + ( + file_reader.metadata().major_version, + file_reader.metadata().minor_version + ), + (0, 3) + ); let schema = file_reader.schema(); assert_eq!(schema.metadata.get("foo").unwrap(), "bar"); @@ -3098,7 +2719,7 @@ mod tests { #[test_log::test(tokio::test)] async fn test_encoded_batch_round_trip( // TODO: Add V2_1 (currently fails) - #[values(LanceFileVersion::V2_0)] version: LanceFileVersion, + #[values(ConcreteFileVersion::V2_0)] version: ConcreteFileVersion, ) { let data = gen_batch() .col("x", array::rand::()) @@ -3115,23 +2736,19 @@ mod tests { buffer_alignment: 64, }; - let encoding_strategy = crate::versions::v2_0::encoding_strategy(); + let encoding_strategy = lance_encoding::array_encoding::ArrayFieldEncodingStrategy::new(); let encoded_batch = encode_batch( &data, lance_schema.clone(), - encoding_strategy.as_ref(), + &encoding_strategy, &encoding_options, ) .await .unwrap(); // Test self described - let bytes = versions::encode_self_described_batch( - ConcreteFileVersion::from(version), - &encoded_batch, - ) - .unwrap(); + let bytes = versions::encode_self_described_batch(version, &encoded_batch).unwrap(); assert_eq!(footer_version(&bytes), (2, 0)); let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap(); @@ -3150,12 +2767,10 @@ mod tests { assert_eq!(data, decoded); // Test mini - let bytes = versions::encode_mini_batch(ConcreteFileVersion::from(version), &encoded_batch) - .unwrap(); + let bytes = versions::encode_mini_batch(version, &encoded_batch).unwrap(); assert_eq!(footer_version(&bytes), (2, 0)); let decoded_batch = - EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref(), LanceFileVersion::V2_0) - .unwrap(); + EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref()).unwrap(); let decoded = decode_batch( &decoded_batch, &FilterExpression::no_filter(), @@ -3173,8 +2788,12 @@ mod tests { #[rstest] #[test_log::test(tokio::test)] async fn test_projection( - #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)] - version: LanceFileVersion, + #[values( + ConcreteFileVersion::V2_0, + ConcreteFileVersion::V2_1, + ConcreteFileVersion::V2_2 + )] + version: ConcreteFileVersion, ) { let fs = FsFixture::default(); @@ -3221,15 +2840,15 @@ mod tests { let projected_schema = written_file.schema.project(&columns).unwrap(); let projection = if use_field_ids { - ReaderProjection::from_field_ids( - file_reader.metadata.version(), + versions::reader_projection_from_field_ids( + file_reader.metadata().version(), &projected_schema, &field_id_mapping, ) .unwrap() } else { - ReaderProjection::from_column_names( - file_reader.metadata.version(), + versions::reader_projection_from_column_names( + file_reader.metadata().version(), &written_file.schema, &columns, ) @@ -3351,8 +2970,8 @@ mod tests { .open_file(&fs.tmp_path, &CachedFileSize::unknown()) .await .unwrap(); - let projection = ReaderProjection::from_column_names( - LanceFileVersion::V2_1, + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, &written_file.schema, &["c10"], ) @@ -3410,8 +3029,8 @@ mod tests { let fs = FsFixture::default(); let written_file = create_wide_direct_file(&fs, 512).await; - let projection = ReaderProjection::from_column_names( - LanceFileVersion::V2_1, + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, &written_file.schema, &["c0"], ) @@ -3588,20 +3207,13 @@ mod tests { async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() { let fs = FsFixture::default(); let written_file = create_wide_fixed_size_list_file(&fs, 512).await; - let projection = ReaderProjection::from_column_names( - LanceFileVersion::V2_1, + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, &written_file.schema, &["c17", "c509"], ) .unwrap(); - assert!(ProjectedFileReader::supports_projection( - &projection, - LanceFileVersion::V2_1 - )); - assert!(!ProjectedFileReader::supports_projection( - &projection, - LanceFileVersion::V2_0 - )); + assert!(projection.prefers_indexed_metadata(512)); assert_lazy_projection_matches_eager_and_reads_metadata_subset( &fs, projection, @@ -3610,12 +3222,44 @@ mod tests { .await; } + #[tokio::test] + async fn test_v2_0_rejects_indexed_metadata_reader() { + let fs = FsFixture::default(); + let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await; + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_0, + &written_file.schema, + &["score"], + ) + .unwrap(); + assert!(projection.prefers_indexed_metadata(100)); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let err = ProjectedFileReader::try_open( + file_scheduler, + Some(projection), + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap_err(); + assert!( + matches!(err, lance_core::Error::NotSupported { .. }), + "expected V2.0 indexed metadata open to fail, got {err:?}" + ); + } + #[tokio::test] async fn test_lazy_reader_nested_projection_compacts_physical_columns() { let fs = FsFixture::default(); let written_file = create_wide_structural_file(&fs, 128).await; - let projection = ReaderProjection::from_column_names( - LanceFileVersion::V2_1, + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, &written_file.schema, &["s97.y", "l4", "s3"], ) @@ -3641,10 +3285,7 @@ mod tests { .any(|indices| indices[0] > indices[1]), "the projection must reorder physical columns to exercise compact remapping" ); - assert!(ProjectedFileReader::supports_projection( - &projection, - LanceFileVersion::V2_1 - )); + assert!(projection.prefers_indexed_metadata(128 * 4)); let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset( &fs, projection, "nested", ) @@ -3692,19 +3333,16 @@ mod tests { #[tokio::test] async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) { let fs = FsFixture::default(); - let written_file = create_some_file(&fs, LanceFileVersion::V2_1).await; + let written_file = create_some_file(&fs, ConcreteFileVersion::V2_1).await; - let ordinary_projection = ReaderProjection::from_column_names( - LanceFileVersion::V2_1, + let ordinary_projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, &written_file.schema, &["location.x"], ) .unwrap(); assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1); - assert!(ProjectedFileReader::supports_projection( - &ordinary_projection, - LanceFileVersion::V2_1 - )); + assert!(ordinary_projection.prefers_indexed_metadata(100)); let file_scheduler = fs .scheduler @@ -3729,10 +3367,7 @@ mod tests { Arc::make_mut(&mut projection.schema).fields[0] .metadata .insert(metadata_key.to_string(), "true".to_string()); - assert!(!ProjectedFileReader::supports_projection( - &projection, - LanceFileVersion::V2_1 - )); + assert!(!projection.prefers_indexed_metadata(100)); let err = ProjectedFileReader::try_open( file_scheduler, @@ -3749,7 +3384,7 @@ mod tests { ); } - // The projection-length validation lives in `FileReadCore`, shared by the + // The projection-length validation lives in `DecodeEngine`, shared by the // eager and the lazy (indexed) metadata providers. The indexed provider loads // only the projected columns and renumbers them 0..N, so this checks that the // renumbered `column_infos`/`column_indices` still line up for the length @@ -3767,8 +3402,7 @@ mod tests { let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap(); let fs = FsFixture::default(); - let mut writer = versions::create_writer( - ConcreteFileVersion::V2_1, + let mut writer = versions::v2_1::create_writer( fs.object_store.create(&fs.tmp_path).await.unwrap(), lance_schema.clone(), FileWriterOptions::default(), @@ -3792,9 +3426,12 @@ mod tests { .unwrap(); let cache = test_cache(); let open_indexed = |names: &[&str]| { - let projection = - ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names) - .unwrap(); + let projection = versions::reader_projection_from_column_names( + ConcreteFileVersion::V2_1, + &lance_schema, + names, + ) + .unwrap(); ProjectedFileReader::try_open( file_scheduler.clone(), Some(projection), @@ -3857,7 +3494,7 @@ mod tests { async fn test_compressing_buffer() { let fs = FsFixture::default(); - let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await; + let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await; let file_scheduler = fs .scheduler .open_file(&fs.tmp_path, &CachedFileSize::unknown()) @@ -3912,7 +3549,7 @@ mod tests { #[tokio::test] async fn test_read_all() { let fs = FsFixture::default(); - let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await; + let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await; let total_rows = data.iter().map(|batch| batch.num_rows()).sum::(); let file_scheduler = fs @@ -3949,8 +3586,12 @@ mod tests { #[rstest] #[tokio::test] async fn test_blocking_take( - #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)] - version: LanceFileVersion, + #[values( + ConcreteFileVersion::V2_0, + ConcreteFileVersion::V2_1, + ConcreteFileVersion::V2_2 + )] + version: ConcreteFileVersion, ) { let fs = FsFixture::default(); let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await; @@ -3963,7 +3604,10 @@ mod tests { .unwrap(); let file_reader = FileReader::try_open( file_scheduler.clone(), - Some(ReaderProjection::from_column_names(version, &schema, &["score"]).unwrap()), + Some( + versions::reader_projection_from_column_names(version, &schema, &["score"]) + .unwrap(), + ), Arc::::default(), &test_cache(), FileReaderOptions::default(), @@ -3994,7 +3638,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_drop_in_progress() { let fs = FsFixture::default(); - let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await; + let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await; let total_rows = data.iter().map(|batch| batch.num_rows()).sum::(); let file_scheduler = fs @@ -4043,7 +3687,7 @@ mod tests { // if the stream was dropped before it finished. let fs = FsFixture::default(); - let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await; + let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await; let total_rows = written_file .data .iter() @@ -4065,11 +3709,11 @@ mod tests { .await .unwrap(); - let projection = - ReaderProjection::from_whole_schema(&written_file.schema, LanceFileVersion::V2_0); - let column_infos = file_reader - .collect_columns_from_projection(&projection) - .unwrap(); + let projection = versions::reader_projection_from_whole_schema( + &written_file.schema, + ConcreteFileVersion::V2_0, + ); + let column_infos = file_reader.metadata().column_infos.clone(); let mut decode_scheduler = DecodeBatchScheduler::try_new( &projection.schema, &projection.column_indices, @@ -4077,7 +3721,7 @@ mod tests { &vec![], total_rows as u64, Arc::::default(), - file_reader.core.scheduler.clone(), + file_reader.scheduler(), test_cache(), &FilterExpression::no_filter(), &DecoderConfig::default(), @@ -4097,14 +3741,14 @@ mod tests { range, &FilterExpression::no_filter(), tx, - file_reader.core.scheduler.clone(), + file_reader.scheduler(), ) } #[tokio::test] async fn test_read_empty_range() { let fs = FsFixture::default(); - create_some_file(&fs, LanceFileVersion::V2_0).await; + create_some_file(&fs, ConcreteFileVersion::V2_0).await; let file_scheduler = fs .scheduler @@ -4162,8 +3806,7 @@ mod tests { )])) .unwrap(); - let mut file_writer = versions::create_writer( - ConcreteFileVersion::V2_1, + let mut file_writer = versions::v2_1::create_writer( fs.object_store.create(&fs.tmp_path).await.unwrap(), lance_schema, FileWriterOptions::default(), @@ -4232,7 +3875,7 @@ mod tests { #[tokio::test] async fn test_read_global_buffer_no_user_buffers() { let fs = FsFixture::default(); - create_some_file(&fs, LanceFileVersion::V2_1).await; + create_some_file(&fs, ConcreteFileVersion::V2_1).await; let file_scheduler = fs .scheduler @@ -4261,12 +3904,12 @@ mod tests { #[tokio::test] async fn test_deep_size_of_includes_column_metadata( #[values( - LanceFileVersion::V2_0, - LanceFileVersion::V2_1, - LanceFileVersion::V2_2, - LanceFileVersion::V2_3 + ConcreteFileVersion::V2_0, + ConcreteFileVersion::V2_1, + ConcreteFileVersion::V2_2, + ConcreteFileVersion::V2_3 )] - version: LanceFileVersion, + version: ConcreteFileVersion, ) { // Regression test: CachedFileMetadata::deep_size_of must account for // column_metadatas and column_infos, otherwise the moka cache weigher @@ -4362,21 +4005,11 @@ mod tests { let run = |dt: DataType, indices: &[u32], lengths: Vec| -> lance_core::Result { let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]); let schema = Schema::try_from(&arrow).unwrap(); - let column_len = |c: usize| Ok(lengths[c]); - let mut cursor = 0usize; - let mut field_lengths = Vec::new(); - for field in &schema.fields { - let rows = validate_field_length( - field, - is_structural, - true, - indices, - &mut cursor, - &column_len, - )?; - field_lengths.push((field.name.as_str(), rows)); + if is_structural { + versions::v2_1::test_projection_length(&schema, indices, &lengths) + } else { + versions::v2_0::test_projection_length(&schema, indices, &lengths) } - verify_uniform_lengths(&field_lengths) }; let struct_ty = || { @@ -4416,24 +4049,14 @@ mod tests { schema: Arc::new(schema), column_indices: vec![0], }; - let column_len = |column: usize| { - assert_eq!(column, 0); - Ok(3) - }; - let mut cursor = 0usize; - - let rows = validate_field_length( - &projection.schema.fields[0], - false, - true, + let rows = versions::v2_0::test_projection_length( + &projection.schema, &projection.column_indices, - &mut cursor, - &column_len, + &[3], ) .unwrap(); assert_eq!(rows, 3); - assert_eq!(cursor, 1); } #[test] @@ -4445,16 +4068,11 @@ mod tests { -> lance_core::Result { let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]); let schema = Schema::try_from(&arrow).unwrap(); - let column_len = |c: usize| Ok(lengths[c]); - let mut cursor = 0usize; - validate_field_length( - &schema.fields[0], - is_structural, - true, - indices, - &mut cursor, - &column_len, - ) + if is_structural { + versions::v2_1::test_projection_length(&schema, indices, &lengths) + } else { + versions::v2_0::test_projection_length(&schema, indices, &lengths) + } }; // A list's items have a different cardinality than its rows; that gap diff --git a/rust/lance-file/src/reader/structural.rs b/rust/lance-file/src/reader/structural.rs new file mode 100644 index 00000000000..18bda6eb933 --- /dev/null +++ b/rust/lance-file/src/reader/structural.rs @@ -0,0 +1,445 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use bytes::Bytes; +use lance_core::{ + Error, Result, + cache::LanceCache, + datatypes::{Field, Schema}, +}; +use lance_encoding::{ + EncodingsIo, + decoder::{ColumnInfo, PageEncoding, PageInfo}, + format::{pb, pb21}, +}; +use prost::{Message, Name}; + +use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, FileMetadataIndex, FileMetadataProvider, FileReader, PreparedProjection, + ReadProjection, ReaderProjection, normalized_column_num_rows, verify_uniform_lengths, + }, + writer::PAGE_BUFFER_ALIGNMENT, +}; + +fn fetch_encoding(encoding: &pbfile::Encoding) -> Result { + match &encoding.location { + Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source( + "Indirect file encodings are not supported".into(), + )), + Some(pbfile::encoding::Location::Direct(encoding)) => { + let envelope = prost_types::Any::decode(Bytes::from(encoding.encoding.clone())) + .map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(), + ) + })?; + envelope.to_msg::().map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding: {error}", M::NAME).into(), + ) + }) + } + Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source( + format!("Missing {} encoding description", M::NAME).into(), + )), + None => Err(Error::invalid_input_source( + format!("Missing {} encoding location", M::NAME).into(), + )), + } +} + +/// Decode structural page syntax before an exact reader validates its grammar. +pub fn decode_page_layout( + column_index: u32, + page_index: usize, + page: &pbfile::column_metadata::Page, +) -> Result { + fetch_encoding(page.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!( + "Column {} page {} is missing its encoding", + column_index, page_index + ) + .into(), + ) + })?) +} + +/// Build normalized page metadata after exact grammar validation. +pub fn build_page_info( + column_index: u32, + page_index: usize, + page: &pbfile::column_metadata::Page, + page_layout: pb21::PageLayout, +) -> Result { + if page.buffer_offsets.len() != page.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} page {} has {} buffer offsets but {} buffer sizes", + column_index, + page_index, + page.buffer_offsets.len(), + page.buffer_sizes.len() + ) + .into(), + )); + } + let buffer_offsets_and_sizes = Arc::from( + page.buffer_offsets + .iter() + .zip(&page.buffer_sizes) + .map(|(offset, size)| { + if offset % PAGE_BUFFER_ALIGNMENT as u64 != 0 { + return Err(Error::invalid_input_source( + format!( + "Column {} page {} buffer offset {} is not aligned to {} bytes", + column_index, page_index, offset, PAGE_BUFFER_ALIGNMENT + ) + .into(), + )); + } + Ok((*offset, *size)) + }) + .collect::>>()?, + ); + Ok(PageInfo { + buffer_offsets_and_sizes, + encoding: PageEncoding::Structural(page_layout), + num_rows: page.length, + priority: page.priority, + }) +} + +/// Finish normalized column metadata after all pages have been validated. +pub fn build_column_info( + column_index: u32, + metadata: &pbfile::ColumnMetadata, + page_infos: Vec, +) -> Result> { + if metadata.buffer_offsets.len() != metadata.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} has {} buffer offsets but {} buffer sizes", + column_index, + metadata.buffer_offsets.len(), + metadata.buffer_sizes.len() + ) + .into(), + )); + } + let buffer_offsets_and_sizes = Arc::from( + metadata + .buffer_offsets + .iter() + .zip(&metadata.buffer_sizes) + .map(|(offset, size)| (*offset, *size)) + .collect::>(), + ); + let encoding: pb::ColumnEncoding = + fetch_encoding(metadata.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Column {} is missing its encoding", column_index).into(), + ) + })?)?; + Ok(Arc::new(ColumnInfo { + index: column_index, + page_infos: Arc::from(page_infos), + buffer_offsets_and_sizes, + encoding, + })) +} + +/// Validate the structural global-buffer alignment contract. +pub fn validate_global_buffers(buffers: &[BufferDescriptor]) -> Result<()> { + for (buffer_index, buffer) in buffers.iter().enumerate() { + if buffer.position % PAGE_BUFFER_ALIGNMENT as u64 != 0 { + return Err(Error::invalid_input_source( + format!( + "Global buffer {} position {} is not aligned to {} bytes", + buffer_index, buffer.position, PAGE_BUFFER_ALIGNMENT + ) + .into(), + )); + } + } + Ok(()) +} + +fn field_column_shape(field: &Field) -> (bool, bool) { + if field.is_blob() || field.is_packed_struct() { + return (true, false); + } + (field.children.is_empty(), !field.children.is_empty()) +} + +/// Count physical columns in the leaf-column layout used by v2.1+. +pub fn physical_column_count(field: &Field) -> usize { + if field.children.is_empty() || field.is_blob() || field.is_packed_struct() { + 1 + } else { + field.children.iter().map(physical_column_count).sum() + } +} + +fn append_physical_fields( + fields: &[Field], + field_ids: &mut Vec, + column_indices: &mut Vec, + next_column: &mut i32, +) { + for field in fields { + if field.children.is_empty() || field.is_blob() || field.is_packed_struct() { + field_ids.push(field.id); + column_indices.push(*next_column); + *next_column += 1; + } else { + append_physical_fields(&field.children, field_ids, column_indices, next_column); + } + } +} + +/// Build persisted field-to-column entries for the v2.1+ leaf layout. +pub fn data_file_columns(schema: &Schema) -> (Vec, Vec) { + let mut field_ids = Vec::new(); + let mut column_indices = Vec::new(); + append_physical_fields(&schema.fields, &mut field_ids, &mut column_indices, &mut 0); + (field_ids, column_indices) +} + +/// Build the field-id lookup for the v2.1+ leaf layout. +pub fn field_id_to_column_index(schema: &Schema) -> BTreeMap { + let (field_ids, column_indices) = data_file_columns(schema); + field_ids + .into_iter() + .zip(column_indices) + .filter_map(|(field_id, column_index)| { + (column_index >= 0).then_some((field_id as u32, column_index as u32)) + }) + .collect() +} + +fn append_field_ids( + fields: &[Field], + field_id_to_column_index: &BTreeMap, + column_indices: &mut Vec, +) { + for field in fields { + let (contributes, recurse) = field_column_shape(field); + if contributes + && let Some(column_index) = field_id_to_column_index.get(&(field.id as u32)).copied() + { + column_indices.push(column_index); + } + if recurse { + append_field_ids(&field.children, field_id_to_column_index, column_indices); + } + } +} + +/// Build the leaf-column projection selected by a v2.1+ exact reader. +pub fn projection_from_field_ids( + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> ReaderProjection { + let mut column_indices = Vec::new(); + append_field_ids( + &schema.fields, + field_id_to_column_index, + &mut column_indices, + ); + ReaderProjection { + schema: Arc::new(schema.clone()), + column_indices, + } +} + +/// Project names using a caller-selected leaf-column mapping. +pub fn projection_from_column_names( + schema: &Schema, + column_names: &[&str], + field_id_to_column_index: &BTreeMap, +) -> Result { + let projected = schema.project(column_names)?; + Ok(projection_from_field_ids( + &projected, + field_id_to_column_index, + )) +} + +fn children_share_parent_length(field: &Field) -> bool { + field.logical_type.is_struct() +} + +fn validate_field_length Result>( + field: &Field, + comparable: bool, + column_indices: &[u32], + cursor: &mut usize, + column_len: &F, +) -> Result { + let (contributes, recurse) = field_column_shape(field); + let mut field_rows = None; + if contributes { + let column = *column_indices.get(*cursor).ok_or_else(|| { + Error::invalid_input(format!( + "projection supplied fewer column indices than its fields require (ran out at field '{}')", + field.name + )) + })?; + *cursor += 1; + field_rows = Some(column_len(column as usize)?); + } + if recurse { + let enforce_children = comparable && children_share_parent_length(field); + for child in &field.children { + let child_rows = + validate_field_length(child, enforce_children, column_indices, cursor, column_len)?; + let expected = *field_rows.get_or_insert(child_rows); + if enforce_children && child_rows != expected { + return Err(Error::invalid_input(format!( + "cannot read field '{}': its children have differing lengths (child '{}' has {} rows, but the field has {}); a struct's children must all have the same length", + field.name, child.name, child_rows, expected + ))); + } + } + } + field_rows.ok_or_else(|| { + Error::invalid_input(format!( + "projected field '{}' maps to no columns", + field.name + )) + }) +} + +/// Determine the normalized logical read length for a structural projection. +pub fn prepared_read_length(prepared: &PreparedProjection) -> Result { + let column_len = |column: usize| { + let info = prepared.column_infos.get(column).ok_or_else(|| { + Error::invalid_input(format!( + "projection references column index {} but only {} columns are available", + column, + prepared.column_infos.len() + )) + })?; + normalized_column_num_rows(info) + }; + projection_length( + &prepared.decoder_projection.schema, + &prepared.decoder_projection.column_indices, + &column_len, + ) +} + +fn projection_length Result>( + schema: &Schema, + column_indices: &[u32], + column_len: &F, +) -> Result { + let mut cursor = 0; + let mut field_lengths = Vec::with_capacity(schema.fields.len()); + for field in &schema.fields { + let rows = validate_field_length(field, true, column_indices, &mut cursor, column_len)?; + field_lengths.push((field.name.as_str(), rows)); + } + if cursor != column_indices.len() { + return Err(Error::invalid_input(format!( + "projection supplied {} column indices but its fields require {}", + column_indices.len(), + cursor + ))); + } + verify_uniform_lengths(&field_lengths) +} + +pub type DecodeColumn = fn(u32, &pbfile::ColumnMetadata) -> Result>; + +#[derive(Debug)] +struct StructuralReadProjection { + decode_column: DecodeColumn, +} + +/// Compose structural projection execution with an exact column grammar. +pub fn read_projection(decode_column: DecodeColumn) -> Arc { + Arc::new(StructuralReadProjection { decode_column }) +} + +#[async_trait] +impl ReadProjection for StructuralReadProjection { + fn validate_indexed( + &self, + projection: &ReaderProjection, + metadata_index: &FileMetadataIndex, + ) -> Result<()> { + FileMetadataProvider::validate_indexed_projection_structure(projection, metadata_index)?; + if FileMetadataProvider::projection_matches_indexed_metadata(projection) { + Ok(()) + } else { + Err(FileMetadataProvider::indexed_projection_error( + projection, + metadata_index, + )) + } + } + + fn read_length(&self, prepared: &PreparedProjection) -> Result { + prepared_read_length(prepared) + } + + async fn prepare( + &self, + metadata_provider: &FileMetadataProvider, + projection: &ReaderProjection, + io: &Arc, + cache: &Arc, + ) -> Result<(PreparedProjection, u64)> { + let prepared = match metadata_provider { + FileMetadataProvider::Full(metadata) => { + FileReader::validate_projection(projection, metadata)?; + PreparedProjection { + column_infos: metadata.column_infos.clone(), + decoder_projection: projection.clone(), + } + } + FileMetadataProvider::Indexed(metadata_index) => { + self.validate_indexed(projection, metadata_index)?; + let column_infos = FileMetadataProvider::load_indexed_column_infos( + metadata_index, + io, + cache, + &projection.column_indices, + self.decode_column, + ) + .await?; + PreparedProjection { + column_infos, + decoder_projection: ReaderProjection { + schema: projection.schema.clone(), + column_indices: (0..projection.column_indices.len()) + .map(|index| index as u32) + .collect(), + }, + } + } + }; + let read_len = self.read_length(&prepared)?; + Ok((prepared, read_len)) + } +} + +#[cfg(test)] +pub fn test_projection_length( + schema: &Schema, + column_indices: &[u32], + column_lengths: &[u64], +) -> Result { + projection_length(schema, column_indices, &|column| { + column_lengths.get(column).copied().ok_or_else(|| { + Error::invalid_input(format!("missing synthetic length for column {column}")) + }) + }) +} diff --git a/rust/lance-file/src/version.rs b/rust/lance-file/src/version.rs index 16f4a7e4b5d..4c53583a835 100644 --- a/rust/lance-file/src/version.rs +++ b/rust/lance-file/src/version.rs @@ -318,7 +318,7 @@ mod tests { crate::determine_file_version(&object_store, &path, Some(footer.len())) .await .unwrap(), - LanceFileVersion::Legacy + ConcreteFileVersion::V1 ); } } diff --git a/rust/lance-file/src/versions/mod.rs b/rust/lance-file/src/versions/mod.rs index 5a02ba79630..410608ee312 100644 --- a/rust/lance-file/src/versions/mod.rs +++ b/rust/lance-file/src/versions/mod.rs @@ -1,17 +1,29 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Exact Lance file-version implementations. +//! Exact file-format composition roots. +//! +//! Each version module lists the mechanisms used by that file format. Callers +//! resolve release selectors before entering this module. Prefer APIs under a +//! concrete module such as [`v2_1`] when the version is statically known. The +//! root functions perform the single exhaustive dispatch for runtime versions. use bytes::Bytes; use lance_core::{Error, Result, datatypes::Schema}; -use lance_encoding::encoder::EncodedBatch; -use lance_io::traits::Writer; +use lance_encoding::{ + decoder::{ColumnInfo, DecoderPlugins, PageInfo}, + encoder::EncodedBatch, +}; +use lance_io::{scheduler::FileScheduler, traits::Writer}; +use std::{collections::BTreeMap, future::Future, sync::Arc}; use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, CachedFileMetadata, FileMetadataIndex, FileMetadataProvider, FileReader, + FileReaderOptions, ProjectedFileReader, RawFileMetadata, ReadProjection, ReaderProjection, + }, version::ConcreteFileVersion, writer::{FileWriter, FileWriterOptions}, }; +use lance_core::cache::LanceCache; pub mod v1; pub mod v2_0; @@ -19,6 +31,265 @@ pub mod v2_1; pub mod v2_2; pub mod v2_3; +pub(crate) fn read_projection(version: ConcreteFileVersion) -> Result> { + match version { + ConcreteFileVersion::V1 => Err(Error::internal( + "current reader composition received Lance v1".to_string(), + )), + ConcreteFileVersion::V2_0 => Ok(v2_0::read_projection()), + ConcreteFileVersion::V2_1 => Ok(v2_1::read_projection()), + ConcreteFileVersion::V2_2 => Ok(v2_2::read_projection()), + ConcreteFileVersion::V2_3 => Ok(v2_3::read_projection()), + } +} + +/// A self-described file reader selected by the exact footer version. +pub enum OpenedFileReader { + /// A v1 file. The persisted footer numbers are retained for diagnostics. + V1 { + /// The major version stored in the footer. + major_version: u16, + /// The minor version stored in the footer. + minor_version: u16, + }, + /// A current-format reader selected from the exact footer identity. + Current(FileReader), +} + +pub(crate) fn finish_metadata( + version: ConcreteFileVersion, + metadata: RawFileMetadata, +) -> Result { + match version { + ConcreteFileVersion::V1 => Err(Error::internal( + "current metadata dispatch received a Lance v1 file".to_string(), + )), + ConcreteFileVersion::V2_0 => v2_0::finish_metadata(metadata), + ConcreteFileVersion::V2_1 => v2_1::finish_metadata(metadata), + ConcreteFileVersion::V2_2 => v2_2::finish_metadata(metadata), + ConcreteFileVersion::V2_3 => v2_3::finish_metadata(metadata), + } +} + +pub(crate) fn finish_metadata_index(index: FileMetadataIndex) -> Result { + match index.version { + ConcreteFileVersion::V1 => Err(Error::version_conflict( + "Attempt to use the Lance current-format reader with a v1 metadata index".to_string(), + 0, + 2, + )), + ConcreteFileVersion::V2_0 => v2_0::finish_metadata_index(index), + ConcreteFileVersion::V2_1 => v2_1::finish_metadata_index(index), + ConcreteFileVersion::V2_2 => v2_2::finish_metadata_index(index), + ConcreteFileVersion::V2_3 => v2_3::finish_metadata_index(index), + } +} + +pub(crate) fn decode_column_metadata( + version: ConcreteFileVersion, + column_metadatas: &[pbfile::ColumnMetadata], +) -> Result>> { + match version { + ConcreteFileVersion::V1 => Err(Error::not_supported( + "self-described batches are not part of the Lance v1 grammar".to_string(), + )), + ConcreteFileVersion::V2_0 => v2_0::decode_column_metadata(column_metadatas), + ConcreteFileVersion::V2_1 => v2_1::decode_column_metadata(column_metadatas), + ConcreteFileVersion::V2_2 => v2_2::decode_column_metadata(column_metadatas), + ConcreteFileVersion::V2_3 => v2_3::decode_column_metadata(column_metadatas), + } +} + +pub(crate) fn validate_global_buffers( + version: ConcreteFileVersion, + buffers: &[BufferDescriptor], +) -> Result<()> { + match version { + ConcreteFileVersion::V1 => Ok(()), + ConcreteFileVersion::V2_0 => v2_0::validate_global_buffers(buffers), + ConcreteFileVersion::V2_1 => v2_1::validate_global_buffers(buffers), + ConcreteFileVersion::V2_2 => v2_2::validate_global_buffers(buffers), + ConcreteFileVersion::V2_3 => v2_3::validate_global_buffers(buffers), + } +} + +pub fn reader_projection_from_field_ids( + version: ConcreteFileVersion, + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> Result { + Ok(match version { + ConcreteFileVersion::V1 => v1::projection_from_field_ids(schema, field_id_to_column_index), + ConcreteFileVersion::V2_0 => { + v2_0::projection_from_field_ids(schema, field_id_to_column_index) + } + ConcreteFileVersion::V2_1 => { + v2_1::projection_from_field_ids(schema, field_id_to_column_index) + } + ConcreteFileVersion::V2_2 => { + v2_2::projection_from_field_ids(schema, field_id_to_column_index) + } + ConcreteFileVersion::V2_3 => { + v2_3::projection_from_field_ids(schema, field_id_to_column_index) + } + }) +} + +pub fn reader_projection_from_whole_schema( + schema: &Schema, + version: ConcreteFileVersion, +) -> ReaderProjection { + match version { + ConcreteFileVersion::V1 => v1::projection_from_whole_schema(schema), + ConcreteFileVersion::V2_0 => v2_0::projection_from_whole_schema(schema), + ConcreteFileVersion::V2_1 => v2_1::projection_from_whole_schema(schema), + ConcreteFileVersion::V2_2 => v2_2::projection_from_whole_schema(schema), + ConcreteFileVersion::V2_3 => v2_3::projection_from_whole_schema(schema), + } +} + +pub fn reader_projection_from_column_names( + version: ConcreteFileVersion, + schema: &Schema, + column_names: &[&str], +) -> Result { + match version { + ConcreteFileVersion::V1 => v1::projection_from_column_names(schema, column_names), + ConcreteFileVersion::V2_0 => v2_0::projection_from_column_names(schema, column_names), + ConcreteFileVersion::V2_1 => v2_1::projection_from_column_names(schema, column_names), + ConcreteFileVersion::V2_2 => v2_2::projection_from_column_names(schema, column_names), + ConcreteFileVersion::V2_3 => v2_3::projection_from_column_names(schema, column_names), + } +} + +/// Count the physical columns represented by one field in an exact grammar. +pub fn physical_column_count( + version: ConcreteFileVersion, + field: &lance_core::datatypes::Field, +) -> usize { + match version { + ConcreteFileVersion::V1 => v1::physical_column_count(field), + ConcreteFileVersion::V2_0 => v2_0::physical_column_count(field), + ConcreteFileVersion::V2_1 => v2_1::physical_column_count(field), + ConcreteFileVersion::V2_2 => v2_2::physical_column_count(field), + ConcreteFileVersion::V2_3 => v2_3::physical_column_count(field), + } +} + +/// Build persisted field-to-column entries for an exact grammar. +pub fn data_file_columns(version: ConcreteFileVersion, schema: &Schema) -> (Vec, Vec) { + match version { + ConcreteFileVersion::V1 => v1::data_file_columns(schema), + ConcreteFileVersion::V2_0 => v2_0::data_file_columns(schema), + ConcreteFileVersion::V2_1 => v2_1::data_file_columns(schema), + ConcreteFileVersion::V2_2 => v2_2::data_file_columns(schema), + ConcreteFileVersion::V2_3 => v2_3::data_file_columns(schema), + } +} + +/// Copy one column's external metadata and buffers according to the exact file +/// grammar. +/// +/// The caller supplies the version-free I/O operation. V2.0 may suppress that +/// operation when a structural header page has already been copied. +pub async fn copy_external_metadata_column( + version: ConcreteFileVersion, + schema: &Schema, + column_index: usize, + has_existing_pages: bool, + copy: Copy, +) -> Result<()> +where + Copy: FnOnce() -> CopyFuture + Send, + CopyFuture: Future> + Send, +{ + match version { + ConcreteFileVersion::V1 => Err(Error::not_supported( + "binary-copy metadata operations are not supported for Lance v1".to_string(), + )), + ConcreteFileVersion::V2_0 => { + if v2_0::should_copy_external_metadata_column(schema, column_index, has_existing_pages) + { + copy().await + } else { + Ok(()) + } + } + ConcreteFileVersion::V2_1 | ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { + copy().await + } + } +} + +/// Normalize one copied column before an exact-version footer is written. +pub fn finalize_external_metadata_column( + version: ConcreteFileVersion, + schema: &Schema, + column_index: usize, + pages: &mut Vec, + num_rows: u64, +) -> Result<()> { + match version { + ConcreteFileVersion::V1 => Err(Error::not_supported( + "binary-copy metadata operations are not supported for Lance v1".to_string(), + )), + ConcreteFileVersion::V2_0 => { + v2_0::finalize_external_metadata_column(schema, column_index, pages, num_rows); + Ok(()) + } + ConcreteFileVersion::V2_1 | ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => Ok(()), + } +} + +/// Open a projected reader while keeping exact metadata-form selection in the +/// file layer. +/// +/// `open_indexed` returns `None` when the loaded index is not selective enough +/// to justify a projected reader. V2.0 never invokes it because indexed +/// metadata is not part of that reader's accepted grammar. +pub async fn open_projected_reader( + version: ConcreteFileVersion, + projection: &ReaderProjection, + prefer_indexed: bool, + open_indexed: OpenIndexed, + open_full: OpenFull, +) -> Result +where + OpenIndexed: FnOnce() -> IndexedFuture + Send, + IndexedFuture: Future>> + Send, + OpenFull: FnOnce() -> FullFuture + Send, + FullFuture: Future> + Send, +{ + match version { + ConcreteFileVersion::V1 => Err(Error::not_supported( + "projected current-format readers cannot open Lance v1 files".to_string(), + )), + ConcreteFileVersion::V2_0 => open_full().await, + ConcreteFileVersion::V2_1 | ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { + if prefer_indexed + && FileMetadataProvider::projection_matches_indexed_metadata(projection) + && let Some(reader) = open_indexed().await? + { + return Ok(reader); + } + open_full().await + } + } +} + +/// Open a self-described file and dispatch to the matching reader. +/// +/// The current-format reader's optimistic tail read is also used for exact +/// version detection, so current files do not pay for a separate footer probe. +pub async fn open_self_described_reader( + scheduler: FileScheduler, + decoder_plugins: Arc, + cache: &LanceCache, + options: FileReaderOptions, +) -> Result { + FileReader::try_open_for_dispatch(scheduler, None, decoder_plugins, cache, options).await +} + /// Create a current-format writer for an exact file version. /// /// V1 uses [`v1::writer::FileWriter`] directly because its manifest provider is diff --git a/rust/lance-file/src/versions/v2_0/mod.rs b/rust/lance-file/src/versions/v2_0/mod.rs index 219218f3ab7..5003345a70c 100644 --- a/rust/lance-file/src/versions/v2_0/mod.rs +++ b/rust/lance-file/src/versions/v2_0/mod.rs @@ -1,27 +1,128 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors +//! Lance v2.0 file composition. -//! Lance v2.0 encoding composition. - -use std::sync::Arc; +use std::{collections::BTreeMap, sync::Arc}; use bytes::Bytes; -use lance_core::{Result, datatypes::Schema}; -use lance_encoding::{ - array_encoding::ArrayFieldEncodingStrategy, - encoder::{EncodedBatch, FieldEncodingStrategy}, +use lance_core::{ + Result, + datatypes::{Field, Schema}, }; +use lance_encoding::{decoder::PageInfo, encoder::EncodedBatch}; use lance_io::traits::Writer as ObjectWriter; -use crate::writer::FileWriterOptions; +use crate::{reader::ReadProjection, writer::FileWriterOptions}; +mod reader; mod writer; +#[cfg(test)] +pub(crate) use reader::test_projection_length; +pub(crate) use reader::{ + decode_column_metadata, finish_metadata, finish_metadata_index, validate_global_buffers, +}; +pub use reader::{ + projection_from_column_names, projection_from_field_ids, projection_from_whole_schema, +}; pub use writer::Writer; -/// Compose the v2.0 field encoding mechanisms. -pub fn encoding_strategy() -> Arc { - Arc::new(ArrayFieldEncodingStrategy::new()) +pub(crate) fn read_projection() -> Arc { + reader::read_projection() +} + +/// Count physical columns represented by a field in a v2.0 footer. +pub fn physical_column_count(field: &Field) -> usize { + if field.is_blob() || field.is_packed_struct() { + 1 + } else { + 1 + field + .children + .iter() + .map(physical_column_count) + .sum::() + } +} + +/// Build persisted field-to-column entries for a v2.0 data file. +pub fn data_file_columns(schema: &Schema) -> (Vec, Vec) { + let mut field_ids = Vec::new(); + let mut column_indices = Vec::new(); + append_physical_fields(&schema.fields, &mut field_ids, &mut column_indices, &mut 0); + (field_ids, column_indices) +} + +pub(super) fn field_id_to_column_index(schema: &Schema) -> BTreeMap { + let (field_ids, column_indices) = data_file_columns(schema); + field_ids + .into_iter() + .zip(column_indices) + .map(|(field_id, column_index)| (field_id as u32, column_index as u32)) + .collect() +} + +fn append_physical_fields( + fields: &[Field], + field_ids: &mut Vec, + column_indices: &mut Vec, + next_column: &mut i32, +) { + for field in fields { + field_ids.push(field.id); + column_indices.push(*next_column); + *next_column += 1; + if !field.is_blob() && !field.is_packed_struct() { + append_physical_fields(&field.children, field_ids, column_indices, next_column); + } + } +} + +fn is_external_metadata_structural_header( + fields: &[Field], + target_column: usize, + next_column: &mut usize, +) -> Option { + for field in fields { + if *next_column == target_column { + return Some(field.logical_type.is_struct() && !field.is_packed_struct()); + } + *next_column += 1; + if !field.is_blob() + && !field.is_packed_struct() + && let Some(is_header) = + is_external_metadata_structural_header(&field.children, target_column, next_column) + { + return Some(is_header); + } + } + None +} + +pub(super) fn should_copy_external_metadata_column( + schema: &Schema, + column_index: usize, + has_existing_pages: bool, +) -> bool { + let mut next_column = 0; + let is_header = + is_external_metadata_structural_header(&schema.fields, column_index, &mut next_column) + .unwrap_or(false); + !is_header || !has_existing_pages +} + +pub(super) fn finalize_external_metadata_column( + schema: &Schema, + column_index: usize, + pages: &mut Vec, + num_rows: u64, +) { + let mut next_column = 0; + let is_header = + is_external_metadata_structural_header(&schema.fields, column_index, &mut next_column) + .unwrap_or(false); + if is_header && !pages.is_empty() { + pages[0].num_rows = num_rows; + pages[0].priority = 0; + pages.truncate(1); + } } /// Create a v2.0 writer with an explicit schema. diff --git a/rust/lance-file/src/versions/v2_0/reader.rs b/rust/lance-file/src/versions/v2_0/reader.rs new file mode 100644 index 00000000000..5a96031ba79 --- /dev/null +++ b/rust/lance-file/src/versions/v2_0/reader.rs @@ -0,0 +1,415 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use bytes::Bytes; +use lance_core::{ + Error, Result, + cache::LanceCache, + datatypes::{Field, Schema}, +}; +use lance_encoding::{ + EncodingsIo, + decoder::{ColumnInfo, PageEncoding, PageInfo}, + format::pb, +}; +use prost::{Message, Name}; + +use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, CachedFileMetadata, FileMetadataIndex, FileMetadataProvider, FileReader, + PreparedProjection, RawFileMetadata, ReadProjection, ReaderProjection, + normalized_column_num_rows, verify_uniform_lengths, + }, + version::ConcreteFileVersion, +}; + +fn fetch_encoding(encoding: &pbfile::Encoding) -> Result { + match &encoding.location { + Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source( + "Indirect file encodings are not supported".into(), + )), + Some(pbfile::encoding::Location::Direct(encoding)) => { + let envelope = prost_types::Any::decode(Bytes::from(encoding.encoding.clone())) + .map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(), + ) + })?; + envelope.to_msg::().map_err(|error| { + Error::invalid_input_source( + format!("Invalid direct {} encoding: {error}", M::NAME).into(), + ) + }) + } + Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source( + format!("Missing {} encoding description", M::NAME).into(), + )), + None => Err(Error::invalid_input_source( + format!("Missing {} encoding location", M::NAME).into(), + )), + } +} + +pub fn decode_column( + column_index: u32, + metadata: &pbfile::ColumnMetadata, +) -> Result> { + let page_infos = metadata + .pages + .iter() + .enumerate() + .map(|(page_index, page)| { + let array_encoding = + fetch_encoding::(page.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!( + "Column {} page {} is missing its encoding", + column_index, page_index + ) + .into(), + ) + })?)?; + if page.buffer_offsets.len() != page.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} page {} has {} buffer offsets but {} buffer sizes", + column_index, + page_index, + page.buffer_offsets.len(), + page.buffer_sizes.len() + ) + .into(), + )); + } + let buffer_offsets_and_sizes = Arc::from( + page.buffer_offsets + .iter() + .zip(&page.buffer_sizes) + .map(|(offset, size)| (*offset, *size)) + .collect::>(), + ); + Ok(PageInfo { + buffer_offsets_and_sizes, + encoding: PageEncoding::Array(array_encoding), + num_rows: page.length, + priority: page.priority, + }) + }) + .collect::>>()?; + + if metadata.buffer_offsets.len() != metadata.buffer_sizes.len() { + return Err(Error::invalid_input_source( + format!( + "Column {} has {} buffer offsets but {} buffer sizes", + column_index, + metadata.buffer_offsets.len(), + metadata.buffer_sizes.len() + ) + .into(), + )); + } + let buffer_offsets_and_sizes = Arc::from( + metadata + .buffer_offsets + .iter() + .zip(&metadata.buffer_sizes) + .map(|(offset, size)| (*offset, *size)) + .collect::>(), + ); + Ok(Arc::new(ColumnInfo { + index: column_index, + page_infos: Arc::from(page_infos), + buffer_offsets_and_sizes, + encoding: fetch_encoding(metadata.encoding.as_ref().ok_or_else(|| { + Error::invalid_input_source( + format!("Column {} is missing its encoding", column_index).into(), + ) + })?)?, + })) +} + +pub fn decode_column_metadata( + column_metadatas: &[pbfile::ColumnMetadata], +) -> Result>> { + column_metadatas + .iter() + .enumerate() + .map(|(column_index, metadata)| { + let column_index = u32::try_from(column_index).map_err(|_| { + Error::invalid_input_source("File has more than u32::MAX columns".into()) + })?; + decode_column(column_index, metadata) + }) + .collect() +} + +pub async fn prepare_projection( + metadata_provider: &FileMetadataProvider, + projection: &ReaderProjection, + _io: &Arc, + _cache: &Arc, +) -> Result { + match metadata_provider { + FileMetadataProvider::Full(metadata) => { + FileReader::validate_projection(projection, metadata)?; + Ok(PreparedProjection { + column_infos: metadata.column_infos.clone(), + decoder_projection: projection.clone(), + }) + } + FileMetadataProvider::Indexed(metadata_index) => { + FileMetadataProvider::validate_indexed_projection_structure( + projection, + metadata_index, + )?; + Err(FileMetadataProvider::indexed_projection_error( + projection, + metadata_index, + )) + } + } +} + +fn field_column_shape(field: &Field) -> (bool, bool) { + if field.is_blob() || field.is_packed_struct() { + return (true, false); + } + (true, !field.children.is_empty()) +} + +fn append_field_ids( + fields: &[Field], + field_id_to_column_index: &BTreeMap, + column_indices: &mut Vec, +) { + for field in fields { + let (contributes, recurse) = field_column_shape(field); + if contributes + && let Some(column_index) = field_id_to_column_index.get(&(field.id as u32)).copied() + { + column_indices.push(column_index); + } + if recurse { + append_field_ids(&field.children, field_id_to_column_index, column_indices); + } + } +} + +pub fn projection_from_field_ids( + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> ReaderProjection { + let mut column_indices = Vec::new(); + append_field_ids( + &schema.fields, + field_id_to_column_index, + &mut column_indices, + ); + ReaderProjection { + schema: Arc::new(schema.clone()), + column_indices, + } +} + +pub fn projection_from_whole_schema(schema: &Schema) -> ReaderProjection { + projection_from_field_ids(schema, &super::field_id_to_column_index(schema)) +} + +pub fn projection_from_column_names( + schema: &Schema, + column_names: &[&str], +) -> Result { + let field_id_to_column_index = super::field_id_to_column_index(schema); + let projected = schema.project(column_names)?; + Ok(projection_from_field_ids( + &projected, + &field_id_to_column_index, + )) +} + +fn children_share_parent_length(field: &Field) -> bool { + field.logical_type.is_struct() +} + +fn validate_field_length Result>( + field: &Field, + comparable: bool, + column_indices: &[u32], + cursor: &mut usize, + column_len: &F, +) -> Result { + let (contributes, recurse) = field_column_shape(field); + let mut field_rows = None; + if contributes { + let column = *column_indices.get(*cursor).ok_or_else(|| { + Error::invalid_input(format!( + "projection supplied fewer column indices than its fields require (ran out at field '{}')", + field.name + )) + })?; + *cursor += 1; + field_rows = Some(column_len(column as usize)?); + } + if recurse { + let enforce_children = comparable && children_share_parent_length(field); + for child in &field.children { + let child_rows = + validate_field_length(child, enforce_children, column_indices, cursor, column_len)?; + let expected = *field_rows.get_or_insert(child_rows); + if enforce_children && child_rows != expected { + return Err(Error::invalid_input(format!( + "cannot read field '{}': its children have differing lengths (child '{}' has {} rows, but the field has {}); a struct's children must all have the same length", + field.name, child.name, child_rows, expected + ))); + } + } + } + field_rows.ok_or_else(|| { + Error::invalid_input(format!( + "projected field '{}' maps to no columns", + field.name + )) + }) +} + +pub fn prepared_read_length(prepared: &PreparedProjection) -> Result { + let column_len = |column: usize| { + let info = prepared.column_infos.get(column).ok_or_else(|| { + Error::invalid_input(format!( + "projection references column index {} but only {} columns are available", + column, + prepared.column_infos.len() + )) + })?; + normalized_column_num_rows(info) + }; + let mut cursor = 0; + let mut field_lengths = Vec::with_capacity(prepared.decoder_projection.schema.fields.len()); + for field in &prepared.decoder_projection.schema.fields { + let rows = validate_field_length( + field, + true, + &prepared.decoder_projection.column_indices, + &mut cursor, + &column_len, + )?; + field_lengths.push((field.name.as_str(), rows)); + } + if cursor != prepared.decoder_projection.column_indices.len() { + return Err(Error::invalid_input(format!( + "projection supplied {} column indices but its fields require {}", + prepared.decoder_projection.column_indices.len(), + cursor + ))); + } + verify_uniform_lengths(&field_lengths) +} + +#[derive(Debug)] +struct V20ReadProjection; + +pub(super) fn read_projection() -> Arc { + Arc::new(V20ReadProjection) +} + +#[async_trait] +impl ReadProjection for V20ReadProjection { + fn validate_indexed( + &self, + projection: &ReaderProjection, + metadata_index: &FileMetadataIndex, + ) -> Result<()> { + FileMetadataProvider::validate_indexed_projection_structure(projection, metadata_index)?; + Err(FileMetadataProvider::indexed_projection_error( + projection, + metadata_index, + )) + } + + fn read_length(&self, prepared: &PreparedProjection) -> Result { + prepared_read_length(prepared) + } + + async fn prepare( + &self, + metadata_provider: &FileMetadataProvider, + projection: &ReaderProjection, + io: &Arc, + cache: &Arc, + ) -> Result<(PreparedProjection, u64)> { + let prepared = prepare_projection(metadata_provider, projection, io, cache).await?; + let read_len = self.read_length(&prepared)?; + Ok((prepared, read_len)) + } +} + +pub fn finish_metadata(raw: RawFileMetadata) -> Result { + if !matches!( + (raw.footer.major_version, raw.footer.minor_version), + (0, 3) | (2, 0) + ) { + return Err(Error::version_conflict( + "Attempt to use the Lance v2.0 reader for a different file version".to_string(), + raw.footer.major_version, + raw.footer.minor_version, + )); + } + let column_infos = decode_column_metadata(&raw.column_metadatas)?; + Ok(CachedFileMetadata { + file_schema: raw.file_schema, + column_metadatas: raw.column_metadatas, + column_infos, + num_rows: raw.num_rows, + file_buffers: raw.file_buffers, + num_data_bytes: raw.num_data_bytes, + num_column_metadata_bytes: raw.num_column_metadata_bytes, + num_global_buffer_bytes: raw.num_global_buffer_bytes, + num_footer_bytes: raw.num_footer_bytes, + major_version: raw.footer.major_version, + minor_version: raw.footer.minor_version, + version: ConcreteFileVersion::V2_0, + file_size_bytes: raw.file_size_bytes, + retained_global_buffers: raw.retained_global_buffers, + }) +} + +pub fn validate_global_buffers(_buffers: &[BufferDescriptor]) -> Result<()> { + Ok(()) +} + +pub fn finish_metadata_index(index: FileMetadataIndex) -> Result { + if index.version == ConcreteFileVersion::V2_0 { + Ok(index) + } else { + let (major, minor) = index.version.to_standard_footer_numbers(); + Err(Error::version_conflict( + "Attempt to use the Lance v2.0 reader for a different metadata index".to_string(), + major, + minor, + )) + } +} + +#[cfg(test)] +pub fn test_projection_length( + schema: &Schema, + column_indices: &[u32], + column_lengths: &[u64], +) -> Result { + let column_len = |column: usize| { + column_lengths.get(column).copied().ok_or_else(|| { + Error::invalid_input(format!("missing synthetic length for column {column}")) + }) + }; + let mut cursor = 0; + let mut field_lengths = Vec::with_capacity(schema.fields.len()); + for field in &schema.fields { + let rows = validate_field_length(field, true, column_indices, &mut cursor, &column_len)?; + field_lengths.push((field.name.as_str(), rows)); + } + verify_uniform_lengths(&field_lengths) +} diff --git a/rust/lance-file/src/versions/v2_1/mod.rs b/rust/lance-file/src/versions/v2_1/mod.rs index 78e2f2ed526..63caea79165 100644 --- a/rust/lance-file/src/versions/v2_1/mod.rs +++ b/rust/lance-file/src/versions/v2_1/mod.rs @@ -1,9 +1,6 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors +//! Lance v2.1 file composition. -//! Lance v2.1 encoding composition. - -use std::sync::Arc; +use std::{collections::BTreeMap, sync::Arc}; use bytes::Bytes; use lance_core::{ @@ -23,13 +20,44 @@ use lance_encoding::{ }; use lance_io::traits::Writer as ObjectWriter; -use crate::writer::FileWriterOptions; +use crate::{ + reader::{ReadProjection, structural}, + writer::FileWriterOptions, +}; mod compression; +mod reader; mod writer; +pub use reader::{ + projection_from_column_names, projection_from_field_ids, projection_from_whole_schema, +}; pub use writer::Writer; +#[cfg(test)] +pub(crate) use reader::test_projection_length; +pub(crate) use reader::{ + decode_column_metadata, finish_metadata, finish_metadata_index, validate_global_buffers, +}; + +pub(crate) fn read_projection() -> Arc { + structural::read_projection(reader::decode_column) +} + +/// Count physical columns represented by a field in a v2.1 footer. +pub fn physical_column_count(field: &Field) -> usize { + structural::physical_column_count(field) +} + +/// Build persisted field-to-column entries for a v2.1 data file. +pub fn data_file_columns(schema: &Schema) -> (Vec, Vec) { + structural::data_file_columns(schema) +} + +pub(super) fn field_id_to_column_index(schema: &Schema) -> BTreeMap { + structural::field_id_to_column_index(schema) +} + #[derive(Debug)] struct FieldStrategy { primitive: PrimitiveFieldEncoding, diff --git a/rust/lance-file/src/versions/v2_1/reader.rs b/rust/lance-file/src/versions/v2_1/reader.rs new file mode 100644 index 00000000000..0c11e5f8976 --- /dev/null +++ b/rust/lance-file/src/versions/v2_1/reader.rs @@ -0,0 +1,304 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::BTreeMap, sync::Arc}; + +use lance_core::{Error, Result, datatypes::Schema}; +use lance_encoding::{decoder::ColumnInfo, format::pb21}; + +use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, CachedFileMetadata, FileMetadataIndex, RawFileMetadata, ReaderProjection, + structural, + }, + version::ConcreteFileVersion, +}; + +fn required<'a, T>(value: Option<&'a T>, label: &str) -> Result<&'a T> { + value.ok_or_else(|| { + Error::invalid_input_source( + format!("Lance v2.1 {label} is missing its nested encoding").into(), + ) + }) +} + +fn validate_compressive_encoding(encoding: &pb21::CompressiveEncoding) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref() { + Some(Compression::Flat(_)) + | Some(Compression::InlineBitpacking(_)) + | Some(Compression::Constant(_)) => Ok(()), + Some(Compression::Variable(variable)) => validate_compressive_encoding(required( + variable.offsets.as_deref(), + "variable offsets", + )?), + Some(Compression::OutOfLineBitpacking(bitpacking)) => { + validate_compressive_encoding(required( + bitpacking.values.as_deref(), + "out-of-line bitpacking values", + )?) + } + Some(Compression::Fsst(fsst)) => { + validate_compressive_encoding(required(fsst.values.as_deref(), "FSST values")?) + } + Some(Compression::Dictionary(dictionary)) => { + validate_compressive_encoding(required( + dictionary.indices.as_deref(), + "dictionary indices", + )?)?; + validate_compressive_encoding(required( + dictionary.items.as_deref(), + "dictionary items", + )?) + } + Some(Compression::Rle(rle)) => { + let values = required(rle.values.as_deref(), "RLE values")?; + let run_lengths = required(rle.run_lengths.as_deref(), "RLE run lengths")?; + let fixed_values = matches!( + values.compression.as_ref(), + Some(Compression::Flat(flat)) + if matches!(flat.bits_per_value, 8 | 16 | 32 | 64) + && flat.data.is_none() + ); + let fixed_u8_lengths = matches!( + run_lengths.compression.as_ref(), + Some(Compression::Flat(flat)) + if flat.bits_per_value == 8 && flat.data.is_none() + ); + if !fixed_values || !fixed_u8_lengths { + return Err(Error::invalid_input_source( + "Lance v2.1 RLE requires flat values and flat u8 run lengths".into(), + )); + } + Ok(()) + } + Some(Compression::ByteStreamSplit(split)) => validate_compressive_encoding(required( + split.values.as_deref(), + "byte-stream-split values", + )?), + Some(Compression::General(general)) => validate_compressive_encoding(required( + general.values.as_deref(), + "general-compression values", + )?), + Some(Compression::FixedSizeList(list)) => validate_compressive_encoding(required( + list.values.as_deref(), + "fixed-size-list values", + )?), + Some(Compression::PackedStruct(packed)) => validate_compressive_encoding(required( + packed.values.as_deref(), + "packed-struct values", + )?), + Some(Compression::VariablePackedStruct(_)) => Err(Error::invalid_input_source( + "Variable packed struct compression is not part of the Lance v2.1 grammar".into(), + )), + None => Err(Error::invalid_input_source( + "Lance v2.1 compressive encoding is missing its compression variant".into(), + )), + } +} + +fn validate_page_layout(layout: &pb21::PageLayout) -> Result<()> { + use pb21::page_layout::Layout; + + match layout.layout.as_ref() { + Some(Layout::MiniBlockLayout(miniblock)) => { + if miniblock.has_large_chunk { + return Err(Error::invalid_input_source( + "Large miniblock chunks are not part of the Lance v2.1 grammar".into(), + )); + } + if let Some(rep) = miniblock.rep_compression.as_ref() { + validate_compressive_encoding(rep)?; + } + if let Some(def) = miniblock.def_compression.as_ref() { + validate_compressive_encoding(def)?; + } + validate_compressive_encoding(required( + miniblock.value_compression.as_ref(), + "miniblock values", + )?)?; + if let Some(dictionary) = miniblock.dictionary.as_ref() { + validate_compressive_encoding(dictionary)?; + } + Ok(()) + } + Some(Layout::FullZipLayout(fullzip)) => validate_compressive_encoding(required( + fullzip.value_compression.as_ref(), + "full-zip values", + )?), + Some(Layout::ConstantLayout(constant)) => { + if constant.inline_value.is_some() { + Err(Error::invalid_input_source( + "Lance v2.1 only accepts the all-null form of constant page layout".into(), + )) + } else if constant.rep_compression.is_some() || constant.def_compression.is_some() { + Err(Error::invalid_input_source( + "Compressed constant-page levels are not part of the Lance v2.1 grammar".into(), + )) + } else { + Ok(()) + } + } + Some(Layout::BlobLayout(blob)) => { + let inner = blob.inner_layout.as_deref().ok_or_else(|| { + Error::invalid_input_source( + "Lance v2.1 blob page layout is missing its inner layout".into(), + ) + })?; + validate_page_layout(inner) + } + Some(Layout::SparseLayout(_)) => Err(Error::invalid_input_source( + "Sparse page layout is not part of the Lance v2.1 grammar".into(), + )), + None => Err(Error::invalid_input_source( + "Lance v2.1 page is missing its page layout".into(), + )), + } +} + +pub fn decode_column( + column_index: u32, + metadata: &pbfile::ColumnMetadata, +) -> Result> { + let page_infos = metadata + .pages + .iter() + .enumerate() + .map(|(page_index, page)| { + let page_layout = structural::decode_page_layout(column_index, page_index, page)?; + validate_page_layout(&page_layout)?; + structural::build_page_info(column_index, page_index, page, page_layout) + }) + .collect::>>()?; + structural::build_column_info(column_index, metadata, page_infos) +} + +pub fn decode_column_metadata( + column_metadatas: &[pbfile::ColumnMetadata], +) -> Result>> { + column_metadatas + .iter() + .enumerate() + .map(|(column_index, metadata)| { + let column_index = u32::try_from(column_index).map_err(|_| { + Error::invalid_input_source("File has more than u32::MAX columns".into()) + })?; + decode_column(column_index, metadata) + }) + .collect() +} + +pub fn projection_from_field_ids( + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> ReaderProjection { + structural::projection_from_field_ids(schema, field_id_to_column_index) +} + +pub fn projection_from_whole_schema(schema: &Schema) -> ReaderProjection { + structural::projection_from_field_ids(schema, &super::field_id_to_column_index(schema)) +} + +pub fn projection_from_column_names( + schema: &Schema, + column_names: &[&str], +) -> Result { + structural::projection_from_column_names( + schema, + column_names, + &super::field_id_to_column_index(schema), + ) +} + +pub fn finish_metadata(raw: RawFileMetadata) -> Result { + if (raw.footer.major_version, raw.footer.minor_version) != (2, 1) { + return Err(Error::version_conflict( + "Attempt to use the Lance v2.1 reader for a different file version".to_string(), + raw.footer.major_version, + raw.footer.minor_version, + )); + } + validate_global_buffers(&raw.file_buffers)?; + let column_infos = decode_column_metadata(&raw.column_metadatas)?; + Ok(CachedFileMetadata { + file_schema: raw.file_schema, + column_metadatas: raw.column_metadatas, + column_infos, + num_rows: raw.num_rows, + file_buffers: raw.file_buffers, + num_data_bytes: raw.num_data_bytes, + num_column_metadata_bytes: raw.num_column_metadata_bytes, + num_global_buffer_bytes: raw.num_global_buffer_bytes, + num_footer_bytes: raw.num_footer_bytes, + major_version: raw.footer.major_version, + minor_version: raw.footer.minor_version, + version: ConcreteFileVersion::V2_1, + file_size_bytes: raw.file_size_bytes, + retained_global_buffers: raw.retained_global_buffers, + }) +} + +pub fn validate_global_buffers(buffers: &[BufferDescriptor]) -> Result<()> { + structural::validate_global_buffers(buffers) +} + +pub fn finish_metadata_index(index: FileMetadataIndex) -> Result { + if index.version != ConcreteFileVersion::V2_1 { + let (major, minor) = index.version.to_standard_footer_numbers(); + return Err(Error::version_conflict( + "Attempt to use the Lance v2.1 reader for a different metadata index".to_string(), + major, + minor, + )); + } + validate_global_buffers(&index.file_buffers)?; + Ok(index) +} + +#[cfg(test)] +pub fn test_projection_length( + schema: &Schema, + column_indices: &[u32], + column_lengths: &[u64], +) -> Result { + structural::test_projection_length(schema, column_indices, column_lengths) +} + +#[cfg(test)] +mod grammar_tests { + use super::*; + use pb21::{ + CompressiveEncoding, FullZipLayout, General, PageLayout, VariablePackedStruct, + compressive_encoding::Compression, page_layout::Layout, + }; + + #[test] + fn rejects_nested_variable_packed_struct() { + let variable_packed = CompressiveEncoding { + compression: Some(Compression::VariablePackedStruct(VariablePackedStruct { + fields: Vec::new(), + })), + }; + let wrapped = CompressiveEncoding { + compression: Some(Compression::General(Box::new(General { + compression: None, + values: Some(Box::new(variable_packed)), + }))), + }; + let layout = PageLayout { + layout: Some(Layout::FullZipLayout(FullZipLayout { + value_compression: Some(wrapped), + ..Default::default() + })), + }; + + let error = validate_page_layout(&layout).unwrap_err(); + assert!( + error + .to_string() + .contains("Variable packed struct compression is not part") + ); + } +} diff --git a/rust/lance-file/src/versions/v2_2/mod.rs b/rust/lance-file/src/versions/v2_2/mod.rs index 569115b6331..af0085509a6 100644 --- a/rust/lance-file/src/versions/v2_2/mod.rs +++ b/rust/lance-file/src/versions/v2_2/mod.rs @@ -1,9 +1,6 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors +//! Lance v2.2 file composition. -//! Lance v2.2 encoding composition. - -use std::sync::Arc; +use std::{collections::BTreeMap, sync::Arc}; use bytes::Bytes; use lance_core::{ @@ -24,13 +21,41 @@ use lance_encoding::{ }; use lance_io::traits::Writer as ObjectWriter; -use crate::writer::FileWriterOptions; +use crate::{ + reader::{ReadProjection, structural}, + writer::FileWriterOptions, +}; mod compression; +mod reader; mod writer; +pub(crate) use reader::{ + decode_column_metadata, finish_metadata, finish_metadata_index, validate_global_buffers, +}; +pub use reader::{ + projection_from_column_names, projection_from_field_ids, projection_from_whole_schema, +}; + +pub(crate) fn read_projection() -> Arc { + structural::read_projection(reader::decode_column) +} pub use writer::Writer; +/// Count physical columns represented by a field in a v2.2 footer. +pub fn physical_column_count(field: &Field) -> usize { + structural::physical_column_count(field) +} + +/// Build persisted field-to-column entries for a v2.2 data file. +pub fn data_file_columns(schema: &Schema) -> (Vec, Vec) { + structural::data_file_columns(schema) +} + +pub(super) fn field_id_to_column_index(schema: &Schema) -> BTreeMap { + structural::field_id_to_column_index(schema) +} + #[derive(Debug)] struct FieldStrategy { primitive: PrimitiveFieldEncoding, diff --git a/rust/lance-file/src/versions/v2_2/reader.rs b/rust/lance-file/src/versions/v2_2/reader.rs new file mode 100644 index 00000000000..81a88720c77 --- /dev/null +++ b/rust/lance-file/src/versions/v2_2/reader.rs @@ -0,0 +1,308 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::BTreeMap, sync::Arc}; + +use lance_core::{Error, Result, datatypes::Schema}; +use lance_encoding::{decoder::ColumnInfo, format::pb21}; + +use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, CachedFileMetadata, FileMetadataIndex, RawFileMetadata, ReaderProjection, + structural, + }, + version::ConcreteFileVersion, +}; + +fn required<'a, T>(value: Option<&'a T>, label: &str) -> Result<&'a T> { + value.ok_or_else(|| { + Error::invalid_input_source( + format!("Lance v2.2 {label} is missing its nested encoding").into(), + ) + }) +} + +fn validate_compressive_encoding(encoding: &pb21::CompressiveEncoding) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref() { + Some(Compression::Flat(_)) + | Some(Compression::InlineBitpacking(_)) + | Some(Compression::Constant(_)) => Ok(()), + Some(Compression::Variable(variable)) => validate_compressive_encoding(required( + variable.offsets.as_deref(), + "variable offsets", + )?), + Some(Compression::OutOfLineBitpacking(bitpacking)) => { + validate_compressive_encoding(required( + bitpacking.values.as_deref(), + "out-of-line bitpacking values", + )?) + } + Some(Compression::Fsst(fsst)) => { + validate_compressive_encoding(required(fsst.values.as_deref(), "FSST values")?) + } + Some(Compression::Dictionary(dictionary)) => { + validate_compressive_encoding(required( + dictionary.indices.as_deref(), + "dictionary indices", + )?)?; + validate_compressive_encoding(required( + dictionary.items.as_deref(), + "dictionary items", + )?) + } + Some(Compression::Rle(rle)) => { + let values = required(rle.values.as_deref(), "RLE values")?; + let run_lengths = required(rle.run_lengths.as_deref(), "RLE run lengths")?; + let fixed_values = matches!( + values.compression.as_ref(), + Some(Compression::Flat(flat)) + if matches!(flat.bits_per_value, 8 | 16 | 32 | 64) + && flat.data.is_none() + ); + let fixed_u8_lengths = matches!( + run_lengths.compression.as_ref(), + Some(Compression::Flat(flat)) + if flat.bits_per_value == 8 && flat.data.is_none() + ); + if !fixed_values || !fixed_u8_lengths { + return Err(Error::invalid_input_source( + "Lance v2.2 RLE requires flat values and flat u8 run lengths".into(), + )); + } + Ok(()) + } + Some(Compression::ByteStreamSplit(split)) => validate_compressive_encoding(required( + split.values.as_deref(), + "byte-stream-split values", + )?), + Some(Compression::General(general)) => validate_compressive_encoding(required( + general.values.as_deref(), + "general-compression values", + )?), + Some(Compression::FixedSizeList(list)) => validate_compressive_encoding(required( + list.values.as_deref(), + "fixed-size-list values", + )?), + Some(Compression::PackedStruct(packed)) => validate_compressive_encoding(required( + packed.values.as_deref(), + "packed-struct values", + )?), + Some(Compression::VariablePackedStruct(packed)) => { + for field in &packed.fields { + validate_compressive_encoding(required( + field.value.as_ref(), + "variable packed-struct field", + )?)?; + } + Ok(()) + } + None => Err(Error::invalid_input_source( + "Lance v2.2 compressive encoding is missing its compression variant".into(), + )), + } +} + +fn validate_page_layout(layout: &pb21::PageLayout) -> Result<()> { + use pb21::page_layout::Layout; + + match layout.layout.as_ref() { + Some(Layout::MiniBlockLayout(miniblock)) => { + if !miniblock.has_large_chunk { + return Err(Error::invalid_input_source( + "Lance v2.2 miniblock pages require the u32 chunk grammar".into(), + )); + } + if let Some(rep) = miniblock.rep_compression.as_ref() { + validate_compressive_encoding(rep)?; + } + if let Some(def) = miniblock.def_compression.as_ref() { + validate_compressive_encoding(def)?; + } + validate_compressive_encoding(required( + miniblock.value_compression.as_ref(), + "miniblock values", + )?)?; + if let Some(dictionary) = miniblock.dictionary.as_ref() { + validate_compressive_encoding(dictionary)?; + } + Ok(()) + } + Some(Layout::FullZipLayout(fullzip)) => validate_compressive_encoding(required( + fullzip.value_compression.as_ref(), + "full-zip values", + )?), + Some(Layout::ConstantLayout(constant)) => { + if let Some(rep) = constant.rep_compression.as_ref() { + validate_compressive_encoding(rep)?; + } + if let Some(def) = constant.def_compression.as_ref() { + validate_compressive_encoding(def)?; + } + Ok(()) + } + Some(Layout::BlobLayout(blob)) => { + let inner = blob.inner_layout.as_deref().ok_or_else(|| { + Error::invalid_input_source( + "Lance v2.2 blob page layout is missing its inner layout".into(), + ) + })?; + validate_page_layout(inner) + } + Some(Layout::SparseLayout(_)) => Err(Error::invalid_input_source( + "Sparse page layout is not part of the Lance v2.2 grammar".into(), + )), + None => Err(Error::invalid_input_source( + "Lance v2.2 page is missing its page layout".into(), + )), + } +} + +pub fn decode_column( + column_index: u32, + metadata: &pbfile::ColumnMetadata, +) -> Result> { + let page_infos = metadata + .pages + .iter() + .enumerate() + .map(|(page_index, page)| { + let page_layout = structural::decode_page_layout(column_index, page_index, page)?; + validate_page_layout(&page_layout)?; + structural::build_page_info(column_index, page_index, page, page_layout) + }) + .collect::>>()?; + structural::build_column_info(column_index, metadata, page_infos) +} + +pub fn decode_column_metadata( + column_metadatas: &[pbfile::ColumnMetadata], +) -> Result>> { + column_metadatas + .iter() + .enumerate() + .map(|(column_index, metadata)| { + let column_index = u32::try_from(column_index).map_err(|_| { + Error::invalid_input_source("File has more than u32::MAX columns".into()) + })?; + decode_column(column_index, metadata) + }) + .collect() +} + +pub fn projection_from_field_ids( + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> ReaderProjection { + structural::projection_from_field_ids(schema, field_id_to_column_index) +} + +pub fn projection_from_whole_schema(schema: &Schema) -> ReaderProjection { + structural::projection_from_field_ids(schema, &super::field_id_to_column_index(schema)) +} + +pub fn projection_from_column_names( + schema: &Schema, + column_names: &[&str], +) -> Result { + structural::projection_from_column_names( + schema, + column_names, + &super::field_id_to_column_index(schema), + ) +} + +pub fn finish_metadata(raw: RawFileMetadata) -> Result { + if (raw.footer.major_version, raw.footer.minor_version) != (2, 2) { + return Err(Error::version_conflict( + "Attempt to use the Lance v2.2 reader for a different file version".to_string(), + raw.footer.major_version, + raw.footer.minor_version, + )); + } + validate_global_buffers(&raw.file_buffers)?; + let column_infos = decode_column_metadata(&raw.column_metadatas)?; + Ok(CachedFileMetadata { + file_schema: raw.file_schema, + column_metadatas: raw.column_metadatas, + column_infos, + num_rows: raw.num_rows, + file_buffers: raw.file_buffers, + num_data_bytes: raw.num_data_bytes, + num_column_metadata_bytes: raw.num_column_metadata_bytes, + num_global_buffer_bytes: raw.num_global_buffer_bytes, + num_footer_bytes: raw.num_footer_bytes, + major_version: raw.footer.major_version, + minor_version: raw.footer.minor_version, + version: ConcreteFileVersion::V2_2, + file_size_bytes: raw.file_size_bytes, + retained_global_buffers: raw.retained_global_buffers, + }) +} + +pub fn validate_global_buffers(buffers: &[BufferDescriptor]) -> Result<()> { + structural::validate_global_buffers(buffers) +} + +pub fn finish_metadata_index(index: FileMetadataIndex) -> Result { + if index.version != ConcreteFileVersion::V2_2 { + let (major, minor) = index.version.to_standard_footer_numbers(); + return Err(Error::version_conflict( + "Attempt to use the Lance v2.2 reader for a different metadata index".to_string(), + major, + minor, + )); + } + validate_global_buffers(&index.file_buffers)?; + Ok(index) +} + +#[cfg(test)] +mod grammar_tests { + use super::*; + use pb21::{ + CompressiveEncoding, Dictionary, Flat, FullZipLayout, PageLayout, Rle, + compressive_encoding::Compression, page_layout::Layout, + }; + + fn flat(bits_per_value: u64) -> CompressiveEncoding { + CompressiveEncoding { + compression: Some(Compression::Flat(Flat { + bits_per_value, + data: None, + })), + } + } + + #[test] + fn rejects_nested_variable_width_rle() { + let rle = CompressiveEncoding { + compression: Some(Compression::Rle(Box::new(Rle { + values: Some(Box::new(flat(32))), + run_lengths: Some(Box::new(flat(16))), + }))), + }; + let dictionary = CompressiveEncoding { + compression: Some(Compression::Dictionary(Box::new(Dictionary { + indices: Some(Box::new(rle)), + items: Some(Box::new(flat(32))), + num_dictionary_items: 1, + }))), + }; + let layout = PageLayout { + layout: Some(Layout::FullZipLayout(FullZipLayout { + value_compression: Some(dictionary), + ..Default::default() + })), + }; + + let error = validate_page_layout(&layout).unwrap_err(); + assert!( + error + .to_string() + .contains("flat values and flat u8 run lengths") + ); + } +} diff --git a/rust/lance-file/src/versions/v2_3/mod.rs b/rust/lance-file/src/versions/v2_3/mod.rs index 6df099e7226..4006ce30de3 100644 --- a/rust/lance-file/src/versions/v2_3/mod.rs +++ b/rust/lance-file/src/versions/v2_3/mod.rs @@ -1,11 +1,11 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors +//! Lance v2.3 file composition. -//! Lance v2.3 encoding composition. - -use std::sync::{ - Arc, - atomic::{AtomicBool, Ordering}, +use std::{ + collections::BTreeMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, }; use bytes::Bytes; @@ -27,15 +27,43 @@ use lance_encoding::{ }; use lance_io::traits::Writer as ObjectWriter; -use crate::writer::FileWriterOptions; +use crate::{ + reader::{ReadProjection, structural}, + writer::FileWriterOptions, +}; mod compression; +mod reader; mod writer; +pub(crate) use reader::{ + decode_column_metadata, finish_metadata, finish_metadata_index, validate_global_buffers, +}; +pub use reader::{ + projection_from_column_names, projection_from_field_ids, projection_from_whole_schema, +}; + +pub(crate) fn read_projection() -> Arc { + structural::read_projection(reader::decode_column) +} pub use writer::Writer; static WARNED_ON_UNSTABLE_FORMAT: AtomicBool = AtomicBool::new(false); +/// Count physical columns represented by a field in a v2.3 footer. +pub fn physical_column_count(field: &Field) -> usize { + structural::physical_column_count(field) +} + +/// Build persisted field-to-column entries for a v2.3 data file. +pub fn data_file_columns(schema: &Schema) -> (Vec, Vec) { + structural::data_file_columns(schema) +} + +pub(super) fn field_id_to_column_index(schema: &Schema) -> BTreeMap { + structural::field_id_to_column_index(schema) +} + #[derive(Debug)] struct FieldStrategy { primitive: PrimitiveFieldEncoding, diff --git a/rust/lance-file/src/versions/v2_3/reader.rs b/rust/lance-file/src/versions/v2_3/reader.rs new file mode 100644 index 00000000000..cdafbd0b9db --- /dev/null +++ b/rust/lance-file/src/versions/v2_3/reader.rs @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::BTreeMap, sync::Arc}; + +use lance_core::{Error, Result, datatypes::Schema}; +use lance_encoding::{decoder::ColumnInfo, format::pb21}; + +use crate::{ + format::pbfile, + reader::{ + BufferDescriptor, CachedFileMetadata, FileMetadataIndex, RawFileMetadata, ReaderProjection, + structural, + }, + version::ConcreteFileVersion, +}; + +fn required<'a, T>(value: Option<&'a T>, label: &str) -> Result<&'a T> { + value.ok_or_else(|| { + Error::invalid_input_source( + format!("Lance v2.3 {label} is missing its nested encoding").into(), + ) + }) +} + +fn validate_compressive_encoding(encoding: &pb21::CompressiveEncoding) -> Result<()> { + use pb21::compressive_encoding::Compression; + + match encoding.compression.as_ref() { + Some(Compression::Flat(_)) + | Some(Compression::InlineBitpacking(_)) + | Some(Compression::Constant(_)) => Ok(()), + Some(Compression::Variable(variable)) => validate_compressive_encoding(required( + variable.offsets.as_deref(), + "variable offsets", + )?), + Some(Compression::OutOfLineBitpacking(bitpacking)) => { + validate_compressive_encoding(required( + bitpacking.values.as_deref(), + "out-of-line bitpacking values", + )?) + } + Some(Compression::Fsst(fsst)) => { + validate_compressive_encoding(required(fsst.values.as_deref(), "FSST values")?) + } + Some(Compression::Dictionary(dictionary)) => { + validate_compressive_encoding(required( + dictionary.indices.as_deref(), + "dictionary indices", + )?)?; + validate_compressive_encoding(required( + dictionary.items.as_deref(), + "dictionary items", + )?) + } + Some(Compression::Rle(rle)) => { + validate_compressive_encoding(required(rle.values.as_deref(), "RLE values")?)?; + validate_compressive_encoding(required(rle.run_lengths.as_deref(), "RLE run lengths")?) + } + Some(Compression::ByteStreamSplit(split)) => validate_compressive_encoding(required( + split.values.as_deref(), + "byte-stream-split values", + )?), + Some(Compression::General(general)) => validate_compressive_encoding(required( + general.values.as_deref(), + "general-compression values", + )?), + Some(Compression::FixedSizeList(list)) => validate_compressive_encoding(required( + list.values.as_deref(), + "fixed-size-list values", + )?), + Some(Compression::PackedStruct(packed)) => validate_compressive_encoding(required( + packed.values.as_deref(), + "packed-struct values", + )?), + Some(Compression::VariablePackedStruct(packed)) => { + for field in &packed.fields { + validate_compressive_encoding(required( + field.value.as_ref(), + "variable packed-struct field", + )?)?; + } + Ok(()) + } + None => Err(Error::invalid_input_source( + "Lance v2.3 compressive encoding is missing its compression variant".into(), + )), + } +} + +fn validate_sparse_positions(positions: Option<&pb21::SparsePositionSet>) -> Result<()> { + use pb21::sparse_position_set::Positions; + + if let Some(Positions::Explicit(encoding)) = + positions.and_then(|positions| positions.positions.as_ref()) + { + validate_compressive_encoding(encoding)?; + } + Ok(()) +} + +fn validate_sparse_validity(validity: Option<&pb21::SparseValiditySet>) -> Result<()> { + if let Some(validity) = validity { + validate_sparse_positions(validity.positions.as_ref())?; + } + Ok(()) +} + +fn validate_sparse_layout(layout: &pb21::SparseLayout) -> Result<()> { + use pb21::{sparse_count_set::Counts, sparse_structural_layer::Layer}; + + if !layout.has_large_chunk { + return Err(Error::invalid_input_source( + "Lance v2.3 sparse pages require the u32 chunk grammar".into(), + )); + } + validate_compressive_encoding(required( + layout.value_compression.as_ref(), + "sparse values", + )?)?; + for layer in &layout.structural_layers { + match layer.layer.as_ref() { + Some(Layer::Validity(validity)) => { + validate_sparse_validity(validity.validity.as_ref())? + } + Some(Layer::List(list)) => { + validate_sparse_positions(list.non_empty_positions.as_ref())?; + if let Some(Counts::Explicit(encoding)) = list + .counts + .as_ref() + .and_then(|counts| counts.counts.as_ref()) + { + validate_compressive_encoding(encoding)?; + } + validate_sparse_validity(list.validity.as_ref())?; + } + Some(Layer::FixedSizeList(list)) => validate_sparse_validity(list.validity.as_ref())?, + None => { + return Err(Error::invalid_input_source( + "Lance v2.3 sparse structural layer is missing its layer variant".into(), + )); + } + } + } + Ok(()) +} + +fn validate_page_layout(layout: &pb21::PageLayout) -> Result<()> { + use pb21::page_layout::Layout; + + match layout.layout.as_ref() { + Some(Layout::MiniBlockLayout(miniblock)) => { + if !miniblock.has_large_chunk { + return Err(Error::invalid_input_source( + "Lance v2.3 miniblock pages require the u32 chunk grammar".into(), + )); + } + if let Some(rep) = miniblock.rep_compression.as_ref() { + validate_compressive_encoding(rep)?; + } + if let Some(def) = miniblock.def_compression.as_ref() { + validate_compressive_encoding(def)?; + } + validate_compressive_encoding(required( + miniblock.value_compression.as_ref(), + "miniblock values", + )?)?; + if let Some(dictionary) = miniblock.dictionary.as_ref() { + validate_compressive_encoding(dictionary)?; + } + Ok(()) + } + Some(Layout::FullZipLayout(fullzip)) => validate_compressive_encoding(required( + fullzip.value_compression.as_ref(), + "full-zip values", + )?), + Some(Layout::ConstantLayout(constant)) => { + if let Some(rep) = constant.rep_compression.as_ref() { + validate_compressive_encoding(rep)?; + } + if let Some(def) = constant.def_compression.as_ref() { + validate_compressive_encoding(def)?; + } + Ok(()) + } + Some(Layout::SparseLayout(sparse)) => validate_sparse_layout(sparse), + Some(Layout::BlobLayout(blob)) => { + let inner = blob.inner_layout.as_deref().ok_or_else(|| { + Error::invalid_input_source( + "Lance v2.3 blob page layout is missing its inner layout".into(), + ) + })?; + validate_page_layout(inner) + } + None => Err(Error::invalid_input_source( + "Lance v2.3 page is missing its page layout".into(), + )), + } +} + +pub fn decode_column( + column_index: u32, + metadata: &pbfile::ColumnMetadata, +) -> Result> { + let page_infos = metadata + .pages + .iter() + .enumerate() + .map(|(page_index, page)| { + let page_layout = structural::decode_page_layout(column_index, page_index, page)?; + validate_page_layout(&page_layout)?; + structural::build_page_info(column_index, page_index, page, page_layout) + }) + .collect::>>()?; + structural::build_column_info(column_index, metadata, page_infos) +} + +pub fn decode_column_metadata( + column_metadatas: &[pbfile::ColumnMetadata], +) -> Result>> { + column_metadatas + .iter() + .enumerate() + .map(|(column_index, metadata)| { + let column_index = u32::try_from(column_index).map_err(|_| { + Error::invalid_input_source("File has more than u32::MAX columns".into()) + })?; + decode_column(column_index, metadata) + }) + .collect() +} + +pub fn projection_from_field_ids( + schema: &Schema, + field_id_to_column_index: &BTreeMap, +) -> ReaderProjection { + structural::projection_from_field_ids(schema, field_id_to_column_index) +} + +pub fn projection_from_whole_schema(schema: &Schema) -> ReaderProjection { + structural::projection_from_field_ids(schema, &super::field_id_to_column_index(schema)) +} + +pub fn projection_from_column_names( + schema: &Schema, + column_names: &[&str], +) -> Result { + structural::projection_from_column_names( + schema, + column_names, + &super::field_id_to_column_index(schema), + ) +} + +pub fn finish_metadata(raw: RawFileMetadata) -> Result { + if (raw.footer.major_version, raw.footer.minor_version) != (2, 3) { + return Err(Error::version_conflict( + "Attempt to use the Lance v2.3 reader for a different file version".to_string(), + raw.footer.major_version, + raw.footer.minor_version, + )); + } + validate_global_buffers(&raw.file_buffers)?; + let column_infos = decode_column_metadata(&raw.column_metadatas)?; + Ok(CachedFileMetadata { + file_schema: raw.file_schema, + column_metadatas: raw.column_metadatas, + column_infos, + num_rows: raw.num_rows, + file_buffers: raw.file_buffers, + num_data_bytes: raw.num_data_bytes, + num_column_metadata_bytes: raw.num_column_metadata_bytes, + num_global_buffer_bytes: raw.num_global_buffer_bytes, + num_footer_bytes: raw.num_footer_bytes, + major_version: raw.footer.major_version, + minor_version: raw.footer.minor_version, + version: ConcreteFileVersion::V2_3, + file_size_bytes: raw.file_size_bytes, + retained_global_buffers: raw.retained_global_buffers, + }) +} + +pub fn validate_global_buffers(buffers: &[BufferDescriptor]) -> Result<()> { + structural::validate_global_buffers(buffers) +} + +pub fn finish_metadata_index(index: FileMetadataIndex) -> Result { + if index.version != ConcreteFileVersion::V2_3 { + let (major, minor) = index.version.to_standard_footer_numbers(); + return Err(Error::version_conflict( + "Attempt to use the Lance v2.3 reader for a different metadata index".to_string(), + major, + minor, + )); + } + validate_global_buffers(&index.file_buffers)?; + Ok(index) +} + +#[cfg(test)] +mod grammar_tests { + use super::*; + use pb21::{ + CompressiveEncoding, Dictionary, Flat, FullZipLayout, PageLayout, Rle, + compressive_encoding::Compression, page_layout::Layout, + }; + + fn flat(bits_per_value: u64) -> CompressiveEncoding { + CompressiveEncoding { + compression: Some(Compression::Flat(Flat { + bits_per_value, + data: None, + })), + } + } + + #[test] + fn accepts_nested_variable_width_rle() { + let rle = CompressiveEncoding { + compression: Some(Compression::Rle(Box::new(Rle { + values: Some(Box::new(flat(32))), + run_lengths: Some(Box::new(flat(16))), + }))), + }; + let dictionary = CompressiveEncoding { + compression: Some(Compression::Dictionary(Box::new(Dictionary { + indices: Some(Box::new(rle)), + items: Some(Box::new(flat(32))), + num_dictionary_items: 1, + }))), + }; + let layout = PageLayout { + layout: Some(Layout::FullZipLayout(FullZipLayout { + value_compression: Some(dictionary), + ..Default::default() + })), + }; + + validate_page_layout(&layout).unwrap(); + } +} diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index 0a7dc3a407f..41b052c6091 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -273,3 +273,7 @@ impl FileWriter { } } } + +#[cfg(test)] +#[path = "writer_tests.rs"] +mod writer_tests; diff --git a/rust/lance-file/src/writer_tests.rs b/rust/lance-file/src/writer_tests.rs new file mode 100644 index 00000000000..c3d69c33f7a --- /dev/null +++ b/rust/lance-file/src/writer_tests.rs @@ -0,0 +1,1887 @@ +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use crate::reader::{FileReader, FileReaderOptions, describe_encoding}; + use crate::testing::FsFixture; + use crate::version::ConcreteFileVersion; + use crate::versions; + use crate::writer::{ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, FileWriter, FileWriterOptions}; + use arrow_array::builder::{Float32Builder, Int32Builder, StringDictionaryBuilder}; + use arrow_array::types::{Float64Type, Int8Type, Int32Type}; + use arrow_array::{ + Array, ArrayRef, Int32Array, LargeBinaryArray, ListArray, RecordBatch, RecordBatchReader, + StringArray, UInt64Array, cast::AsArray, + }; + use arrow_schema::{DataType, Field, Field as ArrowField, Schema, Schema as ArrowSchema}; + use lance_core::cache::LanceCache; + use lance_core::datatypes::Schema as LanceSchema; + use lance_core::utils::tempfile::TempObjFile; + use lance_datagen::{BatchCount, RowCount, array, gen_batch}; + use lance_encoding::compression_config::{CompressionFieldParams, CompressionParams}; + use lance_encoding::decoder::DecoderPlugins; + use lance_io::object_store::ObjectStore; + use lance_io::traits::Writer; + use lance_io::utils::CachedFileSize; + use rstest::rstest; + use tokio::io::AsyncWriteExt; + + fn create_writer( + object_writer: Box, + schema: LanceSchema, + version: ConcreteFileVersion, + options: FileWriterOptions, + ) -> lance_core::Result { + versions::create_writer(version, object_writer, schema, options) + } + + fn create_v2_1_writer_with_compression( + object_writer: Box, + schema: LanceSchema, + options: FileWriterOptions, + compression: CompressionParams, + ) -> lance_core::Result { + versions::v2_1::create_writer_with_compression(object_writer, schema, options, compression) + .map(Into::into) + } + + #[tokio::test] + async fn current_writer_dispatch_rejects_legacy_version() { + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + let object_writer = object_store.create(&path).await.unwrap(); + let Err(error) = versions::create_lazy_writer( + ConcreteFileVersion::V1, + object_writer, + FileWriterOptions::default(), + ) else { + panic!("legacy v1 unexpectedly created a current-format writer"); + }; + assert!(matches!(error, lance_core::Error::NotSupported { .. })); + assert!( + error + .to_string() + .contains("legacy v1 files require an explicit schema and manifest provider") + ); + } + + #[rstest] + #[case::v2_0(ConcreteFileVersion::V2_0)] + #[case::v2_1(ConcreteFileVersion::V2_1)] + #[case::v2_2(ConcreteFileVersion::V2_2)] + #[case::v2_3(ConcreteFileVersion::V2_3)] + #[tokio::test] + async fn version_leaf_writes_exact_standard_footer(#[case] version: ConcreteFileVersion) { + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new( + "value", + DataType::Int32, + true, + )])) + .unwrap(); + let mut writer = create_writer( + object_store.create(&path).await.unwrap(), + schema, + version, + FileWriterOptions::default(), + ) + .unwrap(); + let summary = writer.finish().await.unwrap(); + let footer = object_store + .open(&path) + .await + .unwrap() + .get_range(summary.size_bytes as usize - 8..summary.size_bytes as usize) + .await + .unwrap(); + let actual = ( + u16::from_le_bytes([footer[0], footer[1]]), + u16::from_le_bytes([footer[2], footer[3]]), + ); + assert_eq!(actual, version.to_standard_footer_numbers()); + } + + #[rstest] + #[case::v2_0(ConcreteFileVersion::V2_0)] + #[case::v2_1(ConcreteFileVersion::V2_1)] + #[case::v2_2(ConcreteFileVersion::V2_2)] + #[case::v2_3(ConcreteFileVersion::V2_3)] + fn packed_struct_is_one_physical_column(#[case] version: ConcreteFileVersion) { + let packed = ArrowField::new( + "packed", + DataType::Struct(vec![ArrowField::new("child", DataType::Int32, true)].into()), + true, + ) + .with_metadata(HashMap::from([("packed".to_string(), "true".to_string())])); + let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + packed, + ArrowField::new("tail", DataType::Int32, true), + ])) + .unwrap(); + + let projection = + versions::reader_projection_from_column_names(version, &schema, &["tail"]).unwrap(); + assert_eq!(projection.column_indices, vec![1]); + + let (field_ids, column_indices) = versions::data_file_columns(version, &schema); + assert_eq!(field_ids.len(), 2); + assert_eq!(column_indices, vec![0, 1]); + assert_eq!( + schema + .fields + .iter() + .map(|field| versions::physical_column_count(version, field)) + .sum::(), + 2 + ); + } + + #[rstest] + #[case::v1(ConcreteFileVersion::V1, &[0, 1, 2], &[0, 1, 2])] + #[case::v2_0(ConcreteFileVersion::V2_0, &[0, 1, 2], &[0, 1, 2])] + #[case::v2_1(ConcreteFileVersion::V2_1, &[1, 2], &[0, 1])] + #[case::v2_2(ConcreteFileVersion::V2_2, &[1, 2], &[0, 1])] + #[case::v2_3(ConcreteFileVersion::V2_3, &[1, 2], &[0, 1])] + fn data_file_mapping_tracks_only_version_physical_fields( + #[case] version: ConcreteFileVersion, + #[case] expected_field_ids: &[i32], + #[case] expected_column_indices: &[i32], + ) { + let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new( + "nested", + DataType::Struct(vec![ArrowField::new("child", DataType::Int32, true)].into()), + true, + ), + ArrowField::new("tail", DataType::Int32, true), + ])) + .unwrap(); + + let (field_ids, column_indices) = versions::data_file_columns(version, &schema); + assert_eq!(field_ids, expected_field_ids); + assert_eq!(column_indices, expected_column_indices); + } + + fn compatibility_fixture_batch() -> RecordBatch { + let row_count = 4097; + let ids = Arc::new(Int32Array::from_iter_values(0..row_count)) as ArrayRef; + let names = Arc::new(StringArray::from_iter((0..row_count).map(|index| { + (index % 7 != 0).then(|| format!("value-{index:04}-deterministic-fixture")) + }))) as ArrayRef; + let items = Arc::new(ListArray::from_iter_primitive::( + (0..row_count).map(|index| { + (index % 11 != 0).then(|| { + vec![ + Some(index), + (index % 5 != 0).then_some(index * 2), + Some(index * 3), + ] + }) + }), + )) as ArrayRef; + let mut categories = StringDictionaryBuilder::::new(); + for index in 0..row_count { + if index % 13 == 0 { + categories.append_null(); + } else { + categories + .append(match index % 3 { + 0 => "red", + 1 => "green", + _ => "blue", + }) + .unwrap(); + } + } + let categories = Arc::new(categories.finish()) as ArrayRef; + let blobs = Arc::new(LargeBinaryArray::from_iter_values( + (0..row_count) + .map(|index| format!("blob-{index:04}-deterministic-payload").into_bytes()), + )) as ArrayRef; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("name", DataType::Utf8, true).with_metadata(HashMap::from([( + "lance-encoding:compression".to_string(), + "none".to_string(), + )])), + ArrowField::new( + "items", + DataType::List(Arc::new(ArrowField::new("item", DataType::Int32, true))), + true, + ), + ArrowField::new( + "category", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + ) + .with_metadata(HashMap::from([( + "lance-encoding:dict-values-compression".to_string(), + "none".to_string(), + )])), + ArrowField::new("blob", DataType::LargeBinary, true).with_metadata(HashMap::from([( + "lance-encoding:blob".to_string(), + "true".to_string(), + )])), + ])); + RecordBatch::try_new(schema, vec![ids, names, items, categories, blobs]).unwrap() + } + + fn stable_current_fixture(version: ConcreteFileVersion) -> &'static [u8] { + match version { + ConcreteFileVersion::V1 => unreachable!("v1 uses the legacy writer fixture"), + ConcreteFileVersion::V2_0 => { + include_bytes!("../test_data/exact_versions/v2_0.lance") + } + ConcreteFileVersion::V2_1 => { + include_bytes!("../test_data/exact_versions/v2_1.lance") + } + ConcreteFileVersion::V2_2 => { + include_bytes!("../test_data/exact_versions/v2_2.lance") + } + ConcreteFileVersion::V2_3 => { + unreachable!("v2.3 is unstable and does not have a compatibility fixture") + } + } + } + + fn assert_blob_column_eq(actual: &dyn Array, expected: &dyn Array) { + let actual = actual.as_binary::(); + let expected = expected.as_binary::(); + assert_eq!(actual.len(), expected.len()); + for index in 0..actual.len() { + assert_eq!( + actual.is_null(index), + expected.is_null(index), + "blob validity differs at row {index}" + ); + if actual.is_valid(index) { + assert_eq!( + actual.value(index), + expected.value(index), + "blob payload differs at row {index}" + ); + } + } + } + + fn assert_wire_bytes_equal(actual: &[u8], expected: &[u8]) { + if let Some(offset) = actual + .iter() + .zip(expected) + .position(|(actual, expected)| actual != expected) + { + panic!( + "wire fixture first differs at byte {offset}: actual={}, expected={}", + actual[offset], expected[offset] + ); + } + assert_eq!( + actual.len(), + expected.len(), + "wire fixture length changed after a common {}-byte prefix", + actual.len().min(expected.len()) + ); + } + + #[rstest] + #[case::v2_0(ConcreteFileVersion::V2_0)] + #[case::v2_1(ConcreteFileVersion::V2_1)] + #[case::v2_2(ConcreteFileVersion::V2_2)] + #[tokio::test] + async fn stable_current_writer_is_byte_compatible(#[case] version: ConcreteFileVersion) { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + + let batch = compatibility_fixture_batch(); + let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); + schema.set_dictionary(&batch).unwrap(); + let fs = FsFixture::default(); + let object_writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); + let options = FileWriterOptions { + data_cache_bytes: Some(1), + max_page_bytes: Some(1024), + ..Default::default() + }; + let mut writer: FileWriter = match version { + ConcreteFileVersion::V1 => unreachable!(), + ConcreteFileVersion::V2_0 => { + versions::v2_0::create_writer(object_writer, schema.clone(), options) + .map(Into::into) + } + ConcreteFileVersion::V2_1 => { + versions::v2_1::create_writer(object_writer, schema.clone(), options) + .map(Into::into) + } + ConcreteFileVersion::V2_2 => { + versions::v2_2::create_writer(object_writer, schema.clone(), options) + .map(Into::into) + } + ConcreteFileVersion::V2_3 => unreachable!(), + } + .unwrap(); + for offset in (0..batch.num_rows()).step_by(1024) { + let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); + writer.write_batch(&slice).await.unwrap(); + } + let summary = writer.finish().await.unwrap(); + let actual = fs + .object_store + .open(&fs.tmp_path) + .await + .unwrap() + .get_range(0..summary.size_bytes as usize) + .await + .unwrap(); + let expected = stable_current_fixture(version); + assert_wire_bytes_equal(actual.as_ref(), expected); + + let fixture_fs = FsFixture::default(); + let mut fixture_writer = fixture_fs + .object_store + .create(&fixture_fs.tmp_path) + .await + .unwrap(); + fixture_writer.write_all(expected).await.unwrap(); + Writer::shutdown(fixture_writer.as_mut()).await.unwrap(); + let scheduler = fixture_fs + .scheduler + .open_file( + &fixture_fs.tmp_path, + &CachedFileSize::new(expected.len() as u64), + ) + .await + .unwrap(); + let reader = FileReader::try_open( + scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(reader.metadata().version, version); + assert!( + reader + .metadata() + .column_metadatas + .iter() + .any(|metadata| metadata.pages.len() > 1) + ); + let batches = reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 16, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + batch.num_rows() + ); + assert!( + batches + .iter() + .all(|actual| actual.schema_ref() == batch.schema_ref()) + ); + let mut row_offset = 0; + for actual in &batches { + let expected = batch.slice(row_offset, actual.num_rows()); + assert_blob_column_eq(actual.column(4).as_ref(), expected.column(4).as_ref()); + row_offset += actual.num_rows(); + } + assert_eq!(row_offset, batch.num_rows()); + } + + async fn write_v2_3_fixture(batch: &RecordBatch, schema: &LanceSchema) -> Vec { + let fs = FsFixture::default(); + let object_writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); + let mut writer = versions::v2_3::create_writer( + object_writer, + schema.clone(), + FileWriterOptions { + data_cache_bytes: Some(1), + max_page_bytes: Some(1024), + ..Default::default() + }, + ) + .unwrap(); + for offset in (0..batch.num_rows()).step_by(1024) { + let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); + writer.write_batch(&slice).await.unwrap(); + } + let summary = writer.finish().await.unwrap(); + fs.object_store + .open(&fs.tmp_path) + .await + .unwrap() + .get_range(0..summary.size_bytes as usize) + .await + .unwrap() + .to_vec() + } + + #[tokio::test] + async fn v2_3_writer_emits_current_exact_grammar() { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + + let batch = compatibility_fixture_batch(); + let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); + schema.set_dictionary(&batch).unwrap(); + let first = write_v2_3_fixture(&batch, &schema).await; + let second = write_v2_3_fixture(&batch, &schema).await; + assert_eq!(first, second); + assert_eq!( + &first[first.len() - 8..], + &[2, 0, 3, 0, b'L', b'A', b'N', b'C'] + ); + + let fs = FsFixture::default(); + let mut object_writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); + object_writer.write_all(&first).await.unwrap(); + Writer::shutdown(object_writer.as_mut()).await.unwrap(); + let scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::new(first.len() as u64)) + .await + .unwrap(); + let reader = FileReader::try_open( + scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(reader.metadata().version, ConcreteFileVersion::V2_3); + assert!( + reader + .metadata() + .column_metadatas + .iter() + .any(|metadata| metadata.pages.len() > 1) + ); + let batches = reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 16, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut row_offset = 0; + for actual in &batches { + let expected = batch.slice(row_offset, actual.num_rows()); + assert_blob_column_eq(actual.column(4).as_ref(), expected.column(4).as_ref()); + row_offset += actual.num_rows(); + } + assert_eq!(row_offset, batch.num_rows()); + } + + #[tokio::test] + async fn v1_writer_is_byte_compatible() { + use crate::versions::v1::reader::FileReader as V1Reader; + use crate::versions::v1::writer::{ + FileWriter as V1Writer, FileWriterOptions as V1WriterOptions, NotSelfDescribing, + }; + + let expected = include_bytes!("../test_data/exact_versions/v1.lance"); + let batch = compatibility_fixture_batch(); + let mut schema = LanceSchema::try_from(batch.schema().as_ref()).unwrap(); + schema.set_dictionary(&batch).unwrap(); + let fs = FsFixture::default(); + let mut writer = V1Writer::::try_new( + fs.object_store.as_ref(), + &fs.tmp_path, + schema.clone(), + &V1WriterOptions { + collect_stats_for_fields: Some(Vec::new()), + }, + ) + .await + .unwrap(); + for offset in (0..batch.num_rows()).step_by(1024) { + let slice = batch.slice(offset, (batch.num_rows() - offset).min(1024)); + writer.write(std::slice::from_ref(&slice)).await.unwrap(); + } + let summary = writer.finish().await.unwrap(); + let actual = fs + .object_store + .open(&fs.tmp_path) + .await + .unwrap() + .get_range(0..summary.size_bytes as usize) + .await + .unwrap(); + assert_wire_bytes_equal(actual.as_ref(), expected); + + let fixture_fs = FsFixture::default(); + let mut fixture_writer = fixture_fs + .object_store + .create(&fixture_fs.tmp_path) + .await + .unwrap(); + fixture_writer.write_all(expected).await.unwrap(); + Writer::shutdown(fixture_writer.as_mut()).await.unwrap(); + let reader = V1Reader::try_new( + fixture_fs.object_store.as_ref(), + &fixture_fs.tmp_path, + schema.clone(), + ) + .await + .unwrap(); + let actual_batch = reader + .read_range(0..batch.num_rows(), &schema) + .await + .unwrap(); + assert_eq!(reader.num_batches(), 5); + assert_eq!(actual_batch.num_rows(), batch.num_rows()); + assert_eq!(actual_batch.column(0).to_data(), batch.column(0).to_data()); + assert_eq!(actual_batch.column(1).to_data(), batch.column(1).to_data()); + assert_blob_column_eq(actual_batch.column(4).as_ref(), batch.column(4).as_ref()); + } + + #[tokio::test] + async fn test_basic_write() { + let tmp_path = TempObjFile::default(); + let obj_store = Arc::new(ObjectStore::local()); + + let reader = gen_batch() + .col("score", array::rand::()) + .into_reader_rows(RowCount::from(1000), BatchCount::from(10)); + + let writer = obj_store.create(&tmp_path).await.unwrap(); + + let lance_schema = + lance_core::datatypes::Schema::try_from(reader.schema().as_ref()).unwrap(); + + let mut file_writer = create_writer( + writer, + lance_schema, + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .unwrap(); + + for batch in reader { + file_writer.write_batch(&batch.unwrap()).await.unwrap(); + } + file_writer.add_schema_metadata("foo", "bar"); + file_writer.finish().await.unwrap(); + // Tests asserting the contents of the written file are in reader.rs + } + + #[tokio::test] + async fn test_write_empty() { + let tmp_path = TempObjFile::default(); + let obj_store = Arc::new(ObjectStore::local()); + + let reader = gen_batch() + .col("score", array::rand::()) + .into_reader_rows(RowCount::from(0), BatchCount::from(0)); + + let writer = obj_store.create(&tmp_path).await.unwrap(); + + let lance_schema = + lance_core::datatypes::Schema::try_from(reader.schema().as_ref()).unwrap(); + + let mut file_writer = create_writer( + writer, + lance_schema, + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .unwrap(); + + for batch in reader { + file_writer.write_batch(&batch.unwrap()).await.unwrap(); + } + file_writer.add_schema_metadata("foo", "bar"); + file_writer.finish().await.unwrap(); + } + + // Read a single column back at an explicit range/index set, returning its + // `Int32` values. Reading one column (or an equal-length group) at a time is + // how unequal-length files are consumed: a full scan across columns of + // differing lengths cannot form a single rectangular batch. + async fn read_int32_column( + reader: &FileReader, + schema: &LanceSchema, + version: ConcreteFileVersion, + name: &str, + params: lance_io::ReadBatchParams, + ) -> Vec> { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + + let projection = + versions::reader_projection_from_column_names(version, schema, &[name]).unwrap(); + let batches: Vec = reader + .read_stream_projected(params, 1024, 16, projection, FilterExpression::no_filter()) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + } + + /// A single file may hold columns of differing item counts, written by + /// advancing each column's row counter independently (no shared global + /// counter). + #[rstest] + #[tokio::test] + async fn test_write_columns_unequal_lengths( + #[values(ConcreteFileVersion::V2_0, ConcreteFileVersion::V2_1)] + version: ConcreteFileVersion, + ) { + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + version, + FileWriterOptions::default(), + ) + .unwrap(); + + // Field "a" gets 5 values across two calls (appending), field "b" gets a + // single value, and field "c" is never written (a zero-length column). + let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); + writer.write_column(0, a1).await.unwrap(); + writer.write_column(1, b).await.unwrap(); + let a2: ArrayRef = Arc::new(Int32Array::from(vec![4, 5])); + writer.write_column(0, a2).await.unwrap(); + // An empty array is a no-op whether or not the field already has rows: + // field "a" keeps its 5 rows, field "c" stays a zero-length column. + let empty: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); + writer.write_column(0, empty.clone()).await.unwrap(); + writer.write_column(2, empty).await.unwrap(); + + let summary = writer.finish().await.unwrap(); + // The file's logical length is the longest column. + assert_eq!(summary.num_rows, 5); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // Per-column row counts are recorded in / derivable from file metadata. + assert_eq!(reader.num_rows(), 5); + assert_eq!(reader.column_num_rows(0).unwrap(), 5); + assert_eq!(reader.column_num_rows(1).unwrap(), 1); + assert_eq!(reader.column_num_rows(2).unwrap(), 0); + assert!(reader.column_num_rows(3).is_err()); + + // Each column reads back independently at its own length. + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Range(0..5) + ) + .await, + vec![Some(1), Some(2), Some(3), Some(4), Some(5)], + ); + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "b", + ReadBatchParams::Range(0..1) + ) + .await, + vec![Some(10)], + ); + + // Random access by position within the longer column returns the right + // value even though other columns are shorter. (The take path requires + // strictly increasing indices.) + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Indices(arrow_array::UInt32Array::from(vec![0, 2, 4])), + ) + .await, + vec![Some(1), Some(3), Some(5)], + ); + } + + /// Reading an unequal-length file: + /// - a projection whose columns are equal length full-scans normally; + /// - a full scan across columns of differing length is rejected up front, + /// before any batch is produced (even though a prefix would be rectangular); + /// - a bounded read is valid as long as every projected column covers it; + /// - a single-column `RangeFull` resolves to that column's own length, not + /// the file's (maximum) length. + #[rstest] + #[tokio::test] + async fn test_read_unequal_length_projection( + #[values(ConcreteFileVersion::V2_0, ConcreteFileVersion::V2_1)] + version: ConcreteFileVersion, + ) { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + version, + FileWriterOptions::default(), + ) + .unwrap(); + // "a" and "b" are equal length (5); "c" is shorter (1). + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10]))) + .await + .unwrap(); + writer + .write_column(2, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let read = |names: &'static [&'static str], params: ReadBatchParams| { + let projection = + versions::reader_projection_from_column_names(version, &lance_schema, names) + .unwrap(); + async { + match reader + .read_stream_projected( + params, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + { + Ok(stream) => stream.try_collect::>().await, + Err(e) => Err(e), + } + } + }; + let col_values = |batches: &[RecordBatch], idx: usize| -> Vec> { + batches + .iter() + .flat_map(|b| { + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + }; + + // Equal-length projection [a, b] full-scans into rectangular batches. + let batches = read(&["a", "b"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + assert_eq!( + col_values(&batches, 1), + vec![Some(6), Some(7), Some(8), Some(9), Some(10)] + ); + + // A mismatched-length projection [a, c] (5 vs 1) is rejected before any + // batch is yielded, regardless of the read params — its columns cannot + // be combined into rectangular batches. The error names each column's + // length so the caller can see which column is the odd one out. + let err = read(&["a", "c"], ReadBatchParams::RangeFull) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("a=5") && err.contains("c=1"), + "error should name each column's length, got: {err}" + ); + assert!( + read(&["a", "c"], ReadBatchParams::Range(0..1)) + .await + .is_err(), + "even a common-prefix read of unequal-length columns must error" + ); + + // A single-column RangeFull resolves to that column's own length. + let batches = read(&["c"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(100)]); + let batches = read(&["a"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + + // RangeFrom/RangeTo likewise resolve against the projected column's own + // length rather than the file's longest column. + let batches = read(&["a"], ReadBatchParams::RangeFrom(2..)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(3), Some(4), Some(5)]); + // RangeFrom on the short column "c" resolves to length 1, not 5. + let batches = read(&["c"], ReadBatchParams::RangeFrom(0..)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(100)]); + let batches = read(&["a"], ReadBatchParams::RangeTo(..3)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(1), Some(2), Some(3)]); + // A bound past the projected column's length errors. + assert!( + read(&["a"], ReadBatchParams::RangeTo(..6)).await.is_err(), + "RangeTo past the column length must error" + ); + assert!( + read(&["c"], ReadBatchParams::RangeFrom(2..)).await.is_err(), + "RangeFrom past the column length must error" + ); + } + + /// A struct and a list column each map to multiple physical columns, and a + /// list's item column is longer than its top-level row count. The + /// projection-length check must partition `column_indices` by top-level + /// field and use each field's root column, so an ordinary (rectangular) file + /// with nested columns still reads under the new validation path. + #[rstest] + #[tokio::test] + async fn test_read_nested_columns_under_validation( + #[values(ConcreteFileVersion::V2_0, ConcreteFileVersion::V2_1)] + version: ConcreteFileVersion, + ) { + use arrow_array::types::Int32Type; + use arrow_array::{ListArray, StructArray}; + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let struct_type = DataType::Struct( + vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ] + .into(), + ); + let list_type = DataType::List(Arc::new(ArrowField::new("item", DataType::Int32, true))); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("s", struct_type, true), + ArrowField::new("lst", list_type, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let s: ArrayRef = Arc::new(StructArray::from(vec![ + ( + Arc::new(ArrowField::new("x", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("y", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![11, 21, 31])) as ArrayRef, + ), + ])); + // 3 lists, 6 items: the item column is longer than the top-level rows. + let lst: ArrayRef = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5), Some(6)]), + ])); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![a, s, lst]).unwrap(); + + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + version, + FileWriterOptions::default(), + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // If `validate_field_length` mispartitioned the physical columns, the + // length check would read the wrong root column (e.g. the list's item + // column, length 6) and spuriously reject this rectangular file. + for names in [&["a", "s", "lst"][..], &["a", "lst"][..], &["a", "s"][..]] { + let projection = + versions::reader_projection_from_column_names(version, &lance_schema, names) + .unwrap(); + let batches: Vec = reader + .read_stream_projected( + ReadBatchParams::RangeFull, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, 3, + "projection {names:?} should read 3 top-level rows" + ); + } + } + + /// `write_column` rejects invalid inputs at the API boundary with + /// descriptive errors: a writer without an explicit schema, an + /// out-of-bounds field index, and a null written into a non-nullable field. + #[tokio::test] + async fn test_write_column_validation_errors() { + // A lazy-schema writer cannot infer the schema from a single column. + let fs = FsFixture::default(); + let mut lazy_writer = versions::v2_1::create_lazy_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + FileWriterOptions::default(), + ); + let err = lazy_writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("explicit schema"), + "expected explicit-schema error, got: {err}" + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + // An out-of-bounds field index is rejected, naming the index and count. + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .unwrap(); + let err = writer + .write_column(5, Arc::new(Int32Array::from(vec![1]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains('5') && err.contains('2'), + "expected out-of-bounds error naming index 5 and 2 fields, got: {err}" + ); + + // A null in a non-nullable field ("a") is rejected. + let err = writer + .write_column(0, Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("non-null"), + "expected nullability error, got: {err}" + ); + } + + /// The blocking read path applies the same projection-length validation as + /// the async path: a short single column resolves to its own length, and a + /// mismatched-length projection errors up front. + #[rstest] + #[tokio::test] + async fn test_blocking_read_unequal_length( + #[values(ConcreteFileVersion::V2_0, ConcreteFileVersion::V2_1)] + version: ConcreteFileVersion, + ) { + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + version, + FileWriterOptions::default(), + ) + .unwrap(); + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = Arc::new( + FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(), + ); + + // Single short column: RangeFull resolves to its own length (1). + let proj_c = + versions::reader_projection_from_column_names(version, &lance_schema, &["c"]).unwrap(); + let reader_c = reader.clone(); + let batches = tokio::task::spawn_blocking(move || { + reader_c + .read_stream_projected_blocking( + ReadBatchParams::RangeFull, + 1024, + Some(proj_c), + FilterExpression::no_filter(), + ) + .unwrap() + .collect::, _>>() + .unwrap() + }) + .await + .unwrap(); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 1); + + // A mismatched projection [a, c] errors on the blocking path too. + let proj_ac = + versions::reader_projection_from_column_names(version, &lance_schema, &["a", "c"]) + .unwrap(); + let reader_ac = reader.clone(); + let is_err = tokio::task::spawn_blocking(move || { + reader_ac + .read_stream_projected_blocking( + ReadBatchParams::RangeFull, + 1024, + Some(proj_ac), + FilterExpression::no_filter(), + ) + .is_err() + }) + .await + .unwrap(); + assert!( + is_err, + "blocking full scan across unequal-length columns must error" + ); + } + + /// Files written the ordinary (rectangular) way keep equal column lengths, + /// so the unequal-length support is backwards compatible. + #[tokio::test] + async fn test_write_batch_keeps_equal_lengths() { + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let mut writer = create_writer( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema, + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![4, 5, 6])), + ], + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + let summary = writer.finish().await.unwrap(); + assert_eq!(summary.num_rows, 3); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(reader.column_num_rows(0).unwrap(), 3); + assert_eq!(reader.column_num_rows(1).unwrap(), 3); + } + + #[tokio::test] + async fn test_max_page_bytes_enforced() { + let arrow_field = Field::new("data", DataType::UInt64, false); + let arrow_schema = Schema::new(vec![arrow_field]); + let lance_schema = LanceSchema::try_from(&arrow_schema).unwrap(); + + // 8MiB + let data: Vec = (0..1_000_000).collect(); + let array = UInt64Array::from(data); + let batch = + RecordBatch::try_new(arrow_schema.clone().into(), vec![Arc::new(array)]).unwrap(); + + let options = FileWriterOptions { + max_page_bytes: Some(1024 * 1024), // 1MB + ..Default::default() + }; + + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + let mut writer = create_writer( + object_store.create(&path).await.unwrap(), + lance_schema, + ConcreteFileVersion::V2_0, + options, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + let fs = FsFixture::default(); + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let column_meta = file_reader.metadata(); + + let mut total_page_num: u32 = 0; + for (col_idx, col_metadata) in column_meta.column_metadatas.iter().enumerate() { + assert!( + !col_metadata.pages.is_empty(), + "Column {} has no pages", + col_idx + ); + + for (page_idx, page) in col_metadata.pages.iter().enumerate() { + total_page_num += 1; + let total_size: u64 = page.buffer_sizes.iter().sum(); + assert!( + total_size <= 1024 * 1024, + "Column {} Page {} size {} exceeds 1MB limit", + col_idx, + page_idx, + total_size + ); + } + } + + assert_eq!(total_page_num, 8) + } + + #[tokio::test(flavor = "current_thread")] + async fn test_max_page_bytes_env_var() { + let arrow_field = Field::new("data", DataType::UInt64, false); + let arrow_schema = Schema::new(vec![arrow_field]); + let lance_schema = LanceSchema::try_from(&arrow_schema).unwrap(); + // 4MiB + let data: Vec = (0..500_000).collect(); + let array = UInt64Array::from(data); + let batch = + RecordBatch::try_new(arrow_schema.clone().into(), vec![Arc::new(array)]).unwrap(); + + // 2MiB + unsafe { + std::env::set_var(ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, "2097152"); + } + + let options = FileWriterOptions { + max_page_bytes: None, // enforce env + ..Default::default() + }; + + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + let mut writer = create_writer( + object_store.create(&path).await.unwrap(), + lance_schema.clone(), + ConcreteFileVersion::V2_1, + options, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + let fs = FsFixture::default(); + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + for col_metadata in file_reader.metadata().column_metadatas.iter() { + for page in col_metadata.pages.iter() { + let total_size: u64 = page.buffer_sizes.iter().sum(); + assert!( + total_size <= 2 * 1024 * 1024, + "Page size {} exceeds 2MB limit", + total_size + ); + } + } + + unsafe { + std::env::set_var(ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, ""); + } + } + + #[tokio::test] + async fn test_compression_overrides_end_to_end() { + // Create test schema with different column types + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("customer_id", DataType::Int32, false), + ArrowField::new("product_id", DataType::Int32, false), + ArrowField::new("quantity", DataType::Int32, false), + ArrowField::new("price", DataType::Float32, false), + ArrowField::new("description", DataType::Utf8, false), + ])); + + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + // Create test data with patterns suitable for different compression + let mut customer_ids = Int32Builder::new(); + let mut product_ids = Int32Builder::new(); + let mut quantities = Int32Builder::new(); + let mut prices = Float32Builder::new(); + let mut descriptions = Vec::new(); + + // Generate data with specific patterns: + // - customer_id: highly repetitive (good for RLE) + // - product_id: moderately repetitive (good for RLE) + // - quantity: random values (not good for RLE) + // - price: some repetition + // - description: long strings (good for Zstd) + for i in 0..10000 { + // Customer ID repeats every 100 rows (100 unique customers) + // This creates runs of 100 identical values + customer_ids.append_value(i / 100); + + // Product ID has only 5 unique values with long runs + product_ids.append_value(i / 2000); + + // Quantity is mostly 1 with occasional other values + quantities.append_value(if i % 10 == 0 { 5 } else { 1 }); + + // Price has only 3 unique values + prices.append_value(match i % 3 { + 0 => 9.99, + 1 => 19.99, + _ => 29.99, + }); + + // Descriptions are repetitive but we'll keep them simple + descriptions.push(format!("Product {}", i / 2000)); + } + + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(customer_ids.finish()), + Arc::new(product_ids.finish()), + Arc::new(quantities.finish()), + Arc::new(prices.finish()), + Arc::new(StringArray::from(descriptions)), + ], + ) + .unwrap(); + + // Configure compression parameters + let mut params = CompressionParams::new(); + + // RLE for ID columns (ends with _id) + params.columns.insert( + "*_id".to_string(), + CompressionFieldParams { + rle_threshold: Some(0.5), // Lower threshold to trigger RLE more easily + compression: None, // Will use default compression if any + compression_level: None, + bss: Some(lance_encoding::compression_config::BssMode::Off), // Explicitly disable BSS to ensure RLE is used + minichunk_size: None, + }, + ); + + // For now, we'll skip Zstd compression since it's not imported + // In a real implementation, you could add other compression types here + + // Configure file writer options + let options = FileWriterOptions { + max_page_bytes: Some(64 * 1024), // 64KB pages + ..Default::default() + }; + + // Write the file + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + + let mut writer = create_v2_1_writer_with_compression( + object_store.create(&path).await.unwrap(), + lance_schema.clone(), + options, + params, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + writer.add_schema_metadata("compression_test", "configured_compression"); + writer.finish().await.unwrap(); + + // Now write the same data without compression overrides for comparison + let path_no_compression = TempObjFile::default(); + let default_options = FileWriterOptions { + max_page_bytes: Some(64 * 1024), + ..Default::default() + }; + + let mut writer_no_compression = create_writer( + object_store.create(&path_no_compression).await.unwrap(), + lance_schema.clone(), + ConcreteFileVersion::V2_1, + default_options, + ) + .unwrap(); + + writer_no_compression.write_batch(&batch).await.unwrap(); + writer_no_compression.finish().await.unwrap(); + + // Note: With our current data patterns and RLE compression, the compressed file + // might actually be slightly larger due to compression metadata overhead. + // This is expected and the test is mainly to verify the system works end-to-end. + + // Read back the compressed file and verify data integrity + let fs = FsFixture::default(); + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // Verify metadata + let metadata = file_reader.metadata(); + assert_eq!(metadata.version, ConcreteFileVersion::V2_1); + + let schema = file_reader.schema(); + assert_eq!( + schema.metadata.get("compression_test"), + Some(&"configured_compression".to_string()) + ); + + // Verify the actual encodings used + let column_metadatas = &metadata.column_metadatas; + + // Check customer_id column (index 0) - should use RLE due to our configuration + assert!(!column_metadatas[0].pages.is_empty()); + let customer_id_encoding = describe_encoding(&column_metadatas[0].pages[0]); + assert!( + customer_id_encoding.contains("RLE") || customer_id_encoding.contains("Rle"), + "customer_id column should use RLE encoding due to '*_id' pattern match, but got: {}", + customer_id_encoding + ); + + // Check product_id column (index 1) - should use RLE due to our configuration + assert!(!column_metadatas[1].pages.is_empty()); + let product_id_encoding = describe_encoding(&column_metadatas[1].pages[0]); + assert!( + product_id_encoding.contains("RLE") || product_id_encoding.contains("Rle"), + "product_id column should use RLE encoding due to '*_id' pattern match, but got: {}", + product_id_encoding + ); + } + + #[tokio::test] + async fn test_field_metadata_compression() { + // Test that field metadata compression settings are respected + let mut metadata = HashMap::new(); + metadata.insert( + lance_encoding::constants::COMPRESSION_META_KEY.to_string(), + "zstd".to_string(), + ); + metadata.insert( + lance_encoding::constants::COMPRESSION_LEVEL_META_KEY.to_string(), + "6".to_string(), + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("text", DataType::Utf8, false).with_metadata(metadata.clone()), + ArrowField::new("data", DataType::Int32, false).with_metadata(HashMap::from([( + lance_encoding::constants::COMPRESSION_META_KEY.to_string(), + "none".to_string(), + )])), + ])); + + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + // Create test data + let id_array = Int32Array::from_iter_values(0..1000); + let text_array = StringArray::from_iter_values( + (0..1000).map(|i| format!("test string {} repeated text", i)), + ); + let data_array = Int32Array::from_iter_values((0..1000).map(|i| i * 2)); + + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(id_array), + Arc::new(text_array), + Arc::new(data_array), + ], + ) + .unwrap(); + + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + + // Create encoding strategy that will read from field metadata + let params = CompressionParams::new(); + let options = FileWriterOptions::default(); + let mut writer = create_v2_1_writer_with_compression( + object_store.create(&path).await.unwrap(), + lance_schema.clone(), + options, + params, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + // Read back metadata + let fs = FsFixture::default(); + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let column_metadatas = &file_reader.metadata().column_metadatas; + + // The text column (index 1) should use zstd compression based on metadata + let text_encoding = describe_encoding(&column_metadatas[1].pages[0]); + // For string columns, we expect Binary encoding with zstd compression + assert!( + text_encoding.contains("Zstd"), + "text column should use zstd compression from field metadata, but got: {}", + text_encoding + ); + + // The data column (index 2) should use no compression based on metadata + let data_encoding = describe_encoding(&column_metadatas[2].pages[0]); + // For Int32 columns with "none" compression, we expect Flat encoding without compression + assert!( + data_encoding.contains("Flat") && data_encoding.contains("compression: None"), + "data column should use no compression from field metadata, but got: {}", + data_encoding + ); + } + + #[tokio::test] + async fn test_field_metadata_rle_threshold() { + // Test that RLE threshold from field metadata is respected + let mut metadata = HashMap::new(); + metadata.insert( + lance_encoding::constants::RLE_THRESHOLD_META_KEY.to_string(), + "0.9".to_string(), + ); + // Also set compression to ensure RLE is used + metadata.insert( + lance_encoding::constants::COMPRESSION_META_KEY.to_string(), + "lz4".to_string(), + ); + // Explicitly disable BSS to ensure RLE is tested + metadata.insert( + lance_encoding::constants::BSS_META_KEY.to_string(), + "off".to_string(), + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("status", DataType::Int32, false).with_metadata(metadata), + ])); + + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + // Create data with very high repetition (3 runs for 10000 values = 0.0003 ratio) + let status_array = Int32Array::from_iter_values( + std::iter::repeat_n(200, 8000) + .chain(std::iter::repeat_n(404, 1500)) + .chain(std::iter::repeat_n(500, 500)), + ); + + let batch = + RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(status_array)]).unwrap(); + + let path = TempObjFile::default(); + let object_store = ObjectStore::local(); + + // Create encoding strategy that will read from field metadata + let params = CompressionParams::new(); + let options = FileWriterOptions::default(); + let mut writer = create_v2_1_writer_with_compression( + object_store.create(&path).await.unwrap(), + lance_schema.clone(), + options, + params, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + // Read back and check encoding + let fs = FsFixture::default(); + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let column_metadatas = &file_reader.metadata().column_metadatas; + let status_encoding = describe_encoding(&column_metadatas[0].pages[0]); + assert!( + status_encoding.contains("RLE") || status_encoding.contains("Rle"), + "status column should use RLE encoding due to metadata threshold, but got: {}", + status_encoding + ); + } + + #[tokio::test] + async fn test_large_page_split_on_read() { + use arrow_array::Array; + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + // Test that large pages written with relaxed limits can be split during read + + let arrow_field = ArrowField::new("data", DataType::Binary, false); + let arrow_schema = ArrowSchema::new(vec![arrow_field]); + let lance_schema = LanceSchema::try_from(&arrow_schema).unwrap(); + + // Create a large binary value (40MB) to trigger large page creation + let large_value = vec![42u8; 40 * 1024 * 1024]; + let array = arrow_array::BinaryArray::from(vec![ + Some(large_value.as_slice()), + Some(b"small value"), + ]); + let batch = RecordBatch::try_new(Arc::new(arrow_schema), vec![Arc::new(array)]).unwrap(); + + // Write with relaxed page size limit (128MB) + let options = FileWriterOptions { + max_page_bytes: Some(128 * 1024 * 1024), + ..Default::default() + }; + + let fs = FsFixture::default(); + let path = fs.tmp_path; + + let mut writer = create_writer( + fs.object_store.create(&path).await.unwrap(), + lance_schema.clone(), + ConcreteFileVersion::V2_1, + options, + ) + .unwrap(); + + writer.write_batch(&batch).await.unwrap(); + let write_summary = writer.finish().await.unwrap(); + assert_eq!(write_summary.num_rows, 2); + assert_eq!( + write_summary.size_bytes, + fs.object_store.size(&path).await.unwrap() + ); + + // Read back with split configuration + let file_scheduler = fs + .scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + + // Configure reader to split pages larger than 10MB into chunks + let reader_options = FileReaderOptions { + read_chunk_size: 10 * 1024 * 1024, // 10MB chunks + ..Default::default() + }; + + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + reader_options, + ) + .await + .unwrap(); + + // Read the data back + let stream = file_reader + .read_stream( + ReadBatchParams::RangeFull, + 1024, + 10, // batch_readahead + FilterExpression::no_filter(), + ) + .await + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + + // Verify the data is correctly read despite splitting + let read_array = batches[0].column(0); + let read_binary = read_array + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(read_binary.len(), 2); + assert_eq!(read_binary.value(0).len(), 40 * 1024 * 1024); + assert_eq!(read_binary.value(1), b"small value"); + + // Verify first value matches what we wrote + assert!(read_binary.value(0).iter().all(|&b| b == 42u8)); + } + + fn spill_config() -> (TempObjFile, Arc) { + let spill_path = TempObjFile::default(); + (spill_path, Arc::new(ObjectStore::local())) + } + + fn make_batches(num_batches: i32, num_cols: usize, rows_per_batch: i32) -> Vec { + let fields: Vec<_> = (0..num_cols) + .map(|c| ArrowField::new(format!("c{c}"), DataType::Int32, false)) + .collect(); + let schema = Arc::new(ArrowSchema::new(fields)); + (0..num_batches) + .map(|i| { + let cols: Vec> = (0..num_cols) + .map(|c| { + let start = (i * rows_per_batch + c as i32) * 100; + Arc::new(Int32Array::from_iter_values(start..start + rows_per_batch)) + as Arc + }) + .collect(); + RecordBatch::try_new(schema.clone(), cols).unwrap() + }) + .collect() + } + + async fn write_and_read_batches( + batches: &[RecordBatch], + spill: Option<(Arc, object_store::path::Path)>, + ) -> Vec { + let fs = FsFixture::default(); + let lance_schema = LanceSchema::try_from(batches[0].schema().as_ref()).unwrap(); + let writer = fs.object_store.create(&fs.tmp_path).await.unwrap(); + let mut file_writer = create_writer( + writer, + lance_schema, + ConcreteFileVersion::V2_1, + FileWriterOptions::default(), + ) + .unwrap(); + if let Some((store, path)) = spill { + file_writer = file_writer.with_page_metadata_spill(store, path); + } + for batch in batches { + file_writer.write_batch(batch).await.unwrap(); + } + file_writer.add_schema_metadata("foo", "bar"); + file_writer.finish().await.unwrap(); + + crate::testing::read_lance_file( + &fs, + Arc::::default(), + lance_encoding::decoder::FilterExpression::no_filter(), + ) + .await + } + + #[rstest::rstest] + #[case::multi_col(20, 2, 100)] + #[case::many_batches(50, 2, 100)] + #[tokio::test] + async fn test_page_metadata_spill_roundtrip( + #[case] num_batches: i32, + #[case] num_cols: usize, + #[case] rows_per_batch: i32, + ) { + let batches = make_batches(num_batches, num_cols, rows_per_batch); + let baseline = write_and_read_batches(&batches, None).await; + let (spill_path, spill_store) = spill_config(); + let spilled = + write_and_read_batches(&batches, Some((spill_store, spill_path.as_ref().clone()))) + .await; + assert_eq!(baseline, spilled); + } + + #[tokio::test] + async fn test_page_metadata_spill_many_columns() { + // Many columns forces small per-column buffer limits, exercising mid-write flushing. + let batches = make_batches(10, 500, 100); + let baseline = write_and_read_batches(&batches, None).await; + let (spill_path, spill_store) = spill_config(); + let spilled = + write_and_read_batches(&batches, Some((spill_store, spill_path.as_ref().clone()))) + .await; + assert_eq!(baseline, spilled); + } +} diff --git a/rust/lance-file/test_data/exact_versions/README.md b/rust/lance-file/test_data/exact_versions/README.md index 99438d91ca9..fd63b1284ab 100644 --- a/rust/lance-file/test_data/exact_versions/README.md +++ b/rust/lance-file/test_data/exact_versions/README.md @@ -2,9 +2,9 @@ These files were generated with the writers at baseline commit `3a72f8a61e14613f517dded6816d4bfc77817c93`. The deterministic input batch is -defined by `compatibility_fixture_batch` in `src/compatibility_tests.rs` and -covers primitive, nullable UTF-8, nullable list, nullable dictionary, blob, -multiple input batches, and multiple pages. +defined by `compatibility_fixture_batch` in `src/writer_tests.rs` and covers +primitive, nullable UTF-8, nullable list, nullable dictionary, blob, multiple +input batches, and multiple pages. The baseline generator was run twice in separate processes with a Cargo target directory isolated from the refactored checkout. Both runs produced identical @@ -17,10 +17,10 @@ bytes: | `v2_1.lance` | `3af97ba176b72c7e00a248b4a270a53402a72e594631950f76eb3daab45c50ce` | | `v2_2.lance` | `8298cd9301e657417b0725461345c27cf46515529d2a8b35824be139e3466a14` | -The compatibility tests require each stable writer to reproduce its fixture -byte-for-byte and each reader to open and read the baseline file. V2.3 is -unstable, so it has deterministic current-revision tests instead of a -checked-in compatibility fixture. +The compatibility tests require each refactored stable writer to reproduce its +fixture byte-for-byte and each refactored reader to open and read the baseline +file. V2.3 is unstable, so it has deterministic current-revision tests instead +of a checked-in compatibility fixture. Regenerate these fixtures only from the baseline writer APIs. Files generated with the implementation under test are not independent compatibility evidence. diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index b0b2b9c17c0..5efa4aa54fe 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -13,7 +13,7 @@ use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result, cache::LanceCache}; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; use lance_encoding::version::LanceFileVersion; -use lance_file::reader::{FileReader as CurrentFileReader, FileReaderOptions, ReaderProjection}; +use lance_file::reader::{FileReader as CurrentFileReader, FileReaderOptions}; use lance_file::version::ConcreteFileVersion; use lance_file::versions; use lance_file::versions::v1::reader::FileReader as V1FileReader; @@ -232,13 +232,16 @@ impl IndexReader for CurrentIndexReader { ))); } let projection = if let Some(projection) = projection { - ReaderProjection::from_column_names( + versions::reader_projection_from_column_names( self.0.metadata().version(), self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) + versions::reader_projection_from_whole_schema( + self.0.schema(), + self.0.metadata().version(), + ) }; let batches = self .0 @@ -270,13 +273,16 @@ impl IndexReader for CurrentIndexReader { return empty_batch(); } let projection = if let Some(projection) = projection { - ReaderProjection::from_column_names( + versions::reader_projection_from_column_names( self.0.metadata().version(), self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) + versions::reader_projection_from_whole_schema( + self.0.schema(), + self.0.metadata().version(), + ) }; // `DecodeBatchScheduler::schedule_ranges` requires sorted, // non-overlapping ranges; sort internally and permute the @@ -348,13 +354,16 @@ impl IndexReader for CurrentIndexReader { ))); } let projection = if let Some(projection) = projection { - ReaderProjection::from_column_names( + versions::reader_projection_from_column_names( self.0.metadata().version(), self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) + versions::reader_projection_from_whole_schema( + self.0.schema(), + self.0.metadata().version(), + ) }; self.0 .read_stream_projected( diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 164fbd31268..9aecf2cf226 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -840,7 +840,7 @@ pub async fn merge_partial_vector_auxiliary_files( // Inherit format version from the first shard file if format_version.is_none() { - format_version = Some(meta.version()); + format_version = Some(meta.version().into()); } // Read distance type diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 16da2b25bf6..160923fd04a 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -36,13 +36,13 @@ use lance_core::{ use lance_datafusion::utils::StreamingWriteSource; use lance_encoding::decoder::DecoderPlugins; use lance_file::reader::{ - CachedFileMetadata, FileMetadataIndex, FileReaderOptions, ProjectedFileReader, ReaderProjection, + CachedFileMetadata, FileMetadataIndex, FileReaderOptions, ProjectedFileReader, }; -use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; +use lance_file::version::ConcreteFileVersion; use lance_file::versions::v1::reader::{ FileReader as V1FileReader, read_batch as previous_read_batch, }; -use lance_file::{LanceEncodingsIo, determine_file_version}; +use lance_file::{LanceEncodingsIo, determine_file_version, versions as file_versions}; use lance_io::ReadBatchParams; use lance_io::scheduler::{FileScheduler, ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; @@ -368,7 +368,7 @@ mod v2_adapter { projection: Arc, ) -> BoxFuture<'_, Result> { async move { - let projection = ReaderProjection::from_field_ids( + let projection = file_versions::reader_projection_from_field_ids( self.reader.version(), projection.as_ref(), self.field_id_to_column_idx.as_ref(), @@ -398,7 +398,7 @@ mod v2_adapter { projection: Arc, ) -> BoxFuture<'_, Result> { async move { - let projection = ReaderProjection::from_field_ids( + let projection = file_versions::reader_projection_from_field_ids( self.reader.version(), projection.as_ref(), self.field_id_to_column_idx.as_ref(), @@ -427,7 +427,7 @@ mod v2_adapter { projection: Arc, ) -> BoxFuture<'_, Result> { async move { - let projection = ReaderProjection::from_field_ids( + let projection = file_versions::reader_projection_from_field_ids( self.reader.version(), projection.as_ref(), self.field_id_to_column_idx.as_ref(), @@ -459,7 +459,7 @@ mod v2_adapter { ) -> BoxFuture<'_, Result> { let indices = UInt32Array::from(indices.to_vec()); async move { - let projection = ReaderProjection::from_field_ids( + let projection = file_versions::reader_projection_from_field_ids( self.reader.version(), projection.as_ref(), self.field_id_to_column_idx.as_ref(), @@ -781,7 +781,9 @@ impl FileFragment { let file_version = determine_file_version(dataset.object_store.as_ref(), &filepath, None).await?; - if file_version != dataset.manifest.data_storage_format.lance_file_version()? { + if file_version + != ConcreteFileVersion::from(dataset.manifest.data_storage_format.lance_file_version()?) + { return Err(Error::invalid_input(format!( "File version mismatch. Dataset version: {:?} Fragment version: {:?}", dataset.manifest.data_storage_format.lance_file_version()?, @@ -789,7 +791,7 @@ impl FileFragment { ))); } - if file_version == LanceFileVersion::Legacy { + if file_version == ConcreteFileVersion::V1 { let fragment = Fragment::with_file_legacy( fragment_id as u64, filename, @@ -820,7 +822,7 @@ impl FileFragment { reader .schema() .check_compatible(dataset.schema(), &SchemaCompareOptions::default())?; - let projection = lance_file::reader::ReaderProjection::from_whole_schema( + let projection = file_versions::reader_projection_from_whole_schema( dataset.schema(), reader.metadata().version(), ); @@ -838,7 +840,7 @@ impl FileFragment { filename, dataset.schema().field_ids(), column_indices, - ConcreteFileVersion::from(file_version), + file_version, None, ); Ok(frag) @@ -985,23 +987,6 @@ impl FileFragment { data_file.fields.first().copied().unwrap_or(0) as u32 } - fn should_try_indexed_metadata( - data_file: &DataFile, - projection: &ReaderProjection, - file_version: LanceFileVersion, - ) -> bool { - if !ProjectedFileReader::supports_projection(projection, file_version) { - return false; - } - let total_columns = data_file - .column_indices - .iter() - .filter(|column_index| **column_index >= 0) - .count(); - let selected_columns = projection.column_indices.len(); - selected_columns.saturating_mul(4) < total_columns - } - pub(super) async fn open_reader( &self, data_file: &DataFile, @@ -1126,8 +1111,8 @@ impl FileFragment { } }), )); - let file_version: LanceFileVersion = data_file.file_version()?.into(); - let reader_projection = ReaderProjection::from_field_ids( + let file_version = data_file.file_version()?; + let reader_projection = file_versions::reader_projection_from_field_ids( file_version, schema_per_file.as_ref(), field_id_to_column_idx.as_ref(), @@ -1137,58 +1122,62 @@ impl FileFragment { .clone() .or_else(|| self.dataset.file_reader_options.clone()) .unwrap_or_default(); - let metadata_index = if metadata_mode == MetadataMode::LazyAllowed - && Self::should_try_indexed_metadata( - data_file, - &reader_projection, - file_version, - ) { - let known_schema = self - .metadata - .physical_rows - .map(|num_rows| (data_file_schema.clone(), num_rows as u64)); - let metadata_index = self - .get_file_metadata_index(&file_scheduler, known_schema) - .await?; - if (reader_projection.column_indices.len() as u32).saturating_mul(4) - < metadata_index.num_columns() - { - Some(metadata_index) - } else { - None - } - } else { - None - }; - + 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(); + let known_schema = self + .metadata + .physical_rows + .map(|num_rows| (data_file_schema.clone(), num_rows as u64)); let encodings_io = Arc::new( LanceEncodingsIo::new(file_scheduler.clone()) .with_read_chunk_size(file_reader_options.read_chunk_size), ); - let reader = if let Some(metadata_index) = metadata_index { - ProjectedFileReader::try_open_with_metadata_index( - encodings_io.clone(), - path.clone(), - Some(reader_projection.clone()), - Arc::::default(), - metadata_index, - &metadata_cache, - file_reader_options.clone(), - ) - .await? - } else { - let file_metadata = self.get_file_metadata(&file_scheduler).await?; - ProjectedFileReader::try_open_with_file_metadata( - encodings_io, - path.clone(), - None, - Arc::::default(), - file_metadata, - &metadata_cache, - file_reader_options, - ) - .await? - }; + let reader = file_versions::open_projected_reader( + file_version, + &reader_projection, + prefer_indexed, + || async { + let metadata_index = self + .get_file_metadata_index(&file_scheduler, known_schema) + .await?; + if (reader_projection.column_indices.len() as u32).saturating_mul(4) + >= metadata_index.num_columns() + { + return Ok(None); + } + Ok(Some( + ProjectedFileReader::try_open_with_metadata_index( + encodings_io.clone(), + path.clone(), + Some(reader_projection.clone()), + Arc::::default(), + metadata_index, + &metadata_cache, + file_reader_options.clone(), + ) + .await?, + )) + }, + || async { + let file_metadata = self.get_file_metadata(&file_scheduler).await?; + ProjectedFileReader::try_open_with_file_metadata( + encodings_io.clone(), + path.clone(), + None, + Arc::::default(), + file_metadata, + &metadata_cache, + file_reader_options.clone(), + ) + .await + }, + ) + .await?; let reader = v2_adapter::Reader::new( Arc::new(reader), schema_per_file, @@ -6047,54 +6036,6 @@ mod tests { ); } - #[test] - fn test_indexed_metadata_heuristic_counts_selected_physical_columns() { - let schema = Schema::try_from(&ArrowSchema::new(vec![ - ArrowField::new( - "s", - DataType::Struct( - vec![ - ArrowField::new("x", DataType::Int32, true), - ArrowField::new("y", DataType::Int32, true), - ] - .into(), - ), - true, - ), - ArrowField::new("a", DataType::Int32, true), - ArrowField::new("b", DataType::Int32, true), - ArrowField::new("c", DataType::Int32, true), - ])) - .unwrap(); - let data_file = DataFile { - path: "wide.lance".to_string(), - fields: Arc::from([0, 1, 2, 3, 4, 5]), - column_indices: Arc::from([-1, 0, 1, 2, 3, 4]), - file_major_version: 2, - file_minor_version: 1, - file_size_bytes: CachedFileSize::unknown(), - base_id: None, - }; - - let full_struct = - ReaderProjection::from_column_names(LanceFileVersion::V2_1, &schema, &["s"]).unwrap(); - assert_eq!(full_struct.column_indices.len(), 2); - assert!(!FileFragment::should_try_indexed_metadata( - &data_file, - &full_struct, - LanceFileVersion::V2_1 - )); - - let partial_struct = - ReaderProjection::from_column_names(LanceFileVersion::V2_1, &schema, &["s.x"]).unwrap(); - assert_eq!(partial_struct.column_indices.len(), 1); - assert!(FileFragment::should_try_indexed_metadata( - &data_file, - &partial_struct, - LanceFileVersion::V2_1 - )); - } - #[tokio::test] async fn test_iops_read_small() { // Create a file that has 8 columns. diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index b12a7198e5a..c3a1696014e 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -3936,7 +3936,7 @@ async fn test_index_inherits_dataset_file_version() { // Verify that the index file uses the same version as the dataset assert_eq!( index_reader.metadata().version(), - dataset_version, + dataset_version.into(), "Index file should use the same format version as the dataset" ); @@ -3965,7 +3965,7 @@ async fn test_index_inherits_dataset_file_version() { assert_eq!( aux_reader.metadata().version(), - dataset_version, + dataset_version.into(), "Auxiliary index file should use the same format version as the dataset" ); } @@ -4044,7 +4044,7 @@ async fn test_legacy_dataset_uses_v2_0_for_indexes() { // Verify that the index file uses V2_0 (not legacy) assert_eq!( index_reader.metadata().version(), - LanceFileVersion::V2_0, + LanceFileVersion::V2_0.into(), "Index files should never use legacy format, even for legacy datasets" ); } diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 8c8b62f3858..88c20515981 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -2536,7 +2536,7 @@ async fn write_root_vector_index_from_auxiliary( let arrow_schema = HNSW::schema(); let schema = lance_core::datatypes::Schema::try_from(arrow_schema.as_ref())?; let mut v2_writer = lance_file::versions::create_writer( - lance_file::version::ConcreteFileVersion::from(format_version), + format_version, obj_writer, schema, V2WriterOptions::default(), diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 6b42896f9bd..d79b12fde6e 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -5301,7 +5301,7 @@ mod tests { // Rewrite auxiliary file with PQ codebook inlined into schema metadata. let mut metadata = reader.schema().metadata.clone(); - let projection = lance_file::reader::ReaderProjection::from_whole_schema( + let projection = lance_file::versions::reader_projection_from_whole_schema( reader.schema(), reader.metadata().version(), );